from __future__ import annotations

from dataclasses import dataclass
from typing import Callable, Generic, TypeVar

from .architectures import PolicyDecision
from .art1 import ART1Network, ART1Result
from .backprop import BackpropLayerState, BackpropNetwork, BackpropResult
from .examples import (
    BackpropParityCritic,
    BackpropXorCritic,
    BinaryPairCategorizer,
    BinaryTripleCategorizer,
    CyclicBinaryGenerator,
    CyclicParityGenerator,
    ParityPressurePolicy,
    ParityPressureState,
    ParityPressureTransition,
    XorDemoState,
    XorDemoTransition,
    XorNoveltyPolicy,
    build_parity_pressure_demo,
    build_xor_novelty_demo,
)
from .orchestration import CooperativeSystem, HybridStepMetadata
from .serialization import to_jsonable


StateT = TypeVar("StateT")
CandidateT = TypeVar("CandidateT")
MetadataT = TypeVar("MetadataT")


BinaryPair = tuple[int, int]
BinaryTriple = tuple[int, int, int]


@dataclass(frozen=True)
class DemoDefinition(Generic[StateT, CandidateT, MetadataT]):
    name: str
    build_system: Callable[[], object]
    initial_state: StateT
    alphabet_size: int
    sequence_getter: Callable[[object], list[int]]
    state_encoder: Callable[[StateT], object]
    state_decoder: Callable[[object], StateT]
    candidate_decoder: Callable[[object], CandidateT]
    metadata_decoder: Callable[[object | None], MetadataT | None]
    system_encoder: Callable[[object], dict[str, object]]
    system_decoder: Callable[[dict[str, object]], object]


def _binary_pair_sequence(record: object) -> list[int]:
    return [left * 2 + right for left, right in record.final_state.accepted]  # type: ignore[attr-defined]


def _binary_triple_sequence(record: object) -> list[int]:
    return [
        left * 4 + middle * 2 + right
        for left, middle, right in record.final_state.accepted  # type: ignore[attr-defined]
    ]


def _decode_xor_state(data: object) -> XorDemoState:
    payload = data  # type: ignore[assignment]
    return XorDemoState(
        accepted=tuple(tuple(int(value) for value in pair) for pair in payload["accepted"]),  # type: ignore[index]
        attempts=int(payload["attempts"]),  # type: ignore[index]
    )


def _decode_parity_state(data: object) -> ParityPressureState:
    payload = data  # type: ignore[assignment]
    return ParityPressureState(
        accepted=tuple(
            tuple(int(value) for value in triple)
            for triple in payload["accepted"]  # type: ignore[index]
        ),
        attempts=int(payload["attempts"]),  # type: ignore[index]
    )


def _decode_binary_pair(data: object) -> BinaryPair:
    left, right = data  # type: ignore[misc]
    return (int(left), int(right))


def _decode_binary_triple(data: object) -> BinaryTriple:
    left, middle, right = data  # type: ignore[misc]
    return (int(left), int(middle), int(right))


def _decode_backprop_result(data: object) -> BackpropResult:
    payload = data  # type: ignore[assignment]
    return BackpropResult(
        outputs=tuple(float(value) for value in payload["outputs"]),  # type: ignore[index]
        loss=float(payload["loss"]),  # type: ignore[index]
        layer_states=tuple(
            BackpropLayerState(
                activations=tuple(float(value) for value in layer_state["activations"]),
                deltas=tuple(float(value) for value in layer_state["deltas"]),
                biases=tuple(float(value) for value in layer_state["biases"]),
            )
            for layer_state in payload["layer_states"]  # type: ignore[index]
        ),
    )


def _decode_art1_result(data: object) -> ART1Result:
    payload = data  # type: ignore[assignment]
    return ART1Result(
        winner=int(payload["winner"]),  # type: ignore[index]
        matched=bool(payload["matched"]),  # type: ignore[index]
        new_category=bool(payload["new_category"]),  # type: ignore[index]
        delta_vigilance=bool(payload["delta_vigilance"]),  # type: ignore[index]
        committed_categories=int(payload["committed_categories"]),  # type: ignore[index]
        vigilance=float(payload["vigilance"]),  # type: ignore[index]
        expected_vector=tuple(int(value) for value in payload["expected_vector"]),  # type: ignore[index]
    )


def _decode_policy_decision(data: object) -> PolicyDecision:
    payload = data  # type: ignore[assignment]
    return PolicyDecision(
        accepted=bool(payload["accepted"]),  # type: ignore[index]
        label=str(payload["label"]),  # type: ignore[index]
    )


def _decode_hybrid_metadata(data: object | None) -> HybridStepMetadata[BackpropResult, ART1Result] | None:
    if data is None:
        return None
    payload = data  # type: ignore[assignment]
    return HybridStepMetadata(
        critique=_decode_backprop_result(payload["critique"]),  # type: ignore[index]
        category=_decode_art1_result(payload["category"]),  # type: ignore[index]
        decision=_decode_policy_decision(payload["decision"]),  # type: ignore[index]
    )


def _encode_xor_system(system: object) -> dict[str, object]:
    current = system  # type: ignore[assignment]
    return {
        "critic_network": current.critic.network.to_dict(),
        "categorizer_network": current.categorizer.network.to_dict(),
        "acceptance_threshold": current.policy.acceptance_threshold,
    }


def _decode_xor_system(data: dict[str, object]) -> CooperativeSystem[
    XorDemoState,
    BinaryPair,
    BackpropResult,
    ART1Result,
]:
    return CooperativeSystem(
        generator=CyclicBinaryGenerator(),
        critic=BackpropXorCritic(BackpropNetwork.from_dict(data["critic_network"])),  # type: ignore[arg-type]
        categorizer=BinaryPairCategorizer(ART1Network.from_dict(data["categorizer_network"])),  # type: ignore[arg-type]
        policy=XorNoveltyPolicy(acceptance_threshold=float(data["acceptance_threshold"])),
        transition=XorDemoTransition(),
    )


def _encode_parity_system(system: object) -> dict[str, object]:
    current = system  # type: ignore[assignment]
    return {
        "critic_network": current.critic.network.to_dict(),
        "categorizer_network": current.categorizer.network.to_dict(),
        "acceptance_threshold": current.policy.acceptance_threshold,
    }


def _decode_parity_system(data: dict[str, object]) -> CooperativeSystem[
    ParityPressureState,
    BinaryTriple,
    BackpropResult,
    ART1Result,
]:
    return CooperativeSystem(
        generator=CyclicParityGenerator(),
        critic=BackpropParityCritic(BackpropNetwork.from_dict(data["critic_network"])),  # type: ignore[arg-type]
        categorizer=BinaryTripleCategorizer(ART1Network.from_dict(data["categorizer_network"])),  # type: ignore[arg-type]
        policy=ParityPressurePolicy(acceptance_threshold=float(data["acceptance_threshold"])),
        transition=ParityPressureTransition(),
    )


DEMO_DEFINITIONS: dict[str, DemoDefinition[object, object, object]] = {
    "xor_novelty": DemoDefinition(
        name="xor_novelty",
        build_system=build_xor_novelty_demo,
        initial_state=XorDemoState(),
        alphabet_size=4,
        sequence_getter=_binary_pair_sequence,
        state_encoder=to_jsonable,
        state_decoder=_decode_xor_state,
        candidate_decoder=_decode_binary_pair,
        metadata_decoder=_decode_hybrid_metadata,
        system_encoder=_encode_xor_system,
        system_decoder=_decode_xor_system,
    ),
    "parity_pressure": DemoDefinition(
        name="parity_pressure",
        build_system=build_parity_pressure_demo,
        initial_state=ParityPressureState(),
        alphabet_size=8,
        sequence_getter=_binary_triple_sequence,
        state_encoder=to_jsonable,
        state_decoder=_decode_parity_state,
        candidate_decoder=_decode_binary_triple,
        metadata_decoder=_decode_hybrid_metadata,
        system_encoder=_encode_parity_system,
        system_decoder=_decode_parity_system,
    ),
}


def available_demo_names() -> tuple[str, ...]:
    return tuple(sorted(DEMO_DEFINITIONS))


def get_demo_definition(name: str) -> DemoDefinition[object, object, object]:
    try:
        return DEMO_DEFINITIONS[name]
    except KeyError as exc:
        raise ValueError(f"unknown demo_name {name!r}") from exc
