"""
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: March 29, 2017
"""

import os

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

def read_psf(filename):
    psf = CharmmPsfFile(filename)
    return psf

def read_crd(filename):
    crd = CharmmCrdFile(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):
    extlist = ['rtf', 'prm', 'str']

    parFiles = ()
    for line in open(filename, 'r'):
        if '!' in line: line = line.split('!')[0]
        parfile = line.strip()
        if len(parfile) != 0:
            ext = parfile.lower().split('.')[-1]
            if not ext in extlist: continue
            parFiles += ( parfile, )

    params = CharmmParameterSet( *parFiles )
    return params

def read_box(psf, filename):
    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