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('{}/salt_Uext0.0_mapped.lammpstrj.gz'.format(datadir))
traj = pickleTraj(traj_file)
box = traj.FrameData['BoxL']

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

options0 = yaml.load('./setup/settings.yaml')
options0['box'] = box
ffdef0 = forcefield.ForceField('./setup/ff.yaml')

sysname_prefix = options0['sys_name']
options0['neutralize'] = True
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)

options0['neutralize'] = False
options0['sys_name'] = sysname_prefix + '_elec'
ElecSys = export_sim.create_system(top, options=options0)
export_sim.set_system_settings( ElecSys, options0 )
Elecffs = export_sim.create_forcefield( ElecSys,top,ffdef0, options=options0 )
ElecSys.ForceField.extend(Elecffs)

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=ElecSys)
Opt0.MinReweightFrac = 0.15

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')