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.875
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.30703729455190604,
   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 training in t13.log
weights=list(map(lambda x:tf.constant(x,dtype=tf.float64),
 # f:y$@@:r!awk -v beta=2.875 '/^beta: /{b=$2} p>0{w=w "\n" $0} b==beta&&/^weights: /{p=1;w=$0} p==1&&/]$/{p=0} END{print w}' t13.log
 [0.45615090724163854,
 -0.31097787822669354,
 -0.30507920463515187,
 -0.027893016314395284,
 -0.031378845400177963,
 0.077689083215770949,
 0.075569715367494641,
 0.038699510620482935,
 0.029162385005325472,
 0.0019581497708284694,
 -0.0018231287462758918,
 0.00015888456785728626,
 -0.0028210982286725086,
 0.13124240382350402,
 0.13309785933956725,
 0.017604137564691036,
 0.010907674928860149,
 -0.013780037257168396,
 -0.022445109691812258,
 -0.0045229710423886765,
 -0.0029058196749805151,
 0.0023048449953337728,
 -0.0070235509174246284,
 -0.0014313775421141036,
 0.00081176147554258083,
 -0.014710030999330952,
 -0.010194100966722035,
 0.002744086282626448,
 0.0045756447355585093,
 0.0031292945016411365,
 0.0031592597427928843,
 0.00053880411453796249,
 -0.00058044090213579173,
 0.00095364836258577637,
 -0.0028807214952762316,
 0.0018107008839567691,
 -0.0013583732862177305,
 0.0046931380657292757,
 0.0016671741461710527,
 -0.0031238965035703696,
 -0.0030495300374729362,
 3.7767171335432319e-05,
 0.00034506965785394356,
 -9.8650513910624843e-05,
 0.00084275179037986137,
 0.0012699466261455849,
 0.0012800734726210016,
 0.00078495081260056656,
 -3.6750708339015154e-05,
 -0.31014396639255265,
 -0.3045858543098458,
 -0.010885776010155591,
 -0.015750481987926623,
 0.087259089367838744,
 0.08243283014988155,
 0.040517512492184569,
 0.030525468606565239,
 0.0025872352327758539,
 -0.0027206505719563493,
 -0.00089873373216705352,
 -0.0018318661211866342,
 0.0967308932840898,
 0.095883079309349514,
 0.047763637063773574,
 0.041546863771405255,
 -0.012530825072081196,
 -0.020478495148529022,
 -0.0067227151927674068,
 -0.0052179264725507176,
 0.00418665071041997,
 -0.00771130055753064,
 -0.0013408242290686503,
 0.00065100724836321812,
 -0.040842057940541958,
 -0.03514844539463631,
 -0.025181375323195351,
 -0.023134536637470358,
 0.00242366467545387,
 0.002806728633386199,
 0.00060494371667193494,
 -0.0040390056771061368,
 0.0011595645810642834,
 0.00015374946003506677,
 0.00012011293019308769,
 -0.0021145331363914585,
 0.016401183428638843,
 0.011602504263125767,
 0.0076990960462810717,
 0.0077484140578621538,
 1.1511413473662876e-05,
 0.0011462119410679498,
 -0.0011556563594443477,
 -0.00057730440795531726,
 -0.0018027637615355017,
 -0.0021347460580807263,
 0.00058925948384115634,
 -0.0010558414842687634,
 beta]))
nthmc.run(conf, action, loss, opt, x0, weights=weights, requireInv=True)
