#!/usr/bin/env python

# pmx  Copyright Notice
# ============================
#
# The pmx source code is copyrighted, but you can freely use and
# copy it as long as you don't change or remove any of the copyright
# notices.
#
# ----------------------------------------------------------------------
# pmx is Copyright (C) 2006-2013 by Daniel Seeliger
#
#                        All Rights Reserved
#
# Permission to use, copy, modify, distribute, and distribute modified
# versions of this software and its documentation for any purpose and
# without fee is hereby granted, provided that the above copyright
# notice appear in all copies and that both the copyright notice and
# this permission notice appear in supporting documentation, and that
# the name of Daniel Seeliger not be used in advertising or publicity
# pertaining to distribution of the software without specific, written
# prior permission.
#
# DANIEL SEELIGER DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS
# SOFTWARE, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND
# FITNESS.  IN NO EVENT SHALL DANIEL SEELIGER BE LIABLE FOR ANY
# SPECIAL, INDIRECT OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER
# RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF
# CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
# CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
# ----------------------------------------------------------------------

"""This file contains the Rotation class.
A Rotation instance is built with two vectors.

Examples
--------

    >>> v1 = [1,2,3]
    >>> v2 = [2,3,4]
    >>> r = Rotation(v1,v2)  # create rotation object around v2-v1
    >>> v3 = [4,5,6]
    >>> v3 = r.apply(v3)     # rotate v3 around v2-v1

"""

from numpy import array, linalg, arccos, inner
from . import _pmx as _p


class Rotation:

    def __init__(self, v1, v2):
        """ creates a rotation object
        around the vector v2-v1"""

        self.v1 = array(v1)
        self.v2 = [v2[0], v2[1], v2[2]]
        tmp = array(v2)
        self.vec = tmp-self.v1
        self.m1 = None
        self.m2 = None
        x = 1./linalg.norm(self.vec)
        self.norm_vec = self.vec*x
        self.__rm1()
        self.__rm2()

    def __rm1(self):

        a = self.norm_vec
        self.m1 = [
            [a[0]*a[0], a[0]*a[1], a[0]*a[2]],
            [a[1]*a[0], a[1]*a[1], a[1]*a[2]],
            [a[2]*a[0], a[2]*a[1], a[2]*a[2]]
            ]

    def __rm2(self):

        a = self.norm_vec
        self.m2 = [
            [0.0, -a[2], a[1]],
            [a[2], 0.0, -a[0]],
            [-a[1], a[0], 0.0]
            ]

    def apply(self, v, phi):
        return _p.apply_rotation(self, [v[0], v[1], v[2]], phi)


def vec_ang(v1, v2):
    x1 = linalg.norm(v1)
    x2 = linalg.norm(v2)
    return arccos(inner(v1, v2)/(x1*x2))


def bb_super(mol1, mol2, use_orig_mc_coords=True):
    """ superpose mol2 on mol1"""

    N1, CA1, C1 = mol1.fetchm(['N', 'CA', 'C'])
    N2, CA2, C2 = mol2.fetchm(['N', 'CA', 'C'])

    if ((mol1.resname == 'GLY') or (mol2.resname == 'GLY') or
       (mol2.resname[:2] == 'G2')):
        fit_atoms([N1, CA1, C1], [N2, CA2, C2], mol2.atoms)
    else:
        N1, CA1, C1, CB1 = mol1.fetchm(['N', 'CA', 'C', 'CB'])
        N2, CA2, C2, CB2 = mol2.fetchm(['N', 'CA', 'C', 'CB'])
        fit_atoms([N1, CA1, C1, CB1], [N2, CA2, C2, CB2], mol2.atoms)

    if use_orig_mc_coords:
        atom_set = ['N', 'CA', 'C', 'H', 'O', 'HA', 'HN']
        gly_atom_set = ['N', 'CA', 'C', 'H', 'O', 'HA1', 'HN']
        pro_atom_set = ['N', 'CA', 'C', 'O', 'HA']
        pro_atom_set_forGly = ['N', 'CA', 'C', 'O', 'HA1']

        # for proline it is different
        if mol1.resname=='PRO' or mol2.resname == 'PRO' or mol2.resname[:2] == 'P2':
            atoms1 = mol1.fetchm(pro_atom_set)
            atoms2 = mol2.fetchm(pro_atom_set)
            if mol1.resname=='GLY':
                atoms1 = mol1.fetchm(pro_atom_set_forGly)
            elif mol2.resname=='GLY' or mol2.resname[:2] == 'G2':
                atoms2 = mol2.fetchm(pro_atom_set_forGly)
        else:
            if mol1.resname == 'GLY':
                atoms1 = mol1.fetchm(gly_atom_set)
            else:
                atoms1 = mol1.fetchm(atom_set)

            if mol2.resname == 'GLY' or mol2.resname[:2] == 'G2':
                atoms2 = mol2.fetchm(gly_atom_set)
            else:
                atoms2 = mol2.fetchm(atom_set)

        assert len(atoms1) == len(atoms2), "%s -> %s" % ('-'.join(map(lambda a: a.name, atoms1)), '-'.join(map(lambda a: a.name, atoms2)))
        for atom1, atom2 in zip(atoms1, atoms2):
            atom2.x = atom1.x


def nuc_super(mol1, mol2, name1=None, name2=None):
    """ superpose mol2 on mol1"""

    if name1 is None:
        name1 = mol1.resname[:2]
    if name2 is None:
        name2 = mol2.resname[:2]

    if name1 in ['DT', 'DC', 'RC', 'RU']:
        fit1_atoms = ['C1\'', 'C6', 'N1', 'C2', 'C5', 'N3']
    else:
        fit1_atoms = ['C1\'', 'C8', 'N9', 'C4', 'N7', 'C5']

    if name2 in ['DT', 'DC', 'RC', 'RU']:
        fit2_atoms = ['C1\'', 'C6', 'N1', 'C2', 'C5', 'N3']
    else:
        fit2_atoms = ['C1\'', 'C8', 'N9', 'C4', 'N7', 'C5']

    atoms1 = mol1.fetchm(fit1_atoms)
    atoms2 = mol2.fetchm(fit2_atoms)

    fit_atoms(atoms1, atoms2, mol2.atoms)


def planarity(atom_list):

    coords = list(map(lambda a: a.x, atom_list))
    plan = _p.planarity(coords)
    return plan


def apply_fit_R(atoms, R):

    for atom in atoms:
        x_old = list(map(lambda x: x, atom.x))
        for r in range(3):
            atom.x[r] = 0
            for c in range(3):
                atom.x[r] += R[r][c]*x_old[c]

def center_vector( v ):
    vout = _p.center_vec( v )
    return( vout )

def calc_fit_R( cs1, cs2, m ):
    R = _p.calc_fit_R(cs1, cs2, m)
    return(R)


def fit(model1, model2, atom_names=[]):
    if atom_names:
        subset1 = model1.fetch_atoms(atom_names)
        subset2 = model2.fetch_atoms(atom_names)
        cs1 = list(map(lambda a: a.x, subset1))
        cs2 = list(map(lambda a: a.x, subset2))
    else:
        cs1 = model1.coords()
        cs2 = model2.coords()

    assert(len(cs1) == len(cs2))
    m = list(map(lambda x: 1., cs1))  # dummy array
    v = _p.center_vec(cs1)
    v2 = _p.center_vec(cs2)
    R = _p.calc_fit_R(cs1, cs2, m)
    model2.translate([-v2[0], -v2[1], -v2[2]])
    apply_fit_R(model2.atoms, R)
    model2.translate(v)


def fit_by_ndx(ref, model, ndx1, ndx2):
    crd1 = list(map(lambda i: ref.atoms[i-1].x, ndx1))
    crd2 = list(map(lambda i: model.atoms[i-1].x, ndx2))

    assert(len(crd1) == len(crd2))
    m = list(map(lambda x: 1., crd1))  # dummy array
    v = _p.center_vec(crd1)
    v2 = _p.center_vec(crd2)
    R = _p.calc_fit_R(crd1, crd2, m)
    model.translate([-v2[0], -v2[1], -v2[2]])
    apply_fit_R(model.atoms, R)
    model.translate(v)


def translate_by_ndx(struct, ndx):
    crd = list(map(lambda i: struct.atoms[i-1].x, ndx))
    m = list(map(lambda x: 1., crd))  # FIXME: variable m is not used
    v = _p.center_vec(crd)
    struct.translate([-v[0], -v[1], -v[2]])
    return(v)


def fit_atoms(fit_atoms1, fit_atoms2, rot_atoms2):

    cs1 = list(map(lambda a: a.x, fit_atoms1))
    cs2 = list(map(lambda a: a.x, fit_atoms2))
    assert len(cs1) == len(cs2)
    m = list(map(lambda x: 1., cs1))  # dummy array
    v = _p.center_vec(cs1)
    v2 = _p.center_vec(cs2)
    R = _p.calc_fit_R(cs1, cs2, m)
    for atom in rot_atoms2:
        atom.x[0] -= v2[0]
        atom.x[1] -= v2[1]
        atom.x[2] -= v2[2]
    apply_fit_R(rot_atoms2, R)
    for atom in rot_atoms2:
        atom.x[0] += v[0]
        atom.x[1] += v[1]
        atom.x[2] += v[2]
