#!/usr/bin/env python3
"""Classify a query sample's FFPE damage relative to validated controls.

The classification is made ONLY by comparing measurable end-enrichment
frequencies of the query against the distribution of control samples that
were produced with the same library protocol.  There is no universal
absolute threshold: "low / moderate / high" is defined relative to the
controls.

Metrics used (from ffpe_damage_v2.py --outdir/end_enrichment.csv, window W):
  * C_to_T_R1_5prime_W : C>T frequency in the first W bases of R1 (forward)
  * G_to_A_R2_3prime_W : G>A frequency in the last W bases of R2 (reverse)
  * combined           : mean of the two damage-relevant metrics

Classification per metric (controls mean M, SD S, z = (query - M)/S):
  high      : z >  k      (query clearly above the control distribution)
  moderate  : -k <= z <= k
  low       : z <  -k

Outputs (--out):
  comparison_table.csv
  classification.csv
  control_vs_query.png
"""

import argparse
import os

import pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt


def parse_args():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--controls", nargs="+", required=True,
                    help="dirs with ffpe_damage_v2.py outputs for controls, "
                         "or a file with one dir per line")
    ap.add_argument("--query", required=True,
                    help="dir with ffpe_damage_v2.py outputs for the sample")
    ap.add_argument("--out", default="compare_out")
    ap.add_argument("--window", type=int, default=5,
                    help="end window size (bases from read end)")
    ap.add_argument("--k", type=float, default=2.0,
                    help="number of SDs defining moderate band")
    return ap.parse_args()


def load_samples(dirs):
    paths = []
    for d in dirs:
        if os.path.isfile(d):
            with open(d) as fh:
                paths.extend(line.strip() for line in fh if line.strip())
        else:
            paths.append(d)
    out = {}
    for p in paths:
        csv_path = os.path.join(p, "end_enrichment.csv")
        if not os.path.isfile(csv_path):
            raise FileNotFoundError(f"missing {csv_path}")
        e = pd.read_csv(csv_path)
        parent = os.path.basename(os.path.dirname(os.path.normpath(p)))
        name = parent or os.path.basename(os.path.normpath(p))
        out[name] = e
    return out


def metric_value(e, read, end, col, window):
    m = e[(e.read == read) & (e.end == end) & (e.window_size == window)]
    if m.empty:
        return float("nan")
    return float(m[col].iloc[0])


def build_metrics(e, window):
    return {
        f"C_to_T_R1_5prime_{window}":
            metric_value(e, "R1", "5prime",
                         "C_to_T_frequency_percent", window),
        f"G_to_A_R2_3prime_{window}":
            metric_value(e, "R2", "3prime",
                         "G_to_A_frequency_percent", window),
    }


def classify(z, k):
    if z > k:
        return "high"
    if z < -k:
        return "low"
    return "moderate"


def main():
    args = parse_args()
    os.makedirs(args.out, exist_ok=True)

    controls = load_samples(args.controls)
    query_e = load_samples([args.query])
    query_name = list(query_e)[0]
    query_e = query_e[query_name]

    rows = []
    metric_names = [f"C_to_T_R1_5prime_{args.window}",
                    f"G_to_A_R2_3prime_{args.window}"]
    query_values = {}

    for name, e in controls.items():
        m = build_metrics(e, args.window)
        rows.append({"sample": name, "type": "control", **m,
                     "combined": sum(m.values()) / 2})
    qm = build_metrics(query_e, args.window)
    q_combined = sum(qm.values()) / 2
    rows.append({"sample": query_name, "type": "query", **qm,
                 "combined": q_combined})
    for mn in metric_names:
        query_values[mn] = qm[mn]

    df = pd.DataFrame(rows)

    result_rows = []
    for mn in metric_names + ["combined"]:
        vals = df.loc[df.type == "control", mn].dropna()
        mean = vals.mean()
        sd = vals.std(ddof=1) if len(vals) > 1 else 0.0
        qv = query_values.get(mn, None) if mn in metric_names else q_combined
        if sd == 0 or pd.isna(sd):
            if pd.isna(qv) or qv == mean:
                z = 0.0
            else:
                z = float("inf") if qv > mean else float("-inf")
        else:
            z = (qv - mean) / sd
        cls = classify(z, args.k)
        result_rows.append({
            "metric": mn,
            "control_mean": mean,
            "control_sd": sd,
            "query_value": qv,
            "z_score": z,
            "classification": cls,
        })

    result_df = pd.DataFrame(result_rows)
    result_df.to_csv(os.path.join(args.out, "classification.csv"), index=False)
    df.to_csv(os.path.join(args.out, "comparison_table.csv"), index=False)

    n_metrics = len(metric_names)
    fig, axes = plt.subplots(1, n_metrics, figsize=(6 * n_metrics, 5),
                             squeeze=False)
    for ax, mn in zip(axes[0], metric_names):
        cvals = df.loc[df.type == "control", mn].dropna()
        qv = query_values[mn]
        mean = cvals.mean()
        sd = cvals.std(ddof=1) if len(cvals) > 1 else 0.0
        ax.scatter(range(len(cvals)), cvals, color="#1f77b4", s=60,
                   label="controls")
        ax.axhline(mean, color="#7f7f7f", linestyle="--", label="control mean")
        if sd > 0:
            ax.axhline(mean + args.k * sd, color="#ff7f0e", linestyle=":",
                       label=f"mean ± k·SD (k={args.k:g})")
            ax.axhline(mean - args.k * sd, color="#ff7f0e", linestyle=":")
        ax.scatter([len(cvals)], [qv], color="#d62728", s=90, marker="x",
                   label="query")
        ax.set_xticks([])
        ax.set_ylabel("frequency (%)")
        ax.set_title(mn)
        ax.legend(fontsize=8)
    fig.tight_layout()
    fig.savefig(os.path.join(args.out, "control_vs_query.png"), dpi=300)
    plt.close(fig)

    print(result_df.to_string(index=False))


if __name__ == "__main__":
    main()