##########################################################################################
##############################  IMPORTING USEFUL DEFINITIONS  ############################
# standard libraries imports
import os, sys

# external libraries imports
import wx
import numpy as np

# fullrmc library imports
from fullrmc.Globals import LOGGER, maxint
from fullrmc.Engine import Engine
from fullrmc.Constraints.PairDistributionConstraints import PairDistributionConstraint
from fullrmc.Constraints.DistanceConstraints import InterMolecularDistanceConstraint
from fullrmc.Constraints.BondConstraints import BondConstraint
from fullrmc.Constraints.AngleConstraints import BondsAngleConstraint
from fullrmc.Constraints.ImproperAngleConstraints import ImproperAngleConstraint
from fullrmc.Generators.Translations import TranslationGenerator
from fullrmc.Selectors.RandomSelectors import RandomSelector, SmartRandomSelector


###########################################################################################
####################################  SHUT DOWN LOGGING  ##################################
#LOGGER.set_log_to_stdout_flag(False)
## set very high logging level and force 'move accepted'.
#LOGGER.set_minimum_level(maxint)
#LOGGER.force_log_type_flags("move accepted", stdoutFlag=False, fileFlag=True)
#
###########################################################################################
######################################  CREATE ENGINE  ####################################
## set log files name and turn off stout logging
normalSelLog = "normal"
MLSelLog     = "ML"

# runtime arguments
numberOfSteps = 200000

# set engine variables
pdbPath = "thf.pdb"

def create_engine():
    # create engine
    ENGINE = Engine(path=None)
    ENGINE.set_pdb(pdbPath)
    # initialize constraints
    B_CONSTRAINT   = BondConstraint()
    BA_CONSTRAINT  = BondsAngleConstraint()
    IA_CONSTRAINT  = ImproperAngleConstraint()
    # add constraints
    ENGINE.add_constraints([B_CONSTRAINT])
    B_CONSTRAINT.create_bonds_by_definition( bondsDefinition={"THF": [('O' ,'C1' , 1.20, 1.70),
                                                                      ('O' ,'C4' , 1.20, 1.70),
                                                                      ('C1','C2' , 1.25, 1.90),
                                                                      ('C2','C3' , 1.25, 1.90),
                                                                      ('C3','C4' , 1.25, 1.90),
                                                                      ('C1','H11', 0.88, 1.16),('C1','H12', 0.88, 1.16),
                                                                      ('C2','H21', 0.88, 1.16),('C2','H22', 0.88, 1.16),
                                                                      ('C3','H31', 0.88, 1.16),('C3','H32', 0.88, 1.16),
                                                                      ('C4','H41', 0.88, 1.16),('C4','H42', 0.88, 1.16)] })
    ENGINE.add_constraints([BA_CONSTRAINT])
    BA_CONSTRAINT.create_angles_by_definition( anglesDefinition={"THF": [ ('O'  ,'C1' ,'C4' , 105, 125),
                                                                          ('C1' ,'O'  ,'C2' , 100, 120),
                                                                          ('C4' ,'O'  ,'C3' , 100, 120),
                                                                          ('C2' ,'C1' ,'C3' , 95 , 115),
                                                                          ('C3' ,'C2' ,'C4' , 95 , 115),
                                                                          # H-C-H angle
                                                                          ('C1' ,'H11','H12', 98 , 118),
                                                                          ('C2' ,'H21','H22', 98 , 118),
                                                                          ('C3' ,'H31','H32', 98 , 118),
                                                                          ('C4' ,'H41','H42', 98 , 118),
                                                                          # H-C-O angle
                                                                          ('C1' ,'H11','O'  , 100, 120),
                                                                          ('C1' ,'H12','O'  , 100, 120),
                                                                          ('C4' ,'H41','O'  , 100, 120),
                                                                          ('C4' ,'H42','O'  , 100, 120),
                                                                          # H-C-C
                                                                          ('C1' ,'H11','C2' , 103, 123),
                                                                          ('C1' ,'H12','C2' , 103, 123),
                                                                          ('C2' ,'H21','C1' , 103, 123),
                                                                          ('C2' ,'H21','C3' , 103, 123),
                                                                          ('C2' ,'H22','C1' , 103, 123),
                                                                          ('C2' ,'H22','C3' , 103, 123),
                                                                          ('C3' ,'H31','C2' , 103, 123),
                                                                          ('C3' ,'H31','C4' , 103, 123),
                                                                          ('C3' ,'H32','C2' , 103, 123),
                                                                          ('C3' ,'H32','C4' , 103, 123),
                                                                          ('C4' ,'H41','C3' , 103, 123),
                                                                          ('C4' ,'H42','C3' , 103, 123) ] })
    ENGINE.add_constraints([IA_CONSTRAINT])
    IA_CONSTRAINT.create_angles_by_definition( anglesDefinition={"THF": [ ('C2','O','C1','C4', -15, 15),
                                                                          ('C3','O','C1','C4', -15, 15) ] })
    # initialize constraints data
    ENGINE.initialize_used_constraints()
    # set moves generators
    # set all move generators to Translation with a maximum amplutide of 0.3 A
    [g.set_move_generator(TranslationGenerator(amplitude=0.3)) for g in ENGINE.groups]
    # set randomly 25% of the translation amplitude to 10A
    [g.set_move_generator(TranslationGenerator(amplitude=10.)) for g in ENGINE.groups if np.random.random()>0.25]
    # return engine
    return ENGINE


##########################################################################################
#####################################  DIFFERENT RUNS  ###################################
#################### run normal selector ####################
def normal_run():
    LOGGER.force_log("info", "normal selection %i started... DON'T INTERRUPT"%numberOfSteps, stdout=True, file=False)
    # load fresh engine
    engine = create_engine()
    # delete existing log files
    normalLogs = [fn for fn in next(os.walk("."))[2] if ".log" in fn and normalSelLog in fn]
    [os.remove(l) for l in normalLogs]
    # set log file name
    LOGGER.set_log_file_basename(normalSelLog)
    # set random group selector
    engine.set_group_selector(RandomSelector(engine))
    # run engine
    engine.run(numberOfSteps=numberOfSteps, saveFrequency=2*numberOfSteps, restartPdb=None)
    LOGGER.force_log("info", "normal selection finished", stdout=True, file=False)

############### run machine learning selector ###############
def ML_run():
    LOGGER.force_log("info", "machine learning selection %i started... DON'T INTERRUPT"%numberOfSteps, stdout=True, file=False)
    # load fresh engine
    engine = create_engine()
    # delete existing log files
    MLLogs = [fn for fn in next(os.walk("."))[2] if ".log" in fn and MLSelLog in fn]
    [os.remove(l) for l in MLLogs]
    # set log file name
    LOGGER.set_log_file_basename(MLSelLog)
    # set smart group selector
    engine.set_group_selector(SmartRandomSelector(engine))
    # run engine
    engine.run(numberOfSteps=numberOfSteps, saveFrequency=2*numberOfSteps, restartPdb=None)
    LOGGER.force_log("info", "machine learning selection finished", stdout=True, file=False)

##########################################################################################
#####################################  RUN SIMULATION  ###################################
normal_run()
ML_run()
LOGGER.flush()

##########################################################################################
################################  READ LOGGING FILES DATA ################################
# machine learning logs
MLSelLog = sorted([fn for fn in next(os.walk("."))[2] if ".log" in fn and MLSelLog in fn])
mlGenerated = []
mlAccepted  = []
for log in MLSelLog:
    fd = open(log,'r')
    mlLines = fd.readlines()
    fd.close()
    mlGenerated.extend([float(l.split("Gen:")[1].split("-")[0]) for l in mlLines])
    mlAccepted.extend([float(l.split("Acc:")[1].split("(")[0]) for l in mlLines])
mlGenerated = np.array(mlGenerated)
mlAccepted = np.array(mlAccepted)
mlAccepted = 100.*mlAccepted/mlGenerated
np.savetxt(X=np.transpose([mlGenerated,mlAccepted]),
           fname="MLSelection.dat",
           fmt='%.3f', delimiter="   ",
           header="Generated    Accepted(%)")
# normal logs
normalLogs = sorted([fn for fn in next(os.walk("."))[2] if ".log" in fn and normalSelLog in fn])
nGenerated = []
nAccepted  = []
for log in normalLogs:
    fd = open(log,'r')
    mlLines = fd.readlines()
    fd.close()
    nGenerated.extend([float(l.split("Gen:")[1].split("-")[0]) for l in mlLines])
    nAccepted.extend([float(l.split("Acc:")[1].split("(")[0]) for l in mlLines])
nGenerated = np.array(nGenerated)
nAccepted = np.array(nAccepted)
nAccepted = 100.*nAccepted/nGenerated
np.savetxt(X=np.transpose([nGenerated,nAccepted]),
           fname="traditionalSelection.dat",
           fmt='%.3f', delimiter="   ",
           header="Generated    Accepted(%)")


##########################################################################################
##################################  PLOT LOGGING DATA ####################################
import matplotlib.pyplot as plt
plt.plot(mlGenerated, mlAccepted, 'black',linewidth=3, label="machine learning selection")
plt.plot(nGenerated, nAccepted, 'red', linewidth=3, label="traditional selection")
plt.xlabel("Generated moves")
plt.ylabel("Accepted moves (%)")
plt.legend(frameon=False, loc="upper left")
plt.show()
