# commands/hccnchange.py
import numpy as np
import pandas as pd
import os

from sklearn.metrics import mean_squared_error


def add_hccnchange_subparser(subparsers):
    """
    Add the hccnchange subcommand to the CLI tool.

    This subcommand is used to identify if the CN change from parent to child is correctly identified, using RMSE and ACC.
    """
    parser = subparsers.add_parser(
        "hccnchange",
        help="Identify if the CN change from parent to child is correctly identified, using RMSE and ACC."
    )
    parser.add_argument(
        "--tool-maternal-cna-files", 
        type=str, 
        nargs = "+",
        required=True, 
        metavar="",
        help="Paths to the maternal CNA files generated by different tools. , separated by spaces"
    )
    parser.add_argument(
        "--tool-paternal-cna-files", 
        type=str, 
        nargs = "+",
        required=True, 
        metavar="",
        help="Paths to the paternal CNA files generated by different tools. , separated by spaces"
    )
    parser.add_argument(
        "--tool-names", 
        type=str, 
        nargs = "+",
        required=True, 
        metavar="",
        help="List of tool names corresponding to the classification files"
    )
    parser.add_argument(
        "--changes-file", 
        type=str, 
        required=True, 
        metavar="",
        help="Path to the changes file that contains information about changes in CN status."
    )
    parser.add_argument(
        "--output-dir", 
        type=str, 
        default="./output", 
        metavar="",
        help="Path to the file where the results will be saved, default: './output' "
    )
    return parser


def calculate_predictions_by_haplotype(change_data, predict_data, haplotype_type):

    haplotype_truth = change_data[change_data['Haplotype'] == haplotype_type]
    
    predict_sorted = predict_data.merge(haplotype_truth, left_on='region', right_on='Segment', how='inner')

    if predict_sorted.empty:
        print("Predict sorted is empty.")
        return pd.DataFrame(columns= ['Parent', 'Child', 'Haplotype', 'Type', 'Segment', 'Change', 'Parent_predict_num', 'Child_predict_num'])

    def calculate_mode_for_columns(row, columns_prefix, df):
        columns = [col for col in df.columns if col.startswith(row[columns_prefix])]
        if columns:
            values = df.loc[row.name, columns]
            mode_value = values.mode()
            return mode_value.iloc[0] if not mode_value.empty else None
        else:
            return None

    predict_sorted['Parent_predict_num'] = predict_sorted.apply(
        calculate_mode_for_columns, axis=1, columns_prefix='Parent', df=predict_sorted
    )
    
    predict_sorted['Child_predict_num'] = predict_sorted.apply(
        calculate_mode_for_columns, axis=1, columns_prefix='Child', df=predict_sorted
    )

    columns_to_keep = ['Parent', 'Child', 'Haplotype', 'Type', 'Segment', 'Change', 'Parent_predict_num', 'Child_predict_num']
    predict_sorted_cleaned = predict_sorted[columns_to_keep].copy() 

    predict_sorted_cleaned.loc[:, 'Change_predict'] = (
        predict_sorted_cleaned['Parent_predict_num'].fillna(0).astype(int).astype(str) +
        "->" +
        predict_sorted_cleaned['Child_predict_num'].fillna(0).astype(int).astype(str)
    )

    return predict_sorted_cleaned

def calculate_CN_Change_accuracy(ground_truth,predict_data):
    parent = ground_truth.str.split('->').str[0].apply(pd.to_numeric, errors='coerce')
    parent_predict = predict_data.str.split('->').str[0].apply(pd.to_numeric, errors='coerce')

    child = ground_truth.str.split('->').str[1].apply(pd.to_numeric, errors='coerce')
    child_predict = predict_data.str.split('->').str[1].apply(pd.to_numeric, errors='coerce')
    return np.mean((parent-child) == (parent_predict-child_predict))

def calculate_CN_change_rmse(ground_truth,predict_data):
    parent = ground_truth.str.split('->').str[0].apply(pd.to_numeric, errors='coerce')
    parent_predict = predict_data.str.split('->').str[0].apply(pd.to_numeric, errors='coerce')

    child = ground_truth.str.split('->').str[1].apply(pd.to_numeric, errors='coerce')
    child_predict = predict_data.str.split('->').str[1].apply(pd.to_numeric, errors='coerce')

    rmse = np.sqrt(mean_squared_error(parent-child, parent_predict-child_predict))
    
    return rmse

# python src/main.py hccnchange /
# --tool-maternal-cna-files data/chisel_5x_maternal_cnv.csv /
# --tool-paternal-cna-files data/chisel_5x_paternal_cnv.csv /
# --tool-names chisel_5x --changes-file profile/changes.csv /

def run(args):
    # Validate input
    if (len(args.tool_maternal_cna_files) != len(args.tool_names) ) and (len(args.tool_maternal_cna_files) != len(args.tool_paternal_cna_files)):
        raise ValueError("The number of classification files must match the number of tool names.")
    
    os.makedirs(args.output_dir, exist_ok=True)
    
    change_ground_truth = pd.read_csv(args.changes_file)

    results = []
    for tool_maternal_path,tool_paternal_path, tool_name in zip(args.tool_maternal_cna_files, args.tool_paternal_cna_files,args.tool_names):

  
        predict_maternal = pd.read_csv(tool_maternal_path)
        predict_paternal = pd.read_csv(tool_paternal_path)

        maternal_result = calculate_predictions_by_haplotype(change_ground_truth, predict_maternal, haplotype_type='maternal')
        paternal_result = calculate_predictions_by_haplotype(change_ground_truth, predict_paternal, haplotype_type='paternal')

        combined_result = pd.concat([maternal_result, paternal_result], ignore_index=True)

        for type in combined_result['Type'].unique().tolist():


            rmse = calculate_CN_change_rmse(combined_result[combined_result['Type']==type]['Change'], combined_result[combined_result['Type']==type]['Change_predict'])
            acc = calculate_CN_Change_accuracy(combined_result[combined_result['Type']==type]['Change'], combined_result[combined_result['Type']==type]['Change_predict'])

            results.append({
                "Tool": tool_name,
                "Type": type,
                "RMSE": rmse,
                "ACC": acc
            })
    
    result_df = pd.DataFrame(results)
    output_file = os.path.join(args.output_dir, "evolution_onset_CN_Change.csv")
    result_df.to_csv(output_file, index=False)
    print(f"Results saved to {output_file}")
