import os
import gzip
import argparse
from pathlib import Path
from time import time
from multiprocessing import Pool

import pandas as pd
import numpy as np

pd.options.mode.chained_assignment = None

DST_TRESHHOLD = 20.0


def run(pdb_path, output_path, granularity):
    """Calcuklate all protein descriptors

    Args:
        pdb_path (str): path to PDB
        output_path (str): output file
        granularity (int): granularity of the descriptors
    """    
    proteins = [p for p in pdb_path.glob("**/*.gz") if not os.path.isdir(p)]
    with Pool() as pool:
        results = []
        for protein in proteins:
            result = pool.apply_async(process_protein, (protein, granularity, 3, 7))
            results.append(result)
        print("Processing started")
        t = time()
        data = [n for sublist in [result.get()['data'] for result in results] for n in sublist]
        index = [n for sublist in [result.get()['index'] for result in results] for n in sublist]
        t = time() - t
        df = pd.DataFrame(index=index, data=data)
        df.to_pickle(Path(output_path))
        print(f'Processing took {t:.1f} seconds')



def process_protein(protein, granularity, fstart, fend):
    """Create protein descriptor from file

    Args:
        protein (str): path to protein file
        granularity (int): descriptor granularity
        fstart (_type_): filename protein id start index
        fend (_type_): filename protein id end index

    Returns:
        dict: protein chain id and the descriptor
    """    
    protein_chains = read_and_extract(protein, granularity)
    
    data_list = []
    index_list = []
    for (chain, df) in protein_chains:
        desc = create_descriptor(df, granularity)
        data_list.append(desc)
        index_list.append(f"{protein.stem[fstart:fend].upper()}:{chain}")
    return {'index': index_list, 'data': data_list}


def create_descriptor(chain_df, granularity):
    """Create protein descriptor from extracted data

    Args:
        chain_df (DataFrame): extracted protein data
        granularity (int): granularity of the descriptor
    """    
    def compute_matrix(row):
        dist =  np.linalg.norm(np.array([row['x_x'], row['y_x'], row['z_x']]) - np.array([row['x_y'], row['y_y'], row['z_y']]))
        return (20.0 - dist) / DST_TRESHHOLD if dist <= 20.0 else 0.0
    
    chain_df['key'] = 0
    chain_df = chain_df.sort_values('normalized_rs')
    chain_df = pd.merge(chain_df, chain_df, on='key', how='left')
    chain_df['dist'] = chain_df.apply(lambda row: compute_matrix(row), axis=1)

    chain_df = chain_df.pivot(index='normalized_rs_x', columns='normalized_rs_y', values='dist')
    nparray = chain_df.to_numpy(dtype='float16')
    shape = nparray.shape[0]
    nparray = np.pad(nparray, (0, granularity - shape), "constant")
    nup = nparray[np.triu_indices(nparray.shape[0], k=1)]
    return nup


def read_and_extract(protein_file, granularity):
    """Extract protein descriptor data from PDB gz file

    Args:
        protein_file (str): path to protein file
        granularity (int): descriptor granularity
    """
    def remap(n, min_, max_):
        if max_ - min_ >= granularity:
            return int((n - min_) / (max_ - min_) * (granularity - 1)) + 1
        return n - min_ + 1
    
    df = pd.DataFrame(columns=['atom', 'residue', 'chain', 'residue_sequence', 'x', 'y', 'z'])

    atoms = []
    residues = []
    chains = []
    residue_sequences = []
    xs = []
    ys = []
    zs = []

    with gzip.open(protein_file, 'rt') as file:
        model = True
        for line in file:

            words = line.split()
            if len(words) == 0:
                continue
            if words[0] == "MODEL" and not model and words[1] == '1':
                model = True
            if model and line[0:4] == "ATOM":
                atoms.append(line[12:16].strip())
                residues.append(line[17:20].strip())
                chains.append(line[21:22].strip())
                residue_sequences.append(line[22:27].strip())
                xs.append(line[30:38].strip())
                ys.append(line[38:46].strip())
                zs.append(line[46:54].strip())
            if words[0] == "ENDMDL":
                model = False
                break

    if len(residue_sequences) == 0:
        return []

    coded_residue_sequences = []
    index = 1
    last = residue_sequences[0]
    for rs in residue_sequences:
        if rs == last:
            coded_residue_sequences.append(index)
        else:
            index += 1
            coded_residue_sequences.append(index)
            last = rs

    df = pd.DataFrame({'atom': atoms, 'residue': residues, 'chain': chains, 'residue_sequence': coded_residue_sequences, 'x': xs, 'y': ys, 'z': zs})
    df = df.astype({'residue_sequence': int, 'x': float, 'y': float, 'z': float})
    chains = df['chain'].unique()
    tables = []
    for chain in chains:
        table = df[df["chain"] == chain]
        min_ = np.min(table["residue_sequence"])
        max_ = np.max(table["residue_sequence"])
        table.loc[:, "normalized_rs"] = table.loc[:, "residue_sequence"].apply(lambda x: remap(x, min_, max_))
        table = table.drop(['residue_sequence'], axis=1)
        table = table.groupby(['chain', 'normalized_rs']).mean()
        table = table.reset_index()
        table = table.sort_values(['chain', 'normalized_rs'])
        tables.append((chain, table))
    return tables


if __name__ == "__main__":    
    parser = argparse.ArgumentParser()
    parser.add_argument("--pdb", type=str, required=True)
    parser.add_argument("--output", type=str, required=True)
    parser.add_argument("--granularity", type=int, required=False, default=10)

    args = parser.parse_args()

    input_path = Path(args.pdb)
    output_path = Path(args.output)
    assert input_path.exists()

    run(input_path, output_path, args.granularity)
