"""
Generated by CHARMM-GUI (http://www.charmm-gui.org)

omm_restratins.py

This module contains restraint functions for OpenMM.

Correspondance: jul316@lehigh.edu or wonpil@lehigh.edu
Last update: February 5, 2025
"""

from openmm import *
from openmm.app import *
from openmm.unit import *


def restraints(system, crd, inputs):
    boxlx = system.getDefaultPeriodicBoxVectors()[0][0].value_in_unit(nanometers)
    boxly = system.getDefaultPeriodicBoxVectors()[1][1].value_in_unit(nanometers)
    boxlz = system.getDefaultPeriodicBoxVectors()[2][2].value_in_unit(nanometers)

    if inputs.fc_bb > 0 or inputs.fc_sc > 0:
        # positional restraints for protein
        posresPROT = CustomExternalForce("k*periodicdistance(x, y, z, x0, y0, z0)^2;")
        posresPROT.addPerParticleParameter("k")
        posresPROT.addPerParticleParameter("x0")
        posresPROT.addPerParticleParameter("y0")
        posresPROT.addPerParticleParameter("z0")
        for line in open("restraints/prot_pos.txt", "r"):
            segments = line.strip().split()
            atom1 = int(segments[0])
            state = segments[1]
            xpos = crd.positions[atom1].value_in_unit(nanometers)[0]
            ypos = crd.positions[atom1].value_in_unit(nanometers)[1]
            zpos = crd.positions[atom1].value_in_unit(nanometers)[2]
            if state == "BB" and inputs.fc_bb > 0:
                fc_ppos = inputs.fc_bb
                posresPROT.addParticle(atom1, [fc_ppos, xpos, ypos, zpos])
            if state == "SC" and inputs.fc_sc > 0:
                fc_ppos = inputs.fc_sc
                posresPROT.addParticle(atom1, [fc_ppos, xpos, ypos, zpos])
        system.addForce(posresPROT)

    if inputs.fc_mpos > 0:
        # positional restraints for micelle lipid head group
        posresMICE = CustomExternalForce("k*periodicdistance(x, y, z, x0, y0, z0)^2;")
        posresMICE.addGlobalParameter("k", inputs.fc_mpos)
        posresMICE.addPerParticleParameter("x0")
        posresMICE.addPerParticleParameter("y0")
        posresMICE.addPerParticleParameter("z0")
        for line in open("restraints/lipid_pos.txt", "r"):
            segments = line.strip().split()
            atom1 = int(segments[0])
            xpos = crd.positions[atom1].value_in_unit(nanometers)[0]
            ypos = crd.positions[atom1].value_in_unit(nanometers)[1]
            zpos = crd.positions[atom1].value_in_unit(nanometers)[2]
            posresMICE.addParticle(atom1, [xpos, ypos, zpos])
        system.addForce(posresMICE)

    if inputs.fc_lpos > 0:
        # positional restraints for bilayer lipid head group
        posresMEMB = CustomExternalForce("k*periodicdistance(0, 0, z, 0, 0, z0)^2;")
        posresMEMB.addGlobalParameter("k", inputs.fc_lpos)
        posresMEMB.addPerParticleParameter("z0")
        for line in open("restraints/lipid_pos.txt", "r"):
            segments = line.strip().split()
            atom1 = int(segments[0])
            zpos = crd.positions[atom1].value_in_unit(nanometers)[2]
            posresMEMB.addParticle(atom1, [zpos])
        system.addForce(posresMEMB)

    if inputs.fc_hmmm > 0:
        posresHMMM = CustomCentroidBondForce(1, "k*(z1 - z0)^2")
        posresHMMM.addGlobalParameter("k", inputs.fc_hmmm)
        posresHMMM.addPerBondParameter("z0")

        # add groups and bonds
        for i, line in enumerate(open("restraints/hmmm_pos.txt", "r")):
            segments = line.strip().split()
            atoms = list(map(int, segments[:-1]))
            zpos = float(segments[-1]) + boxlz / 2.0
            posresHMMM.addGroup(atoms)
            posresHMMM.addBond([i], [zpos])
        posresHMMM.setUsesPeriodicBoundaryConditions(True)
        system.addForce(posresHMMM)

    if inputs.fc_dcle > 0:
        posresDCLE = CustomCentroidBondForce(
            1,
            "k*dr^2; \
                                                dr=max(0, r-rfb); \
                                                r=abs(z1 - z0);",
        )
        posresDCLE.addGlobalParameter("k", inputs.fc_dcle)
        posresDCLE.addGlobalParameter("rfb", inputs.fbres_rfb)
        posresDCLE.addGlobalParameter("z0", boxlz / 2.0)

        # add groups and bonds
        for i, line in enumerate(open("restraints/dcle_pos.txt", "r")):
            atoms = list(map(int, line.strip().split()))
            posresDCLE.addGroup(atoms)
            posresDCLE.addBond([i])
        posresDCLE.setUsesPeriodicBoundaryConditions(True)
        system.addForce(posresDCLE)

    if inputs.fc_ldih > 0:
        # dihedral restraints for lipids
        dihresMEMB = CustomTorsionForce(
            "fc_ldih*max(0, abs(diff+wrap) - rwidth)^2; \
                                         wrap = 2*pi*(step(-diff-pi)-step(diff-pi)); \
                                         diff = theta - rtheta0; \
                                         rtheta0 = theta0*pi/180; \
                                         rwidth = width*pi/180;"
        )
        dihresMEMB.addGlobalParameter("fc_ldih", inputs.fc_ldih)
        dihresMEMB.addGlobalParameter("pi", 3.141592653589793)
        dihresMEMB.addPerTorsionParameter("width")
        dihresMEMB.addPerTorsionParameter("theta0")
        for line in open("restraints/dihe.txt", "r"):
            segments = line.strip().split()
            atom1 = int(segments[0])
            atom2 = int(segments[1])
            atom3 = int(segments[2])
            atom4 = int(segments[3])
            theta0 = float(segments[4])
            width = float(segments[5])
            dihresMEMB.addTorsion(atom1, atom2, atom3, atom4, [width, theta0])
        system.addForce(dihresMEMB)

    if inputs.fc_cdih > 0:
        # dihedral restraints for carbohydrates
        dihresCARB = CustomTorsionForce(
            "fc_cdih*max(0, abs(diff+wrap) - rwidth)^2; \
                                         wrap = 2*pi*(step(-diff-pi)-step(diff-pi)); \
                                         diff = theta - rtheta0; \
                                         rtheta0 = theta0*pi/180; \
                                         rwidth = width*pi/180;"
        )
        dihresCARB.addGlobalParameter("fc_cdih", inputs.fc_cdih)
        dihresCARB.addGlobalParameter("pi", 3.141592653589793)
        dihresCARB.addPerTorsionParameter("width")
        dihresCARB.addPerTorsionParameter("theta0")
        if os.path.isfile("restraints/carbohydrate_restraint.dat"):
            for line in open("restraints/carbohydrate_restraint.dat", "r"):
                segments = line.strip().split()
                atom1 = int(segments[0])
                atom2 = int(segments[1])
                atom3 = int(segments[2])
                atom4 = int(segments[3])
                theta0 = float(segments[4])
                width = float(segments[5])
                dihresCARB.addTorsion(atom1, atom2, atom3, atom4, [width, theta0])
        if os.path.isfile("restraints/detergent_carbohydrate_restraint.dat"):
            for line in open("restraints/detergent_carbohydrate_restraint.dat", "r"):
                segments = line.strip().split()
                atom1 = int(segments[0])
                atom2 = int(segments[1])
                atom3 = int(segments[2])
                atom4 = int(segments[3])
                theta0 = float(segments[4])
                width = float(segments[5])
                dihresCARB.addTorsion(atom1, atom2, atom3, atom4, [width, theta0])
        system.addForce(dihresCARB)

    return system