"""
    This file is part of the TrIPP software
    (https://github.com/fornililab/TrIPP).
    Copyright (c) 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 pandas as pd 
import logging
import numpy as np
import re

logger = logging.getLogger(__name__)

def sort_key(label):
    left, chain = label.split(":")    
    match = re.match(r"([A-Za-z\+\-]+)(\d*)", left)
    resname, number = match.group(1), int(match.group(2)) 
    
    return (chain, resname, number)
    
def output_df(output_directory, output_prefix, data, save_disulphide_pka, extract_buriedness_data, predictor):
    """
    Combines the results from all worker processes and saves the pKa and buriedness (buried ratio) values to CSV files.
    Parameters
    ----------
    output_directory : str
        The path to the directory where output files will be saved.
    output_prefix : str
        Output file prefix.
    data : list
        A list of data from all workers, where each worker's data is a dictionary
        generated by extract_pka_buriedness_data
    save_disulphide_pka : bool
        Whether to save pKa values for cysteines in disulphide bonds.
        Only valid for PROPKA predictor, this parameter will be ignored if
        pKAI/pKAI+ predictor is used.
    extract_buriedness_data : bool
        Whether to extract and save buriedness data.
        Only valid for PROPKA predictor, this parameter will be ignored if
        pKAI/pKAI+ predictor is used.
    predictor : str
        The pKa predictor used.
    Returns
    -------
    disulphide_cys_col : list or None
        A list of residue identifiers for disulphide-bonded cysteines if save_disulphide_pka is True, 
        otherwise None.
    """
    pka_full_data = []
    buriedness_full_data = []
    time_list = []
    for data_worker in data:
        for pka_buriedness_resname_worker_dict in data_worker:
            for time, pka_buriedness_resname_dict in pka_buriedness_resname_worker_dict.items():
                pka_full_data.append(pka_buriedness_resname_dict['pka_list'])
                buriedness_full_data.append(pka_buriedness_resname_dict['buriedness_list'])
                time_list.append(time)
    pka_df = pd.DataFrame(pka_full_data, columns=pka_buriedness_resname_dict['residue_identifier_list'])
    pka_df.insert(0,'Time [ps]',time_list)
    
    try:
        pka_df = pka_df.astype(float).sort_values('Time [ps]')
    except ValueError:
        contains_star = pka_df.applymap(lambda x: '*' in x if isinstance(x, str) else False)
        coupled_residues = ", ".join(list(pka_df.loc[:, contains_star.any()].columns))
        logging.info(f"""PROPKA detected that your system has the following coupled residues:
{coupled_residues}
run -d in the optargs to consider alternative state.
""")
        pka_df = pka_df.applymap(lambda x: x.replace('*','') if isinstance(x, str) else x)
        pka_df = pka_df.astype(float).sort_values('Time [ps]')
        
    if extract_buriedness_data and predictor == 'propka':
        buriedness_df = pd.DataFrame(buriedness_full_data, columns=pka_buriedness_resname_dict['residue_identifier_list'])
        buriedness_df.insert(0,'Time [ps]',time_list)
        buriedness_df['Time [ps]'] = buriedness_df['Time [ps]'].astype(float)
        buriedness_df.sort_values('Time [ps]', inplace=True)

    if not save_disulphide_pka and predictor == 'propka':
        disulphide_cys_col = pka_df.columns[(pka_df == 99.99).any()].to_list()
        pka_df.drop(disulphide_cys_col, axis=1, inplace=True)
        if extract_buriedness_data:
            buriedness_df.drop(disulphide_cys_col, axis=1, inplace=True)
    else:
        disulphide_cys_col = None
        
    cols = [pka_df.columns[0]] + sorted(pka_df.columns[1:], key=sort_key)
    pka_df = pka_df[cols]    
    pka_df.to_csv(f'{output_directory}/{output_prefix}_{predictor}_pka.csv', index=False)
    
    if extract_buriedness_data and predictor == 'propka':
        cols = [buriedness_df.columns[0]] + sorted(buriedness_df.columns[1:], key=sort_key)
        buriedness_df = buriedness_df[cols]
        buriedness_df.to_csv(f'{output_directory}/{output_prefix}_{predictor}_buriedness.csv', index=False)

    return disulphide_cys_col