import rosetta.core.scoring.methods as methods
import rosetta

@rosetta.EnergyMethod()
class testScoreMethod(methods.WholeStructureEnergy):
    def __init__(self):
        methods.WholeStructureEnergy.__init__(self, self.creator())

    def finalize_total_energy(self, pose, efunc, emap):

        #Calculate what score should be. Below is just demo scoring method that should force the protein to crunch at position 0,0,0.
        score = 0
        for i in range (1, pose.total_residue()):
            pos = pose.residue(i).xyz('CA')
            score += abs(pos[0]) + abs(pos[1]) + abs(pos[2])


        emap.set(self.scoreType, score)