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

# external libraries imports
import numpy as np

# fullrmc library imports
from fullrmc.Globals import LOGGER, maxint
from fullrmc.Engine import Engine
from fullrmc.Core.MoveGenerator import MoveGeneratorCollector
from fullrmc.Constraints.DistanceConstraints import InterMolecularDistanceConstraint
from fullrmc.Generators.Rotations import RotationGenerator
from fullrmc.Generators.Translations import TranslationTowardsCenterGenerator


##########################################################################################
##################################  SHUT DOWN LOGGING  ###################################
LOGGER.set_minimum_level(maxint, stdoutFlag=True, fileFlag=True)


##########################################################################################
#####################################  CREATE ENGINE  ####################################
pdbPath = "system.pdb"
ENGINE = Engine(path=None)
ENGINE.set_pdb(pdbPath)
# add inter-molecular distance constraint
EMD_CONSTRAINT = InterMolecularDistanceConstraint(defaultDistance=1.75)
ENGINE.add_constraints([EMD_CONSTRAINT])

##########################################################################################
####################################  DIFFERENT RUNS  ####################################
def move_towards():
    # set only one molecule group
    ENGINE.set_groups_as_molecules()
    secMolIdxs = ENGINE.groups[1].indexes
    ENGINE.set_groups( ENGINE.groups[0] )
    # set move generator
    for g in ENGINE.groups:
        t = TranslationTowardsCenterGenerator(center={'indexes': secMolIdxs},amplitude=0.15, angle=90)
        r = RotationGenerator(amplitude=10)
        mg = MoveGeneratorCollector(collection=[t,r], randomize=True, weights=[(0,1),(1,5)])
        g.set_move_generator(mg)
    # set runtime parameters
    nsteps = 1000
    xyzFrequency = 1
    # run engine
    xyzPath="trajectory.xyz"
    if os.path.isfile(xyzPath): os.remove(xyzPath)
    ENGINE.run(numberOfSteps=nsteps, saveFrequency=2*nsteps, xyzFrequency=xyzFrequency, xyzPath=xyzPath, restartPdb=None)



##########################################################################################
#####################################  RUN SIMULATION  ###################################
move_towards()


##########################################################################################
##################################  VISUALIZE SIMULATION  ################################
ENGINE.set_pdb(pdbPath)
ENGINE.visualize( commands = ["trajectory.xyz"],
                  boxWidth=0, bgColor='white',
                  representationParams='CPK 1.0 0.2 50 50')
