#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Oct  2 10:32:56 2019

@author: zahramazhar
"""

import os
os.chdir('/Users/zahramazhar/Desktop/PyRosetta4/')

"""Importing different classes/protocols from pyrosetta module"""

import random
import string
from pyrosetta.rosetta.core import *
from pyrosetta.rosetta.protocols.loops.loop_closure.ccd import *
from rosetta.protocols.loops.loop_closure.kinematic_closure import * 
from pyrosetta.rosetta.protocols.loops.loop_mover.refine import *
from pyrosetta.rosetta.core.pose import *
from pyrosetta.rosetta.core.fragment import *
from pyrosetta.rosetta.core.io.pdb import*
from pyrosetta.rosetta.protocols.grafting import *
from pyrosetta.rosetta.protocols.moves import *
from pyrosetta.rosetta.protocols.simple_moves import *
from pyrosetta.rosetta.protocols.loops import *
from pyrosetta import *
pyrosetta.init()


    
"""pose from PDB of interest"""
starting_pose=pose_from_file('full_model_4115_3.pdb')
pose=starting_pose
pA=Pose()
pA.assign(pose)

"""sending pose to PyMOL"""
from pyrosetta import PyMOLMover
pym = PyMOLMover()

residue_list=['A','R','N','D','B','C','E','Q','Z','G','H','I','L','K','M','F','P','S','T','W','Y','V']
for residue in residue_list:
        sequence_list=random.sample(residue_list,7)
        sequence="".join(sequence_list)
    
        non_residue_sequence1= sequence.count('CA')
        if non_residue_sequence1 == 0:
            print(sequence)
        else: 
            exit
print("Our sequence is fine and we can move forward")


"""To generate random filenames for randomly generated fragments"""
i= random.sample(range(0,10000),1)

f1= 'frag7_' + str(i) +'.pdb'
filename1=f1

"""produces a pose from the given sequence with standard dihedral angles"""
"""6-mer fragment generated"""
  
iposeA=pose_from_sequence(sequence, "fa_standard")
pyrosetta.rosetta.core.io.pdb.dump_pdb(iposeA, filename1)
iposeB=pose_from_file(filename1)


"""setting up the fragment mover"""
"""Use Protocol from AnchoredGraftMover under rosetta.protocols.grafting"""
"""AnchorGraftmover information includes (cut start res, cut end res, pose to be inserted)"""
moverA=AnchoredGraftMover(76,131,iposeB)

"""Applied the two AnchorGraftMover information in separate poses"""
moverA.apply(pA)

"""sending pose to PyMOL"""
from pyrosetta import PyMOLMover
pym = PyMOLMover()
pym.apply(pA)
 

"""To generate random filenames for pose containing fragment insertions """
j= random.sample(range(10001,20000),1)

f2= '4115_7mer_' + str(j) +'.pdb'
filename2= f2 

"""Dumps newly generated pose after fragment iinisertion as a pdb"""
pyrosetta.rosetta.core.io.pdb.dump_pdb(pA,filename2)

"""to start the loop modeling"""
pose2=pose_from_pdb(filename2)

"""Pose to be used for refinement cycle"""
native_pose= pose_from_pdb(filename2)

   
"""Assigning changed pose'pA' to pose2"""
pose2=Pose()
pose2=pA
    
"""Defined fold tree for the rest of the script to use"""
ft=FoldTree()
"""ft.add_edge represents the edge of the loop"""
"""ft.add_edge(1,125,-1) represents (start,end,)"""

"""Loop1"""
ft.add_edge(1,74,-1)
ft.add_edge(74,79,-1)
ft.add_edge(74,84,1)
ft.add_edge(84,80,-1)
ft.add_edge(84,175,-1)

print(ft)

"""Check foldtree that has been constructed"""
ft.check_fold_tree()
pose.fold_tree(ft)
pym.apply(pose2)

"""makes sure that all states are saved from beginning to end in pymol"""
pym.keep_history(True)

"""Movemap defined: which part of the protein will be perturbed"""
pym.keep_history(True)
movemap= MoveMap()
movemap.set_bb_true_range(75,84)
movemap.set_chi_true_range(75,84)

"""kT defines temperature and n_moves defines the number of attempts by the specialized mover i.e. smallmover, shearmover etc."""
kT=1
n_moves=1

"""Defined for Movemap: the small mover will only perturb φ or ψ of a random residue by a random small angle"""
smallmover= pyrosetta.rosetta.protocols.simple_moves.SmallMover(movemap,kT,n_moves)
backbone_angle_max= 10
smallmover.angle_max(backbone_angle_max)
smallmover.angle_max('E',5)
smallmover.angle_max('H',10)
smallmover.angle_max('L',20)
smallmover.apply(pose2)
    
"""shear mover perturbs φ of a random residue by a small angle and ψ of the same """
""" residue by the same small angle of opposite sign"""
shearmover= pyrosetta.rosetta.protocols.simple_moves.ShearMover(movemap,kT,n_moves)
shearmover.angle_max(backbone_angle_max)
shearmover.angle_max('L',20)
shearmover.angle_max('H',10)
shearmover.angle_max('E',5)
shearmover.apply(pose2)


"""Loop definitions loop_=Loop(start,end, first midpoint defined)"""
loop1=Loop(76,83,80)
ccd=CCDLoopClosureMover(loop1,movemap)

"""To observe changes in PyMOL created by the Minmover"""
#observer= AddPyMOLObserver(pose2, True)


"""Before running the CCD algorithm, it is important to convert the residues around the cut site of the loop to “cut-point variants”"""
add_single_cutpoint_variant(pose2,loop1)

"""To add multiple loops to the Loops object"""
loops=Loops()
loops.add_loop(loop1)

"""Defining a Kinematic Mover which determines the torsion angles for the pivot residues by solving closure equations"""
kic_mover = KinematicMover()

"""Loop definitions loop_=Loop(start, first midpoint defined,end)"""
kic_mover.set_pivots(76,80,83)
kic_mover.apply(pose2)
    
"""Define score function to use for Refinement protocol"""
scorefxn= get_fa_scorefxn()
print(scorefxn)


"""This Refine_KIC loopmover operates in full atom representation of protein and makes small perturbations from a starting conformation to refine a loop"""
kic_refine = LoopMover_Refine_KIC(loops)

"""To generate random filenames for pose containing fragment insertions """
k= random.sample(range(20001,30000),1)

f3= '4115_6mer_' + str(k) +'.pdb'
filename3= f3


"""Defining Jobdistributor"""
jd=PyJobDistributor(filename3,1,scorefxn)
jd.native_pose=native_pose

while not jd.job_complete:
    kic_refine.apply(pose2)
    jd.output_decoy(pose2)
    """To get all sequences in one file"""
    with open ('Loop_sequence_PDB_name.txt','w') as L:
        L.write(sequence)
        L.close()