from openmm.app import Modeller, PME, HBonds, ForceField, PDBFile
from openmm import NonbondedForce
from openmm.unit import nanometer
from openff.interchange import Interchange

for lipid, name in zip(["popc", "pope", "dppc"], ["POP", "POP", "DPP"]):
    print(lipid)
    pdb = PDBFile("memb_protein_trimer.pdb")

    m = Modeller(pdb.topology, pdb.positions)
    ff = ForceField("amber14-all.xml", "amber14/tip3pfb.xml")
    m.addHydrogens(ff)
    attempts_left = 5
    while attempts_left > 0:
        try:
            m.addMembrane(ff, lipidType=lipid.upper())
            break
        except Exception as e:
            print(e)
            attempts_left -= 1

    print(len(list(m.topology.residues())))
    names = []
    for i, r in enumerate(m.topology.residues()):
        if r.name not in names:
            names.append(r.name)
        if r.name == name:
            print(i)
            found = i
            break
    print(names)
    m.delete(list(m.topology.residues())[:found])
    m.delete(list(m.topology.residues())[1:])
    system = ff.createSystem(
        m.topology,
        nonbondedMethod=PME,
        nonbondedCutoff=1.0 * nanometer,
    )
    PDBFile.writeFile(
        m.topology, m.positions, open(f"test_memb_{lipid}.pdb", "w")
    )

    for force in system.getForces():
        if isinstance(force, NonbondedForce):
            n_p = force.getNumParticles()
            print(n_p)
            charges = [
                force.getParticleParameters(i)[0]._value for i in range(n_p)
            ]
            print(charges)
            print(sum(charges))

    i = Interchange.from_openmm(
        system, m.topology, m.positions, system.getDefaultPeriodicBoxVectors()
    )
    off_ch = [
        x.parameters["charge"].m
        for x in i.collections["Electrostatics"].potentials.values()
    ]
    print(off_ch)
    print(sum(off_ch))
    i.to_gromacs(lipid, _merge_atom_types=True)
    break
    # i.to_pdb(f"{lipid}.pdb")
    # i.topology.molecule(0).to_file(f"{lipid}.sdf")