"""Script 09: Run ablation experiments on the Proposed model.

Three variants (see src/ablation.py for full design notes):

  A1  -- remove text-type metadata (mask TF-IDF slice; keep numeric).
  A2  -- remove numeric metadata    (mask numeric slice; keep TF-IDF).
  A3  -- replace Cross-Attention fusion with [text;meta] concat + Linear.

For each variant the script:
  1. Trains the corresponding ablated model on train/val (same hyperparameters
     as 04_train_proposed.py).
  2. Evaluates on the held-out test set, producing per-aspect F1 / accuracy.
  3. Writes:
        checkpoints/<variant_subdir>/best.pt
        checkpoints/<variant_subdir>/history.json
        reports/ablation_<VARIANT>.json     (per-aspect detail)

After all requested variants are done, the script writes a summary table to
`reports/ablation_summary.json` that includes the Proposed model's per-aspect
numbers (read from reports/per_aspect_proposed.json if present) so the
ablation deltas can be read off directly.

Usage
-----
    # Run all three ablations
    python scripts/09_ablation.py

    # Run a single variant
    python scripts/09_ablation.py --variant A1

    # Skip training and only evaluate (requires checkpoints to exist already)
    python scripts/09_ablation.py --eval_only

Examples of the expected output table (reports/ablation_summary.json):
    {
      "Proposed":            {"mean_macro_f1": 0.71, "per_aspect": {...}},
      "A1_no_text_meta":     {"mean_macro_f1": 0.68, "per_aspect": {...}},
      "A2_no_numeric_meta":  {"mean_macro_f1": 0.70, "per_aspect": {...}},
      "A3_concat_fusion":    {"mean_macro_f1": 0.69, "per_aspect": {...}}
    }
"""
import argparse
import json
import logging
import sys
from pathlib import Path

import pandas as pd

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from src.utils import setup_logging
from src import config as cfg
from src.meta_encoder import MetaEncoder
from src.trainer import train_meta_acsa
from src.evaluator import evaluate_per_aspect
from src.ablation import (
    ABLATION_VARIANTS, MaskedMetaEncoder,
    train_concat_fusion_acsa,
    load_meta_acsa_for_ablation, load_concat_fusion_acsa,
    predict_per_aspect_for_ablation,
)


logger = logging.getLogger(__name__)


def _build_meta_for_variant(variant: str, base_encoder: MetaEncoder):
    """Return the (possibly masked) meta encoder used by the variant."""
    spec = ABLATION_VARIANTS[variant]
    if spec["meta_mask"] is None:
        return base_encoder
    return MaskedMetaEncoder(base=base_encoder, mask=spec["meta_mask"])


def _train_variant(variant: str, train_df, val_df, base_encoder, args):
    spec = ABLATION_VARIANTS[variant]
    meta_enc = _build_meta_for_variant(variant, base_encoder)
    output_dir = cfg.CHECKPOINT_DIR / spec["checkpoint_subdir"]

    print(f"\n>>> Training ablation {variant} ({spec['name']})")
    print(f"    {spec['description']}")
    print(f"    checkpoint -> {output_dir}")

    if spec["fusion"] == "cross_attention":
        # A1/A2 reuse the proposed training loop with a masked encoder.
        train_meta_acsa(
            train_df=train_df, val_df=val_df, meta_encoder=meta_enc,
            bert_name=args.bert_name, epochs=args.epochs,
            batch_size=args.batch_size,
            lr_bert=args.lr_bert, lr_heads=args.lr_heads,
            use_class_weights=not args.no_class_weights,
            output_dir=output_dir,
        )
    elif spec["fusion"] == "concat":
        train_concat_fusion_acsa(
            train_df=train_df, val_df=val_df, meta_encoder=meta_enc,
            bert_name=args.bert_name, epochs=args.epochs,
            batch_size=args.batch_size,
            lr_bert=args.lr_bert, lr_heads=args.lr_heads,
            use_class_weights=not args.no_class_weights,
            output_dir=output_dir,
        )
    else:
        raise ValueError(f"Unknown fusion type: {spec['fusion']}")


def _evaluate_variant(variant: str, test_df, base_encoder):
    spec = ABLATION_VARIANTS[variant]
    ckpt_dir = cfg.CHECKPOINT_DIR / spec["checkpoint_subdir"]
    if not (ckpt_dir / "best.pt").exists():
        print(f"[warn] checkpoint not found for {variant}: {ckpt_dir/'best.pt'}  -- skipping")
        return None

    meta_enc = _build_meta_for_variant(variant, base_encoder)

    print(f"\n>>> Evaluating ablation {variant} ({spec['name']}) on test set")
    if spec["fusion"] == "cross_attention":
        model, tok, device = load_meta_acsa_for_ablation(ckpt_dir, meta_enc)
    else:
        model, tok, device = load_concat_fusion_acsa(ckpt_dir, meta_enc)

    preds, labels = predict_per_aspect_for_ablation(model, tok, test_df, meta_enc, device)
    detail = evaluate_per_aspect(preds, labels)

    detail["variant"] = variant
    detail["variant_name"] = spec["name"]
    detail["description"] = spec["description"]

    out_path = cfg.REPORT_DIR / f"ablation_{variant}.json"
    with open(out_path, "w") as f:
        json.dump(detail, f, indent=2)
    print(f"    per-aspect detail -> {out_path}")
    print(f"    mean macro-F1 = {detail['overall']['mean_macro_f1']:.4f}  "
          f"mean acc = {detail['overall']['mean_accuracy']:.4f}")
    return detail


def _load_proposed_detail():
    """Read previously saved per_aspect_proposed.json so the summary can
    show the Proposed baseline alongside the ablations.
    """
    p = cfg.REPORT_DIR / "per_aspect_proposed.json"
    if not p.exists():
        return None
    with open(p) as f:
        return json.load(f)


def _summarise(details):
    """Build a compact summary table dict from per-variant detail dicts."""
    summary = {}
    for tag, d in details.items():
        if d is None:
            summary[tag] = None
            continue
        per_aspect = {a: {"macro_f1": m["macro_f1"], "accuracy": m["accuracy"]}
                      for a, m in d["per_aspect"].items()}
        summary[tag] = {
            "mean_macro_f1": d["overall"]["mean_macro_f1"],
            "mean_accuracy": d["overall"]["mean_accuracy"],
            "per_aspect": per_aspect,
        }
    return summary


def _print_summary_table(summary):
    """Pretty-print a per-aspect F1 table across the ablations + Proposed."""
    tags = [t for t in summary if summary[t] is not None]
    if not tags:
        print("[info] No results to summarise.")
        return
    # Header
    aspects = list(cfg.ASPECTS)
    col_w = 12
    header = "Aspect".ljust(col_w) + "".join(t.ljust(col_w) for t in tags)
    print("\n" + "=" * len(header))
    print("ABLATION SUMMARY  -- per-aspect macro-F1")
    print("=" * len(header))
    print(header)
    print("-" * len(header))
    for a in aspects:
        row = a.ljust(col_w)
        for t in tags:
            f1 = summary[t]["per_aspect"].get(a, {}).get("macro_f1")
            row += (f"{f1:.4f}".ljust(col_w) if f1 is not None else " --".ljust(col_w))
        print(row)
    print("-" * len(header))
    mean_row = "MEAN_F1".ljust(col_w)
    for t in tags:
        mean_row += f"{summary[t]['mean_macro_f1']:.4f}".ljust(col_w)
    print(mean_row)
    mean_acc_row = "MEAN_ACC".ljust(col_w)
    for t in tags:
        mean_acc_row += f"{summary[t]['mean_accuracy']:.4f}".ljust(col_w)
    print(mean_acc_row)
    print("=" * len(header))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--variant", choices=["A1", "A2", "A3", "all"], default="all",
                        help="Which ablation variant to run (default: all).")
    parser.add_argument("--eval_only", action="store_true",
                        help="Skip training; only evaluate existing checkpoints.")
    parser.add_argument("--bert_name", default=cfg.BERT_MODEL_NAME)
    parser.add_argument("--epochs", type=int, default=cfg.DEFAULT_EPOCHS)
    parser.add_argument("--batch_size", type=int, default=cfg.DEFAULT_BATCH_SIZE)
    parser.add_argument("--lr_bert", type=float, default=cfg.DEFAULT_LR_BERT)
    parser.add_argument("--lr_heads", type=float, default=cfg.DEFAULT_LR_HEADS)
    parser.add_argument("--no_class_weights", action="store_true")
    args = parser.parse_args()

    setup_logging()

    # Sanity-check inputs
    if not cfg.META_ENCODER_PATH.exists():
        raise RuntimeError(f"MetaEncoder not found at {cfg.META_ENCODER_PATH}. "
                           "Run scripts/03_label.py first.")
    train_df = pd.read_parquet(cfg.TRAIN_PATH)
    val_df = pd.read_parquet(cfg.VAL_PATH)
    test_df = pd.read_parquet(cfg.TEST_PATH).reset_index(drop=True)

    expected = [f"aspect_{a}" for a in cfg.ASPECTS]
    missing = [c for c in expected if c not in train_df.columns]
    if missing:
        raise RuntimeError(f"Missing aspect label columns: {missing}. "
                           "Run scripts/03_label.py first.")

    base_encoder = MetaEncoder.load()
    print(f"Loaded base MetaEncoder (total_dim={base_encoder.total_dim}, "
          f"tfidf={base_encoder.tfidf_dim}, numeric={base_encoder.num_dim})")

    variants = ["A1", "A2", "A3"] if args.variant == "all" else [args.variant]

    # 1. Train (unless --eval_only)
    if not args.eval_only:
        for v in variants:
            _train_variant(v, train_df, val_df, base_encoder, args)

    # 2. Evaluate
    details = {}
    for v in variants:
        details[ABLATION_VARIANTS[v]["name"]] = _evaluate_variant(v, test_df, base_encoder)

    # 3. Pull in Proposed for direct comparison, then write summary
    proposed = _load_proposed_detail()
    if proposed is not None:
        details = {"Proposed": proposed, **details}
    else:
        print("[info] reports/per_aspect_proposed.json not found; "
              "summary will not include the Proposed baseline.")

    summary = _summarise(details)
    _print_summary_table(summary)

    out_path = cfg.REPORT_DIR / "ablation_summary.json"
    with open(out_path, "w") as f:
        json.dump(summary, f, indent=2)
    print(f"\nSummary -> {out_path}")


if __name__ == "__main__":
    main()