#!/usr/bin/env python3
"""
Process the peg transfer user study CSV and print paper-style summary tables.

This script reads `source_data/peg_transfer_data.csv` directly and prints results for:
- Overall
- Junior Surgeons
- Senior Surgeons


"""

import argparse
import csv
import math
import warnings
from pathlib import Path
from typing import Dict, Iterable, List, Sequence, Tuple

import numpy as np
from scipy.stats import shapiro, t as t_dist, ttest_rel, wilcoxon


DEFAULT_PLATFORM_ORDER = ["Manual", "Humanoid", "da Vinci"]
DEFAULT_COMPARISON_PAIRS = [
    ("da Vinci", "Humanoid"),
    ("da Vinci", "Manual"),
    ("Humanoid", "Manual"),
]
DEFAULT_PLATFORM_ALIAS = {
    "manual": "Manual",
    "humanoid": "Humanoid",
    "da vinci": "da Vinci",
    "davinci": "da Vinci",
}
DEFAULT_GROUP_ORDER = ["Overall", "Junior Surgeons", "Senior Surgeons"]
DEFAULT_GROUP_ALIAS = {
    "junior": "Junior Surgeons",
    "junior surgeon": "Junior Surgeons",
    "junior surgeons": "Junior Surgeons",
    "senior": "Senior Surgeons",
    "senior surgeon": "Senior Surgeons",
    "senior surgeons": "Senior Surgeons",
}


def normalize_platform_name(value: str, alias: Dict[str, str]) -> str:
    if not isinstance(value, str):
        return value
    key = value.strip().lower()
    return alias.get(key, value.strip())


def normalize_group_name(value: str) -> str:
    if not isinstance(value, str):
        return "Unspecified"
    key = value.strip().lower()
    if not key:
        return "Unspecified"
    return DEFAULT_GROUP_ALIAS.get(key, value.strip())


def parse_float(value: str) -> float:
    if value is None:
        return np.nan
    text = str(value).strip()
    if not text:
        return np.nan
    return float(text)


def clean_occurrence_label(name: str) -> str:
    cleaned = name.strip()
    lower = cleaned.lower()
    aliases = {
        "drops": "Drops",
        "collision": "Collision",
        "failed pick up": "Failed Pick Up",
    }
    return aliases.get(lower, cleaned)


def detect_error_structure(
    fieldnames: Sequence[str],
    required_columns: Sequence[str],
) -> Tuple[List[str], List[str], List[str], str, str]:
    normalized = {name: name.strip().lower() for name in fieldnames}
    required_set = {name.strip().lower() for name in required_columns}

    exact_total = [name for name, lower in normalized.items() if lower in {"total error", "error total"}]
    legacy_error = [name for name, lower in normalized.items() if lower == "error"]

    excluded = set(required_set)
    excluded.update(name.strip().lower() for name in exact_total)
    excluded.update(name.strip().lower() for name in legacy_error)

    occurrence_columns = [
        name
        for name in fieldnames
        if normalized[name] not in excluded
    ]

    if exact_total:
        if occurrence_columns:
            joined = ", ".join(repr(name) for name in occurrence_columns)
            return exact_total, occurrence_columns, legacy_error, "Total Error", (
                f"using column {exact_total[0]!r} with row-wise fallback to summed components {joined}"
            )
        return exact_total, occurrence_columns, legacy_error, "Total Error", f"using column {exact_total[0]!r}"

    if occurrence_columns:
        joined = ", ".join(repr(name) for name in occurrence_columns)
        return [], occurrence_columns, legacy_error, "Total Error", f"summing columns {joined}"

    if legacy_error:
        return [], occurrence_columns, legacy_error, "Error", f"using legacy column {legacy_error[0]!r}"

    raise ValueError(
        "Could not find an error metric column. Expected 'Total Error', component error columns, or 'Error'."
    )


def compute_error_value(
    row: Dict[str, str],
    total_error_columns: Sequence[str],
    occurrence_columns: Sequence[str],
    legacy_error_columns: Sequence[str],
) -> float:
    total_values = [parse_float(row.get(column, "")) for column in total_error_columns]
    total_present = [value for value in total_values if not np.isnan(value)]
    if total_present:
        return float(total_present[0])

    occurrence_values = [parse_float(row.get(column, "")) for column in occurrence_columns]
    occurrence_present = [value for value in occurrence_values if not np.isnan(value)]
    if occurrence_present:
        return float(sum(occurrence_present))

    legacy_values = [parse_float(row.get(column, "")) for column in legacy_error_columns]
    legacy_present = [value for value in legacy_values if not np.isnan(value)]
    if legacy_present:
        return float(legacy_present[0])

    return np.nan


def mean_std(values: Sequence[float]) -> Tuple[float, float]:
    arr = np.asarray(values, dtype=float)
    arr = arr[~np.isnan(arr)]
    if len(arr) == 0:
        return np.nan, np.nan
    if len(arr) == 1:
        return float(arr[0]), 0.0
    return float(np.mean(arr)), float(np.std(arr, ddof=1))


def fmt_mean_std(mean: float, std: float, ndigits: int) -> str:
    if np.isnan(mean):
        return "--"
    return f"{mean:.{ndigits}f} +/- {std:.{ndigits}f}"


def fmt_p(value: float, sci_thresh: float = 1e-3) -> str:
    if value is None or (isinstance(value, float) and np.isnan(value)):
        return "--"
    if value < sci_thresh:
        return f"{value:.2e}"
    return f"{value:.3f}"


def fmt_ci(mean_diff: float, ci_low: float, ci_high: float, ndigits: int = 2) -> str:
    if np.isnan(mean_diff) or np.isnan(ci_low) or np.isnan(ci_high):
        return "--"
    return f"{mean_diff:.{ndigits}f} [{ci_low:.{ndigits}f}, {ci_high:.{ndigits}f}]"


def fmt_d(value: float, ndigits: int = 2) -> str:
    if value is None or (isinstance(value, float) and np.isnan(value)):
        return "--"
    return f"{abs(value):.{ndigits}f}"


def load_and_compute_metrics(
    csv_path: str,
    platform_order: Sequence[str],
    w_time: float,
) -> Tuple[List[Dict[str, object]], str, str, List[str]]:
    with open(csv_path, newline="", encoding="utf-8-sig") as handle:
        reader = csv.DictReader(handle)
        raw_rows = list(reader)

    participant_column = "Participant"
    if raw_rows and participant_column not in raw_rows[0] and "User Name" in raw_rows[0]:
        participant_column = "User Name"

    required = [participant_column, "Notes", "Modality", "Duration (s)"]
    if not raw_rows:
        raise ValueError("The CSV is empty.")

    missing = [name for name in required if name not in raw_rows[0]]
    if missing:
        raise ValueError(f"Missing required CSV columns: {missing}")

    total_error_columns, occurrence_columns, legacy_error_columns, error_label, error_source_description = detect_error_structure(
        list(raw_rows[0].keys()),
        required,
    )
    occurrence_labels = [clean_occurrence_label(name) for name in occurrence_columns]

    cleaned_rows: List[Dict[str, object]] = []
    valid_platforms = set(platform_order)
    for row in raw_rows:
        participant = str(row[participant_column]).strip()
        platform = normalize_platform_name(row["Modality"], DEFAULT_PLATFORM_ALIAS)
        duration = parse_float(row["Duration (s)"])
        total_error = compute_error_value(row, total_error_columns, occurrence_columns, legacy_error_columns)

        if not participant or platform not in valid_platforms:
            continue
        if np.isnan(duration) or np.isnan(total_error):
            continue

        occurrence_values = {}
        for column, label in zip(occurrence_columns, occurrence_labels):
            occurrence_values[label] = parse_float(row.get(column, ""))

        cleaned_rows.append(
            {
                "Participant": participant,
                "Group": normalize_group_name(row.get("Notes", "")),
                "Platform": platform,
                "TimePerTrialSec": duration,
                "TotalError": total_error,
                **occurrence_values,
            }
        )

    if not cleaned_rows:
        raise ValueError("No valid rows found after filtering missing duration/error values.")

    times = np.asarray([float(row["TimePerTrialSec"]) for row in cleaned_rows], dtype=float)
    errors = np.asarray([float(row["TotalError"]) for row in cleaned_rows], dtype=float)

    t_min = float(np.min(times))
    t_max = float(np.max(times))
    e_min = float(np.min(errors))
    e_max = float(np.max(errors))

    t_range = max(t_max - t_min, 1e-9)
    e_range = max(e_max - e_min, 1e-9)

    for row in cleaned_rows:
        t_norm = (float(row["TimePerTrialSec"]) - t_min) / t_range
        e_norm = (float(row["TotalError"]) - e_min) / e_range
        row["BalancedFLS"] = 100.0 * (1.0 - (w_time * t_norm + (1.0 - w_time) * e_norm))

    return cleaned_rows, error_label, error_source_description, occurrence_labels


def build_participant_platform_map(
    rows: Sequence[Dict[str, object]],
    value_key: str,
) -> Dict[str, Dict[str, float]]:
    participant_platform_values: Dict[str, Dict[str, float]] = {}
    for row in rows:
        participant = str(row["Participant"])
        platform = str(row["Platform"])
        participant_platform_values.setdefault(participant, {})[platform] = float(row[value_key])
    return participant_platform_values


def paired_stats(
    participant_platform_values: Dict[str, Dict[str, float]],
    platform_a: str,
    platform_b: str,
    alpha: float = 0.05,
) -> Dict[str, float]:
    paired_a: List[float] = []
    paired_b: List[float] = []

    for platform_values in participant_platform_values.values():
        if platform_a in platform_values and platform_b in platform_values:
            paired_a.append(float(platform_values[platform_a]))
            paired_b.append(float(platform_values[platform_b]))

    n = len(paired_a)
    if n < 2:
        return {
            "n": n,
            "mean_diff": np.nan,
            "ci_low": np.nan,
            "ci_high": np.nan,
            "p_ttest": np.nan,
            "cohens_dz": np.nan,
            "shapiro_p": np.nan,
            "p_wilcoxon": np.nan,
        }

    a = np.asarray(paired_a, dtype=float)
    b = np.asarray(paired_b, dtype=float)
    diffs = a - b

    mean_diff = float(np.mean(diffs))
    sd_diff = float(np.std(diffs, ddof=1))
    se = sd_diff / math.sqrt(n) if sd_diff > 0 else 0.0

    if sd_diff > 0:
        tcrit = float(t_dist.ppf(1 - alpha / 2, df=n - 1))
        ci_low = mean_diff - tcrit * se
        ci_high = mean_diff + tcrit * se
    else:
        ci_low = mean_diff
        ci_high = mean_diff

    if np.allclose(diffs, 0):
        p_ttest = 1.0
        p_wilcoxon = 1.0
    else:
        try:
            _, p_ttest = ttest_rel(a, b, alternative="two-sided")
            p_ttest = float(p_ttest)
        except Exception:
            p_ttest = np.nan

        try:
            with warnings.catch_warnings():
                warnings.simplefilter("ignore", UserWarning)
                p_wilcoxon = float(wilcoxon(diffs, alternative="two-sided", zero_method="pratt").pvalue)
        except Exception:
            p_wilcoxon = np.nan

    shapiro_p = float(shapiro(diffs).pvalue) if n >= 3 else np.nan
    cohens_dz = (mean_diff / sd_diff) if sd_diff > 0 else np.nan

    return {
        "n": n,
        "mean_diff": mean_diff,
        "ci_low": float(ci_low),
        "ci_high": float(ci_high),
        "p_ttest": p_ttest,
        "cohens_dz": float(cohens_dz) if not np.isnan(cohens_dz) else np.nan,
        "shapiro_p": shapiro_p,
        "p_wilcoxon": p_wilcoxon,
    }


def choose_test(stats: Dict[str, float], normality_alpha: float) -> Tuple[str, str]:
    shapiro_p = stats["shapiro_p"]
    if not np.isnan(shapiro_p) and shapiro_p < normality_alpha:
        return "w", fmt_p(stats["p_wilcoxon"])
    return "t", fmt_p(stats["p_ttest"])


def build_summary_rows(
    rows: Sequence[Dict[str, object]],
    platform_order: Sequence[str],
    error_label: str,
) -> List[Dict[str, object]]:
    summary_rows: List[Dict[str, object]] = []
    for platform in platform_order:
        platform_rows = [row for row in rows if row["Platform"] == platform]
        time_mean, time_std = mean_std([float(row["TimePerTrialSec"]) for row in platform_rows])
        err_mean, err_std = mean_std([float(row["TotalError"]) for row in platform_rows])
        fls_mean, fls_std = mean_std([float(row["BalancedFLS"]) for row in platform_rows])

        summary_rows.append(
            {
                "Platform": platform,
                "Time (s)": fmt_mean_std(time_mean, time_std, 1),
                error_label: fmt_mean_std(err_mean, err_std, 2),
                "Balanced FLS (0-100)": fmt_mean_std(fls_mean, fls_std, 2),
            }
        )

    return summary_rows


def build_occurrence_summary_rows(
    rows: Sequence[Dict[str, object]],
    platform_order: Sequence[str],
    occurrence_labels: Sequence[str],
) -> List[Dict[str, object]]:
    occurrence_rows: List[Dict[str, object]] = []
    for platform in platform_order:
        platform_rows = [row for row in rows if row["Platform"] == platform]
        summary_row: Dict[str, object] = {"Platform": platform}
        for label in occurrence_labels:
            values = [float(row[label]) for row in platform_rows if label in row and not np.isnan(float(row[label]))]
            mean_value, std_value = mean_std(values)
            summary_row[label] = fmt_mean_std(mean_value, std_value, 2)
        occurrence_rows.append(summary_row)
    return occurrence_rows


def build_pairwise_rows(
    rows: Sequence[Dict[str, object]],
    comparison_pairs: Sequence[Tuple[str, str]],
    normality_alpha: float,
    error_label: str,
) -> Tuple[List[Dict[str, object]], List[Dict[str, object]]]:
    metric_maps = {
        "Time (s)": build_participant_platform_map(rows, "TimePerTrialSec"),
        error_label: build_participant_platform_map(rows, "TotalError"),
        "Balanced FLS": build_participant_platform_map(rows, "BalancedFLS"),
    }

    pairwise_rows: List[Dict[str, object]] = []
    diagnostic_rows: List[Dict[str, object]] = []

    for metric_name, participant_values in metric_maps.items():
        for first, second in comparison_pairs:
            stats = paired_stats(participant_values, first, second)
            selected_test, selected_p = choose_test(stats, normality_alpha)
            pairwise_rows.append(
                {
                    "Metric": metric_name,
                    "Comparison": f"{first} - {second}",
                    "Delta [95% CI]": fmt_ci(stats["mean_diff"], stats["ci_low"], stats["ci_high"], 2),
                    "p": selected_p,
                    "d_z": fmt_d(stats["cohens_dz"], 2),
                    "Test": selected_test,
                }
            )
            diagnostic_rows.append(
                {
                    "Metric": metric_name,
                    "Comparison": f"{first} - {second}",
                    "n": stats["n"],
                    "Shapiro p": fmt_p(stats["shapiro_p"]),
                    "p_ttest": fmt_p(stats["p_ttest"]),
                    "p_wilcoxon": fmt_p(stats["p_wilcoxon"]),
                }
            )

    return pairwise_rows, diagnostic_rows


def write_csv(path: str, rows: Sequence[Dict[str, object]]) -> None:
    if not rows:
        return
    output_path = Path(path)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    with output_path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(rows[0].keys()))
        writer.writeheader()
        writer.writerows(rows)


def format_table(rows: Sequence[Dict[str, object]], columns: Sequence[str]) -> str:
    if not rows:
        return "(Empty)"

    widths = []
    for column in columns:
        width = len(column)
        for row in rows:
            width = max(width, len(str(row.get(column, ""))))
        widths.append(width)

    def build_line(values: Iterable[str]) -> str:
        return "  ".join(str(value).ljust(width) for value, width in zip(values, widths))

    header = build_line(columns)
    separator = "  ".join("-" * width for width in widths)
    body = [build_line([row.get(column, "") for column in columns]) for row in rows]
    return "\n".join([header, separator, *body])


def print_notes(normality_alpha: float, error_label: str, error_source_description: str) -> None:
    threshold = f"{normality_alpha:.2f}".rstrip("0").rstrip(".")
    print(
        "\nNotes: Delta is the mean paired difference (first minus second) with a 95% confidence interval (CI). "
        "d_z is Cohen's d for paired samples computed on paired differences. "
        "Test indicates a paired two-sided t-test (t) or Wilcoxon signed-rank test (w), "
        f"selected using the Shapiro-Wilk normality check on paired differences (p >= {threshold} -> t; otherwise w)."
    )
    print(f"Error metric: {error_label} ({error_source_description}). Balanced FLS uses this error metric.")


def print_summary_table(title: str, rows: Sequence[Dict[str, object]], error_label: str) -> None:
    print(f"\n{title}\n")
    print(format_table(rows, ["Platform", "Time (s)", error_label, "Balanced FLS (0-100)"]))


def print_occurrence_summary_table(
    title: str,
    rows: Sequence[Dict[str, object]],
    occurrence_labels: Sequence[str],
) -> None:
    print(f"\n{title}\n")
    print(format_table(rows, ["Platform", *occurrence_labels]))


def print_pairwise_table(title: str, rows: Sequence[Dict[str, object]]) -> None:
    print(f"\n{title}\n")
    print(format_table(rows, ["Metric", "Comparison", "Delta [95% CI]", "p", "d_z", "Test"]))


def print_diagnostics_table(title: str, rows: Sequence[Dict[str, object]]) -> None:
    print(f"\n{title}\n")
    print(format_table(rows, ["Metric", "Comparison", "n", "Shapiro p", "p_ttest", "p_wilcoxon"]))


def ordered_groups(rows: Sequence[Dict[str, object]]) -> List[str]:
    present = {str(row["Group"]) for row in rows}
    ordered = [group for group in DEFAULT_GROUP_ORDER if group == "Overall" or group in present]
    ordered.extend(group for group in sorted(present) if group not in ordered)
    return ordered


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--csv", type=str, default="source_data/peg_transfer_data.csv")
    parser.add_argument("--platform_order", type=str, default="Manual,Humanoid,da Vinci")
    parser.add_argument("--w_time", type=float, default=0.5, help="Balanced FLS weight for normalized time.")
    parser.add_argument("--normality_alpha", type=float, default=0.05, help="Normality threshold used to choose t versus Wilcoxon.")
    parser.add_argument("--out_summary_csv", type=str, default="", help="If set, write the summary rows to this CSV path.")
    parser.add_argument("--out_occurrence_csv", type=str, default="", help="If set, write the error-occurrence summary rows to this CSV path.")
    parser.add_argument("--out_pairwise_csv", type=str, default="", help="If set, write the paired-comparison rows to this CSV path.")
    parser.add_argument("--out_diag_csv", type=str, default="", help="If set, write the diagnostics rows to this CSV path.")
    parser.add_argument("--show_diagnostics", action="store_true", help="If set, print the diagnostics table for each group.")
    args = parser.parse_args()

    platform_order = [item.strip() for item in args.platform_order.split(",") if item.strip()]
    comparison_pairs = [
        (first, second)
        for first, second in DEFAULT_COMPARISON_PAIRS
        if first in platform_order and second in platform_order
    ]

    rows, error_label, error_source_description, occurrence_labels = load_and_compute_metrics(
        args.csv,
        platform_order,
        args.w_time,
    )

    summary_export_rows: List[Dict[str, object]] = []
    occurrence_export_rows: List[Dict[str, object]] = []
    pairwise_export_rows: List[Dict[str, object]] = []
    diagnostic_export_rows: List[Dict[str, object]] = []

    print_notes(args.normality_alpha, error_label, error_source_description)

    group_rows_map: Dict[str, List[Dict[str, object]]] = {"Overall": list(rows)}
    for row in rows:
        group_rows_map.setdefault(str(row["Group"]), []).append(row)

    groups_to_print = ordered_groups(rows)

    for index, group_name in enumerate(groups_to_print, start=1):
        current_rows = group_rows_map[group_name]
        summary_rows = build_summary_rows(current_rows, platform_order, error_label)
        occurrence_rows = build_occurrence_summary_rows(current_rows, platform_order, occurrence_labels)
        pairwise_rows, diagnostic_rows = build_pairwise_rows(current_rows, comparison_pairs, args.normality_alpha, error_label)

        summary_export_rows.extend({"Group": group_name, **row} for row in summary_rows)
        occurrence_export_rows.extend({"Group": group_name, **row} for row in occurrence_rows)
        pairwise_export_rows.extend({"Group": group_name, **row} for row in pairwise_rows)
        diagnostic_export_rows.extend({"Group": group_name, **row} for row in diagnostic_rows)

        print_summary_table(f"(D{index}) FLS peg transfer: performance summary ({group_name})", summary_rows, error_label)
        print_pairwise_table(f"(E{index}) FLS peg transfer: paired comparison statistics ({group_name})", pairwise_rows)
        if occurrence_labels:
            print_occurrence_summary_table(
                f"(F{index}) FLS peg transfer: error type occurrences ({group_name})",
                occurrence_rows,
                occurrence_labels,
            )

        if args.show_diagnostics:
            print_diagnostics_table(f"(Diag {index}) FLS peg transfer: diagnostics ({group_name})", diagnostic_rows)

    if args.out_summary_csv:
        write_csv(args.out_summary_csv, summary_export_rows)
        print(f"\nSaved summary CSV: {args.out_summary_csv}")

    if args.out_occurrence_csv:
        write_csv(args.out_occurrence_csv, occurrence_export_rows)
        print(f"Saved occurrence CSV: {args.out_occurrence_csv}")

    if args.out_pairwise_csv:
        write_csv(args.out_pairwise_csv, pairwise_export_rows)
        print(f"Saved paired-comparison CSV: {args.out_pairwise_csv}")

    if args.out_diag_csv:
        write_csv(args.out_diag_csv, diagnostic_export_rows)
        print(f"Saved diagnostics CSV: {args.out_diag_csv}")


if __name__ == "__main__":
    main()