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=4E-5)

# from the end of t17.log
weights=list(map(lambda x:tf.constant(x,dtype=tf.float64), [
 0.20168195491865906,
 -0.528247418923954,
 -0.49710260600030143,
 0.0945204044191504,
 0.09317345752005862,
 0.373457812709759,
 0.38402251350965216,
 0.20330475108684667,
 0.1974260639550408,
 -0.0005484199341088154,
 -0.007712243762308879,
 0.007018458862729426,
 0.003181895225064764,
 0.35444648930771805,
 0.3544879387200731,
 0.12378436024017102,
 0.11768257436613305,
 0.007462452015615455,
 0.0011981519562075723,
 -0.023990092220123837,
 -0.033099400535327644,
 -0.0019624173352754083,
 0.0036939791353533275,
 -0.0024644617695774905,
 0.0071291330524728466,
 -0.04701458995544631,
 -0.052234734106293576,
 -0.07821405007781404,
 -0.064049755998159,
 -0.05002797596863401,
 -0.05627941110521701,
 0.0024106599285186455,
 -0.001168529157892512,
 0.00020135282908499303,
 -0.0036232976564688764,
 0.0017652237101491006,
 -0.002511445933480341,
 0.020308702526657518,
 -0.009471319288351265,
 -0.010352499498503242,
 -0.0071223878673138084,
 0.01600378910175528,
 0.010316603088210611,
 0.0015558543455745644,
 -0.00023526136219970906,
 -0.00024308294962199122,
 0.001608028095559756,
 0.0009177938211149464,
 0.000639715577251889,
 -0.49088334452703275,
 -0.46590247383859384,
 0.14155475150864258,
 0.14151535273027077,
 0.3808405252930217,
 0.37305684927596355,
 0.17817079084556417,
 0.17447603207198292,
 -0.005198573982589358,
 -0.004969897077922295,
 0.005212867786665684,
 0.0018505675821249507,
 0.32379567176490287,
 0.31831869615920394,
 0.2671422622608067,
 0.2322044600676465,
 0.04735447593433718,
 0.03762883031201281,
 -0.017807762577113084,
 -0.027124222209042884,
 -0.0005900092798843302,
 0.003345612962214431,
 -0.002918904020330505,
 0.0008924424538644182,
 -0.06083297831111033,
 -0.05495779335480908,
 -0.040323259900219066,
 -0.04200241504776445,
 -0.01955375639434779,
 -0.019309803206697736,
 0.004732135309307072,
 0.0060610832202624716,
 0.0018258852508407383,
 1.1690984178396865e-05,
 0.0033162976621228814,
 0.0029257798842658484,
 0.02552174055222456,
 0.012513577756078278,
 0.014663805983868341,
 0.002739948651423212,
 0.004411251399144746,
 0.005609517242742677,
 -0.0005191857670020354,
 -0.002363055270551431,
 0.0009098290717334516,
 0.0009092985346618658,
 0.001475716088810317,
 -0.0005894767761914989,
 beta]))
nthmc.run(conf, action, loss, opt, x0, weights=weights, requireInv=True)
