from backbone.ramachandran import RamachandranSampler
from rotamers.dunbrack import DunbrackRotamerLibrary
from geometry.backbone import GEO, ORI
from geometry.nerf import nerf
from protein.residue import ResidueName, Residue

rs = RamachandranSampler()
rl = DunbrackRotamerLibrary()

ghost_residue = Residue(ResidueName.GLY)
backbone_atoms = ["N", "CA", "C", "O"]
for atom in backbone_atoms:
    ghost_residue.backbone["atoms"][atom] = ORI[atom]

primary_sequence = ["PHE", "SER", "THR"]

residues = [ ghost_residue ]

# First 2 residues of the sequence
for residue_index, residue_name in enumerate(primary_sequence[:2]):
    residue = Residue(ResidueName(residue_name))
    residue.backbone["phi"], residue.backbone["psi"] = rs.sample(residue.name.value)
    
    residue.backbone["atoms"]["N"] = nerf(
        residues[residue_index].backbone["atoms"]["N"],
        residues[residue_index].backbone["atoms"]["CA"],
        residues[residue_index].backbone["atoms"]["C"],
        GEO["C_N_length"],
        GEO["CA_C_N_angle"],
        residues[residue_index].backbone["psi"],
        )
    
    residue.backbone["atoms"]["CA"] = nerf(
        residues[residue_index].backbone["atoms"]["CA"],
        residues[residue_index].backbone["atoms"]["C"],
        residue.backbone["atoms"]["N"],
        GEO["N_CA_length"],
        GEO["C_N_CA_angle"],
        180,
        )
    
    residue.backbone["atoms"]["C"] = nerf(
        residues[residue_index].backbone["atoms"]["C"],
        residue.backbone["atoms"]["N"],
        residue.backbone["atoms"]["CA"],
        GEO["CA_C_length"],
        GEO["N_CA_C_angle"],
        residue.backbone["phi"],
        )
    
    residue.backbone["atoms"]["O"] = nerf(
        residue.backbone["atoms"]["N"],
        residue.backbone["atoms"]["CA"],
        residue.backbone["atoms"]["C"],
        GEO["C_O_length"],
        GEO["CA_C_O_angle"],
        residue.backbone["psi"] + 180,
        )

    residues.append(residue)

    residue.sidechain["chis"] =  rl.rotamer_params(residue.name.value, residue.backbone["phi"], residue.backbone["psi"])

    print(residue.summary())

# Generate a tripeptide backbone
# Populate the sidechain of the middle residue