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

conf = nthmc.Conf(nbatch=2048, nepoch=1, nstepEpoch=64, nstepMixing=64, initDt=0.4, refreshOpt=False, nthr=8)
beta=1.625
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.3737090001953306,
   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=5E-5)

weights=list(map(lambda x:tf.constant(x,dtype=tf.float64), [
 0.36690518142636847,
 -0.2350499626351434,
 -0.23236403789535023,
 0.039628683046730964,
 0.03955621731120355,
 0.12369040331513048,
 0.11711523730587103,
 0.04507532025789665,
 0.04748881864279912,
 0.005880398058223657,
 0.0013283811439883362,
 0.002950447074563798,
 0.0021982511255234077,
 0.1031983396449511,
 0.10277780224660524,
 0.05638098165459163,
 0.05210973590514252,
 0.005902264675774449,
 0.004238864534277979,
 -0.0004355354737754071,
 -0.00031552476038428486,
 0.0015308187003705831,
 0.0005955326790122433,
 -0.0004019655118871211,
 -0.0013044933130967665,
 -0.028915697148967284,
 -0.026239863446361707,
 -0.010452949800621177,
 -0.009696469967752343,
 -0.0014091052542447389,
 -0.0023235177106805424,
 -0.0001456935540350357,
 -1.6339233156804906e-05,
 -0.0012448222067823794,
 0.0003516721774415144,
 -1.87272967392963e-06,
 -0.00021449536182441552,
 0.007341464129446544,
 0.004777938841718168,
 -0.0013232431741844676,
 -0.0015235219520319287,
 0.0010002842985669354,
 0.0003972128665523457,
 0.00045491260290411083,
 0.00032914370170055774,
 -0.00044756655412584286,
 -0.00043505509262282055,
 -0.00017650833299237448,
 -0.0004297015820867862,
 -0.23051934077436112,
 -0.22760454747006717,
 0.04763876070940241,
 0.04671847081339848,
 0.12406463990455231,
 0.11734170732460082,
 0.04248090822795133,
 0.045992698671530205,
 0.005967444349188979,
 0.001642507281982355,
 0.002492935079953428,
 0.0022128069908070504,
 0.08955653724256905,
 0.0887033052197043,
 0.07369496477613433,
 0.06980509011643392,
 0.006055001807475078,
 0.004734065300928894,
 -0.0009436276700335873,
 -0.0013310770503804947,
 0.001598979539050544,
 0.0005766265500856805,
 -0.0004717235074499824,
 -0.0016059553608090908,
 -0.034101020249122906,
 -0.031221906572674797,
 -0.024075133220092467,
 -0.024467421219070647,
 -0.0025258204260913183,
 -0.0032701721818881088,
 -0.0008839831878879536,
 -0.0005257237023484202,
 -0.0011157254271515103,
 1.9712509381617236e-05,
 0.00014132614817972966,
 0.0004621553248295664,
 0.011159669698624244,
 0.008626380789992392,
 0.005283771575548108,
 0.005118675244144138,
 -3.619299983114965e-05,
 0.0003997980011772861,
 -0.00014591149180966093,
 9.919004256654938e-05,
 -0.00015596066385126177,
 -0.0005090096158742788,
 -0.0001347939094261849,
 0.00010623428278517314,
 beta]))
nthmc.run(conf, action, loss, opt, x0, weights=weights, requireInv=True)
