#! /usr/bin/env python

# MIT License
#
# Copyright (c) 2022 Shintaro Minami
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import numpy as np
from einops import repeat, rearrange

CONST_Q1Q2 = 0.084
CONST_F = 332
DEFAULT_CUTOFF = -0.5
DEFAULT_MARGIN = 1.0


def _unfold(a: np.ndarray, window: int, axis: int):
    idx = np.arange(window)[:, None] + np.arange(a.shape[axis] - window + 1)[None, :]
    unfolded = np.take(a, idx, axis=axis)
    return np.moveaxis(unfolded, axis - 1, -1)


def _check_input(coord):
    org_shape = coord.shape
    assert (len(org_shape) == 3) or (len(org_shape) == 4), (
        "Shape of input tensor should be [batch, L, atom, xyz] or [L, atom, xyz]"
    )
    coord = repeat(coord, "... -> b ...", b=1) if len(org_shape) == 3 else coord
    return coord, org_shape


def _get_hydrogen_atom_position(coord: np.ndarray) -> np.ndarray:
    # A little bit lazy (but should be OK) definition of H position here.
    vec_cn = coord[:, 1:, 0] - coord[:, :-1, 2]
    vec_cn = vec_cn / np.linalg.norm(vec_cn, axis=-1, keepdims=True)
    vec_can = coord[:, 1:, 0] - coord[:, 1:, 1]
    vec_can = vec_can / np.linalg.norm(vec_can, axis=-1, keepdims=True)
    vec_nh = vec_cn + vec_can
    vec_nh = vec_nh / np.linalg.norm(vec_nh, axis=-1, keepdims=True)
    return coord[:, 1:, 0] + 1.01 * vec_nh


def get_hbond_map(
    coord: np.ndarray,
    donor_mask: np.ndarray = None,
    cutoff: float = DEFAULT_CUTOFF,
    margin: float = DEFAULT_MARGIN,
    return_e: bool = False,
) -> np.ndarray:
    # check input
    coord, org_shape = _check_input(coord)
    b, l, a, _ = coord.shape
    # add pseudo-H atom if not available
    assert (a == 4) or (a == 5), "Number of atoms should be 4 (N,CA,C,O) or 5 (N,CA,C,O,H)"
    h = coord[:, 1:, 4] if a == 5 else _get_hydrogen_atom_position(coord)
    # distance matrix
    nmap = repeat(coord[:, 1:, 0], "... m c -> ... m n c", n=l - 1)
    hmap = repeat(h, "... m c -> ... m n c", n=l - 1)
    cmap = repeat(coord[:, 0:-1, 2], "... n c -> ... m n c", m=l - 1)
    omap = repeat(coord[:, 0:-1, 3], "... n c -> ... m n c", m=l - 1)
    d_on = np.linalg.norm(omap - nmap, axis=-1)
    d_ch = np.linalg.norm(cmap - hmap, axis=-1)
    d_oh = np.linalg.norm(omap - hmap, axis=-1)
    d_cn = np.linalg.norm(cmap - nmap, axis=-1)
    # electrostatic interaction energy
    e = np.pad(CONST_Q1Q2 * (1.0 / d_on + 1.0 / d_ch - 1.0 / d_oh - 1.0 / d_cn) * CONST_F, [[0, 0], [1, 0], [0, 1]])
    if return_e:
        return e
    # mask for local pairs (i,i), (i,i+1), (i,i+2)
    local_mask = ~np.eye(l, dtype=bool)
    local_mask *= ~np.diag(np.ones(l - 1, dtype=bool), k=-1)
    local_mask *= ~np.diag(np.ones(l - 2, dtype=bool), k=-2)
    # mask for donor H absence (Proline)
    donor_mask = np.array(donor_mask).astype(float) if donor_mask is not None else np.ones(l, dtype=float)
    donor_mask = repeat(donor_mask, "l1 -> l1 l2", l2=l)
    # hydrogen bond map (continuous value extension of original definition)
    hbond_map = np.clip(cutoff - margin - e, a_min=-margin, a_max=margin)
    hbond_map = (np.sin(hbond_map / margin * np.pi / 2) + 1.0) / 2
    hbond_map = hbond_map * repeat(local_mask, "l1 l2 -> b l1 l2", b=b)
    hbond_map = hbond_map * repeat(donor_mask, "l1 l2 -> b l1 l2", b=b)
    # return h-bond map
    hbond_map = np.squeeze(hbond_map, axis=0) if len(org_shape) == 3 else hbond_map
    return hbond_map


def assign(coord: np.ndarray, donor_mask: np.ndarray = None) -> np.ndarray:
    # check input
    coord, org_shape = _check_input(coord)
    # get hydrogen bond map
    hbmap = get_hbond_map(coord, donor_mask=donor_mask)
    hbmap = rearrange(hbmap, "... l1 l2 -> ... l2 l1")  # convert into "i:C=O, j:N-H" form
    # identify turn 3, 4, 5
    turn3 = np.diagonal(hbmap, axis1=-2, axis2=-1, offset=3) > 0.0
    turn4 = np.diagonal(hbmap, axis1=-2, axis2=-1, offset=4) > 0.0
    turn5 = np.diagonal(hbmap, axis1=-2, axis2=-1, offset=5) > 0.0
    # assignment of helical sses
    h3 = np.pad(turn3[:, :-1] * turn3[:, 1:], [[0, 0], [1, 3]])
    h4 = np.pad(turn4[:, :-1] * turn4[:, 1:], [[0, 0], [1, 4]])
    h5 = np.pad(turn5[:, :-1] * turn5[:, 1:], [[0, 0], [1, 5]])
    # helix4 first
    helix4 = h4 + np.roll(h4, 1, 1) + np.roll(h4, 2, 1) + np.roll(h4, 3, 1)
    h3 = h3 * ~np.roll(helix4, -1, 1) * ~helix4  # helix4 is higher prioritized
    h5 = h5 * ~np.roll(helix4, -1, 1) * ~helix4  # helix4 is higher prioritized
    helix3 = h3 + np.roll(h3, 1, 1) + np.roll(h3, 2, 1)
    helix5 = h5 + np.roll(h5, 1, 1) + np.roll(h5, 2, 1) + np.roll(h5, 3, 1) + np.roll(h5, 4, 1)
    # identify bridge
    unfoldmap = _unfold(_unfold(hbmap, 3, -2), 3, -2) > 0.0
    unfoldmap_rev = rearrange(unfoldmap, "b l1 l2 ... -> b l2 l1 ...")
    p_bridge = (unfoldmap[:, :, :, 0, 1] * unfoldmap_rev[:, :, :, 1, 2]) + (
        unfoldmap_rev[:, :, :, 0, 1] * unfoldmap[:, :, :, 1, 2]
    )
    p_bridge = np.pad(p_bridge, [[0, 0], [1, 1], [1, 1]])
    a_bridge = (unfoldmap[:, :, :, 1, 1] * unfoldmap_rev[:, :, :, 1, 1]) + (
        unfoldmap[:, :, :, 0, 2] * unfoldmap_rev[:, :, :, 0, 2]
    )
    a_bridge = np.pad(a_bridge, [[0, 0], [1, 1], [1, 1]])
    # ladder
    ladder = (p_bridge + a_bridge).sum(-1) > 0
    # H, E, L of C3
    helix = (helix3 + helix4 + helix5) > 0
    strand = ladder
    loop = ~helix * ~strand
    onehot = np.stack([loop, helix, strand], axis=-1)
    onehot = rearrange(onehot, "1 ... -> ...") if len(org_shape) == 3 else onehot
    return onehot


atomnum = {" N  ": 0, " CA ": 1, " C  ": 2, " O  ": 3}


# Credits: Shintaro Minami, adapted from pydssp to enable filtering by chain id
def read_pdbtext_with_checking(pdbstring: str, return_sequence: bool = False, chain_id=None):
    lines = pdbstring.split("\n")
    coords, atoms, resid_old, resname_old, sequence, check = [], None, None, None, [], []
    for l in lines:
        if l.startswith("ATOM"):
            if chain_id is not None and l[21] != chain_id:
                continue
            iatom = atomnum.get(l[12:16], None)
            resid = l[21:26]
            resname = l[17:20]
            if resid != resid_old:
                if atoms is not None:
                    coords.append(atoms)
                    sequence.append(resname_old)
                    check.append(atom_check)
                atoms, resid_old, resname_old, atom_check = [], resid, resname, []
            if iatom is not None:
                xyz = [float(l[30:38]), float(l[38:46]), float(l[46:54])]
                atoms.append(xyz)
                atom_check.append(iatom)
    if atoms is not None:
        coords.append(atoms)
        sequence.append(resname_old)
        check.append(atom_check)
    coords = np.array(coords)
    sequence = np.array(sequence)
    # check
    assert len(coords.shape) == 3, "Some required atoms [N,CA,C,O] are missing in the input PDB file"
    check = np.array(check)
    assert np.all(check[:, 0] == 0), (
        "Order of PDB line may be broken. It's required to be N->CA->C->O w/o any duplicate or lack"
    )
    assert np.all(check[:, 1] == 1), (
        "Order of PDB line may be broken. It's required to be N->CA->C->O w/o any duplicate or lack"
    )
    assert np.all(check[:, 2] == 2), (
        "Order of PDB line may be broken. It's required to be N->CA->C->O w/o any duplicate or lack"
    )
    assert np.all(check[:, 3] == 3), (
        "Order of PDB line may be broken. It's required to be N->CA->C->O w/o any duplicate or lack"
    )
    # output
    if return_sequence:
        return coords, sequence
    else:
        return coords
