#!/usr/bin/env python3
"""Build the NAR measured-DMD evidence figure and audit tables."""

from __future__ import annotations

import hashlib
import json
import os
from pathlib import Path

os.environ.setdefault("MPLCONFIGDIR", "temporary_workspace/matplotlib")
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
import numpy as np
import pandas as pd


ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "results" / "nar_measured_dmd_evidence_v01"

BLUE = "#2A6F97"
TEAL = "#248277"
VERMILION = "#C84C3A"
GOLD = "#D39B2A"
INK = "#20262E"
MID = "#69737D"
LIGHT = "#E8EDF1"
PALE = "#F6F8FA"


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for block in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def panel_label(ax: plt.Axes, label: str, title: str) -> None:
    ax.text(
        -0.08,
        1.06,
        label,
        transform=ax.transAxes,
        fontsize=11,
        fontweight="bold",
        color=INK,
        va="top",
    )
    ax.text(
        0.0,
        1.06,
        title,
        transform=ax.transAxes,
        fontsize=9.5,
        fontweight="bold",
        color=INK,
        va="top",
    )


def build_figure() -> tuple[pd.DataFrame, pd.DataFrame]:
    induction = pd.read_csv(
        ROOT
        / "results/gse233606_single_cell_dmd_induction_v01/condition_gene_effects.tsv.gz",
        sep="\t",
    )
    induction_metrics = pd.read_csv(
        ROOT
        / "results/gse233606_single_cell_dmd_induction_v01/unselected_disease_axis_metrics.tsv",
        sep="\t",
    ).iloc[0]
    induction = induction[induction["eligible_unselected_axis_gene"]].copy()

    external_consensus = pd.read_csv(
        ROOT
        / "results/external_dmd_pathway_reversal_v01/"
        "external_dmd_cross_dataset_pathway_consensus.tsv",
        sep="\t",
    )
    external_summary = json.loads(
        (
            ROOT
            / "results/external_dmd_pathway_reversal_v01/"
            "external_dmd_pathway_reversal_summary.json"
        ).read_text()
    )
    correction = pd.read_csv(
        ROOT
        / "results/dmd_causal_perturbation_truth_v01/"
        "gse272233_correction_vector_metrics.tsv",
        sep="\t",
    ).set_index("background")
    reactome_summary = json.loads(
        (
            ROOT
            / "results/gse272233_replicate_aware_reactome_v01/"
            "gse272233_replicate_aware_reactome_summary.json"
        ).read_text()
    )
    organoid = pd.read_csv(
        ROOT
        / "results/gse277637_external_dmd_organoid_validation_v01/"
        "external_validation_metrics.tsv",
        sep="\t",
    )

    effect_cols = {
        "PRJNA772047": "disease_score_effect.PRJNA772047",
        "PRJNA1218493": "disease_score_effect.PRJNA1218493",
        "GSE156497": "disease_score_effect.GSE156497",
    }
    complete = external_consensus.dropna(subset=list(effect_cols.values())).copy()
    agreements = []
    reference = np.sign(complete[effect_cols["PRJNA772047"]])
    for dataset in ["PRJNA1218493", "GSE156497"]:
        agreements.append(
            {
                "dataset": dataset,
                "agreement": float(
                    np.mean(reference == np.sign(complete[effect_cols[dataset]]))
                ),
                "n_pathways": int(len(complete)),
                "evidence_tier": (
                    "direction only, 2 DMD / 1 Normal"
                    if dataset == "PRJNA1218493"
                    else "direction only, mouse 1 / 1"
                ),
            }
        )
    all_three_fraction = float(complete["direction_consistent_all_3"].mean())
    agreements.append(
        {
            "dataset": "All three",
            "agreement": all_three_fraction,
            "n_pathways": int(len(complete)),
            "evidence_tier": "same sign in all three datasets",
        }
    )
    agreement_frame = pd.DataFrame(agreements)

    organoid_rows = organoid[
        (organoid["scope"] == "REACTOME_PATHWAY")
        & organoid["left_axis"].isin(["DMD1_minus_WT", "DMD2_minus_WT", "DMD3_minus_WT"])
        & (organoid["right_axis"] == "GSE272233_DISEASE_MEDIAN")
    ].copy()
    organoid_rows["line"] = organoid_rows["left_axis"].str.replace("_minus_WT", "")
    organoid_rows = organoid_rows.set_index("line").loc[["DMD1", "DMD2", "DMD3"]]

    plt.rcParams.update(
        {
            "font.family": "DejaVu Sans",
            "font.size": 7.2,
            "axes.labelcolor": INK,
            "axes.edgecolor": "#AEB7BF",
            "xtick.color": MID,
            "ytick.color": MID,
            "axes.titlecolor": INK,
            "pdf.fonttype": 42,
            "svg.fonttype": "none",
        }
    )
    fig = plt.figure(figsize=(7.2, 7.5), facecolor="white")
    grid = fig.add_gridspec(
        2, 2, left=0.085, right=0.975, top=0.88, bottom=0.16, wspace=0.32, hspace=0.74
    )
    ax_a = fig.add_subplot(grid[0, 0])
    ax_b = fig.add_subplot(grid[0, 1])
    ax_c = fig.add_subplot(grid[1, 0])
    ax_d = fig.add_subplot(grid[1, 1])

    fig.suptitle(
        "Measured DMD evidence and experimental boundaries",
        x=0.53,
        y=0.965,
        fontsize=11.5,
        fontweight="bold",
        color=INK,
    )
    fig.text(
        0.53,
        0.938,
        "Engineered induction -> patient muscle -> CRISPR correction -> organoid stress test",
        ha="center",
        va="center",
        fontsize=7.8,
        color=MID,
    )

    # Panel A: unbiased gene-axis concordance.
    x = induction["crispr_dmd_minus_healthy"].to_numpy()
    y = induction["patient_dmd_minus_healthy"].to_numpy()
    concordant = induction["direction_concordant"].to_numpy(dtype=bool)
    limit = float(np.quantile(np.abs(np.concatenate([x, y])), 0.995))
    limit = max(limit, 0.2)
    ax_a.scatter(
        x[~concordant],
        y[~concordant],
        s=2.0,
        c="#BAC2C9",
        alpha=0.35,
        linewidths=0,
        rasterized=True,
    )
    ax_a.scatter(
        x[concordant],
        y[concordant],
        s=2.0,
        c=TEAL,
        alpha=0.28,
        linewidths=0,
        rasterized=True,
    )
    ax_a.axhline(0, color="#C8CFD5", lw=0.6)
    ax_a.axvline(0, color="#C8CFD5", lw=0.6)
    ax_a.plot([-limit, limit], [-limit, limit], color=VERMILION, lw=0.9, ls="--")
    ax_a.set_xlim(-limit, limit)
    ax_a.set_ylim(-limit, limit)
    ax_a.set_xlabel("Engineered DMD - healthy\nmean log1p(CP10K)")
    ax_a.set_ylabel("Patient DMD - healthy\nmean log1p(CP10K)")
    ax_a.text(
        0.04,
        0.95,
        f"n = {int(induction_metrics['n_genes']):,} genes\n"
        f"Spearman rho = {induction_metrics['spearman_rho']:.3f}\n"
        f"direction agreement = {100 * induction_metrics['direction_agreement_fraction']:.1f}%",
        transform=ax_a.transAxes,
        va="top",
        color=INK,
        bbox={"facecolor": "white", "edgecolor": LIGHT, "boxstyle": "square,pad=0.35"},
    )
    panel_label(ax_a, "A", "Induction -> patient replication")

    # Panel B: external sample-level pathway direction transfer.
    labels = ["PRJNA1218493", "GSE156497", "All three"]
    values = agreement_frame.set_index("dataset").loc[labels, "agreement"].to_numpy() * 100
    ypos = np.arange(len(labels))
    ax_b.barh(ypos, values, color=[BLUE, GOLD, TEAL], height=0.56)
    ax_b.axvline(50, color=VERMILION, lw=0.8, ls="--")
    ax_b.set_xlim(0, 100)
    ax_b.set_yticks(ypos, ["PRJNA1218493\n2 DMD / 1 Normal", "GSE156497\nmouse 1 / 1", "All 3 datasets\nsame direction"])
    ax_b.invert_yaxis()
    ax_b.set_xlabel("Pathway direction agreement (%)")
    for y_pos, value in zip(ypos, values):
        ax_b.text(value + 1.2, y_pos, f"{value:.1f}%", va="center", color=INK, fontsize=7.2)
    ax_b.text(
        0.02,
        -0.25,
        f"PRJNA772047 camera: {external_summary['results']['prjna772047_camera_supported_disease_pathways']} / "
        f"{external_summary['results']['pathways_tested_prjna772047']} pathways",
        transform=ax_b.transAxes,
        color=INK,
        fontsize=7.0,
    )
    ax_b.text(
        0.02,
        -0.35,
        "Formal inference: PRJNA772047 only (3 DMD / 2 Normal)",
        transform=ax_b.transAxes,
        color=MID,
        fontsize=6.7,
    )
    panel_label(ax_b, "B", "Patient-muscle replication")

    # Panel C: correction response and pathway closure by background.
    backgrounds = ["dup2", "dup2_9", "dup8_9"]
    x_pos = np.arange(3)
    reversal_percent = correction.loc[backgrounds, "significant_reversal_fraction"].to_numpy() * 100
    cosine_percent = correction.loc[backgrounds, "correction_to_ideal_cosine"].to_numpy() * 100
    ax_c.bar(x_pos, reversal_percent, width=0.62, color=[BLUE, TEAL, GOLD], alpha=0.88)
    ax_c.plot(x_pos, cosine_percent, color=VERMILION, marker="o", lw=1.2, ms=4.5)
    ax_c.set_ylim(0, 100)
    ax_c.set_xticks(x_pos, ["dup2", "dup2-9", "dup8-9"])
    ax_c.set_ylabel("Gene reversal or ideal-rescue cosine (%)")
    camera_counts = external_summary["results"]["camera_supported_independent_reversal_by_background"]
    strict_counts = reactome_summary["results"]["background_specific_dual_reversal_counts"]
    for idx, background in enumerate(backgrounds):
        ax_c.text(
            idx,
            reversal_percent[idx] + 3.0,
            f"{reversal_percent[idx]:.1f}%",
            ha="center",
            color=INK,
            fontsize=7.1,
        )
        ax_c.text(
            idx,
            4.0,
            f"external camera: {int(camera_counts[background])}\nwithin-bg dual: {int(strict_counts[background])}",
            ha="center",
            va="bottom",
            color=INK,
            fontsize=6.2,
        )
    legend = [
        Line2D([0], [0], color=BLUE, lw=6, label="significant gene reversal"),
        Line2D([0], [0], color=VERMILION, marker="o", lw=1.2, label="ideal-rescue cosine"),
    ]
    ax_c.legend(
        handles=legend,
        loc="upper left",
        bbox_to_anchor=(0.0, 0.97),
        frameon=False,
        fontsize=6.2,
    )
    ax_c.text(
        0.99,
        0.97,
        "cross-background dual: 0",
        transform=ax_c.transAxes,
        ha="right",
        va="top",
        color=VERMILION,
        fontsize=6.5,
        fontweight="bold",
    )
    panel_label(ax_c, "C", "Correction depends on genotype")

    # Panel D: organoid heterogeneity.
    lines = ["DMD1", "DMD2", "DMD3"]
    organoid_agreement = organoid_rows.loc[lines, "direction_agreement_to_ideal"].to_numpy() * 100
    d_y = np.arange(3)
    ax_d.barh(d_y, organoid_agreement, color=[BLUE, TEAL, GOLD], height=0.52)
    ax_d.axvline(50, color=VERMILION, lw=0.8, ls="--")
    ax_d.set_xlim(40, 75)
    ax_d.set_yticks(d_y, lines)
    ax_d.invert_yaxis()
    ax_d.set_xlabel("Agreement with GSE272233 disease pathways (%)")
    for y_pos, value in zip(d_y, organoid_agreement):
        ax_d.text(value + 0.6, y_pos, f"{value:.1f}%", va="center", color=INK, fontsize=7.1)
    panel_label(ax_d, "D", "Organoid heterogeneity")

    for ax in [ax_a, ax_b, ax_c, ax_d]:
        ax.spines["top"].set_visible(False)
        ax.spines["right"].set_visible(False)
        ax.grid(axis="y", color="#EDF1F4", linewidth=0.55, zorder=0)
        ax.set_axisbelow(True)

    fig.text(
        0.53,
        0.045,
        "MEASURED EVIDENCE    DMD induction    |    patient replication    |    "
        "3 correction backgrounds    |    organoid heterogeneity\n"
        "INFERENCE    descriptive axis    |    sample competition    |    "
        "genotype-specific reversal    |    shared-WT stress test",
        ha="center",
        va="center",
        color=INK,
        fontsize=6.4,
        fontweight="bold",
        bbox={"facecolor": PALE, "edgecolor": "#C9D2D9", "boxstyle": "square,pad=0.55"},
    )

    OUT.mkdir(parents=True, exist_ok=True)
    fig.savefig(OUT / "figure8_measured_dmd_evidence.png", dpi=600, facecolor="white")
    fig.savefig(OUT / "figure8_measured_dmd_evidence.pdf", facecolor="white")
    fig.savefig(OUT / "figure8_measured_dmd_evidence.svg", facecolor="white")
    fig.savefig(
        OUT / "figure8_measured_dmd_evidence_600dpi.tiff",
        dpi=600,
        facecolor="white",
        pil_kwargs={"compression": "tiff_lzw"},
    )
    plt.close(fig)

    source_rows = [
        {
            "panel": "A",
            "metric": "unselected_gene_axis_spearman",
            "value": induction_metrics["spearman_rho"],
            "n": int(induction_metrics["n_genes"]),
            "source": "GSE233606",
            "inference": "descriptive; one line per condition",
        },
        {
            "panel": "A",
            "metric": "unselected_gene_direction_agreement",
            "value": induction_metrics["direction_agreement_fraction"],
            "n": int(induction_metrics["n_genes"]),
            "source": "GSE233606",
            "inference": "descriptive; no cell-level P values",
        },
        {
            "panel": "B",
            "metric": "prjna772047_camera_supported_pathways",
            "value": external_summary["results"]["prjna772047_camera_supported_disease_pathways"],
            "n": external_summary["results"]["pathways_tested_prjna772047"],
            "source": "PRJNA772047",
            "inference": "camera, 3 DMD versus 2 Normal donors/samples",
        },
    ]
    for row in agreement_frame.to_dict("records"):
        source_rows.append(
            {
                "panel": "B",
                "metric": f"pathway_direction_agreement_{row['dataset']}",
                "value": row["agreement"],
                "n": row["n_pathways"],
                "source": row["dataset"],
                "inference": row["evidence_tier"],
            }
        )
    for background in backgrounds:
        source_rows.extend(
            [
                {
                    "panel": "C",
                    "metric": "significant_gene_reversal_fraction",
                    "value": correction.loc[background, "significant_reversal_fraction"],
                    "n": int(correction.loc[background, "n_both_fdr_lt_0_05"]),
                    "source": f"GSE272233:{background}",
                    "inference": "measured correction response",
                },
                {
                    "panel": "C",
                    "metric": "external_camera_reversal_count",
                    "value": camera_counts[background],
                    "n": external_summary["results"]["pathways_tested_prjna772047"],
                    "source": f"PRJNA772047->{background}",
                    "inference": "camera-supported independent direction reversal",
                },
                {
                    "panel": "C",
                    "metric": "within_background_dual_reversal_count",
                    "value": strict_counts[background],
                    "n": reactome_summary["design"]["tested_reactome_pathways"],
                    "source": f"GSE272233:{background}",
                    "inference": "camera and sample-score FDR below 0.05",
                },
            ]
        )
    for line in lines:
        source_rows.append(
            {
                "panel": "D",
                "metric": "organoid_pathway_direction_agreement_to_gse272233_disease",
                "value": organoid_rows.loc[line, "direction_agreement_to_ideal"],
                "n": int(organoid_rows.loc[line, "n_overlap"]),
                "source": f"GSE277637:{line}",
                "inference": "descriptive; shared WT reference",
            }
        )
    return pd.DataFrame(source_rows), agreement_frame


def build_tables(source_data: pd.DataFrame) -> None:
    source_data.to_csv(OUT / "figure8_source_data.tsv", sep="\t", index=False)
    dataset_rows = [
        {
            "dataset": "GSE233606",
            "species": "human",
            "modality": "single-cell RNA-seq",
            "groups": "healthy; patient DMD; engineered DMD",
            "biological_units": "one line per condition",
            "cells_or_samples": "3566 cells",
            "analysis_role": "measured DMD induction and patient-line replication",
            "formal_inference": "no",
            "claim_ceiling": "descriptive unselected gene-axis concordance",
        },
        {
            "dataset": "PRJNA772047",
            "species": "human",
            "modality": "single-nucleus RNA-seq pseudobulk",
            "groups": "3 DMD; 2 Normal",
            "biological_units": "5 reported donors/samples",
            "cells_or_samples": "8776 nuclei; 5 pseudobulks",
            "analysis_role": "independent patient-muscle pathway replication",
            "formal_inference": "yes, small-n camera/limma",
            "claim_ceiling": "sample-level pathway association",
        },
        {
            "dataset": "PRJNA1218493",
            "species": "human",
            "modality": "single-cell/nucleus RNA-seq pseudobulk",
            "groups": "2 DMD; 1 Normal",
            "biological_units": "3 reported samples",
            "cells_or_samples": "7439 cells; 3 pseudobulks",
            "analysis_role": "external disease-direction sensitivity",
            "formal_inference": "no, one control",
            "claim_ceiling": "directional sensitivity only",
        },
        {
            "dataset": "GSE156497",
            "species": "mouse",
            "modality": "single-cell RNA-seq pseudobulk",
            "groups": "1 DMD-like; 1 WT",
            "biological_units": "2 reported samples",
            "cells_or_samples": "11581 cells; 2 pseudobulks",
            "analysis_role": "cross-species disease-direction sensitivity",
            "formal_inference": "no",
            "claim_ceiling": "descriptive direction only",
        },
        {
            "dataset": "GSE272233",
            "species": "human",
            "modality": "bulk RNA-seq",
            "groups": "WT; 3 DMD duplication backgrounds; corrected clones",
            "biological_units": "3 reported biological repeats per group",
            "cells_or_samples": "21 samples",
            "analysis_role": "measured DMD-locus CRISPR correction",
            "formal_inference": "yes, camera/limma",
            "claim_ceiling": "background-specific correction response",
        },
        {
            "dataset": "GSE277637",
            "species": "human",
            "modality": "organoid single-cell RNA-seq pseudobulk",
            "groups": "3 DMD patient lines; 1 WT line",
            "biological_units": "4 lines with one shared WT",
            "cells_or_samples": "10480 cells; 4 pseudobulks",
            "analysis_role": "model-system heterogeneity stress test",
            "formal_inference": "no",
            "claim_ceiling": "line-specific descriptive effects",
        },
    ]
    pd.DataFrame(dataset_rows).to_csv(
        OUT / "table5_dmd_dataset_statistical_units.tsv", sep="\t", index=False
    )

    candidates = pd.read_csv(
        ROOT
        / "submission_nar_regular/stage1_v17_nar_dra_author_review_candidate/"
        "assets/tables/supplementary_table_s1_candidate_evidence_to_experiment_v26.tsv",
        sep="\t",
    )
    boundary = candidates[["gene", "action", "highest_supported_level", "gap_code", "gap_label"]].copy()
    boundary["measured_dmd_candidate_perturbation_response"] = 0
    boundary["candidate_level_dmd_pathway_prediction"] = 0
    boundary["current_claim"] = "governed hypothesis; no DMD candidate response truth"
    boundary.to_csv(
        OUT / "supplementary_table_candidate_dmd_truth_boundary.tsv",
        sep="\t",
        index=False,
    )

    summary = {
        "schema": "nmd-vcell-nar-measured-dmd-evidence/1.0",
        "figure": "figure8_measured_dmd_evidence",
        "candidate_count": int(len(boundary)),
        "measured_candidate_gene_dmd_perturbations": int(
            boundary["measured_dmd_candidate_perturbation_response"].sum()
        ),
        "candidate_level_dmd_pathway_predictions": int(
            boundary["candidate_level_dmd_pathway_prediction"].sum()
        ),
        "claim_boundary": (
            "Measured DMD-gene induction and correction are distinct from candidate-gene "
            "responses. All 21 candidates remain unmeasured in DMD perturbation context."
        ),
    }
    (OUT / "nar_measured_dmd_evidence_summary.json").write_text(
        json.dumps(summary, indent=2) + "\n", encoding="utf-8"
    )

    report = """# NAR measured DMD evidence module

This module joins four measured or externally observed axes without converting any
candidate into a DMD perturbation claim:

1. GSE233606 engineered DMD induction and patient-line replication.
2. Sample-level external pathway direction checks, with formal camera inference only for PRJNA772047.
3. GSE272233 DMD-locus CRISPR correction across three mutation backgrounds.
4. GSE277637 organoid-line heterogeneity using one shared WT reference.

The 21 governed candidates still have 0 measured DMD-context perturbation responses
and 0 candidate-level DMD pathway predictions.
"""
    (OUT / "NAR_MEASURED_DMD_EVIDENCE_REPORT.md").write_text(report, encoding="utf-8")


def write_manifest() -> None:
    paths = sorted(path for path in OUT.iterdir() if path.name != "manifest.tsv")
    rows = [
        {"file": path.name, "size_bytes": path.stat().st_size, "sha256": sha256(path)}
        for path in paths
    ]
    pd.DataFrame(rows).to_csv(OUT / "manifest.tsv", sep="\t", index=False)


def main() -> None:
    OUT.mkdir(parents=True, exist_ok=True)
    source_data, _ = build_figure()
    build_tables(source_data)
    write_manifest()
    print(f"Wrote NAR measured DMD evidence module to {OUT}")


if __name__ == "__main__":
    main()
