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

## external libraries imports
import numpy as np
from pdbparser.pdbparser import pdbparser
from pdbparser.Utilities.Modify import set_records_attribute_values

## fullrmc imports
from fullrmc.Engine import Engine
from fullrmc.Core.Group import EmptyGroup
from fullrmc.Constraints.PairDistributionConstraints import PairDistributionConstraint
from fullrmc.Constraints.DistanceConstraints import InterMolecularDistanceConstraint
from fullrmc.Generators.Translations import TranslationGenerator
from fullrmc.Generators.Swaps import SwapPositionsGenerator


#  ####################################################################################  #
#  ############################# DECLARE USEFUL VARIABLES #############################  #
# dirname
try:
    DIR_PATH = os.path.dirname( os.path.realpath(__file__) )
except:
    DIR_PATH = ''

# files name
grFileName     = "pdf.exp"
pdbFileName    = "system.pdb"
engineFileName = "system.rmc"
NCORES         = 1
FRESH_START    = False

# engine variables
experimentalDataPath = os.path.join(DIR_PATH, grFileName)
structurePdbPath     = os.path.join(DIR_PATH, pdbFileName)
engineSavePath       = os.path.join(DIR_PATH, engineFileName)



#  ####################################################################################  #
#  ################################### CREATE ENGINE ##################################  #
ENGINE = Engine(path=None)
if not ENGINE.is_engine(engineSavePath) or FRESH_START:
    ENGINE = Engine(path=engineSavePath, freshStart=True)
    ENGINE.set_pdb(structurePdbPath)
    ## create and add pair distribution constraint
    PDF_CONSTRAINT = PairDistributionConstraint(experimentalData=experimentalDataPath, weighting="atomicNumber")
    ENGINE.add_constraints([PDF_CONSTRAINT])
    ## create and add intermolecular distances constraint
    EMD_CONSTRAINT = InterMolecularDistanceConstraint()
    ENGINE.add_constraints([EMD_CONSTRAINT])
    EMD_CONSTRAINT.set_type_definition("element")
    EMD_CONSTRAINT.set_pairs_distance([('Co','Co',2.00),('Co','Mn',2.00),('Co','Ni',2.00),('Co','Li',2.00),('Co','O' ,1.7),
                                       ('Mn','Mn',2.00),('Mn','Li',2.00),('Mn','Ni',2.00),('Mn','O' ,1.7),
                                       ('Ni','Ni',2.00),('Ni','Li',2.00),('Ni','O' ,1.7),
                                       ('Li','Li',2.40),('Li' ,'O',1.8),
                                       ('O' ,'O' ,1.20),])
    ## save engine
    ENGINE.save()
else:
    ENGINE = ENGINE.load(engineSavePath)
    ## unpack constraints before fitting in case tweaking is needed
    PDF_CONSTRAINT, EMD_CONSTRAINT = ENGINE.constraints

#  ####################################################################################  #
#  ############################### DEFINE DIFFERENT RUNS ##############################  #
def normal_run(numberOfSteps=100000, saveFrequency=10000):
    ## reset groups as atoms
    ENGINE.set_groups_as_atoms()
    ## run engine
    ENGINE.run(numberOfSteps=numberOfSteps, saveFrequency=saveFrequency)

def swaps_run(numberOfSteps=100000, saveFrequency=10000):
    ## reset groups as atoms
    ENGINE.set_groups_as_atoms()
    ## build swap lists
    ALL_ELEMENTS = ENGINE.get_original_data('allElements')
    liSwaps = [[idx] for idx in range(len(ALL_ELEMENTS)) if ALL_ELEMENTS[idx]=='li' or ALL_ELEMENTS[idx]=='Li']
    meSwaps = [[idx] for idx in range(len(ALL_ELEMENTS)) if ALL_ELEMENTS[idx] in ('co','ni','mn') or ALL_ELEMENTS[idx] in ('Co','Ni','Mn')]
    ## set swap generators
    for g in ENGINE.groups:
        idx   = g.indexes[0]
        elIdx = ALL_ELEMENTS[idx]
        if elIdx in ('li','Li'):
            SPG=SwapPositionsGenerator(swapList=meSwaps)
            g.set_move_generator(SPG)
        elif elIdx in ('co','Co','mn','Mn','ni','Ni'):
            SPG=SwapPositionsGenerator(swapList=liSwaps)
            g.set_move_generator(SPG)
        else:
            g.set_move_generator( TranslationGenerator(amplitude=0.05) )
    ## run engine
    ENGINE.run(numberOfSteps=numberOfSteps, saveFrequency=saveFrequency)

def removes_run(numberOfSteps=100, saveFrequency=100):
    ## compute indexes lists
    ALL_ELEMENTS = ENGINE.get_original_data('allElements')
    oIndexes  = [idx for idx in range(len(ALL_ELEMENTS)) if ALL_ELEMENTS[idx]=='o'  or ALL_ELEMENTS[idx]=='O']
    liIndexes = [idx for idx in range(len(ALL_ELEMENTS)) if ALL_ELEMENTS[idx]=='li' or ALL_ELEMENTS[idx]=='Li']
    meIndexes = [idx for idx in range(len(ALL_ELEMENTS)) if ALL_ELEMENTS[idx] in ('co','ni','mn') or ALL_ELEMENTS[idx] in ('Co','Ni','Mn')]
    ## create empty group to remove oxygen.
    ## By default EmptyGroup move generator is AtomsRemoveGenerator with its atomsList
    ## None which means it will remove any atom from system.
    RO  = EmptyGroup()
    RO.moveGenerator.set_maximum_collected(20)
    RO.moveGenerator.set_atoms_list(oIndexes)
    ## create empty group to remove lithium
    RLi = EmptyGroup()
    RLi.moveGenerator.set_maximum_collected(500)
    RLi.moveGenerator.set_atoms_list(liIndexes)
    ## create empty group to remove cobalt
    RCo = EmptyGroup()
    RCo.moveGenerator.set_maximum_collected(500)
    RCo.moveGenerator.set_atoms_list(meIndexes)
    ## set groups to engine
    ENGINE.set_groups([RO,RLi,RCo])
    ## run engine
    ENGINE.run(numberOfSteps=numberOfSteps, saveFrequency=saveFrequency)


#  ####################################################################################  #
#  ################################## RUN SIMULATION ##################################  #
## run normal
normal_run(numberOfSteps=100000, saveFrequency=50000)
PDF_CONSTRAINT.set_adjust_scale_factor((10,0.7,1.3))
normal_run(numberOfSteps=250000, saveFrequency=50000)

## add first swapping frame
if not ENGINE.is_frame('swap_1'):
    ENGINE.add_frame('swap_1')
## use swap_1 frame and run swapping
ENGINE.set_used_frame('swap_1')
swaps_run(numberOfSteps=200000, saveFrequency=50000)
normal_run(numberOfSteps=250000, saveFrequency=50000)

## add first removing frame
if not ENGINE.is_frame('removes_2'):
    ENGINE.add_frame('removes_2')
## use removes_2 frame and run some removes
ENGINE.set_used_frame('removes_2')
for _ in range(20):
    # remove as little as possible then try to refine
    removes_run(numberOfSteps=100, saveFrequency=100)
    normal_run(numberOfSteps=100000, saveFrequency=10000)
    swaps_run(numberOfSteps=50000, saveFrequency=10000)





##########################################################################################
###################################### CALL plot.py ######################################
os.system("%s %s"%(sys.executable, os.path.join(DIR_PATH, 'plot.py')))
