#!/usr/bin/env python3

from Bio.PDB.PDBParser import PDBParser
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import matplotlib
from mpl_toolkits.mplot3d import Axes3D
from collections import OrderedDict
import sys, math, os

try:
    from lib import pdb_resdepth
except:
    import pdb_resdepth

def read_pdb(filename):
    """
        READs PDB FILE with BioPython
        INPUT:
            PDB path/filename(str)
        OUTPUT:
            structure(BioPython Structure object)
    """
    parser = PDBParser(PERMISSIVE=1)
    pdbid = filename.split(".")[0]
    structure = parser.get_structure(pdbid, filename)
    return(structure)


def pdb_data_extractor(structure, chainId, depth_dict, depthCutoff, resScale, method):
    """
    COMPUTES RESIDUE POSITIONS AND DEPTH from Structure object
    INPUT:
        structure(BioPython Structure object)
        chainId(str) chain for which to compute residue positions and depth
        depth_dict(dict) Dictionary of residue depth generated by calculate_resdepth()
        depthCutoff(float) depth Cutoff for surface determination
        resScale(str) resolution scale (residue or atom)
    OUPUT:
        df(pandas dataframe) positions of CB/CA and RESIDUE DEPTH (to add)
            x(list) x position of residues (CA or CB)
            y(list) y position of residues (CA or CB)
            z(list) z position of residues (CA or CB)
            resdepth(list) residue depth, 1 if at surface, else -1
    """
    amino = [
        "ALA",
        "ILE",
        "LEU",
        "PRO",
        "VAL",
        "PHE",
        "TRP",
        "TYR",
        "ASP",
        "GLU",
        "ARG",
        "HIS",
        "LYS",
        "SER",
        "THR",
        "CYS",
        "MET",
        "ASN",
        "GLN"]
    x = []
    y = []
    z = []
    resname = []
    resdepth = []
    for model in structure:
        for chain in model:
            if chain.get_id() == chainId :
                for residue in chain:
                    if "GLY" in residue.get_resname():
                        thisDepth = pdb_resdepth.resdepth_to_fft(residue = residue, cutoff = depthCutoff, mydict = depth_dict, method = method)
                        thisResname = residue.get_resname()+str(residue.get_full_id()[3][1])+residue.get_full_id()[2]

                        for atom in residue:
                            if resScale == "residue":
                                if atom.get_name() == "CA":
                                    posx = int(atom.get_coord()[0])
                                    x.append(posx)
                                    posy = int(atom.get_coord()[1])
                                    y.append(posy)
                                    posz = int(atom.get_coord()[2])
                                    z.append(posz)
                                    resdepth.append(thisDepth)
                                    resname.append(thisResname)
                                    break
                            elif resScale == "atom":
                                posx = int(atom.get_coord()[0])
                                x.append(posx)
                                posy = int(atom.get_coord()[1])
                                y.append(posy)
                                posz = int(atom.get_coord()[2])
                                z.append(posz)
                                resdepth.append(thisDepth)
                                resname.append(thisResname)

                    elif residue.get_resname() in amino:
                        thisDepth = pdb_resdepth.resdepth_to_fft(residue = residue, cutoff = depthCutoff, mydict = depth_dict, method = method)
                        thisResname = residue.get_resname()+str(residue.get_full_id()[3][1])+residue.get_full_id()[2]

                        for atom in residue:
                            if resScale == "residue":
                                if atom.get_name() == "CB":
                                    posx = int(atom.get_coord()[0])
                                    x.append(posx)
                                    posy = int(atom.get_coord()[1])
                                    y.append(posy)
                                    posz = int(atom.get_coord()[2])
                                    z.append(posz)
                                    resdepth.append(thisDepth)
                                    resname.append(thisResname)
                                    break
                            elif resScale == "atom":
                                posx = int(atom.get_coord()[0])
                                x.append(posx)
                                posy = int(atom.get_coord()[1])
                                y.append(posy)
                                posz = int(atom.get_coord()[2])
                                z.append(posz)
                                resdepth.append(thisDepth)
                                resname.append(thisResname)

    df = pd.DataFrame(x, columns = ["x"], index = resname)
    df["y"] = y
    df["z"] = z
    df["resdepth"] = resdepth
    return(df)


def interface_data_extractor(interface, depth_dict, depthCutoff, resScale):
    """
    COMPUTES RESIDUE POSITIONS AND DEPTH from Structure object
    INPUT:
        structure(BioPython Structure object)
        chainId(str) chain for which to compute residue positions and depth
        depth_dict(dict) Dictionary of residue depth generated by calculate_resdepth()
        depthCutoff(float) depth Cutoff for surface determination
        resScale(str) resolution scale (residue or atom)
    OUPUT:
        df(pandas dataframe) positions of CB/CA and RESIDUE DEPTH (to add)
            x(list) x position of residues (CA or CB)
            y(list) y position of residues (CA or CB)
            z(list) z position of residues (CA or CB)
            resdepth(list) residue depth, 1 if at surface, else -1
    """
    amino = [
        "ALA",
        "ILE",
        "LEU",
        "PRO",
        "VAL",
        "PHE",
        "TRP",
        "TYR",
        "ASP",
        "GLU",
        "ARG",
        "HIS",
        "LYS",
        "SER",
        "THR",
        "CYS",
        "MET",
        "ASN",
        "GLN"]
    x = []
    y = []
    z = []
    resname = []
    resdepth = []
    for residue in interface:
        if "GLY" in residue.get_resname():
            thisDepth = pdb_resdepth.resdepth_to_fft(residue = residue, cutoff = depthCutoff, mydict = depth_dict)
            thisResname = residue.get_resname()+str(residue.get_full_id()[3][1])+residue.get_full_id()[2]

            for atom in residue:
                if resScale == "residue":
                    if atom.get_name() == "CA":
                        posx = int(atom.get_coord()[0])
                        x.append(posx)
                        posy = int(atom.get_coord()[1])
                        y.append(posy)
                        posz = int(atom.get_coord()[2])
                        z.append(posz)
                        resdepth.append(thisDepth)
                        resname.append(thisResname)
                        break
                elif resScale == "atom":
                    posx = int(atom.get_coord()[0])
                    x.append(posx)
                    posy = int(atom.get_coord()[1])
                    y.append(posy)
                    posz = int(atom.get_coord()[2])
                    z.append(posz)
                    resdepth.append(thisDepth)
                    resname.append(thisResname)

        elif residue.get_resname() in amino:
            thisDepth = pdb_resdepth.resdepth_to_fft(residue = residue, cutoff = depthCutoff, mydict = depth_dict)
            thisResname = residue.get_resname()+str(residue.get_full_id()[3][1])+residue.get_full_id()[2]

            for atom in residue:
                if resScale == "residue":
                    if atom.get_name() == "CB":
                        posx = int(atom.get_coord()[0])
                        x.append(posx)
                        posy = int(atom.get_coord()[1])
                        y.append(posy)
                        posz = int(atom.get_coord()[2])
                        z.append(posz)
                        resdepth.append(thisDepth)
                        resname.append(thisResname)
                        break
                elif resScale == "atom":
                    posx = int(atom.get_coord()[0])
                    x.append(posx)
                    posy = int(atom.get_coord()[1])
                    y.append(posy)
                    posz = int(atom.get_coord()[2])
                    z.append(posz)
                    resdepth.append(thisDepth)
                    resname.append(thisResname)

    df = pd.DataFrame(x, columns = ["x"], index = resname)
    df["y"] = y
    df["z"] = z
    df["resdepth"] = resdepth
    return(df)
