#Script to calculate the free energy change
# upon removing a set of host-guest distance
# restraints and applying standard state conditions
# @authors: Stefano Bosisio and Julien Michel

import os,sys, random
import math
from math import pi, cos, sin
from Sire.Tools.OpenMMMD import *
from Sire.Tools import Parameter, resolveParameters
from Sire.Maths import *
# Python dependencies

try:
    mdtraj = Sire.try_import("mdtraj")
except:
    pass

try:
    shutil = Sire.try_import("shutil")
except:
    pass

try:
    numpy = Sire.try_import("numpy")
except:
    pass

# Constant for conversion
NM_TO_ANG = 10.0


trajfile = Parameter("trajfile", "traj000000001.dcd",
                    """File name of the trajectory to process.""")

stepframe = Parameter("step_frame",1,"""The number of frames to step
 to between two succcessive evaluations.""")
simfile  = Parameter("simfile", "sim.cfg", """ Configuration file with distance restraints dictionary""")
topfile = Parameter("topfile", "SYSTEM.top",
                    """File name of the topology file containing the system to be simulated.""")
buffer = Parameter("buffer", 5.0, """Buffer to be added to coordinates to create domain of integration""")
delta_trans = Parameter("dtrans", 0.20, """Step size for translational volume elements, in Angstrom""")
nor
ient = Parameter("norient", 4, """Number of orientations per [0,2pi] Euler Angles interval""")
temperature = Parameter("temperature",298 * kelvin,"""The temperature of the system""")

verbose = Parameter("verbose", False, """Print debug output""")


def averageCoordinates(restr_dict):
    r"""This function computes the average x,y and z coordinates for host atoms
    Parameters
    ----------
    restr_dict : dictionary
                 restr_dict[lig_idx,host_idx] = ([req,D,K],[x0,y0,z0],[x1,y1,z1]...)
                 where x0 is the x coordinate at frame 0 and so on
    Returns
    ----------
    restr_dict : dictionary
                 restr_dict[lig_idx,host_idx] = ([req,D,K],[avgx,avgy,avgz])
                 where avgx,avgy and avgz are the average coordinates
    """
    #restr_dict[idx]=([req,K,D],[coords])
    #Calculation of the mean coordinate for every atoms
    for pairs in restr_dict:
        coords1 = restr_dict[pairs][1] #here the list of all the coords
        coords2 = restr_dict[pairs][2]
        coord1_avg = [ 0.0, 0.0, 0.0 ]
        coord2_avg = [ 0.0, 0.0, 0.0 ]
        for x in range(0,len(coords1)):
            val1 = coords1[x]
            val2 = coords2[x]
            coord1_avg[0] += (val1[0] - coord1_avg[0])/(x+1)
            coord1_avg[1] += (val1[1] - coord1_avg[1])/(x+1)
            coord1_avg[2] += (val1[2] - coord1_avg[2])/(x+1)
            coord2_avg[0] += (val2[0] - coord2_avg[0])/(x+1)
            coord2_avg[1] += (val2[1] - coord2_avg[1])/(x+1)
            coord2_avg[2] += (val2[2] - coord2_avg[2])/(x+1)

        #Substitution of values and reset each avg coords
        # Note that mdtraj coordinates are in nm, but potential parameters
        # in angstrom
        restr_dict[pairs][1] = [ coord1_avg[0] * NM_TO_ANG,
                                 coord1_avg[1] * NM_TO_ANG,
                                 coord1_avg[2] * NM_TO_ANG ]
        restr_dict[pairs][2] = [ coord2_avg[0] * NM_TO_ANG,
                                 coord2_avg[1] * NM_TO_ANG,
                                 coord2_avg[2] * NM_TO_ANG ]

    return restr_dict

def defineIntegrationDomain(restr_dict):
    r"""Definition of the integration domain starting from host coordiantes
    Parameters
    ----------
    restr_dict : dictionary
                 restr_dict[lig_idx,host_idx] = ([req,D,K],[avg_lig_x, avg_lig_y, avg_lig_z], [avg_hostx,avg_hosty,avg_hostz])
                 where avgx,avgy and avgz are the average coordinates of the ligand and host atoms defined by lig_idx and host_idx
    Returns
    ----------
    space : 3D array
            space = [(-x,-y,-z)(x,y,z)]
            where the coordinates define the minimum and maximum coordinates
            of the integration domain

    """
    max_x = -999999
    max_y = -999999
    max_z = -999999
    min_x = +999999
    min_y = +999999
    min_z = +999999

    for pairs in restr_dict:
        val = restr_dict[pairs][2]
        if val[0] > max_x:
            max_x = val[0]
        if val[0] < min_x:
            min_x = val[0]
        if val[1] > max_y:
            max_y = val[1]
        if val[1] < min_y :
            min_y = val[1]
        if val[2] > max_z:
            max_z = val[2]
        if val[2] < min_z :
            min_z = val[2]
    #print(max_x,max_y,max_z,min_x,min_y,min_z)
    # Adding a buffer region
    max_x += buffer.val
    max_y += buffer.val
    max_z += buffer.val
    min_x -= buffer.val
    min_y -= buffer.val
    min_z -= buffer.val
    space = [(min_x,min_y, min_z), (max_x,max_y,max_z)]
    return space

def genOrientations(restr_dict, norientations=5):
    r"""Generates a set of orientations for guest atoms
    The coordinates are in a frame of reference centered on the COG
    of the guest atoms.
    Norient orientations are generated by multiplying the input coordinates
    by a series of rotation matrices, each corresponding to a rotation along
    Euler angles theta, phi and psi.

    Parameters
    ----------
    restr_dict : dictionary
                 restr_dict[lig_idx,host_idx] = ([req,D,K],[avg_lig_x, avg_lig_y, avg_lig_z], [avg_hostx,avg_hosty,avg_hostz])
                 where avgx,avgy and avgz are the average coordinates of the ligand and host atoms defined by lig_idx and host_idx
    Returns
    ----------
    orientations : array of norient**3 * n_guest atom 3D coordinates
    """

    # First pass: work out COG of guest atoms
    guest_cog = [0.0, 0.0, 0.0]
    guest_indices = []
    for pairs in restr_dict:
        if not (pairs[0] in guest_indices):
            guest_indices.append(pairs[0])
        guest_cog[0] += restr_dict[pairs][1][0]*(1/len(restr_dict))
        guest_cog[1] += restr_dict[pairs][1][1]*(1/len(restr_dict))
        guest_cog[2] += restr_dict[pairs][1][2]*(1/len(restr_dict))
    # if only one guest atom then return null coordinates
    body = []
    if len(guest_indices) < 2:
        print ("Restraints apply to a single guest atom, isotropic case, collapsing orientations.")
        coords = []
        for pairs in restr_dict:
            coords.append( [0.0, 0.0, 0.0] )
        body.append(coords)
        return body
    # Second pass: Subtract COG to get COG centered coordinates
    body = []
    for pairs in restr_dict:
        new_x = restr_dict[pairs][1][0] - guest_cog[0]
        new_y = restr_dict[pairs][1][1] - guest_cog[1]
        new_z = restr_dict[pairs][1][2] - guest_cog[2]
        body.append( Vector(new_x, new_y, new_z) )

    # Now work out set of rotations along Euler Angles
    PI = math.pi
    TWOPI = 2*PI
    orientations = []
    for x in range(0,norientations):
        phi = (x*TWOPI)/norientations
        for y in range(0,int(norientations/2)):
            theta = (y*PI)/(norientations/2)
            for z in range(0,norientations):
                psi = (z*TWOPI)/norientations
                rot00 = cos(phi)*cos(psi)-cos(theta)*sin(phi)*sin(psi)
                rot10 = sin(phi)*cos(psi)+cos(theta)*cos(phi)*sin(psi)
                rot20 = sin(theta)*sin(psi)
                rot01 = -cos(phi)*sin(psi)-cos(theta)*sin(phi)*cos(psi)
                rot11 = -sin(phi)*sin(psi)+cos(theta)*cos(phi)*cos(psi)
                rot21 = sin(theta)*cos(phi)
                rot02 = sin(theta)*sin(phi)
                rot12 = -sin(theta)*cos(phi)
                rot22 = cos(theta)
                rotmat = Matrix(rot00,rot01,rot02,\
                                rot10,rot11,rot12,\
                                rot20,rot21,rot22)
                rotvecs = []
                #import pdb; pdb.set_trace()
                for vec in body:
                    rotvec = rotmat*vec
                    rotvecs.append(rotvec)
                orientations.append(rotvecs)
    return orientations

@resolveParameters
def run():
    try:
        host = os.environ['HOSTNAME']
    except KeyError:
        host = "unknow