"""Standard GRM ablation matrix as configuration metadata."""

from __future__ import annotations

from dataclasses import dataclass, replace

from graph.evaluation.evaluator import evaluate
from graph.model import GraphModel
from graph.synthetic.generator import ReasoningExample


@dataclass(frozen=True)
class Ablation:
    name: str
    description: str


ABLATIONS = (
    Ablation("sage", "relation-aware GraphSAGE production encoder"),
    Ablation("gat", "relation-aware GAT with retained edge attention"),
    Ablation(
        "without-pat", "GraphSAGE plus reasoner without path-aware representations"
    ),
    Ablation("pat-without-query", "PAT with query conditioning disabled"),
    Ablation(
        "without-query-conditioning",
        "all query conditioning disabled in the PAT ablation",
    ),
    Ablation("greedy", "beam width one"),
    Ablation("without-verifier", "neural scores without deterministic evidence gate"),
    Ablation(
        "without-path-supervision", "path and PAT supervision weights set to zero"
    ),
    Ablation("untyped", "type features removed"),
    Ablation("random-negatives", "hard-negative generation disabled"),
    Ablation("synthetic-only", "no approved product-structure adaptation"),
)


def run_ablation_suite(
    examples: list[ReasoningExample],
    model: GraphModel,
    *,
    names: tuple[str, ...] | None = None,
) -> dict[str, dict[str, object]]:
    """Run the executable ablation matrix and return JSON-safe reports.

    Ablations that change architecture are instantiated explicitly. Shared
    parameter names are copied where shapes permit; the report records that
    status so an untrained architectural comparison cannot be mistaken for a
    trained-model result.
    """
    requested = names or tuple(item.name for item in ABLATIONS)
    reports: dict[str, dict[str, object]] = {}
    for name in requested:
        ablated, reason, trained = _ablate(model, name)
        if name in {"greedy", "without-verifier", "without-pat", "untyped"}:
            if name == "greedy":
                predictor = lambda graph, query, active=ablated: active.reason(
                    graph, query, beam_width=1
                )
            elif name == "without-verifier":
                predictor = lambda graph, query, active=ablated: active.reason(
                    graph, query, verify=False
                )
            elif name == "without-pat":
                predictor = lambda graph, query, active=ablated: active.reason(
                    graph, query, use_pat=False
                )
            else:
                predictor = lambda graph, query, active=ablated: active.reason(
                    graph, query, use_types=False
                )
        else:
            predictor = lambda graph, query, active=ablated: active.reason(graph, query)
        report = evaluate(examples, baseline=name, model=ablated, predictor=predictor)
        reports[name] = {
            "description": next(
                (item.description for item in ABLATIONS if item.name == name),
                "custom ablation",
            ),
            "trained_weights": trained,
            "construction": reason,
            "metrics": report.to_dict(),
        }
    return reports


def _ablate(model: GraphModel, name: str) -> tuple[GraphModel, str, bool]:
    if name in {
        "sage",
        "without-pat",
        "greedy",
        "without-verifier",
        "untyped",
        "random-negatives",
        "synthetic-only",
    }:
        return model, "shared model weights with runtime ablation", True
    if name == "gat":
        config = replace(model.config, graph_encoder="gat", use_gat=True)
        ablated = GraphModel(config)
        _copy_compatible(model, ablated)
        return (
            ablated,
            "GAT architecture instantiated; compatible weights copied",
            False,
        )
    if name in {"pat-without-query", "without-query-conditioning"}:
        pat = replace(model.config.pat, query_conditioning=False)
        config = replace(model.config, pat=pat)
        ablated = GraphModel(config)
        _copy_compatible(model, ablated)
        return ablated, "PAT query gate disabled; compatible weights copied", False
    if name == "without-path-supervision":
        loss = replace(model.config.loss, path=0.0, pat=0.0)
        config = replace(model.config, loss=loss)
        ablated = GraphModel(config)
        _copy_compatible(model, ablated)
        return (
            ablated,
            "path/PAT objective weights disabled; compatible weights copied",
            False,
        )
    raise ValueError(f"unknown ablation {name!r}")


def _copy_compatible(source: GraphModel, target: GraphModel) -> None:
    source_state = source.state_dict()
    target_state = target.state_dict()
    compatible = {
        key: value
        for key, value in source_state.items()
        if key in target_state and value.shape == target_state[key].shape
    }
    target_state.update(compatible)
    target.load_state_dict(target_state)