#!/usr/bin/python3

import random , math , subprocess , time
from pyrosetta import *
from pyrosetta.toolbox import *
from rosetta.core.fragment import *
from rosetta.protocols.simple_moves import *
init()

#Probability To Make A Choice
def Decision(score_before , score_after):
	''' Metropolis Criterion, P = 1 to accept all structures when final score is lower than starting score, P closer to 0 = accept less and less but sometimes accept to escape local energy minima '''
	''' Returns a string with either Accept or Reject '''
	E = score_after - score_before		#Takes the difference between the score before and after a change in a structure
	if E >= 0:				#If the score after change is larger than the score before change have low P value
		negE = E * -1
		kt = 1
		e = math.e
		P = e**(negE/kt)
	elif E < 0:				#If the score after change is larger than the score before change have P of 1
		P = 1
	if (random.random() < P) == True:	#Randomly accept of reject the score difference. Probability to accept increases as P gets closer to 1
		return('Accept')
	else:
		return('Reject')
#------------------------------------------------------------------------------------------------------
structure = pose_from_pdb('structure.pdb')				#Pose a PDB structure
sequence = structure.sequence()						#Get sequence from PDB structure as string
pose = pose_from_sequence(sequence)					#Use the sequence string to generate the main pose
size = pose.total_residue()						#Get total residue number of the structure

fragset = ConstantLengthFragSet(3)					#Call the Fragment loader for 3-mer
fragset.read_fragment_file('aat000_03_05.200_v1_3')			#Load 3-mer fragment file

pose_copy = Pose()
minValue = None
for samples in range(1,100):						#Monte Carlo number of simulations

	pose_copy.assign(pose)

	#0 - Fragment Insertion
	movemap = MoveMap()						#Load Map Mover
	movemap.set_bb(True)						#A Move Map specifies which degrees of freedom are allowed to change in the pose when the mover is applied (in this case, all backbone torsion angles)
	mover_3mer = ClassicFragmentMover(fragset, movemap)		#Load Fragment Mover
	mover_3mer.apply(pose_copy)					#Apply 3-mer mover

	#1 - Random Trial Move
	scorefxn = get_fa_scorefxn()
	score_before = scorefxn(pose_copy)				#Find score of the structure before changing it
	start = random.randint(1 , size)				#Monte Carlo random residue number as start position
	choice = random.randint(0 , 1)					#Setup to randomly choose weather to change Phi or Psi angles
	if choice == 0:
		#Choose Phi Angle
		phi_angle_now = pose_copy.phi(start)			#Get current Phi angle
		phi_angle_new = random.gauss(phi_angle_now , 25)	#Randomly move current angel by any angel within 25 degrees, gauss mean = current angle and standard deviation = 25
		pose_copy.set_phi(start , phi_angle_new)		#Set Phi angle to new value
	else:
		#Choose Psi Angle
		psi_angle_now = pose_copy.psi(start)			#Get current Psi angle
		psi_angle_new = random.gauss(psi_angle_now , 25)	#Randomly move current angel by any angel within 25 degrees. gauss mean = current angle and standard deviation = 25
		pose_copy.set_psi(start , psi_angle_new)		#Set Psi angle to new value

	#2 - Score Structure
	score_after = scorefxn(pose_copy)				#Find score of the structure after changing it

	#3 - Accept/Reject						#Metropolis Criterion - P = 1 to accept all the time when final score is lower than starting score, P closer to 0 = accept less and less but sometimes to escape local energy minima
	if Decision(score_before , score_after) == 'Accept':
		print('[+] Accept')
		pose.assign(pose_copy)					#Accept by updating the original pose
		if minValue == None:
			minValue = score_after
		else:
			minValue = min(minValue, score_after)
		continue
	else:
		print('[-] Reject')
		continue						#Revert back by not changing the original pose (do nothing), top of the script always uses the original pose

print('Score: ' , minValue)