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


def add_hconsetacc_subparser(subparsers):
    """
    Add the hconsetacc subcommand to the CLI tool.

    This subcommand is used to identify whether CNA events are correctly detected at their first branch of evolution, using ACC.
    """
    parser = subparsers.add_parser(
        "hconsetacc",
        help="Identify whether CNA events are correctly detected at their first branch of evolution, using 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_Evolution_Onset_accuracy(ground_truth,predict_data):
    return np.mean(predict_data == ground_truth)


##python src/main.py hconsetacc 
# --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():


            acc = calculate_Evolution_Onset_accuracy(combined_result[combined_result['Type']==type]['Change'], combined_result[combined_result['Type']==type]['Change_predict'])

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