import sys, os, copy
print('python path:\n    {}'.format(sys.path))
import sim, pickleTraj
import subprocess
from datetime import datetime
#import my special md-manager utilities
from utility import *
print('... execution date: {}'.format(datetime.now()))
print(context.get_devinfo(sim))
print('... cwd:\t{}'.format(os.getcwd()))

import numpy as np
from collections import OrderedDict
from collections import namedtuple
units = sim.units.DimensionlessUnits

sim.export.omm.platformName = 'OpenCL'
sim.export.omm.device = 0 #-1 is default, let openmm choose its own platform. otherwise is GPU device #
sim.export.omm.NPairPotentialKnots = 500 #number of points used to spline-interpolate the potential
sim.export.omm.InnerCutoff = 0.001 #0.001 is default. Note that a small value is not necessary, like in the lammps export, because the omm export interpolates to zero
sim.srel.optimizetrajomm.OpenMMStepsMin = 0 #number of steps to minimize structure, 0 is default
sim.srel.optimizetrajomm.OpenMMDelTempFiles = False #False is Default
sim.export.omm.UseTabulated = True
sim.srel.base.n_process = 1
#print('NUM PROCESSES: {}'.format(sim.srel.base.n_process))

sim.export.omm.VVerbose = False
sim.srel.base.VVerbose = False
sim.srel.penalty.VVerbose = False

UseWPenalty = False
StageCoefs = []
MaxIter=None
SteepestIter=0

datadir='./'
forcefield_file = None #initial force field file
traj_file = os.path.abspath('{}/dodecane-water_mapped.lammpstrj.gz'.format(datadir))
traj = pickleTraj(traj_file)
box = traj.FrameData['BoxL']

print('===== Defining system and simulation =====')
top = topologify.Topology('./setup/top_dodecane-water.yaml')

options0 = yaml.load('./setup/settings_dodecane-water.yaml')
options0['box'] = box
ffdef0 = forcefield.ForceField('./setup/ff_dodecane-water.yaml')

sysname_prefix = options0['sys_name']
Sys0 = export_sim.create_system(top, options=options0)
export_sim.set_system_settings( Sys0, options0 )
ffs0 = export_sim.create_forcefield( Sys0,top,ffdef0, options=options0 )
Sys0.ForceField.extend(ffs0)

export_sim.load_system(Sys0, forcefield_file)

Opt0 = optimizer.create_optimizer(Sys0, traj, md_engine=options0['run']['md_engine'],
    steps_equil=options0['run']['steps_equil'],
    steps_prod=options0['run']['steps_prod'],
    steps_stride=options0['run']['steps_stride'],
    ElecSys=None)
Opt0.MinReweightFrac = 0.15

print('=== Adding measures ===')
Dist2 = sim.measure.distang.Distance2(Sys0,sim.atomselect.PolyFilter(Filters=[Sys0.World.SiteTypes[0],Sys0.World.SiteTypes[5]],Intra=True), Name='tail_ree2')
Sys0.Measures.append(Dist2)
Dist2_target = 1.01
sim.srel.penalty.VVerbose = False
StageCoefs = [100.0,1.0e3,1.0e4]
UDist = Opt0.AddPenalty(dist2, dist2_target, MeasureScale = 1., Coef=1.e-80)
print('Opt penalties: {}'.format(Opt0.Penalties))
Opt0.UpdateMode = 2 #0=both, 1=srel only, 2=bias only

print('===== RUNNING LJG SREL =====')

Opt0.RunConjugateGradient(MaxIter=MaxIter)
#sav/back up files and results to separate directory
dname='ljg0'
FilePrefix = Opt0.FilePrefix
TempFilePrefix = Opt0.TempFilePrefix
try:
  os.mkdir(dname)
  os.system('cp {}* {}'.format(FilePrefix,dname))
  os.system('cp {}* {}'.format(TempFilePrefix,dname))
except:
  print('file moving commands failed')