#!/usr/bin/env python3
"""Shared optimizer constructor for backbone-body learning-rate multiplier experiments.

Splits trainable parameters into two groups:
  - ``backbone_body``: parameters whose name starts with ``backbone.body.``
  - ``detection``:    everything else (FPN, RPN, ROI heads, predictors)

The body group receives ``base_lr * multiplier``; the detection group keeps
``base_lr``.  Both groups must be non-empty and together must cover every
trainable parameter exactly once.
"""

from __future__ import annotations

import json
import hashlib
import math
from pathlib import Path
from typing import Any, Dict, List, Tuple

import torch

BODY_PREFIX = "backbone.body."
GROUP_BODY = "backbone_body"
GROUP_DETECTION = "detection"
_EPS = 1e-12


def validate_multiplier(multiplier: float) -> float:
    """Return *multiplier* if valid, else raise ``ValueError``."""
    m = float(multiplier)
    if not math.isfinite(m):
        raise ValueError(f"backbone-lr-multiplier must be finite, got {multiplier}")
    if m <= 0.0:
        raise ValueError(f"backbone-lr-multiplier must be positive, got {multiplier}")
    return m


def validate_base_lr(base_lr: float) -> float:
    """Return *base_lr* if it is finite and positive."""
    lr = float(base_lr)
    if not math.isfinite(lr) or lr <= 0.0:
        raise ValueError(f"base_lr must be finite and positive, got {base_lr}")
    return lr


def split_trainable_parameters(
    model: torch.nn.Module,
    *,
    require_both_groups: bool = True,
) -> Tuple[List[torch.nn.Parameter], List[torch.nn.Parameter]]:
    """Split *model*'s trainable parameters into (body, detection) lists.

    Guarantees:
      * no parameter appears in both lists;
      * the union covers every trainable parameter exactly once;
      * both lists are non-empty.
    """
    body: List[torch.nn.Parameter] = []
    detection: List[torch.nn.Parameter] = []
    for name, param in model.named_parameters():
        if not param.requires_grad:
            continue
        if name.startswith(BODY_PREFIX):
            body.append(param)
        else:
            detection.append(param)

    if require_both_groups and not body:
        raise RuntimeError("No trainable parameters found under 'backbone.body.*'")
    if require_both_groups and not detection:
        raise RuntimeError("No trainable parameters found outside 'backbone.body.*'")
    if not body and not detection:
        raise RuntimeError("Model has no trainable parameters")

    # Sanity: no overlap (should be impossible given the prefix check, but be explicit)
    body_ids = {id(p) for p in body}
    det_ids = {id(p) for p in detection}
    overlap = body_ids & det_ids
    if overlap:
        raise RuntimeError(f"Parameter(s) appear in both groups: {overlap}")

    trainable_ids = {id(p) for p in model.parameters() if p.requires_grad}
    grouped_ids = body_ids | det_ids
    if grouped_ids != trainable_ids:
        missing = trainable_ids - grouped_ids
        extra = grouped_ids - trainable_ids
        raise RuntimeError(
            "Optimizer groups do not exactly cover trainable parameters: "
            f"missing={len(missing)}, extra={len(extra)}"
        )

    return body, detection


def build_adaptation_optimizer(
    model: torch.nn.Module,
    base_lr: float,
    multiplier: float,
    momentum: float = 0.9,
    weight_decay: float = 5e-4,
    require_both_groups: bool = True,
) -> torch.optim.SGD:
    """Create an SGD optimizer with separate LR for backbone body vs detection.

    Parameters
    ----------
    model : torch.nn.Module
        The Mask R-CNN model (already on device).
    base_lr : float
        Learning rate for the detection group (and body when multiplier=1).
    multiplier : float
        Multiplier applied to *base_lr* for the ``backbone.body.*`` group.
    momentum, weight_decay : float
        Passed through to SGD.

    Returns
    -------
    torch.optim.SGD
        Optimizer with two named param groups.
    """
    base_lr = validate_base_lr(base_lr)
    multiplier = validate_multiplier(multiplier)
    body_params, det_params = split_trainable_parameters(
        model, require_both_groups=require_both_groups
    )

    param_groups = []
    if body_params:
        param_groups.append({
            "params": body_params,
            "lr": base_lr * multiplier,
            "group_name": GROUP_BODY,
        })
    if det_params:
        param_groups.append({
            "params": det_params,
            "lr": base_lr,
            "group_name": GROUP_DETECTION,
        })
    return torch.optim.SGD(
        param_groups,
        lr=base_lr,  # fallback (overridden per-group)
        momentum=momentum,
        weight_decay=weight_decay,
    )


def optimizer_group_summary(optimizer: torch.optim.Optimizer) -> Dict[str, Any]:
    """Return a JSON-serialisable summary of optimizer param groups."""
    groups: Dict[str, Any] = {}
    for g in optimizer.param_groups:
        name = g.get("group_name", "unnamed")
        if name in groups:
            raise RuntimeError(f"Duplicate optimizer group_name: {name}")
        n_tensors = len(g["params"])
        n_scalars = sum(p.numel() for p in g["params"])
        groups[name] = {
            "lr": g["lr"],
            "num_tensors": n_tensors,
            "num_scalars": n_scalars,
        }
    return groups


# ---------------------------------------------------------------------------
# Parameter-adaptation statistics
# ---------------------------------------------------------------------------


def snapshot_trainable_parameters(model: torch.nn.Module) -> Dict[str, torch.Tensor]:
    """Return a CPU copy of every trainable parameter (name → tensor)."""
    return {
        name: param.detach().cpu().clone()
        for name, param in model.named_parameters()
        if param.requires_grad
    }


def compute_parameter_adaptation(
    initial: Dict[str, torch.Tensor],
    model: torch.nn.Module,
    optimizer: torch.optim.Optimizer,
) -> Dict[str, Any]:
    """Compute per-group and overall L2 adaptation statistics.

    Parameters
    ----------
    initial : dict
        Output of :func:`snapshot_trainable_parameters` taken *before* training.
    model : torch.nn.Module
        The model *after* training (on any device).
    optimizer : torch.optim.Optimizer
        The optimizer used during training (for effective LR lookup).

    Returns
    -------
    dict
        JSON-serialisable statistics with keys ``backbone_body``,
        ``detection``, and ``all_trainable``.
    """
    # Build effective-LR lookup from optimizer
    lr_by_group: Dict[str, float] = {}
    for g in optimizer.param_groups:
        lr_by_group[g.get("group_name", "unnamed")] = g["lr"]

    current_names = {
        name for name, param in model.named_parameters() if param.requires_grad
    }
    initial_names = set(initial)
    if current_names != initial_names:
        raise RuntimeError(
            "Initial snapshot does not match current trainable parameters: "
            f"missing={sorted(current_names - initial_names)}, "
            f"extra={sorted(initial_names - current_names)}"
        )

    # Accumulate in float64 to avoid precision loss on large models
    stats: Dict[str, Dict[str, float]] = {
        GROUP_BODY: {"initial_l2": 0.0, "final_l2": 0.0, "delta_l2": 0.0, "num_tensors": 0, "num_scalars": 0},
        GROUP_DETECTION: {"initial_l2": 0.0, "final_l2": 0.0, "delta_l2": 0.0, "num_tensors": 0, "num_scalars": 0},
    }

    for name, param in model.named_parameters():
        if not param.requires_grad:
            continue
        if name not in initial:
            raise RuntimeError(f"Parameter '{name}' missing from initial snapshot")
        group = GROUP_BODY if name.startswith(BODY_PREFIX) else GROUP_DETECTION
        init_t = initial[name].to(torch.float64)
        final_t = param.detach().cpu().to(torch.float64)
        stats[group]["initial_l2"] += float((init_t ** 2).sum())
        stats[group]["final_l2"] += float((final_t ** 2).sum())
        stats[group]["delta_l2"] += float(((final_t - init_t) ** 2).sum())
        stats[group]["num_tensors"] += 1
        stats[group]["num_scalars"] += param.numel()

    # Take square roots and compute relative delta
    result: Dict[str, Any] = {}
    for group_name in (GROUP_BODY, GROUP_DETECTION):
        s = stats[group_name]
        initial_l2 = math.sqrt(s["initial_l2"])
        final_l2 = math.sqrt(s["final_l2"])
        delta_l2 = math.sqrt(s["delta_l2"])
        relative = delta_l2 / max(initial_l2, _EPS)
        result[group_name] = {
            "initial_l2": initial_l2,
            "final_l2": final_l2,
            "delta_l2": delta_l2,
            "relative_l2_delta": relative,
            "num_tensors": s["num_tensors"],
            "num_scalars": s["num_scalars"],
            "effective_lr": lr_by_group.get(group_name),
        }

    # All-trainable aggregate
    all_init = sum(stats[g]["initial_l2"] for g in (GROUP_BODY, GROUP_DETECTION))
    all_final = sum(stats[g]["final_l2"] for g in (GROUP_BODY, GROUP_DETECTION))
    all_delta = sum(stats[g]["delta_l2"] for g in (GROUP_BODY, GROUP_DETECTION))
    all_init_l2 = math.sqrt(all_init)
    all_final_l2 = math.sqrt(all_final)
    all_delta_l2 = math.sqrt(all_delta)
    result["all_trainable"] = {
        "initial_l2": all_init_l2,
        "final_l2": all_final_l2,
        "delta_l2": all_delta_l2,
        "relative_l2_delta": all_delta_l2 / max(all_init_l2, _EPS),
        "num_tensors": sum(stats[g]["num_tensors"] for g in (GROUP_BODY, GROUP_DETECTION)),
        "num_scalars": sum(stats[g]["num_scalars"] for g in (GROUP_BODY, GROUP_DETECTION)),
        "effective_lr": None,  # mixed
    }

    # Validate finiteness
    for group_name, group_stats in result.items():
        for key, value in group_stats.items():
            if isinstance(value, float) and not math.isfinite(value):
                raise RuntimeError(f"Non-finite value in parameter_adaptation[{group_name}][{key}]: {value}")

    return result


def write_parameter_adaptation(
    output_dir: Path,
    initial: Dict[str, torch.Tensor],
    model: torch.nn.Module,
    optimizer: torch.optim.Optimizer,
) -> Path:
    """Compute adaptation stats and write ``parameter_adaptation.json``."""
    stats = compute_parameter_adaptation(initial, model, optimizer)
    path = output_dir / "parameter_adaptation.json"
    path.write_text(json.dumps(stats, indent=2), encoding="utf-8")
    return path


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