# -*- coding: utf-8 -*-
"""
    __author__ = 'Zehra Sarica'
    __email__ = ['sarica16@itu.edu.tr','zehraacar559@gmail.com']
"""

import logging
import os

import pandas as pd
from biopandas.pdb import PandasPdb

from rinpy import utils
from rinpy.constants import *
from rinpy.log_util import log_time
from rinpy.utils import create_folder_not_exists, PdbRecord, write_ppdb_to_pdb_file, CentralityType

PDB_FILE = 'pdb_file'
OUTPUT_FILE = 'output_file'
KEY = 'key'


class CentralityPdbMapper:
    """ Creates centrality hub PDBs by replacing b-factor value with centrality score of the given centrality type.
    Here, the given pdb is regenerated by replacing b-factor value with centrality score of each centrality type
    This allows to visualize centrality scores on the pdb via PyMol or other tools supporting pdb format.

    Attributes
    ----------
    pdb_name : str, default: None
        PDB ID which must be same with file name, such as 4OBE (4OBE.pdb).
    destination_output_path : str, default: output, this path must be provided.
        Location of the output results in which all files for each PDB file process are stored in this path.
    calculation_options: dict, default: None
        The calculation options as the input for the whole process.
        Example:
        {
            'remove_hydrogen': {
                'is_checked': True,
                'value': 0
            },
            'betweenness': {
                'is_checked': True,
                'value': 5
            },
            'closeness': {
                'is_checked': False,
                'value': 10
            },
            'degree': {
                'is_checked': False,
                'value': 10
            },
            'cluster_number': {
                'is_checked': None,
                'value': 4
            },
            'cutoff': {
                'is_checked': True,
                'value': 4.5
            },
        }
        """

    def __init__(self, graph, pdb_name=None, destination_output_path=None, calculation_options=None):
        if destination_output_path is None:
            raise ValueError('You must provide an output path to proceed.')
        self.graph = graph
        self.pdb_name = pdb_name
        self.destination_output_path = destination_output_path
        create_folder_not_exists(self.destination_output_path)
        self.residue_average_pdb_file_path = os.path.join(
            os.path.join(str(self.destination_output_path), self.pdb_name),
            f'{self.pdb_name}_{RESIDUE_AVERAGE_PDB_FILE}')
        self.calculation_options = calculation_options
        self.ppdb = utils.get_ppdb(pdb_file_path=self.residue_average_pdb_file_path)

    @log_time("Pdb Mapper Process")
    def process(self):
        logging.info(f'Centrality PDB Mapper has been started...')
        base_file_path = os.path.join(self.destination_output_path, self.pdb_name)
        columns = [ATOM_NUMBER, CENTRALITY_SCORE, RESIDUE_NAME, CHAIN_ID, RESIDUE_NUMBER, INSERTION, SEGMENT_ID]
        dtypes = {ATOM_NUMBER: int, CENTRALITY_SCORE: float, RESIDUE_NAME: str, CHAIN_ID: str, RESIDUE_NUMBER: int,
                  INSERTION: str, SEGMENT_ID: str}
        sort_keys = [ATOM_NUMBER]

        centrality_types = {
            CentralityType.BET.display_name: {
                OUTPUT_FILE: CENTRALITY_CSV_TEMPLATE.format(type=CentralityType.BET.display_name),
                PDB_FILE: CENTRALITY_PDB_TEMPLATE.format(type=CentralityType.BET.display_name)
            },
            CentralityType.CLOS.display_name: {
                OUTPUT_FILE: CENTRALITY_CSV_TEMPLATE.format(type=CentralityType.CLOS.display_name),
                PDB_FILE: CENTRALITY_PDB_TEMPLATE.format(type=CentralityType.CLOS.display_name)
            },
            CentralityType.DEG.display_name: {
                OUTPUT_FILE: CENTRALITY_CSV_TEMPLATE.format(type=CentralityType.DEG.display_name),
                PDB_FILE: CENTRALITY_PDB_TEMPLATE.format(type=CentralityType.DEG.display_name)
            }
        }

        for centrality, files in centrality_types.items():
            option = self.calculation_options.get(centrality, {})
            if not option.get(IS_CHECKED, False):
                continue

            output_file = os.path.join(str(base_file_path), f"{self.pdb_name}_{files[OUTPUT_FILE]}")
            df = utils.get_df(output_file_path=output_file,
                              columns=columns,
                              dtypes=dtypes,
                              sort_keys=sort_keys)
            hub_file_path = os.path.join(self.destination_output_path, self.pdb_name,
                                         f"{self.pdb_name}_{files[PDB_FILE]}")
            self._write_to_pdb(df, hub_file_path)
            self._write_to_original_pdb_file(df, files[PDB_FILE])
        logging.info(f'Centrality Hub analyses has been done...')

    def _write_to_pdb(self, out_df, out_filename):
        logging.info(f'Starting creation of {out_filename}.')

        atom_df = utils.get_atom_pdb_df(ppdb=self.ppdb).copy()
        out_df = out_df.copy()

        out_df[ATOM_NUMBER] = out_df[ATOM_NUMBER].astype(int)
        atom_df[ATOM_NUMBER] = atom_df[ATOM_NUMBER].astype(int)

        merged_df = atom_df.merge(out_df[[ATOM_NUMBER, CENTRALITY_SCORE]], on=ATOM_NUMBER, how='left')
        merged_df[B_FACTOR] = merged_df[CENTRALITY_SCORE].apply(
            lambda x: float(TWO_DECIMAL_FLOAT_FORMATTER.format(x)) if pd.notna(x) else merged_df[B_FACTOR])

        merged_df[BLANK_4] = ''
        merged_df[OCCUPANCY] = merged_df[OCCUPANCY].fillna(0)

        clean_df = merged_df.copy()
        for col in [CENTRALITY_SCORE]:
            if col in clean_df.columns:
                clean_df = clean_df.drop(columns=[col])

        ppdb = PandasPdb()
        ppdb.df[PdbRecord.ATOM.name] = clean_df
        write_ppdb_to_pdb_file(ppdb=ppdb, out_filename=out_filename)

        logging.info(f'Finished creating {out_filename}.')

    def _write_to_original_pdb_file(self, out_df, out_filename):
        hub_file_path = os.path.join(self.destination_output_path, self.pdb_name,
                                     f"{self.pdb_name}_original_{out_filename}")

        nodes = self.graph.nodes()

        residue_numbers = [nodes[residue_index][RESIDUE_NUMBER] for residue_index in out_df[ATOM_NUMBER]]
        chain_ids = [nodes[residue_index][CHAIN_ID] for residue_index in out_df[ATOM_NUMBER]]
        insertions = [nodes[residue_index][INSERTION] for residue_index in out_df[ATOM_NUMBER]]
        residue_names = [nodes[residue_index][RESIDUE_NAME] for residue_index in out_df[ATOM_NUMBER]]
        segment_ids = [nodes[residue_index][SEGMENT_ID] for residue_index in out_df[ATOM_NUMBER]]

        centrality_df = out_df.copy()

        centrality_df[RESIDUE_NUMBER] = residue_numbers
        centrality_df[CHAIN_ID] = chain_ids
        centrality_df[INSERTION] = insertions
        centrality_df[RESIDUE_NAME] = residue_names
        centrality_df[SEGMENT_ID] = segment_ids

        centrality_df[KEY] = list(
            zip(centrality_df[RESIDUE_NAME], centrality_df[CHAIN_ID], centrality_df[RESIDUE_NUMBER],
                centrality_df[INSERTION], centrality_df[SEGMENT_ID]))

        residue_to_centrality = centrality_df.set_index(KEY)[CENTRALITY_SCORE].to_dict()

        base_pdb_df = utils.get_base_pdb_df(residue_to=residue_to_centrality,
                                            destination_output_path=self.destination_output_path,
                                            pdb_name=self.pdb_name)

        utils.write_ppdb_to_pdb_file_atom_and_hetatom(ppdb=base_pdb_df, out_filename=hub_file_path)


def main():
    pass


if __name__ == '__main__':
    main()
