"""
Bootstrap the 100% data points without retraining.

For each model in {unet, segformer_b0}:
  1. Copy the existing best checkpoint to checkpoints/{model}_100_best.pth
  2. Parse the existing training_logs.txt into per-epoch JSON entries
     (loss / dice / iou / pixel_acc — but NOT mIoU, which the old trainer
     never logged).
  3. Run a single no-grad pass over the full val set with the new metrics
     code so we get a comparable mIoU/IoU/Dice/PixelAcc number to put on
     the data-share-vs-final scaling chart.

Source checkpoints:
    pv_panel_models/unet_model/checkpoints/best_model.pth     →  unet_100_best.pth
    pv_panel_models/vit_model/checkpoints/best_model.pth      →  segformer_b0_100_best.pth

Usage:
    python bootstrap_100.py                # full bootstrap (copy + parse + val recompute)
    python bootstrap_100.py --skip-val     # copy + parse only (no GPU work)

The val recompute is inference-only — it does not modify the checkpoints.
"""
import argparse
import json
import re
import shutil
import time
from pathlib import Path

import torch
from torch.utils.data import DataLoader

from dataset import SubsetSolarPanelDataset
from metrics import SegMetrics
from models import MODEL_REGISTRY

THIS_DIR = Path(__file__).resolve().parent
REPO_ROOT = THIS_DIR.parents[1]

VAL_IMG = REPO_ROOT / "final_data" / "val" / "images"
VAL_MSK = REPO_ROOT / "final_data" / "val" / "masks"

LOG_DIR = THIS_DIR / "logs"
CKPT_DIR = THIS_DIR / "checkpoints"

# (model_name, source_checkpoint, source_text_log)
SOURCES = [
    (
        "unet",
        REPO_ROOT / "pv_panel_models" / "unet_model" / "checkpoints" / "best_model.pth",
        REPO_ROOT / "pv_panel_models" / "unet_model" / "checkpoints" / "training_logs.txt",
    ),
    (
        "segformer_b0",
        REPO_ROOT / "pv_panel_models" / "vit_model" / "checkpoints" / "best_model.pth",
        REPO_ROOT / "pv_panel_models" / "vit_model" / "checkpoints" / "training_logs.txt",
    ),
]

EPOCH_HEADER = re.compile(r"^Epoch\s+(\d+)\s*/\s*(\d+)\s*$")
TRAIN_LINE = re.compile(
    r"Train Loss:\s*([0-9.]+).*?Acc:\s*([0-9.]+).*?Prec:\s*([0-9.]+).*?"
    r"Rec:\s*([0-9.]+).*?Dice:\s*([0-9.]+).*?IoU:\s*([0-9.]+)"
)
VAL_LINE = re.compile(
    r"Val Loss:\s*([0-9.]+).*?Acc:\s*([0-9.]+).*?Prec:\s*([0-9.]+).*?"
    r"Rec:\s*([0-9.]+).*?Dice:\s*([0-9.]+).*?IoU:\s*([0-9.]+)"
)


def parse_text_log(log_path: Path):
    """Parse pv_panel_models text log → list of per-epoch dicts.

    Old logs only contain {loss, accuracy, precision, recall, dice, iou}.
    mIoU was never computed there, so it is recorded as None.
    """
    if not log_path.is_file():
        raise FileNotFoundError(f"text log not found: {log_path}")

    text = log_path.read_text()
    epochs = []
    current = None

    for raw in text.splitlines():
        line = raw.strip()
        m = EPOCH_HEADER.match(line)
        if m:
            if current is not None and "train_loss" in current and "val_loss" in current:
                epochs.append(current)
            current = {"epoch": int(m.group(1))}
            continue

        if current is None:
            continue

        m = TRAIN_LINE.search(line)
        if m:
            loss, acc, _prec, _rec, dice, iou = map(float, m.groups())
            current.update({
                "train_loss": loss,
                "train_pixel_acc": acc,
                "train_dice": dice,
                "train_iou": iou,
                "train_miou": None,
            })
            continue

        m = VAL_LINE.search(line)
        if m:
            loss, acc, _prec, _rec, dice, iou = map(float, m.groups())
            current.update({
                "val_loss": loss,
                "val_pixel_acc": acc,
                "val_dice": dice,
                "val_iou": iou,
                "val_miou": None,
            })

    if current is not None and "train_loss" in current and "val_loss" in current:
        epochs.append(current)

    if not epochs:
        raise RuntimeError(f"no epochs parsed from {log_path}")

    return epochs


@torch.no_grad()
def recompute_val_metrics(model_name: str, ckpt_path: Path, device: str):
    """One forward pass over the full val set with the new metrics code."""
    model_fn = MODEL_REGISTRY[model_name]
    model, _ = model_fn()

    state = torch.load(ckpt_path, map_location=device, weights_only=False)
    model.load_state_dict(state["model_state_dict"])
    model.to(device).eval()

    val_set = SubsetSolarPanelDataset(
        VAL_IMG, VAL_MSK, file_list=None, image_size=128, augment=False,
    )
    val_loader = DataLoader(val_set, batch_size=16, shuffle=False, num_workers=4, pin_memory=True)

    metrics = SegMetrics()
    for images, masks in val_loader:
        images = images.to(device, non_blocking=True)
        masks = masks.to(device, non_blocking=True)
        outputs = model(images)
        metrics.update(outputs, masks)

    return metrics.compute()


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--skip-val", action="store_true",
                   help="skip the no-grad val pass (faster, but the scaling chart "
                        "loses the new-definition mIoU/Dice/IoU at 100%%)")
    return p.parse_args()


def main():
    args = parse_args()

    LOG_DIR.mkdir(parents=True, exist_ok=True)
    CKPT_DIR.mkdir(parents=True, exist_ok=True)

    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"[bootstrap] device={device}  skip_val={args.skip_val}")

    for model_name, src_ckpt, src_log in SOURCES:
        print(f"\n── {model_name} @ 100% ────────────────────────────────")
        if not src_ckpt.is_file():
            print(f"  ✗ source checkpoint missing: {src_ckpt}")
            continue
        if not src_log.is_file():
            print(f"  ✗ source text log missing: {src_log}")
            continue

        # 1. Copy checkpoint
        dst_ckpt = CKPT_DIR / f"{model_name}_100_best.pth"
        shutil.copy2(src_ckpt, dst_ckpt)
        print(f"  ✓ copied checkpoint → {dst_ckpt.name}")

        # 2. Parse training_logs.txt
        epochs = parse_text_log(src_log)
        print(f"  ✓ parsed {len(epochs)} epochs from {src_log.name}")

        history = {
            "model": model_name,
            "share": 100,
            "n_train": 5325,
            "n_val": 1331,
            "epochs": epochs,
            "bootstrapped_from": str(src_ckpt.relative_to(REPO_ROOT)),
            "metric_caveat": (
                "Per-epoch metrics parsed from existing training_logs.txt "
                "(per-batch averaging). mIoU was not logged in the old "
                "trainer and is null per epoch. The 'recomputed_val_metrics' "
                "field below holds new-definition (global confusion matrix) "
                "values comparable to the 25%/50% runs."
            ),
            "recomputed_val_metrics": None,
            "val_recompute_seconds": None,
            "best_val_dice": max(e["val_dice"] for e in epochs),
            "bootstrap_time_iso": time.strftime("%Y-%m-%dT%H:%M:%S"),
        }

        # 3. Optional val recompute
        if not args.skip_val:
            print("  · running val pass with new metrics code…", end=" ", flush=True)
            t0 = time.time()
            new_val = recompute_val_metrics(model_name, dst_ckpt, device)
            dt = time.time() - t0
            history["recomputed_val_metrics"] = new_val
            history["val_recompute_seconds"] = dt
            print(
                f"done in {dt:.1f}s  "
                f"dice={new_val['dice']:.4f}  iou={new_val['iou']:.4f}  "
                f"miou={new_val['miou']:.4f}  pixel_acc={new_val['pixel_acc']:.4f}"
            )

        log_path = LOG_DIR / f"{model_name}_100.json"
        with open(log_path, "w") as f:
            json.dump(history, f, indent=2)
        print(f"  ✓ wrote log → {log_path.name}")

    print("\n[bootstrap] done.")


if __name__ == "__main__":
    main()