"""Deterministic and neural baseline adapters used by the evaluator."""

from __future__ import annotations

import math
from collections.abc import Callable, Mapping
from dataclasses import dataclass, replace
from itertools import islice

import torch

from graph.graph import GraphStore
from graph.model import GraphModel
from graph.types import (
    EntityAnswer,
    GraphPath,
    GraphQuery,
    PathStep,
    ReasoningResult,
    ReasoningStatus,
)
from graph.verification import detect_conflicting_paths, verify_path

#: Every baseline is a function from a graph and a query to a typed result.
Predictor = Callable[[GraphStore, GraphQuery], ReasoningResult]


class BaselineUnavailableError(RuntimeError):
    """Raised when a named upstream baseline has no supplied checkpoint/adapter."""


#: Paths enumerated from each start entity.  A cap on the whole search lets the
#: first start's dense neighbourhood starve every later start of paths.
PATHS_PER_START = 8192


def _bounded_paths(graph: GraphStore, starts, max_hops: int):
    for start in starts:
        yield from islice(graph.iter_paths((start,), max_hops), PATHS_PER_START)


def heuristic_graph_search(graph: GraphStore, query: GraphQuery) -> ReasoningResult:
    paths = tuple(
        path
        for path in _bounded_paths(
            graph, query.start_entity_ids or query.start_entities, query.max_hops
        )
        if query.min_hops <= path.hop_count <= query.max_hops
    )
    paths = tuple(verify_path(path, graph, query) for path in paths)
    valid = tuple(path for path in paths if path.verified)
    answers = tuple(
        EntityAnswer(endpoint, 1.0)
        for endpoint in dict.fromkeys(
            path.endpoint_id for path in valid if path.endpoint_id
        )
    )
    relation_answers = tuple(
        dict.fromkeys(path.steps[-1].relation_id for path in valid if path.steps)
    )
    boolean_answer = bool(answers) if query.answer_mode == "boolean" else None
    aggregate_value = (
        len(answers) if query.answer_mode in {"count", "aggregate"} else None
    )
    entity_answers = answers if query.answer_mode == "entity" else ()
    conflicts = tuple(
        path for group in detect_conflicting_paths(valid, graph) for path in group
    )
    status = (
        ReasoningStatus.SUPPORTED
        if valid
        else (
            ReasoningStatus.UNSUPPORTED
            if query.target_entity_ids
            else ReasoningStatus.INSUFFICIENT_EVIDENCE
        )
    )
    if conflicts and query.operation in {"verify", "boolean", "is_supported"}:
        status = ReasoningStatus.CONTRADICTED
    if (
        query.target_entity_ids
        and valid
        and not (set(query.target_entity_ids) & {path.endpoint_id for path in valid})
    ):
        status = ReasoningStatus.UNSUPPORTED
    if query.answer_mode == "boolean" and status != ReasoningStatus.SUPPORTED:
        boolean_answer = (
            False
            if status in {ReasoningStatus.UNSUPPORTED, ReasoningStatus.CONTRADICTED}
            else None
        )
    entity_answers = (
        answers
        if query.answer_mode == "entity" and status == ReasoningStatus.SUPPORTED
        else ()
    )
    return ReasoningResult(
        status=status,
        answers=entity_answers,
        supporting_paths=valid
        if query.return_paths and status == ReasoningStatus.SUPPORTED
        else (),
        conflicting_paths=conflicts if query.return_paths else (),
        confidence=1.0 if valid else 0.0,
        calibrated=False,
        model_version="heuristic-graph-search",
        graph_version=graph.graph_version,
        reasoning_steps=max((path.hop_count for path in valid), default=0),
        relation_answers=relation_answers if query.answer_mode == "relation" else (),
        boolean_answer=boolean_answer,
        aggregate_value=aggregate_value,
    )


def adamic_adar_search(graph: GraphStore, query: GraphQuery) -> ReasoningResult:
    """Use local common-neighbor evidence to rank otherwise valid graph paths."""
    starts = query.start_entity_ids or query.start_entities
    neighbor_sets = {
        node: {edge.target_id for edge in graph.outgoing(node)} for node in graph.nodes
    }
    paths = tuple(
        path
        for path in _bounded_paths(graph, starts, query.max_hops)
        if query.min_hops <= path.hop_count <= query.max_hops
    )
    scored: list[GraphPath] = []
    for path in paths:
        checked = verify_path(path, graph, query)
        if not checked.verified:
            continue
        score = 0.0
        if path.start_id and path.endpoint_id:
            common = neighbor_sets.get(path.start_id, set()) & neighbor_sets.get(
                path.endpoint_id, set()
            )
            score = sum(
                1.0 / max(1.0, math.log1p(len(neighbor_sets.get(node, ()))))
                for node in common
            )
        scored.append(replace(checked, score=score))
    scored.sort(key=lambda path: path.score, reverse=True)
    answers = tuple(
        EntityAnswer(endpoint, path.score)
        for endpoint, path in (
            (path.endpoint_id, path) for path in scored if path.endpoint_id
        )
    )
    relation_answers = tuple(
        dict.fromkeys(path.steps[-1].relation_id for path in scored if path.steps)
    )
    boolean_answer = bool(answers) if query.answer_mode == "boolean" else None
    aggregate_value = (
        len(answers) if query.answer_mode in {"count", "aggregate"} else None
    )
    conflicts = tuple(
        path
        for group in detect_conflicting_paths(tuple(scored), graph)
        for path in group
    )
    status = (
        ReasoningStatus.SUPPORTED
        if scored
        else (
            ReasoningStatus.UNSUPPORTED
            if query.target_entity_ids
            else ReasoningStatus.INSUFFICIENT_EVIDENCE
        )
    )
    if conflicts and query.operation in {"verify", "boolean", "is_supported"}:
        status = ReasoningStatus.CONTRADICTED
    if (
        query.target_entity_ids
        and scored
        and not (set(query.target_entity_ids) & {path.endpoint_id for path in scored})
    ):
        status = ReasoningStatus.UNSUPPORTED
    if query.answer_mode == "boolean" and status != ReasoningStatus.SUPPORTED:
        boolean_answer = (
            False
            if status in {ReasoningStatus.UNSUPPORTED, ReasoningStatus.CONTRADICTED}
            else None
        )
    entity_answers = (
        answers
        if query.answer_mode == "entity" and status == ReasoningStatus.SUPPORTED
        else ()
    )
    return ReasoningResult(
        status=status,
        answers=entity_answers,
        supporting_paths=tuple(scored)
        if query.return_paths and status == ReasoningStatus.SUPPORTED
        else (),
        conflicting_paths=conflicts if query.return_paths else (),
        confidence=1.0 if scored else 0.0,
        calibrated=False,
        model_version="adamic-adar",
        graph_version=graph.graph_version,
        reasoning_steps=max((path.hop_count for path in scored), default=0),
        relation_answers=relation_answers if query.answer_mode == "relation" else (),
        boolean_answer=boolean_answer,
        aggregate_value=aggregate_value,
    )


def neural_link_predictor(
    graph: GraphStore, query: GraphQuery, model: GraphModel, *, gat: bool = False
) -> ReasoningResult:
    """One-hop GraphSAGE/GAT link-prediction baseline, separate from GRM search."""
    expected_encoder = "gat" if gat else "sage"
    actual_encoder = (
        "gat" if model.config.graph_encoder == "gat" or model.config.use_gat else "sage"
    )
    if actual_encoder != expected_encoder:
        raise BaselineUnavailableError(
            f"{expected_encoder} baseline requires a model configured with graph_encoder={expected_encoder!r}; got {actual_encoder!r}"
        )
    example = model.graph_to_example(graph)
    # ``graph_to_example`` builds CPU tensors; the encoder lives wherever the
    # model was placed.  Without this move the baseline works on CPU and raises a
    # device mismatch on GPU, which is how it was actually found.
    device = next(model.parameters()).device
    encoded = model.graph_encoder(example.graph.to(device))
    starts = set(query.start_entity_ids or query.start_entities)
    candidates: list[GraphPath] = []
    for edge in graph.edges:
        if edge.source_id not in starts:
            continue
        source_index = example.node_external_ids.index(edge.source_id)
        target_index = example.node_external_ids.index(edge.target_id)
        score = float(
            (encoded.node_states[source_index] * encoded.node_states[target_index])
            .sum()
            .detach()
            .cpu()
        )
        candidate = GraphPath(
            (
                PathStep(
                    edge.source_id,
                    edge.relation_id,
                    edge.target_id,
                    edge.source_type,
                    edge.target_type,
                    edge.attributes,
                ),
            ),
            score=score,
        )
        checked = verify_path(candidate, graph, query)
        if checked.verified:
            candidates.append(checked)
    candidates.sort(key=lambda path: path.score, reverse=True)
    answers = tuple(
        EntityAnswer(path.endpoint_id, path.score)
        for path in candidates
        if path.endpoint_id
    )
    return ReasoningResult(
        status=ReasoningStatus.SUPPORTED
        if candidates
        else ReasoningStatus.INSUFFICIENT_EVIDENCE,
        answers=answers if query.answer_mode == "entity" else (),
        supporting_paths=tuple(candidates[: model.config.answer_top_k])
        if query.return_paths
        else (),
        conflicting_paths=(),
        confidence=float(
            torch.sigmoid(
                torch.tensor(candidates[0].score if candidates else -4.0)
            ).item()
        ),
        calibrated=False,
        model_version=f"{expected_encoder}-link-predictor",
        graph_version=graph.graph_version,
        reasoning_steps=1 if candidates else 0,
    )


@dataclass(frozen=True)
class ReasonOverrides:
    """Typed keyword overrides for :meth:`GraphModel.reason`.

    A typed record rather than a ``dict[str, object]``: the untyped mapping could
    not be splatted into ``reason`` without defeating the type checker, and the
    core pipeline is meant to avoid ``Any`` (PLAN section 51).
    """

    use_pat: bool | None = None
    verify: bool | None = None
    select: bool = True
    query_conditioned: bool = True
    use_types: bool = True
    beam_width: int | None = None

    def describe(self) -> dict[str, str]:
        """Human-readable summary of the fields this ablation changes."""
        default = ReasonOverrides()
        return {
            field: str(getattr(self, field))
            for field in (
                "use_pat",
                "verify",
                "select",
                "query_conditioned",
                "use_types",
                "beam_width",
            )
            if getattr(self, field) != getattr(default, field)
        }


#: Runtime ablations change how one fixed set of weights is used.  Ablations that
#: require *different weights* (no path supervision, no PAT supervision) are not
#: listed here: they are produced by training a variant and are dispatched
#: through an explicit predictor instead.  Returning an unmodified
#: ``model.reason`` under an ablation name, as the previous implementation did
#: for ``graph-without-query-conditioning`` and ``graph-without-path-supervision``,
#: reports a comparison that never ran.
RUNTIME_ABLATIONS: Mapping[str, ReasonOverrides] = {
    "graph": ReasonOverrides(),
    "graph-without-pat": ReasonOverrides(use_pat=False),
    "graph-without-verifier": ReasonOverrides(verify=False),
    "graph-without-selection": ReasonOverrides(select=False),
    "graph-without-query-conditioning": ReasonOverrides(query_conditioned=False),
    "graph-untyped": ReasonOverrides(use_types=False),
    "graph-greedy": ReasonOverrides(beam_width=1),
}

#: Ablations that can only be evaluated against separately trained weights.
TRAINING_ABLATIONS: frozenset[str] = frozenset(
    {"graph-without-path-supervision", "graph-without-pat-supervision"}
)

#: The model was renamed from "grm" to "graph".  Evaluation artifacts produced before
#: the rename key their rows by the old names, so those names keep resolving rather
#: than becoming "unknown baseline" errors when an old report is re-read or compared.
DEPRECATED_BASELINE_NAMES: Mapping[str, str] = {
    "grm": "graph",
    "grm-3b": "graph",
    "grm-full": "graph",
    **{
        f"grm-{suffix}": f"graph-{suffix}"
        for suffix in (
            "without-pat",
            "without-verifier",
            "without-selection",
            "without-query-conditioning",
            "untyped",
            "greedy",
            "without-path-supervision",
            "without-pat-supervision",
        )
    },
}


def apply_overrides(
    model: GraphModel, graph: GraphStore, query: GraphQuery, overrides: ReasonOverrides
) -> ReasoningResult:
    """Call ``reason`` with an ablation's overrides, fully typed."""
    return model.reason(
        graph,
        query,
        use_pat=overrides.use_pat,
        verify=overrides.verify,
        beam_width=overrides.beam_width,
        use_types=overrides.use_types,
        query_conditioned=overrides.query_conditioned,
        select=overrides.select,
    )


def make_baseline(
    name: str,
    model: GraphModel | None = None,
    upstream: Predictor | None = None,
):
    """Resolve a baseline name to a predictor, or fail loudly.

    ``upstream`` supplies the released G-Reasoner-34M predictor.  It is injected
    rather than constructed here because it needs the downloaded checkpoint, a text
    embedding model and ``torch_geometric``, none of which should become an import
    dependency of the evaluation package.
    """
    if name == "heuristic":
        return heuristic_graph_search
    if name == "adamic_adar":
        return adamic_adar_search
    if name == "g-reasoner-34m":
        if upstream is None:
            raise BaselineUnavailableError(
                "g-reasoner-34m requires the released upstream checkpoint and a text "
                "embedder; build one with graph.evaluation.upstream_baseline and pass "
                "it as upstream="
            )
        return upstream
    if name == "g-reasoner-34m-finetuned":
        raise BaselineUnavailableError(
            "g-reasoner-34m-finetuned requires fine-tuning the released checkpoint on "
            "this task; the released weights alone are the zero-shot baseline"
        )
    if name in TRAINING_ABLATIONS:
        raise BaselineUnavailableError(
            f"{name} is a training-time ablation: evaluate it against weights trained "
            f"with that objective disabled, then pass them as an explicit predictor"
        )
    if name in {"graphsage", "gat"}:
        if model is None:
            raise ValueError(f"baseline {name} requires a configured model")
        return lambda graph, query: neural_link_predictor(
            graph, query, model, gat=name == "gat"
        )
    if model is None:
        raise ValueError(f"baseline {name} requires a GRM model adapter")
    # Pre-rename names stay readable so old artifacts do not become errors.
    key = DEPRECATED_BASELINE_NAMES.get(name, name)
    if key in RUNTIME_ABLATIONS:
        overrides = RUNTIME_ABLATIONS[key]
        return lambda graph, query: apply_overrides(model, graph, query, overrides)
    raise ValueError(f"unknown baseline {name!r}; known: {sorted(RUNTIME_ABLATIONS)}")