import datetime
import IMP
import IMP.core
import IMP.algebra
import IMP.atom
import IMP.container

import IMP.pmi.restraints.crosslinking
import IMP.pmi.restraints.stereochemistry
import IMP.pmi.restraints.em
import IMP.pmi.restraints.basic
import IMP.pmi.representation
import IMP.pmi.tools
import IMP.pmi.samplers
import IMP.pmi.output
import IMP.pmi.macros
import IMP.pmi.topology

import os
import sys

datadirectory = "../data/"
topology_file = datadirectory+"topology.txt"


num_frames = 500
num_mc_steps = 50
bead_max_trans = 5.00
rb_max_trans = 5.00	
rb_max_rot = 0

#van szerkezet es mozog
#homer coiled coil egy rigid body h egyutt mozogjon (3)
rigid_bodies = [["PDZ_1"],["PDZ_2"],["PDZ_3"],["SH3"],["GK"],
				]
#egy molekulaban levo domainek
#homer tobb molekulat akarunk egyutt mogatni, ezeket lanconkent deklaralhattuk, szal itt a lancokat egynek kell venni			
super_rigid_bodies = [ ["PDZ_1","PDZ_2","PDZ_3","SH3","GK"],
					  ]

starttime = datetime.datetime.now()	
				  
m = IMP.Model()
topology = IMP.pmi.topology.TopologyReader(topology_file)
domains = topology.get_components()

bm = IMP.pmi.macros.BuildModel(m,
                    component_topologies=domains,
                    list_of_rigid_bodies=rigid_bodies,
					list_of_super_rigid_bodies=super_rigid_bodies)

					
representation = bm.get_representation()

for nc,component in enumerate(domains):
    name = component.name
    sel = IMP.atom.Selection(representation.prot,molecule=name)
    ps = sel.get_selected_particles()
    clr = IMP.display.get_rgb_color(float(nc)/len(domains))
    for p in ps:
        if not IMP.display.Colored.get_is_setup(p):
            IMP.display.Colored.setup_particle(p,clr)
        else:
            IMP.display.Colored(p).set_color(clr)



#representation.shuffle_configuration(50)

representation.set_rigid_bodies_max_rot(rb_max_rot)
representation.set_floppy_bodies_max_trans(bead_max_trans)
representation.set_rigid_bodies_max_trans(rb_max_trans)


outputobjects = [] # reporter objects (for stat files)
sampleobjects = [] # sampling objects





# add colors to the components
for nc,component in enumerate(domains):
    name = component.name
    sel = IMP.atom.Selection(representation.prot,molecule=name)
    ps = sel.get_selected_particles()
    clr = IMP.display.get_rgb_color(float(nc)/len(domains))
    for p in ps:
        if not IMP.display.Colored.get_is_setup(p):
            IMP.display.Colored.setup_particle(p,clr)
        else:
            IMP.display.Colored(p).set_color(clr)




outputobjects = [] # reporter objects (for stat files)
sampleobjects = [] # sampling objects

outputobjects.append(representation)
sampleobjects.append(representation)


ev = IMP.pmi.restraints.stereochemistry.ExcludedVolumeSphere(representation, 
															resolution=8,
															kappa=1.0)
ev.add_to_model()
outputobjects.append(ev)


# Crosslinks - dataset 1 fejlecei az interactions.csvnek
columnmap={}
columnmap["Protein1"]="prot1"
columnmap["Protein2"]="prot2"
columnmap["Residue1"]="res1"
columnmap["Residue2"]="res2"
columnmap["IDScore"]=None


xl1 = IMP.pmi.restraints.crosslinking.ISDCrossLinkMS(representation,
                                   datadirectory+'interactions.csv', #most interactions.csv
                                   length=15.0, #interakcio kornyezete
                                   slope=0.05, #kozeledes gyorsasaga
                                   columnmapping=columnmap,
                                   resolution=10.0, #beam
                                   label="Chen", #interaction type?
                                   csvfile=True)
xl1.add_to_model()

sampleobjects.append(xl1)
outputobjects.append(xl1)


#floppy_bodies flexibilis szerkezetnelkuli csak szekvencia van
representation.optimize_floppy_bodies(10)


#--------------------------
# Monte-Carlo Sampling
#--------------------------
# This object defines all components to be sampled as well as the sampling protocol
mc1=IMP.pmi.macros.ReplicaExchange0(m,
                                    representation,
                                    monte_carlo_sample_objects=sampleobjects,
                                    output_objects=outputobjects,
									crosslink_restraints=[xl1], #chimera kimenethez csv file nev
                                    monte_carlo_temperature=1.0,
                                    simulated_annealing=True,
                                    simulated_annealing_minimum_temperature=1.0,
                                    simulated_annealing_maximum_temperature=2.5,
                                    simulated_annealing_minimum_temperature_nframes=200,
                                    simulated_annealing_maximum_temperature_nframes=20,
                                    replica_exchange_minimum_temperature=1.0,
                                    replica_exchange_maximum_temperature=2.5,
                                    number_of_best_scoring_models=20,
                                    monte_carlo_steps=num_mc_steps,
                                    number_of_frames=num_frames,
                                    global_output_directory="output")

# Start Sampling
mc1.execute_macro()

endtime = datetime.datetime.now()

file = open("simulation_data.txt","w") 


file.write("\n")
file.write("Start: ")
file.write(str(starttime))
file.write("End: ")
file.write(str(endtime))
file.write("Elapsed: ")
file.write(str(endtime-starttime))

file.close()

print("\n")
print("Start: ")
print(starttime)
print("End: ")
print(endtime)
print("Elapsed: ")
print(endtime-starttime)