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

omm_readparams.py

This module is for reading coordinates and parameters in OpenMM.

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

import os, json

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


def read_top(filename, fftype="CHARMM"):
    if fftype == "CHARMM":
        top = CharmmPsfFile(filename)
    elif fftype == "AMBER":
        top = AmberPrmtopFile(filename)
    return top


def read_crd(filename, fftype="CHARMM"):
    if fftype == "CHARMM":
        crd = CharmmCrdFile(filename)
    elif fftype == "AMBER":
        crd = AmberInpcrdFile(filename)
    return crd


def read_charmm_rst(filename):
    charmm_rst = CharmmRstFile(filename)

    for i, line in enumerate(charmm_rst.header):
        line = line.strip()
        words = line.split()
        if len(line) != 0:
            if words[0] == "CRYSTAL" or words[0] == "!CRYSTAL":
                line1 = charmm_rst.header[i + 1]
                line2 = charmm_rst.header[i + 2]
                boxlx = Vec3(float(line1.split()[0].replace("D", "E")), 0.0, 0.0)
                boxly = Vec3(0.0, float(line1.split()[2].replace("D", "E")), 0.0)
                boxlz = Vec3(0.0, 0.0, float(line2.split()[2].replace("D", "E")))
                box = (boxlx, boxly, boxlz)
                break

    positions = charmm_rst.positionsold
    new_positions = []

    for position in positions:
        oldx = position[0] / angstrom
        oldy = position[1] / angstrom
        oldz = position[2] / angstrom

        newx = oldx + boxlx[0] / 2.0
        newy = oldy + boxly[1] / 2.0
        newz = oldz + boxlz[2] / 2.0

        new_positions.append(Vec3(newx, newy, newz))

    charmm_rst.box = Quantity(box, angstroms)
    charmm_rst.positions = Quantity(new_positions, angstroms)

    return charmm_rst


def read_params(filename):
    parFiles = ()
    for line in open(filename, "r"):
        if "!" in line:
            line = line.split("!")[0]
        parfile = line.strip()
        if len(parfile) != 0:
            parFiles += (parfile,)

    params = CharmmParameterSet(*parFiles)
    return params


def read_box(psf, filename):
    try:
        sysinfo = json.load(open(filename, "r"))
        boxlx, boxly, boxlz = map(float, sysinfo["dimensions"][:3])
    except:
        for line in open(filename, "r"):
            segments = line.split("=")
            if segments[0].strip() == "BOXLX":
                boxlx = float(segments[1])
            if segments[0].strip() == "BOXLY":
                boxly = float(segments[1])
            if segments[0].strip() == "BOXLZ":
                boxlz = float(segments[1])
    psf.setBox(boxlx * angstroms, boxly * angstroms, boxlz * angstroms)
    return psf


def gen_box(psf, crd):
    coords = crd.positions

    min_crds = [coords[0][0], coords[0][1], coords[0][2]]
    max_crds = [coords[0][0], coords[0][1], coords[0][2]]

    for coord in coords:
        min_crds[0] = min(min_crds[0], coord[0])
        min_crds[1] = min(min_crds[1], coord[1])
        min_crds[2] = min(min_crds[2], coord[2])
        max_crds[0] = max(max_crds[0], coord[0])
        max_crds[1] = max(max_crds[1], coord[1])
        max_crds[2] = max(max_crds[2], coord[2])

    boxlx = max_crds[0] - min_crds[0]
    boxly = max_crds[1] - min_crds[1]
    boxlz = max_crds[2] - min_crds[2]

    psf.setBox(boxlx, boxly, boxlz)
    return psf