"""
    @release_date  : $release_date
    @version       : $release_version
    @author        : Christos Matsingos, Ka Fu Man 
    
    This file is part of the TrIPP software
    (https://github.com/fornililab/TrIPP).
    Copyright (c) 2024 Christos Matsingos, Ka Fu Man and Arianna Fornili.

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, version 3.

    This program is distributed in the hope that it will be useful, but
    WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
    General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program. If not, see <http://www.gnu.org/licenses/>.
"""


import MDAnalysis as mda
import numpy as np
import pandas as pd 
from tripp._pymol_template_ import gen_pymol_template 
from tripp._model_pka_values_ import model_pka_values 
from tripp._create_mda_universe_ import create_mda_universe 
import logging

class Visualization: 


    """ 
    Class to process and colour-map residue properties on a reference structure using PyMOL.
    
    Parameters:
    ----------
    topology_file: str
        The path of the topology file used for visualization
    pKa_file: str or list, default=None
        The path or a list of paths for the CSV file(s) containing the pKa values to be visualized.
    correlation_file: str, default=None
        The path of the CSV file generated by PCProjectionScreening() containing the correlation 
        coefficient between a given property (PC1 projection in the tutorial) and the pKa value 
        of each ionisable residue.
    start: int, default=0
        Start time (from pKa_file). The average pKa or pKa difference values will be calculated over the frames 
        with simulation times from start to end. Not applicable when correlation values are visualized.
    end: int, default=-1
        End time (from pKa_file). The average pKa or pKa difference values will be calculated over the frames 
        with simulation times from start to end. Not applicable when correlation values are visualized.
    """


    def __init__(self, 
                 topology_file,
                 pka_file=None,
                 correlation_file=None,
                 start = 0, 
                 end = -1):
        self.logger = logging.getLogger()
        if self.logger.hasHandlers:
            for handler in self.logger.handlers[:]:
                self.logger.removeHandler(handler)
        logging.Formatter('%(message)s')
        self.logger.setLevel(logging.WARNING)
        
        self.topology_file = topology_file
        if correlation_file != None:
            self.correlation_file = pd.read_csv(correlation_file)
        self.u = create_mda_universe(topology_file=self.topology_file, trajectory_file=None) 
        
        if pka_file:
            if type(pka_file) == list: 
                pka_file_l = [] 
                for file in pka_file: 
                    pka_file_l.append(pd.read_csv(file)) 
                pka_values = pd.concat(pka_file_l) 
            
            else: 
                pka_values = pd.read_csv(pka_file)
            # If user wants to use part of the trajectory to calculate mean/difference_to_model_value.
            self.start = start
            if end == -1:
                self.end = pka_values['Time [ps]'].iloc[end]
            else:
                self.end = end
            self.pka_values = pka_values[pka_values['Time [ps]'].isin(np.arange(self.start,self.end+1))]
            
            # Reorder the values_df so that C- and N+ go to the end of the dataframe.
            N_C_cols = [col for col in self.pka_values.columns if 'C-' in col or 'N+' in col]
            Rest_cols = [col for col in self.pka_values.columns if col not in N_C_cols]
            self.pka_values = self.pka_values[Rest_cols + N_C_cols]
            
    def gen_pse(self, 
                pymol_path, 
                output_directory,
                output_prefix,
                coloring_method='mean', 
                lower_limit=0, 
                upper_limit=14,
                correlation_threshold=None,
                color_palette='red_white_blue'): 
        """
        Generates a PyMOL session file.
        
        Parameters
        ----------
        pymol_path: str
            Path to PyMOL. The script will spawn a subprocess shell to run a 
            python script in PyMOL.
        output_directory : str
            The directory where the output files are saved.
        output_prefix: str
            Prefix of the output PyMOL file (.pse). The prefix will be combined 
            with the coloring_method ('mean', 'difference_to_model_value' or 'correlation') 
            to give pse_output_filename.
        coloring_method: str, int, or float, default 'mean'
            Determines the property visualised as colour gradient. 
            If set to 'mean', the average pKa value calculated over all frames is used. 
            If set to 'difference_to_model_value', the difference between the PROPKA predicted 
            and model pKa value is used. 
            If a specific simulation time is indicated, the pKa values predicted for the 
            corresponding frames are visualised. 
            If set to 'correlation', the correlation coefficient between a time-dependent property
            and the pKa time evolution will be visualised. In this case a correlation_file will need
            to be provided and specified in the Visualization class.
        lower_limit: int or float, default=0
            Lower bound of the color scale in the PyMOL session. Residues with values
            below this threshold are assigned the gradient’s minimum color.
        upper_limit: int or float, default=14
            Upper bound of the color scale in the PyMOL session. Residues with values
            above this threshold are assigned the gradient’s maximum color.
        correlation_threshold: float, default=None
            If the coloring_method is set to 'correlation', this parameter can be used to
            filter the correlation values. Only residues with an absolute correlation value
            above the threshold will be included in the PyMOL session.
        color_palette: str, default='red_white_blue'
            Colour palette. The default is set to 'red_white_blue'. See PyMOL spectrum
            for allowed colour palettes (a three-colour palette is recommended).
        """

        # Calculation of values depending on colouring method
        if coloring_method == 'mean':
            del self.pka_values['Time [ps]']
            values_df = self.pka_values.mean(axis=0).to_frame(name='pKa').reset_index(names='Residue')
            tempfactors_output_topology_file = f"{output_directory}/{output_prefix}_mean.pdb"

        elif coloring_method == 'difference_to_model_value':
            del self.pka_values['Time [ps]']
            pka_values_mean = self.pka_values.mean(axis=0)
            for residue in pka_values_mean.keys():
                if 'N+' in residue:
                    pka_values_mean[residue] = pka_values_mean[residue] - model_pka_values['NTR']
                elif 'C-' in residue:
                    pka_values_mean[residue] = pka_values_mean[residue] - model_pka_values['CTR']
                else:
                    pka_values_mean[residue] = pka_values_mean[residue] - model_pka_values[residue[0:3]]
            values_df = pka_values_mean.to_frame(name='pKa').reset_index(names='Residue')
            tempfactors_output_topology_file = f"{output_directory}/{output_prefix}_difference_to_model_value.pdb"


        elif coloring_method == 'correlation':
            correlation_df = self.correlation_file
            if correlation_threshold is not None:
                correlation_df = correlation_df[abs(correlation_df['Correlation']) > correlation_threshold]
            values_df = correlation_df[['Residue', 'Correlation']]
            tempfactors_output_topology_file = f"{output_directory}/{output_prefix}_correlation.pdb"

        elif type(coloring_method) == int or type(coloring_method) == float:

            if type(self.pka_file) == list and len(self.pka_file) != 1:
                raise Exception('Please select only one CSV file when selecting a specific time step.')

            else:
                self.pka_values = self.pka_values[self.pka_values['Time [ps]'] == coloring_method]
                del self.pka_values['Time [ps]']
                values_df = self.pka_values.mean(axis=0).to_frame(name='pKa').reset_index(names='Residue')
                tempfactors_output_topology_file = f"{output_directory}/{output_prefix}_time{coloring_method}.pdb"
        else:
            raise SyntaxError(f"{coloring_method} is not allowed for coloring_method, allowed values are 'mean', 'difference_to_model_value', or 'correlation'.")
        
        # Unpacking the residues and values (pka or correlation) from values_df.
        # Looping through them to assign the value onto the tempfactor of ionisable
        # residues. The topology_file with the tempfactor is written as pdb and a PyMOL
        # session is generated as .pse.
        residue_identifiers, values = (columns for _, columns in values_df.items())
        for residue_identifier, value in zip(residue_identifiers, values):
            residue = residue_identifier.split(':')[0]
            chain = residue_identifier.split(':')[-1]
            if 'N+' in residue or 'C-' in residue:
                resid = int(residue[2:])
            else:
                resid = int(residue[3:])

            rounded_predicted_pka = round(value, 2)
            if 'N+' in residue:
                # Only tempfactor of NtermCap_atom_name of the first residue will be written
                ag = self.u.select_atoms(f'((around 2 (chainID {chain} and resid {resid} and backbone and type N)) and chainID {chain} and resid {resid}) and type H or (chainID {chain} and resid {resid} and backbone and type N)')
            elif 'C-' in residue:
                # Only tempfactor of CtermCap_atom_name of the last residue will be written
                ag = self.u.select_atoms(f'((around 2 (chainID {chain} and resid {resid} and backbone and type C and not name CA)) and chainID {chain} and resid {resid}) and type O or (chainID {chain} and resid {resid} and backbone and type C and not name CA)')
            else:
                # All atoms in the ionisable residue will be written
                ag = self.u.select_atoms(f'(chainID {chain} and resid {resid} and not backbone) or (chainID {chain} and resid {resid} and name CA)')
            ag.tempfactors = np.full(ag.tempfactors.shape,rounded_predicted_pka)
        ag = self.u.select_atoms('all')
        ag.write(tempfactors_output_topology_file)
        pse_output_filename = tempfactors_output_topology_file.replace('.pdb', '.pse')
        gen_pymol_template(tempfactors_output_topology_file,
                           pymol_path, 
                           pse_output_filename, 
                           values_df, 
                           lower_limit, 
                           upper_limit, 
                           color_palette)