#!/usr/bin/env python3
"""Validate the synthetic Rosie PoC and render reproducible result figures.

Chart contracts
---------------
1. Evidence score
   Question: Can a high sequencing-evidence score override a failed safety filter?
   Takeaway: No. SYN004 ranks second by score but fails the matched-normal filter.
   Form: horizontal bar chart; status labels and rejection text duplicate color.

2. SYN001 peptide windows
   Question: How does one amino-acid substitution yield nine distinct 9-mers?
   Takeaway: Moving the window places the same mutation at positions 9 through 1.
   Form: aligned sequence-window diagram; an outline and bold letter mark mutation.

3. Sensitivity scenarios
   Question: What does relaxing one filter add to this five-row synthetic input?
   Takeaway: Each relaxation adds one variant and nine peptide windows, but loses a
   different kind of evidence.
   Form: horizontal bar chart; exact counts and admitted IDs are printed on bars.
"""

from __future__ import annotations

import csv
from pathlib import Path

import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt
from matplotlib.patches import Patch, Rectangle

from shortlist import evidence_score, mutant_peptides, read_variants, rejection_reason


ROOT = Path(__file__).resolve().parent
INPUT_PATH = ROOT / "demo_variants.tsv"
OUTPUT_DIR = ROOT / "outputs"

AUDIT_PATH = OUTPUT_DIR / "filter_audit.tsv"
WINDOWS_PATH = OUTPUT_DIR / "syn001_peptide_windows.tsv"
SENSITIVITY_PATH = OUTPUT_DIR / "sensitivity.tsv"

SCORE_FIGURE_PATH = OUTPUT_DIR / "evidence-score.png"
WINDOWS_FIGURE_PATH = OUTPUT_DIR / "syn001-peptide-windows.png"
SENSITIVITY_FIGURE_PATH = OUTPUT_DIR / "sensitivity.png"

TEAL = "#00A6A6"
TEAL_DARK = "#007C7C"
ORANGE = "#E97932"
INK = "#242424"
MUTED = "#777777"
GRID = "#DFDFDF"
LIGHT_TEAL = "#D9F2F2"


def write_tsv(path: Path, fieldnames: list[str], rows: list[dict[str, object]]) -> None:
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(
            handle,
            fieldnames=fieldnames,
            delimiter="\t",
            lineterminator="\n",
        )
        writer.writeheader()
        writer.writerows(rows)


def validate_input(rows: list[dict[str, str]]) -> None:
    expected_columns = {
        "variant_id",
        "gene",
        "protein_change",
        "wildtype_sequence",
        "mutant_sequence",
        "mutation_index",
        "tumor_dna_vaf",
        "normal_dna_vaf",
        "tumor_rna_alt_reads",
        "tumor_tpm",
    }
    if len(rows) != 5:
        raise AssertionError(f"Expected 5 variants, found {len(rows)}")
    if set(rows[0]) != expected_columns:
        raise AssertionError("Input columns differ from the expected schema")
    if len({row["variant_id"] for row in rows}) != len(rows):
        raise AssertionError("variant_id must be unique")

    for row in rows:
        wildtype = row["wildtype_sequence"]
        mutant = row["mutant_sequence"]
        mutation_index = int(row["mutation_index"])
        differences = [
            index
            for index, (wildtype_aa, mutant_aa) in enumerate(zip(wildtype, mutant))
            if wildtype_aa != mutant_aa
        ]
        if len(wildtype) != len(mutant):
            raise AssertionError(f'{row["variant_id"]}: sequence lengths differ')
        if not 0 <= mutation_index < len(mutant):
            raise AssertionError(f'{row["variant_id"]}: mutation_index is out of range')
        if differences != [mutation_index]:
            raise AssertionError(
                f'{row["variant_id"]}: expected one difference at mutation_index'
            )


def make_audit(rows: list[dict[str, str]]) -> list[dict[str, object]]:
    audit_rows = []
    for row in rows:
        reasons = rejection_reason(row)
        peptides = list(mutant_peptides(row)) if not reasons else []
        audit_rows.append(
            {
                "variant_id": row["variant_id"],
                "gene": row["gene"],
                "normal_dna_pass": float(row["normal_dna_vaf"]) <= 0.02,
                "tumor_dna_pass": float(row["tumor_dna_vaf"]) >= 0.10,
                "mutant_rna_pass": int(row["tumor_rna_alt_reads"]) >= 3,
                "tumor_tpm_pass": float(row["tumor_tpm"]) >= 1.0,
                "final_status": "PASS" if not reasons else "REJECT",
                "rejection_reason": reasons,
                "evidence_score": f"{evidence_score(row):.6f}",
                "peptide_count": len(peptides),
            }
        )
    write_tsv(AUDIT_PATH, list(audit_rows[0]), audit_rows)
    return audit_rows


def make_syn001_windows(rows: list[dict[str, str]]) -> list[dict[str, object]]:
    syn001 = next(row for row in rows if row["variant_id"] == "SYN001")
    mutation_index = int(syn001["mutation_index"])
    window_rows = []
    for start, wildtype_peptide, mutant_peptide in mutant_peptides(syn001):
        window_rows.append(
            {
                "variant_id": "SYN001",
                "peptide_start": start + 1,
                "wildtype_peptide": wildtype_peptide,
                "mutant_peptide": mutant_peptide,
                "mutation_position_in_peptide": mutation_index - start + 1,
            }
        )
    write_tsv(WINDOWS_PATH, list(window_rows[0]), window_rows)
    return window_rows


def passes_thresholds(row: dict[str, str], scenario: str) -> bool:
    normal_max = None if scenario == "no_normal_filter" else 0.02
    tumor_min = 0.05 if scenario == "tumor_vaf_relaxed" else 0.10
    rna_min = 1 if scenario == "rna_relaxed" else 3

    return (
        (normal_max is None or float(row["normal_dna_vaf"]) <= normal_max)
        and float(row["tumor_dna_vaf"]) >= tumor_min
        and int(row["tumor_rna_alt_reads"]) >= rna_min
        and float(row["tumor_tpm"]) >= 1.0
    )


def make_sensitivity(rows: list[dict[str, str]]) -> list[dict[str, object]]:
    scenarios = [
        "baseline",
        "rna_relaxed",
        "no_normal_filter",
        "tumor_vaf_relaxed",
    ]
    sensitivity_rows = []
    baseline_ids: set[str] = set()
    for scenario in scenarios:
        passing_rows = [row for row in rows if passes_thresholds(row, scenario)]
        passing_ids = {row["variant_id"] for row in passing_rows}
        if scenario == "baseline":
            baseline_ids = passing_ids
        peptide_count = sum(len(list(mutant_peptides(row))) for row in passing_rows)
        sensitivity_rows.append(
            {
                "scenario": scenario,
                "passing_variant_ids": ",".join(row["variant_id"] for row in passing_rows),
                "new_vs_baseline": ",".join(sorted(passing_ids - baseline_ids)),
                "variant_count": len(passing_rows),
                "peptide_count": peptide_count,
            }
        )
    write_tsv(SENSITIVITY_PATH, list(sensitivity_rows[0]), sensitivity_rows)
    return sensitivity_rows


def plot_evidence_score(rows: list[dict[str, str]], audit_rows: list[dict[str, object]]) -> None:
    status_by_id = {row["variant_id"]: row for row in audit_rows}
    plot_rows = sorted(rows, key=evidence_score)
    scores = [evidence_score(row) for row in plot_rows]
    colors = [
        TEAL if status_by_id[row["variant_id"]]["final_status"] == "PASS" else "#B7B7B7"
        for row in plot_rows
    ]

    fig, ax = plt.subplots(figsize=(10, 5.8))
    bars = ax.barh(
        [row["variant_id"] for row in plot_rows],
        scores,
        color=colors,
        edgecolor=INK,
        linewidth=0.7,
    )
    for bar, row, score in zip(bars, plot_rows, scores):
        audit = status_by_id[row["variant_id"]]
        status = audit["final_status"]
        reason_in_english = {
            "normal DNA에도 보임": "present in normal DNA",
            "tumor DNA VAF 부족": "tumor DNA VAF too low",
            "mutant RNA read 부족": "mutant RNA reads too low",
            "tumor 발현량 부족": "tumor expression too low",
        }
        detail = (
            "PASS"
            if status == "PASS"
            else f'REJECT: {reason_in_english[str(audit["rejection_reason"])]}'
        )
        ax.text(
            score + 0.015,
            bar.get_y() + bar.get_height() / 2,
            f"{score:.3f}  {detail}",
            va="center",
            fontsize=10,
            color=INK,
        )
    ax.set_xlim(0, 1.24)
    ax.set_xlabel("Sequencing-evidence score (PoC rule, 0 to 1)")
    fig.suptitle(
        "A high evidence score does not override a failed filter",
        x=0.125,
        y=0.98,
        ha="left",
        fontsize=15,
        weight="bold",
    )
    ax.set_title(
        "SYN004 ranks second by score but appears in matched-normal DNA",
        loc="left",
        color=MUTED,
        fontsize=10,
        pad=10,
    )
    ax.xaxis.grid(True, color=GRID, linewidth=0.7)
    ax.set_axisbelow(True)
    ax.spines[["top", "right", "left"]].set_visible(False)
    ax.tick_params(axis="y", length=0)
    ax.legend(
        handles=[
            Patch(facecolor=TEAL, edgecolor=INK, label="PASS"),
            Patch(facecolor="#B7B7B7", edgecolor=INK, label="REJECT"),
        ],
        frameon=False,
        loc="lower right",
    )
    fig.tight_layout(rect=[0, 0, 1, 0.92])
    fig.savefig(SCORE_FIGURE_PATH, dpi=180, bbox_inches="tight")
    plt.close(fig)


def plot_syn001_windows(rows: list[dict[str, str]], window_rows: list[dict[str, object]]) -> None:
    syn001 = next(row for row in rows if row["variant_id"] == "SYN001")
    sequence = syn001["mutant_sequence"]
    mutation_index = int(syn001["mutation_index"])

    fig, ax = plt.subplots(figsize=(12.5, 7.2))
    ax.set_xlim(-4.6, len(sequence) + 0.2)
    ax.set_ylim(-0.9, len(window_rows) + 1.9)
    ax.invert_yaxis()
    ax.axis("off")

    cell_width = 0.92
    cell_height = 0.68
    for index, amino_acid in enumerate(sequence):
        is_mutation = index == mutation_index
        ax.add_patch(
            Rectangle(
                (index + 0.04, 0.1),
                cell_width,
                cell_height,
                facecolor=ORANGE if is_mutation else "#F2F2F2",
                edgecolor=INK if is_mutation else "white",
                linewidth=1.8 if is_mutation else 0.5,
            )
        )
        ax.text(
            index + 0.5,
            0.44,
            amino_acid,
            ha="center",
            va="center",
            family="monospace",
            weight="bold" if is_mutation else "normal",
            color="white" if is_mutation else INK,
        )
        ax.text(index + 0.5, -0.15, str(index + 1), ha="center", fontsize=8, color=MUTED)

    for row_index, window in enumerate(window_rows, start=1):
        start = int(window["peptide_start"]) - 1
        mutation_position = int(window["mutation_position_in_peptide"])
        ax.text(
            -0.15,
            row_index + 0.44,
            f"start {start + 1:>2}  |  mutation pos {mutation_position}",
            ha="right",
            va="center",
            family="monospace",
            fontsize=9,
            color=MUTED,
        )
        for offset, amino_acid in enumerate(str(window["mutant_peptide"])):
            sequence_index = start + offset
            is_mutation = offset + 1 == mutation_position
            ax.add_patch(
                Rectangle(
                    (sequence_index + 0.04, row_index + 0.1),
                    cell_width,
                    cell_height,
                    facecolor=ORANGE if is_mutation else LIGHT_TEAL,
                    edgecolor=INK if is_mutation else "white",
                    linewidth=1.8 if is_mutation else 0.5,
                )
            )
            ax.text(
                sequence_index + 0.5,
                row_index + 0.44,
                amino_acid,
                ha="center",
                va="center",
                family="monospace",
                weight="bold" if is_mutation else "normal",
                color="white" if is_mutation else INK,
            )
    ax.text(
        -4.6,
        -0.75,
        "One substitution produces nine distinct 9-mer windows",
        fontsize=15,
        weight="bold",
        color=INK,
    )
    ax.text(
        -4.6,
        -0.3,
        "SYN001 mutant sequence; K9V is outlined and moves from peptide position 9 to 1",
        fontsize=10,
        color=MUTED,
    )
    fig.tight_layout()
    fig.savefig(WINDOWS_FIGURE_PATH, dpi=180, bbox_inches="tight")
    plt.close(fig)


def plot_sensitivity(sensitivity_rows: list[dict[str, object]]) -> None:
    labels = {
        "baseline": "Baseline",
        "rna_relaxed": "RNA reads ≥ 1",
        "no_normal_filter": "No normal-DNA filter",
        "tumor_vaf_relaxed": "Tumor DNA VAF ≥ 0.05",
    }
    plot_rows = list(reversed(sensitivity_rows))
    values = [int(row["peptide_count"]) for row in plot_rows]
    colors = [TEAL_DARK if row["scenario"] == "baseline" else ORANGE for row in plot_rows]

    fig, ax = plt.subplots(figsize=(10, 5.5))
    bars = ax.barh(
        [labels[str(row["scenario"])] for row in plot_rows],
        values,
        color=colors,
        edgecolor=INK,
        linewidth=0.7,
    )
    for bar, row in zip(bars, plot_rows):
        new_id = str(row["new_vs_baseline"])
        addition = "baseline" if not new_id else f"adds {new_id}"
        ax.text(
            int(row["peptide_count"]) + 0.5,
            bar.get_y() + bar.get_height() / 2,
            f'{row["variant_count"]} variants, {row["peptide_count"]} peptides ({addition})',
            va="center",
            fontsize=10,
            color=INK,
        )
    ax.set_xlim(0, 42)
    ax.set_xlabel("Generated 9-mer peptide windows")
    fig.suptitle(
        "Relaxing one filter adds one synthetic variant and nine peptides",
        x=0.125,
        y=0.98,
        ha="left",
        fontsize=15,
        weight="bold",
    )
    ax.set_title(
        "Candidate quantity rises equally, but each scenario gives up different evidence",
        loc="left",
        color=MUTED,
        fontsize=10,
        pad=10,
    )
    ax.xaxis.grid(True, color=GRID, linewidth=0.7)
    ax.set_axisbelow(True)
    ax.spines[["top", "right", "left"]].set_visible(False)
    ax.tick_params(axis="y", length=0)
    fig.tight_layout(rect=[0, 0, 1, 0.92])
    fig.savefig(SENSITIVITY_FIGURE_PATH, dpi=180, bbox_inches="tight")
    plt.close(fig)


def verify_expected_results(
    audit_rows: list[dict[str, object]],
    window_rows: list[dict[str, object]],
    sensitivity_rows: list[dict[str, object]],
) -> None:
    passing_ids = {
        str(row["variant_id"])
        for row in audit_rows
        if row["final_status"] == "PASS"
    }
    if passing_ids != {"SYN001", "SYN002"}:
        raise AssertionError(f"Unexpected baseline PASS set: {passing_ids}")
    if sum(int(row["peptide_count"]) for row in audit_rows) != 18:
        raise AssertionError("Baseline must produce exactly 18 peptides")
    if len(window_rows) != 9:
        raise AssertionError("SYN001 must produce exactly nine 9-mers")
    if [int(row["mutation_position_in_peptide"]) for row in window_rows] != list(
        range(9, 0, -1)
    ):
        raise AssertionError("SYN001 mutation positions must run from 9 to 1")

    expected_sensitivity = {
        "baseline": ({"SYN001", "SYN002"}, 18),
        "rna_relaxed": ({"SYN001", "SYN002", "SYN003"}, 27),
        "no_normal_filter": ({"SYN001", "SYN002", "SYN004"}, 27),
        "tumor_vaf_relaxed": ({"SYN001", "SYN002", "SYN005"}, 27),
    }
    for row in sensitivity_rows:
        scenario = str(row["scenario"])
        ids = set(str(row["passing_variant_ids"]).split(","))
        expected_ids, expected_peptides = expected_sensitivity[scenario]
        if ids != expected_ids or int(row["peptide_count"]) != expected_peptides:
            raise AssertionError(f"Unexpected result for {scenario}: {row}")


def main() -> None:
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    rows = list(read_variants(INPUT_PATH))
    validate_input(rows)
    audit_rows = make_audit(rows)
    window_rows = make_syn001_windows(rows)
    sensitivity_rows = make_sensitivity(rows)
    verify_expected_results(audit_rows, window_rows, sensitivity_rows)

    plot_evidence_score(rows, audit_rows)
    plot_syn001_windows(rows, window_rows)
    plot_sensitivity(sensitivity_rows)

    print("Validated 5 synthetic variants")
    print("Baseline PASS: SYN001, SYN002; 18 peptide windows")
    print("Sensitivity: each one-filter relaxation adds 1 variant and 9 peptides")
    for output_path in [
        AUDIT_PATH,
        WINDOWS_PATH,
        SENSITIVITY_PATH,
        SCORE_FIGURE_PATH,
        WINDOWS_FIGURE_PATH,
        SENSITIVITY_FIGURE_PATH,
    ]:
        print(f"WROTE\t{output_path.relative_to(ROOT)}")


if __name__ == "__main__":
    main()