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

omm_rewrap.py

This module is for re-wrapping coordinates for CHARMM DOMDEC output coordinates.

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

from __future__ import print_function
import os
from math import *

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


def rewrap(simulation):
    bonds = simulation.topology.bonds()
    positions = simulation.context.getState(getPositions=True).getPositions()
    box = simulation.context.getState().getPeriodicBoxVectors()
    boxlx = box[0][0] / angstrom
    boxly = box[1][1] / angstrom
    boxlz = box[2][2] / angstrom

    min_crds = [
        positions[0][0] / angstrom,
        positions[0][1] / angstrom,
        positions[0][2] / angstrom,
    ]
    max_crds = [
        positions[0][0] / angstrom,
        positions[0][1] / angstrom,
        positions[0][2] / angstrom,
    ]

    for position in positions:
        min_crds[0] = min(min_crds[0], position[0] / angstrom)
        min_crds[1] = min(min_crds[1], position[1] / angstrom)
        min_crds[2] = min(min_crds[2], position[2] / angstrom)
        max_crds[0] = max(max_crds[0], position[0] / angstrom)
        max_crds[1] = max(max_crds[1], position[1] / angstrom)
        max_crds[2] = max(max_crds[2], position[2] / angstrom)

    xcen = (max_crds[0] + min_crds[0]) / 2.0
    ycen = (max_crds[1] + min_crds[1]) / 2.0
    zcen = (max_crds[2] + min_crds[2]) / 2.0

    for bond in bonds:
        atom1 = bond[0]
        atom2 = bond[1]
        atom1id = atom1.index
        atom2id = atom2.index
        res1 = atom1.residue
        res2 = atom2.residue

        x1, y1, z1 = positions[atom1id]
        x2, y2, z2 = positions[atom2id]
        dx = fabs(x1 / angstrom - x2 / angstrom)
        dy = fabs(y1 / angstrom - y2 / angstrom)
        dz = fabs(z1 / angstrom - z2 / angstrom)

        if dx > boxlx / 2 or dy > boxly / 2 or dz > boxlz / 2:
            for atom in res2.atoms():
                oldx = positions[atom.index][0] / angstrom
                oldy = positions[atom.index][1] / angstrom
                oldz = positions[atom.index][2] / angstrom
                if dx > boxlx / 2.0:
                    if oldx < xcen:
                        newx = oldx + boxlx
                    else:
                        newx = oldx - boxlx
                else:
                    newx = oldx
                if dy > boxly / 2.0:
                    if oldy < ycen:
                        newy = oldy + boxly
                    else:
                        newy = oldy - boxly
                else:
                    newy = oldy
                if dz > boxlz / 2.0:
                    if oldz < zcen:
                        newz = oldz + boxlz
                    else:
                        newz = oldz - boxlz
                else:
                    newz = oldz

                new_position = Vec3(newx, newy, newz)
                positions[atom.index] = Quantity(new_position, angstroms)

    simulation.context.setPositions(positions)

    return simulation