import tensorflow as tf
import tensorflow.keras as tk
import nthmc

conf = nthmc.Conf(nbatch=2048, nepoch=1, nstepEpoch=256, nstepMixing=64, initDt=0.4, refreshOpt=False, nthr=8)
beta=2.25
nthmc.setup(conf)

action = nthmc.OneD(beta=beta, transform=nthmc.Ident())
loss = nthmc.LossFun(action, cCosDiff=1.0, cTopoDiff=1.0, dHmin=0.5, topoFourierN=1)
x0 = action.initState(conf.nbatch)
weights=list(map(lambda x:tf.constant(x,dtype=tf.float64),
  [0.34406691292594005,
   beta]))
tf.print('-------- pretrain mixing --------')
x0 = nthmc.runInfer(conf, action, loss, weights, x0, detail=False)
tf.print('-------- done pretrain mixing --------')

conf.nstepEpoch=4096
conf.nstepMixing=0
action = nthmc.OneD(beta=beta, transform=nthmc.TransformChain([
  nthmc.OneDNeighbor(mask='even'), nthmc.OneDNeighbor(mask='odd'),
  nthmc.OneDNeighbor(mask='even',distance=2), nthmc.OneDNeighbor(mask='odd',distance=2),
  nthmc.OneDNeighbor(mask='even',distance=4), nthmc.OneDNeighbor(mask='odd',distance=4),
  nthmc.OneDNeighbor(mask='even',distance=8), nthmc.OneDNeighbor(mask='odd',distance=8),
  nthmc.OneDNeighbor(mask='even',distance=16), nthmc.OneDNeighbor(mask='odd',distance=16),
  nthmc.OneDNeighbor(mask='even',distance=32), nthmc.OneDNeighbor(mask='odd',distance=32),
  nthmc.OneDNeighbor(mask='even',order=2), nthmc.OneDNeighbor(mask='odd',order=2),
  nthmc.OneDNeighbor(mask='even',order=2,distance=2), nthmc.OneDNeighbor(mask='odd',order=2,distance=2),
  nthmc.OneDNeighbor(mask='even',order=2,distance=4), nthmc.OneDNeighbor(mask='odd',order=2,distance=4),
  nthmc.OneDNeighbor(mask='even',order=2,distance=8), nthmc.OneDNeighbor(mask='odd',order=2,distance=8),
  nthmc.OneDNeighbor(mask='even',order=2,distance=16), nthmc.OneDNeighbor(mask='odd',order=2,distance=16),
  nthmc.OneDNeighbor(mask='even',order=2,distance=32), nthmc.OneDNeighbor(mask='odd',order=2,distance=32),
  nthmc.OneDNeighbor(mask='even',order=3), nthmc.OneDNeighbor(mask='odd',order=3),
  nthmc.OneDNeighbor(mask='even',order=3,distance=2), nthmc.OneDNeighbor(mask='odd',order=3,distance=2),
  nthmc.OneDNeighbor(mask='even',order=3,distance=4), nthmc.OneDNeighbor(mask='odd',order=3,distance=4),
  nthmc.OneDNeighbor(mask='even',order=3,distance=8), nthmc.OneDNeighbor(mask='odd',order=3,distance=8),
  nthmc.OneDNeighbor(mask='even',order=3,distance=16), nthmc.OneDNeighbor(mask='odd',order=3,distance=16),
  nthmc.OneDNeighbor(mask='even',order=3,distance=32), nthmc.OneDNeighbor(mask='odd',order=3,distance=32),
  nthmc.OneDNeighbor(mask='even',order=4), nthmc.OneDNeighbor(mask='odd',order=4),
  nthmc.OneDNeighbor(mask='even',order=4,distance=2), nthmc.OneDNeighbor(mask='odd',order=4,distance=2),
  nthmc.OneDNeighbor(mask='even',order=4,distance=4), nthmc.OneDNeighbor(mask='odd',order=4,distance=4),
  nthmc.OneDNeighbor(mask='even',order=4,distance=8), nthmc.OneDNeighbor(mask='odd',order=4,distance=8),
  nthmc.OneDNeighbor(mask='even',order=4,distance=16), nthmc.OneDNeighbor(mask='odd',order=4,distance=16),
  nthmc.OneDNeighbor(mask='even',order=4,distance=32), nthmc.OneDNeighbor(mask='odd',order=4,distance=32),
  nthmc.OneDNeighbor(mask='even'), nthmc.OneDNeighbor(mask='odd'),
  nthmc.OneDNeighbor(mask='even',distance=2), nthmc.OneDNeighbor(mask='odd',distance=2),
  nthmc.OneDNeighbor(mask='even',distance=4), nthmc.OneDNeighbor(mask='odd',distance=4),
  nthmc.OneDNeighbor(mask='even',distance=8), nthmc.OneDNeighbor(mask='odd',distance=8),
  nthmc.OneDNeighbor(mask='even',distance=16), nthmc.OneDNeighbor(mask='odd',distance=16),
  nthmc.OneDNeighbor(mask='even',distance=32), nthmc.OneDNeighbor(mask='odd',distance=32),
  nthmc.OneDNeighbor(mask='even',order=2), nthmc.OneDNeighbor(mask='odd',order=2),
  nthmc.OneDNeighbor(mask='even',order=2,distance=2), nthmc.OneDNeighbor(mask='odd',order=2,distance=2),
  nthmc.OneDNeighbor(mask='even',order=2,distance=4), nthmc.OneDNeighbor(mask='odd',order=2,distance=4),
  nthmc.OneDNeighbor(mask='even',order=2,distance=8), nthmc.OneDNeighbor(mask='odd',order=2,distance=8),
  nthmc.OneDNeighbor(mask='even',order=2,distance=16), nthmc.OneDNeighbor(mask='odd',order=2,distance=16),
  nthmc.OneDNeighbor(mask='even',order=2,distance=32), nthmc.OneDNeighbor(mask='odd',order=2,distance=32),
  nthmc.OneDNeighbor(mask='even',order=3), nthmc.OneDNeighbor(mask='odd',order=3),
  nthmc.OneDNeighbor(mask='even',order=3,distance=2), nthmc.OneDNeighbor(mask='odd',order=3,distance=2),
  nthmc.OneDNeighbor(mask='even',order=3,distance=4), nthmc.OneDNeighbor(mask='odd',order=3,distance=4),
  nthmc.OneDNeighbor(mask='even',order=3,distance=8), nthmc.OneDNeighbor(mask='odd',order=3,distance=8),
  nthmc.OneDNeighbor(mask='even',order=3,distance=16), nthmc.OneDNeighbor(mask='odd',order=3,distance=16),
  nthmc.OneDNeighbor(mask='even',order=3,distance=32), nthmc.OneDNeighbor(mask='odd',order=3,distance=32),
  nthmc.OneDNeighbor(mask='even',order=4), nthmc.OneDNeighbor(mask='odd',order=4),
  nthmc.OneDNeighbor(mask='even',order=4,distance=2), nthmc.OneDNeighbor(mask='odd',order=4,distance=2),
  nthmc.OneDNeighbor(mask='even',order=4,distance=4), nthmc.OneDNeighbor(mask='odd',order=4,distance=4),
  nthmc.OneDNeighbor(mask='even',order=4,distance=8), nthmc.OneDNeighbor(mask='odd',order=4,distance=8),
  nthmc.OneDNeighbor(mask='even',order=4,distance=16), nthmc.OneDNeighbor(mask='odd',order=4,distance=16),
  nthmc.OneDNeighbor(mask='even',order=4,distance=32), nthmc.OneDNeighbor(mask='odd',order=4,distance=32),
]))
loss = nthmc.LossFun(action, cCosDiff=1.0, cTopoDiff=1.0, dHmin=0.5, topoFourierN=1)
opt = tk.optimizers.Adam(learning_rate=1E-4)

# from the end of t16.log
weights=list(map(lambda x:tf.constant(x,dtype=tf.float64), [
 0.27467031569717887,
 -0.35079525228697017,
 -0.3463004172411952,
 0.08642702199768838,
 0.09291043721720227,
 0.2250623460821537,
 0.21856133022867158,
 0.0866841925841807,
 0.0915280114570428,
 0.004317996316240299,
 -0.0020768690146567973,
 -0.0026930881470839125,
 0.0008428096389416606,
 0.20576432902458863,
 0.20387440184716182,
 0.09477415494278256,
 0.09278395006364205,
 0.027261343846169267,
 0.025317061861749923,
 0.0025558109447300975,
 0.00011587519368271666,
 0.00020262041594415814,
 0.0001471120215669548,
 -0.0011822219968058566,
 -0.0005116630631441119,
 -0.04059539696342997,
 -0.03967354697693106,
 -0.022191054178163458,
 -0.02039186123588369,
 -0.013214337666031521,
 -0.013721250067491505,
 0.0013132162229792892,
 -0.0008986215564624871,
 -0.0019147647080815012,
 -0.002911776022339967,
 0.001026344348351285,
 0.0007105730323132924,
 0.00707242379752812,
 0.0007786031681717592,
 -0.0018400361199394529,
 -0.00322085308073767,
 0.0014898401798478214,
 0.0017817604637083498,
 0.000513064704705146,
 -0.0005834196597080542,
 0.0001492800189452559,
 0.0003836187479534708,
 0.00012953683316717222,
 0.00010091230153452416,
 -0.33277841543664677,
 -0.32929846989472394,
 0.10552781165238931,
 0.11085623521274739,
 0.22205804115732578,
 0.2157969787561753,
 0.07917064282826275,
 0.08338665473737698,
 0.004159324147506869,
 -0.0015816850023714166,
 -0.0037297327525413182,
 0.001017107822169157,
 0.1887051260407191,
 0.18465763015939937,
 0.1415955662617319,
 0.13940160433726745,
 0.027124643533201886,
 0.025064491731788,
 0.0010426744501802643,
 -0.0009217291058926577,
 0.00021604526039150265,
 0.00034203491430328755,
 -2.7916500875223948e-05,
 -0.0006780944446978524,
 -0.03383972393252698,
 -0.031224928767883867,
 -0.035394987598204995,
 -0.03599968661932616,
 -0.011722749416250355,
 -0.011666885540960657,
 0.0009174390714426379,
 -0.00044215695287379693,
 -0.0005842603445760292,
 -0.002572194452607748,
 0.0011823673303244884,
 0.0013164376122312897,
 0.010669744620610579,
 0.006033375801065042,
 0.007044829792370666,
 0.0025156648400946253,
 0.0020781739549382288,
 0.003069839401052712,
 -0.00010264771919958654,
 0.00031438332504875007,
 -0.0007835652977912688,
 0.0007698351478308178,
 -0.000603759792144621,
 6.098626973537799e-05,
 beta]))
nthmc.run(conf, action, loss, opt, x0, weights=weights, requireInv=True)
