"""Produce comprehensive statistics on the current healthy-brain corpus."""
import pandas as pd
from pathlib import Path
import textwrap

ROOT = Path("/home/MRI-DataSet")
df = pd.read_csv(ROOT / "manifest.csv")

out = []
out.append("=" * 72)
out.append("HEALTHY-BRAIN TRAINING CORPUS — FINAL SNAPSHOT")
out.append("=" * 72)
out.append(f"Manifest:                     {ROOT/'manifest.csv'}")
out.append(f"Golden-0-to-25 (primary):     {ROOT/'Golden-0-to-25'}/")
out.append(f"Golden-25plus (fine-tune):    {ROOT/'Golden-25plus'}/")
out.append("")

# Overall
total = len(df)
healthy = (df["healthy"] == True).sum()
with_age = df["age_years"].notna().sum()
age025 = ((df["age_years"].notna()) & (df["age_years"] < 25)).sum()
age25p = ((df["age_years"].notna()) & (df["age_years"] >= 25)).sum()
no_age = df["age_years"].isna().sum()

out.append(f"Total scans:                  {total}")
out.append(f"Labelled healthy:             {healthy} ({100*healthy/total:.1f}%)")
out.append(f"With known age:               {with_age} ({100*with_age/total:.1f}%)")
out.append(f"  age <  25 (primary):        {age025}")
out.append(f"  age >= 25 (pretrain):       {age25p}")
out.append(f"Unknown age (to impute):      {no_age}")
out.append("")

# Age stats
if with_age:
    ages = df["age_years"].dropna()
    out.append("Age distribution (years) over dated scans:")
    out.append(f"  min={ages.min():.1f}  median={ages.median():.1f}  max={ages.max():.1f}")
    out.append(f"  mean={ages.mean():.1f} ± {ages.std():.1f}")
    out.append("")

# Split breakdown
if "split" in df.columns:
    out.append("-" * 72)
    out.append("SPLIT BREAKDOWN")
    out.append("-" * 72)
    for split, grp in df.groupby("split"):
        n = len(grp)
        h = (grp["healthy"] == True).sum()
        ages = grp["age_years"].dropna()
        rng = f"{ages.min():.1f}-{ages.max():.1f}" if len(ages) else "n/a"
        out.append(f"  {split:<28}  {n:>5} scans   (healthy={h})  range={rng}")
    out.append("")

# Per-dataset breakdown
out.append("-" * 72)
out.append("PER-DATASET BREAKDOWN")
out.append("-" * 72)
out.append(f"{'dataset':<28}{'scans':>7}{'healthy':>9}{'age<25':>8}{'age>=25':>9}"
           f"{'age range':>18}")
out.append("-" * 72)
for ds, grp in df.groupby("dataset"):
    n = len(grp)
    h = (grp["healthy"] == True).sum()
    a = ((grp["age_years"].notna()) & (grp["age_years"] < 25)).sum()
    b = ((grp["age_years"].notna()) & (grp["age_years"] >= 25)).sum()
    ages = grp["age_years"].dropna()
    if len(ages):
        rng = f"{ages.min():.1f}-{ages.max():.1f}"
    else:
        rng = "n/a"
    out.append(f"{ds:<28}{n:>7}{h:>9}{a:>8}{b:>9}{rng:>18}")
out.append("-" * 72)
out.append("")

# Age histogram in 5-year bins
out.append("AGE HISTOGRAM (5-year bins, across all dated scans)")
ages = df["age_years"].dropna()
if len(ages):
    max_age = int(ages.max()) + 5
    edges = list(range(0, max_age + 1, 5))
    counts, _ = pd.cut(ages, bins=edges, right=False, include_lowest=True).value_counts(sort=False), None
    for interval, n in counts.items():
        bar = "#" * int(40 * n / counts.max())
        out.append(f"  [{interval.left:>3}, {interval.right:>3})  {n:>5}  {bar}")
    out.append("")

# Split targets
out.append("=" * 72)
out.append("TARGET SPLITS FOR TRAINING")
out.append("=" * 72)
out.append("")
out.append(textwrap.dedent(f"""\
    Primary training (0-25 years)   : {age025:>5} scans
    Pretraining pool (25+ years)    : {age25p:>5} scans
    Grand total labelled            : {age025+age25p:>5} scans

    Strategy (per user):
      1. Train EfficientNet-B3 multi-view on Training_0to25/   (pediatric + young adult)
      2. Ablation: pretrain on Training_0to25 + Pretrain_25plus
         (full lifespan), then fine-tune on 0-25
      3. Compare pediatric MAE between the two strategies in thesis
"""))

text = "\n".join(out)
print(text)
(ROOT / "STATISTICS.md").write_text(text)
(ROOT / "STATISTICS.txt").write_text(text)
print(f"\n-> wrote {ROOT/'STATISTICS.md'} and {ROOT/'STATISTICS.txt'}")