#!/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 module comtains some functions to build protein
structures from scratch.

Most important:
- build_chain

Usage:
    >>> from pmx.builder import *
    >>> ch = build_chain('AAAAAA')    # builds an ordinary polyalanine
    >>> ch.write('polyala.pdb')       # store file
    >>> ch = build_chain('FGHRTCV',hydrogens = False) # build chain without H
    >>> ch = build_chain('FGHRTCV',ss='HHHHHHH') build as helix
    >>> ch = build_chain('FGHRTCV',dihedrals = ((phi1,psi1,omega1),(...)))
    build chain with defined dihedral angles
"""

import sys
import copy as cp
from numpy import array, pi, linalg, cross
from . import library
from .geometry import Rotation
from .chain import Chain
from .atom import Atom
from .model import Model
from .molecule import Molecule
from . import geometry


def add_bp(m, strand=None, bRNA=False):
    if strand:
        N = len(strand)/2
        if bRNA:
            N = len(strand)
    else:
        N = 1

    r = Rotation([0, 0, 0], [0, 0, 1])
    phi = 0.0
    phi = N*36*pi/180.
    for atom in m.atoms:
        atom.x = r.apply(atom.x, phi)
    for atom in m.atoms:
        atom.x[2] += N*3.4


def make_3ter(r):
    # we add a proton at O3'
    c3, o3 = r.fetchm(['C3\'', 'O3\''])
    v = array(o3.x)-array(c3.x)
    normed = v*1./linalg.norm(v)
    newpos = array(o3.x) + normed
    a = Atom(name='H3T', x=newpos)
    r.append(a)


def make_5ter(r):
    del r['O1P']
    del r['O2P']
    # we add a proton at O5'
    a = r.fetch_atoms('P')[0]
    a.name = 'H5T'
    h5t, o5 = r.fetchm(['H5T', 'O5\''])
    v = array(h5t.x)-array(o5.x)
    normed = v*1./linalg.norm(v)
    newpos = array(o5.x) + normed
    h5t.x = newpos


def build_dna_strand(seq):

    dic = library.pmx_data_file('bp.pkl')

    seq = seq.lower()
    ss = []
    for x in seq:
        ss.append(x)
    new = []
    new.extend(dic[ss[0]].residues)
    ss.pop(0)
    while ss:
        newbp = ss.pop(0)
        newm = dic[newbp].copy()
        add_bp(newm, new)
        new.extend(newm.residues)
    chA = []
    chB = []
    a = 1
    b = 1
    for r in new:
        if r.chain_id == 'B':
            r.set_resid(a)
            a += 1
            chA.append(r)
        else:
            chB.append(r)
    chB.reverse()
    for r in chB:
        r.set_resid(b)
        b += 1
    mm = Model(residues=chA+chB)
    mm.chains[1].set_chain_id('C')
    mm.chains[0].set_chain_id('A')
    mm.chains[1].set_chain_id('B')
    for chain in mm.chains:
        r = chain.residues[0]
        r.set_resname(r.resname+'5')
        make_5ter(r)
        r = chain.residues[-1]
        r.set_resname(r.resname+'3')
        make_3ter(r)
    return mm


def build_rna_strand(seq):

    dic = library.pmx_data_file('rna.pkl')

    seq = seq.lower()
    ss = []
    for x in seq:
        ss.append(x)
    new = []
    new.extend(dic[ss[0]].residues)
    ss.pop(0)
    while ss:
        newbp = ss.pop(0)
        newm = dic[newbp].copy()
        add_bp(newm, new, True)
        new.extend(newm.residues)
    chA = []
    a = 1
    for r in new:
        r.set_resid(a)
        a += 1
        chA.append(r)
    mm = Model(residues=chA)
    mm.chains[0].set_chain_id('A')
    for chain in mm.chains:
        r = chain.residues[0]
        r.set_resname(r.resname+'5')
        make_5ter(r)
        r = chain.residues[-1]
        r.set_resname(r.resname+'3')
        make_3ter(r)
    return mm


def get_fragments():
    dic = library.pmx_data_file('fragments.pkl')
    n = len(dic.keys())
    print("pmx__> # Fragments loaded: %d" % n, file=sys.stderr)
    return dic


def read_pdb_with_connect(f):
    m = Model().read(f)
    m.nm2a()
    mol = m.residues[0]
    l = open(f).readlines()
    conn = []
    for line in l:
        if line[0:6] == 'CONECT':
            conn.append([int(x) for x in line[6:].split()])
    for lst in conn:
        if len(lst) > 1:
            atoms = mol.get_by_id(lst)
            a0 = atoms[0]
            for atom in atoms[1:]:
                a0.bonds.append(atom)
    return mol


def write_pdb_with_connect(mol, f, n=1):
    if not hasattr(f, "write"):
        fp = open(f, "w")
    else:
        fp = f
    print("MODEL%5d" % n, file=fp)
    for atom in mol.atoms:
        print('{}'.format(atom), file=fp)
    for atom in mol.atoms:
        s = "CONECT%5d" % atom.id
        for a in atom.bonds:
            s += '%5d' % a.id
        print('{}'.format(s), file=fp)


def attach_group(atom, mol):
    master = atom.molecule

    bb = atom.bonds[0]
    R = mol.fetch_atoms('R#')[0]
    bR = R.bonds[0]
    diff = array(R.x)-array(bb.x)
    if bR.name[0] == 'C':
        l = 1.54
    elif bR.name[0] == 'S':
        l = 1.8
    elif bR.name[0] == 'O':
        l = 1.4
    else:
        l = 1.5
    for a in mol.atoms:
        a.x = array(a.x)-diff
    ang = R.angle(bR, atom)
    v1 = array(atom.x)-array(bb.x)
    v2 = array(bR.x)-array(R.x)
    cr = cross(v1, v2)
    x = 1./linalg.norm(cr)
    cr = cr*x
    cr = R.x + cr
    r = Rotation(cr, R.x)
    for a in mol.atoms:
        a.x = r.apply(a.x, ang)
    vec = array(bb.x) - array(bR.x)
    x = 1./linalg.norm(vec)
    for a in mol.atoms:
        a.x += vec*(1-x*l)
    for a in master.atoms:
        if atom in a.bonds:
            a.bonds.remove(atom)
    master.remove_atom(atom)
    bR.bonds.remove(R)
    mol.remove_atom(R)
    bR.bonds.append(bb)
    bb.bonds.append(bR)
    for a in mol.atoms:
        master.append(a)


def _is_mainchain( a ):
    if a.name=='N':
        return(True)
    if a.name=='H' or a.name=='HN' or a.name=='H1' or a.name=='H2' or a.name=='H3':
        return(True)
    if a.name=='CA':
        return(True)
    if a.name=='C':
        return(True)
    if a.name=='O' or a.name=='OXT' or a.name=='OX' or a.name=='OT':
        return(True)
    return(False)

def make_stereoisomer_residue( m ):
    mtmp = cp.deepcopy( m )
    # fit
    N1, CA1, C1 = m.fetchm(['N', 'CA', 'C'])
    N2, CA2, C2 = mtmp.fetchm(['N', 'CA', 'C'])
    geometry.fit_atoms( [N1,CA1,C1], [C2,CA2,N2], mtmp.atoms )
    for atmp,a in zip(mtmp.atoms,m.atoms):
        if _is_mainchain( a )==False:
            a.x = atmp.x
    return(m)

def make_residue(key, hydrogens=True, chirality='L'):
    """ returns a molecule object with
    default geometry"""

    if key not in library._aacids:
        raise(KeyError, "Residue %s not known" % key)
    m = Molecule()
    m.unity = 'A'
    m.resname = key
    for i, entry in enumerate(library._aacids[key]):
        if hydrogens is False and entry[0][0] == 'H':
            continue
        else:
            a = Atom()
            a.id = i+1
            a.name = entry[0]
            a.symbol = a.name[0]
            a.x = entry[1]
            a.occ = 1.
            a.resname = key
            a.m = library._atommass[a.symbol]
            a.unity = 'A'
            m.atoms.append(a)

    if chirality=='D':
        m = make_stereoisomer_residue( m )
    return m


def build_chain(sequence, dihedrals=None,
                hydrogens=True, ss=None, chain_id='X', bCyclic=False, bCap=False,
                chirality=None):

    # chirality 
    if chirality==None:
        chirality = ''
        for s in sequence:
            chirality = chirality+'L'

    # if cyclic, form dihedrals
    n = len(sequence)
    if bCyclic==True:
        dihedrals = []
        fudge = float(n)/5.0 # empirical fudge to have some space between the termini
        dphi = 2.0*360.0/float(n+fudge)
        if bCap==True: # leave space for caps
            dphi = 2.0*360.0/float(n+fudge+1)
        dpsi = 0.0
        dom = 180.0
        sign = -1.0
        for i in range(0,n):
            dihedrals.append( (sign*dphi,-1.0*sign*dpsi,dom) )
            sign = -1.0*sign

    ch = Chain()
    if dihedrals is None:
        if ss is None:
            dihedrals = []
            for i in range(len(sequence)):
                dihedrals.append(library._extended)
        else:
            dihedrals = []
            for i, s in enumerate(ss):
                if s == '-':
                    dihedrals.append(library._extended)
                elif s == 'H':
                    dihedrals.append(library._helix)
                elif s == 'E':
                    dihedrals.append(library._anti_beta)

    rl = []
    start = make_residue(library._aacids_dic[sequence[0]], hydrogens=hydrogens, chirality=chirality[0])
    start.set_resid(1)
    set_psi(start, dihedrals[0][1])
    rl.append(start)
    for i in range(1,n):
        aa = sequence[i]
        ld = chirality[i]
        phi, psi, om = dihedrals[i]
        phi2, psi2, om2 = dihedrals[i-1]
        mol = attach_aminoacid(start, library._aacids_ext_amber[aa],
                               hydrogens=hydrogens,
                               phi=phi, psi=psi, omega=om2,chirality=ld)
        mol.set_resid(i+1)
        set_phi(start, mol, dihedrals[i][0])
        start = mol
        rl.append(mol)
    al = []
    for r in rl:
        al += r.atoms
    ch.unity = 'A'
    ch.atoms = al
    ch.renumber_atoms()
    ch.residues = rl
    ch.set_chain_id(chain_id)
    for atom in ch.atoms:
        atom.unity = 'A'

    if bCap==True:
        ch.add_nterm_cap()
        ch.add_cterm_cap()

    return ch


def set_phi(mol1, mol2, phi):
    # select c, n, ca, c
    C = mol1.fetchm(['C'])[0]
    N, CA, C2 = mol2.fetchm(['N', 'CA', 'C'])
    dih = C.dihedral(N, CA, C2)
    phi = pi/180*phi
    delta = phi-dih

    # rotation is around N -> CA
    r = Rotation(N.x, CA.x)

    for atom in mol2.atoms:
        if atom.name not in ['N', 'H']:
            atom.x = r.apply(atom.x, delta)


def set_psi(mol, phi):

    # select n, c, ca, o
    N, CA, C, O = mol.fetchm(['N', 'CA', 'C', 'O'])
    dih = N.dihedral(CA, C, O)  # -pi
    phi = pi/180*phi - pi  # psi is defined with N'
    delta = phi - dih
    r = Rotation(CA.x, C.x)

    # only O is affected
    O.x = r.apply(O.x, delta)


def set_omega(mol1, mol2, phi):

    # select c, n, ca, c
    CA, C = mol1.fetchm(['CA', 'C'])
    N, CA2, C2 = mol2.fetchm(['N', 'CA', 'C'])
    dih = CA.dihedral(C, N, CA2)  # -pi
    phi = pi/180*phi
    delta = phi-dih

    # rotation is around C -> N
    r = Rotation(C.x, N.x)
    for atom in mol2.atoms:
        if atom.name != 'N':
            atom.x = r.apply(atom.x, delta)


def attach_aminoacid(mol, resname, hydrogens=True,
                     omega=180., phi=-139., psi=135., chirality='L'):

    new = make_residue(resname, hydrogens=hydrogens, chirality=chirality)
    # we need Ca, C and O from mol
    Ca, C, O = mol.fetchm(['CA', 'C', 'O'])
    # and N from new
    N, CA2 = new.fetchm(['N', 'CA'])

    l = 1.33  # length of peptide bond

    v = array(C.x)-array(Ca.x)
    x = 1./linalg.norm(v)
    vec = v*x*l
    newpos = array(C.x)+vec
    t = array(N.x)-newpos

    # shift to new position
    for a in new.atoms:
        a.x = array(a.x)-t

    # gen rotation vector
    # this rotation yields the correct
    # O-C-N angle
    v1 = array(C.x)-array(Ca.x)
    v2 = array(O.x)-array(C.x)
    cr = cross(v1, v2)
    x = 1./linalg.norm(cr)
    cr = cr*x
    cr = C.x + cr

    # do first rotation
    r = Rotation(cr, C.x)
    for atom in new.atoms:
        atom.x = r.apply(atom.x, pi/3.)

    # do next rotation
    # here we correct the C-N-Ca angle
    v1 = array(N.x)-array(C.x)
    v2 = array(CA2.x)-array(N.x)
    cr = cross(v1, v2)
    x = 1./linalg.norm(cr)
    cr = cr*x
    cr = N.x + cr
    an = N.angle(C, CA2)
    delta = an - 120*pi/180
    # do second rotation
    r = Rotation(cr, N.x)
    for atom in new.atoms:
        if atom.name != 'N':
            atom.x = r.apply(atom.x, -delta)

    # correct H if we have one
    H = new.fetch('H')
    if H:
        H = H[0]
        # we place the H-atom along the
        # N-Ca vector.
        # then we move it
        v1 = array(CA2.x)-array(N.x)
        x = 1./linalg.norm(v1)
        H.x = N.x + v1*x
        H.x = r.apply(H.x, 2*pi/3.)

    set_omega(mol, new, omega)
    set_phi(mol, new, phi)
    set_psi(new, psi)

    return new
