# commands/hcPhasing.py

import pandas as pd
import os

def add_hcPhasing_subparser(subparsers):
    """
    Add the hcPhasing subcommand to the CLI tool.

    This subcommand is used to calculate the mismatch error and switch error.
    
    """
    parser = subparsers.add_parser(
        "hcPhasing",
        help="Calculate Mismatch Error and Switch Error."
    )
    parser.add_argument(
        "--tool-maternal-cna-files", 
        type=str, 
        nargs = "+",
        required=True, 
        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, 
        help="Paths to the paternal CNA files generated by different tools. , separated by spaces"
    )
    parser.add_argument(
        "--ground-truth-maternal-file", 
        type=str, 
        required=True, 
        help="Path to the ground truth maternal file."
    )
    parser.add_argument(
        "--ground-truth-paternal-file", 
        type=str, 
        required=True, 
        help="Path to the ground truth paternal file."
    )
    parser.add_argument(
        "--tool-names", 
        type=str, 
        nargs = "+",
        required=True, 
        help="List of tool names corresponding to the classification files"
    )
    parser.add_argument(
        "--output-dir", 
        type=str, 
        default="./output", 
        help="Path to the file where the results will be saved, default: './output' "
    )
    return parser

def generate_binary_df(tool_maternal_df,tool_paternal_df):
    """
    Generate binary DataFrames for maternal and paternal data.
    """
    maternal_binary_df = (tool_maternal_df > tool_paternal_df).astype(int)
    maternal_binary_df = maternal_binary_df.mask(tool_maternal_df == tool_paternal_df, -1)
    maternal_binary_df = maternal_binary_df.mask(tool_maternal_df.isna() | tool_paternal_df.isna())

    paternal_binary_df = (tool_paternal_df > tool_maternal_df).astype(int)
    paternal_binary_df = paternal_binary_df.mask(tool_maternal_df == tool_paternal_df, -1)
    paternal_binary_df = paternal_binary_df.mask(tool_maternal_df.isna() | tool_paternal_df.isna())

    return maternal_binary_df, paternal_binary_df


def evaluate_mismatch_and_switch_error(ground_maternal_binary_df,ground_paternal_binary_df,maternal_binary_df,paternal_binary_df, tool_name="tool_name"):

    ground_maternal_binary_df = ground_maternal_binary_df.dropna(axis=1, how='all')
    ground_paternal_binary_df = ground_paternal_binary_df.dropna(axis=1, how='all')

    common_index = ground_maternal_binary_df.index.intersection(maternal_binary_df.index)
    common_columns = ground_maternal_binary_df.columns.intersection(maternal_binary_df.columns)

    ground_maternal_binary_df = ground_maternal_binary_df.loc[common_index, common_columns]
    ground_paternal_binary_df = ground_paternal_binary_df.loc[common_index, common_columns]

    maternal_binary_df = maternal_binary_df.loc[common_index, common_columns]
    paternal_binary_df = paternal_binary_df.loc[common_index, common_columns]

    column_prefixes = maternal_binary_df.columns.str.extract(r"(^[^_]+)_")[0].unique()

    result_list = []

    for prefix in column_prefixes:

        group_columns = [col for col in maternal_binary_df.columns if col.startswith(f"{prefix}_")]

        gm = ground_maternal_binary_df[group_columns]
        gp = ground_paternal_binary_df[group_columns]
        m = maternal_binary_df[group_columns]
        p = paternal_binary_df[group_columns]

        maternal_mask = gm != -1
        paternal_mask = gp != -1

        maternal_diff_count = (m != gm)[maternal_mask].sum().sum()
        paternal_diff_count = (p != gp)[paternal_mask].sum().sum()

        total_maternal = maternal_mask.sum().sum()
        total_paternal = paternal_mask.sum().sum()


        if maternal_diff_count != paternal_diff_count:
            print(f"{tool_name} mismatch maternal and paternal result not same!")

        result = {
            "tool_name": tool_name,
            "cell_group": prefix,
            "mismatch_count": maternal_diff_count,
            "total": total_maternal,
            "mismatch_ratio": maternal_diff_count / total_maternal if total_maternal > 0 else None,
        }

        #-----------------switch error-----------------

        n_row, n_col =  gm.shape

        total = 0
        maternal_diff_count = 0
        paternal_diff_count = 0
        for row in range(0,n_row):
            for col in range(0,n_col-1):
                if gm.iloc[row,col] == -1 or gm.iloc[row,col+1] == -1:
                    continue
                else:
                    total +=1
                    if gm.iloc[row,col] != m.iloc[row,col] or gm.iloc[row,col+1] != m.iloc[row,col+1]:
                        maternal_diff_count +=1

                    if gp.iloc[row,col] != p.iloc[row,col] or gp.iloc[row,col+1] != p.iloc[row,col+1]:
                        paternal_diff_count +=1


        if maternal_diff_count != paternal_diff_count:
            print(f"{tool_name} switch error maternal and paternal result not same!")

        result.update({
            "switch_error_count": maternal_diff_count,
            "total_switch_compare_count": total,
            "switch_error_ratio": maternal_diff_count / total if total > 0 else None,
        })

        result_list.append(result)
            
    
    return result_list

# python src/main.py hcPhasing --tool-maternal-cna-files data/chisel_5x_maternal_cnv.csv 
# --tool-paternal-cna-files data/chisel_5x_paternal_cnv.csv 
# --ground-truth-maternal-file data/ground_truth_maternal_cnv.csv 
# --ground-truth-paternal-file data/ground_truth_paternal_cnv.csv --tool-names chisel_5x

def run(args):
    # Validate input
    if len(args.tool_maternal_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)

    ground_maternal_df = pd.read_csv(args.ground_truth_maternal_file, index_col=0)
    ground_paternal_df = pd.read_csv(args.ground_truth_paternal_file, index_col=0)

    ground_maternal_binary_df, ground_paternal_binary_df = generate_binary_df(ground_maternal_df, ground_paternal_df)

    results = []
    for maternal_path,paternal_path, tool in zip(args.tool_maternal_cna_files,args.tool_paternal_cna_files, args.tool_names):
        tool_maternal_df = pd.read_csv(maternal_path, index_col=0)
        tool_paternal_df = pd.read_csv(paternal_path, index_col=0)

        maternal_binary_df, paternal_binary_df = generate_binary_df(tool_maternal_df, tool_paternal_df)

        result_list = evaluate_mismatch_and_switch_error(ground_maternal_binary_df, ground_paternal_binary_df,maternal_binary_df, paternal_binary_df, tool)
        
        results.extend(result_list)
    
    result_df = pd.DataFrame(results)
    output_file = os.path.join(args.output_dir, "hcPhasing.csv")
    result_df.to_csv(output_file, index=False)
    print(f"Results saved to {output_file}")
