# commands/cndetect.py

import pandas as pd
import numpy as np
import os

from sklearn.isotonic import spearmanr
from sklearn.metrics import mean_squared_error


def add_cndetect_subparser(subparsers):
    """
    Add the cndetect subcommand to the CLI tool.

    This subcommand is used to calculate the RMSE,ACC,SCC of copy number (CN) status at the bin level.
    """
    parser = subparsers.add_parser(
        "cndetect",
        help="Calculate the RMSE,ACC,SCC of CN status at the bin level."
    )
    parser.add_argument(
        "--tool-cna-files", 
        type=str, 
        nargs = "+",
        required=True, 
        metavar="",
        help="Paths to the CNA files generated by different tools. "
    )
    parser.add_argument(
        "--cna-profile-file", 
        type=str, 
        required=True, 
        metavar="",
        help="Path to the ground truth CNA profile file."
    )
    parser.add_argument(
        "--tool-names", 
        type=str, 
        nargs = "+",
        required=True, 
        metavar="",
        help="List of tool names corresponding to the classification files"
    )
    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_accuracy(all_cell_data, ground_truth):
    return np.mean(all_cell_data == ground_truth)

def calculate_rmse(all_cell_data, ground_truth):
    all_cell_data_sum = all_cell_data.apply(lambda x : x.str.split('|').str[0]).apply(pd.to_numeric, errors = "coerce") 
    + all_cell_data.apply(lambda x : x.str.split('|').str[1]).apply(pd.to_numeric, errors = "coerce")
    ground_truth_sum = ground_truth.apply(lambda x : x.str.split('|').str[0]).apply(pd.to_numeric, errors = "coerce") 
    + ground_truth.apply(lambda x : x.str.split('|').str[1]).apply(pd.to_numeric, errors = "coerce")

    return np.sqrt(mean_squared_error(ground_truth_sum ,all_cell_data_sum))

def evaluate_haplotype_predictions(haplotype, all_cell_data_path, ground_truth_path):
    """
    Evaluate the performance of a haplotype prediction tool, including RMSE, Spearman correlation coefficient, and accuracy.

    Parameters:
    haplotype (str): Specifies the haplotype mode, can be 'combined' or other values.
    all_cell_data_path (str): Path to the CSV file containing predicted data.
    ground_truth_path (str): Path to the CSV file containing ground truth data.

    Returns:
    tuple: A tuple containing the following metrics:
        - rmse (float): Root Mean Square Error (RMSE).
        - scc (float): Spearman correlation coefficient.
        - acc (float): Accuracy.
    """
    all_cell_data = pd.read_csv(all_cell_data_path, index_col=0)
    ground_truth = pd.read_csv(ground_truth_path, index_col=0)
    
    # Align rows (index) by keeping only the common indices
    common_index = all_cell_data.index.intersection(ground_truth.index)
    all_cell_data_aligned = all_cell_data.loc[common_index]
    ground_truth_aligned = ground_truth.loc[common_index]

    # Align columns by keeping only the common column names
    common_columns = all_cell_data_aligned.columns.intersection(ground_truth_aligned.columns)
    all_cell_data_aligned = all_cell_data_aligned[common_columns]
    ground_truth_aligned = ground_truth_aligned[common_columns]

    if haplotype == 'combined':
        rmse = calculate_rmse(ground_truth_aligned, all_cell_data_aligned)
    else:
        rmse = np.sqrt(mean_squared_error(ground_truth_aligned, all_cell_data_aligned))

    scc, _ = spearmanr(ground_truth_aligned.values.flatten(), all_cell_data_aligned.values.flatten())

    acc = calculate_accuracy(all_cell_data_aligned.values, ground_truth_aligned.values)
    
    return rmse, scc, acc


##  python src/main.py cndetect 
    # --tool-cna-files data/chisel_2x_500kb_combined_cnv.csv data/chisel_5x_500kb_combined_cnv.csv 
    # --cna-profile-file data/profile_combined_cnv_500kb.csv 
    # --tool-names chisel_2x_500kb chisel_5x_500kb

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

    results = []
    for tool_cna_path, tool in zip(args.tool_cna_files, args.tool_names):
        all_cell_data_path = tool_cna_path
        ground_truth_path = args.cna_profile_file
        haplotype = "combined"

        rmse, scc, acc = evaluate_haplotype_predictions(haplotype, all_cell_data_path, ground_truth_path)
        
        results.append({
            "Tool":  tool,
            "RMSE": rmse,
            "SCC": scc,
            "ACC": acc
        })
    
    result_df = pd.DataFrame(results)
    output_file = os.path.join(args.output_dir, "bin_level_results.csv")
    result_df.to_csv(output_file, index=False)
    print(f"Results saved to {output_file}")
