import testScoreMethod as t
from rosetta import *
init()

pose = Pose()
pose_from_pdb(pose, 'SomeProtein.pdb')

starting_pose = Pose()
starting_pose.assign(pose)

scorefxn = ScoreFunction()

test_score = t.testScoreMethod.scoreType
scorefxn.set_weight(test_score, 1.0) #Scoring function with only my custom coring method

movemap = MoveMap()
movemap.set_bb(True)

kT = 1.0
smallmoves = 3
shearmoves = 7
backbone_angle_max = 7

smallmover = SmallMover(movemap, kT, smallmoves)
smallmover.angle_max(backbone_angle_max)

shearmover = ShearMover(movemap, kT, shearmoves)
shearmover.angle_max(backbone_angle_max)

minmover = MinMover()
minmover.movemap(movemap)
minmover.score_function(scorefxn)

to_pack = standard_packer_task(starting_pose)
to_pack.restrict_to_repacking()
to_pack.or_include_current(True)
packmover = PackRotamersMover(scorefxn, to_pack)

combined_mover = SequenceMover()
#combined_mover.add_mover(smallmover)
#combined_mover.add_mover(shearmover)
combined_mover.add_mover(minmover)
combined_mover.add_mover(packmover)

mc = MonteCarlo(pose, scorefxn, kT)

trial = TrialMover(combined_mover, mc)

cycles = 12

refinement = RepeatMover(trial, cycles)

jobs = 1

jd = PyJobDistributor('Refined', jobs, scorefxn)
jd.native_pose = starting_pose

counter = 0
scores = []

while not jd.job_complete:
    pose.assign(starting_pose)
    counter += 1
    pose.pdb_info().name('Refined' + str(counter))

    mc.reset(pose)
    refinement.apply(pose)

    mc.recover_low(pose)
    jd.output_decoy(pose)
    scores.append(scorefxn(pose))

print 'Starting Pose: ' + str(scorefxn(starting_pose))
print 'Refined Poses: ' + str(scores)