from __future__ import annotations

from dataclasses import dataclass
import random

from .architectures import PolicyDecision
from .art1 import ART1Network, ART1Params, ART1Result
from .backprop import BackpropNetwork, BackpropResult
from .orchestration import CooperativeSystem


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


@dataclass(frozen=True)
class XorDemoState:
    accepted: tuple[BinaryPair, ...] = ()
    attempts: int = 0


class CyclicBinaryGenerator:
    _patterns: tuple[BinaryPair, ...] = (
        (0, 0),
        (0, 1),
        (1, 0),
        (1, 1),
    )

    def generate(self, state: XorDemoState) -> BinaryPair:
        return self._patterns[state.attempts % len(self._patterns)]


class BackpropXorCritic:
    def __init__(self, network: BackpropNetwork) -> None:
        self.network = network

    def critique(self, state: XorDemoState, candidate: BinaryPair) -> BackpropResult:
        return self.network.predict(tuple(float(value) for value in candidate))


class BinaryPairCategorizer:
    def __init__(self, network: ART1Network) -> None:
        self.network = network

    def categorize(self, state: XorDemoState, candidate: BinaryPair) -> ART1Result:
        return self.network.categorize(candidate)


class XorNoveltyPolicy:
    def __init__(self, *, acceptance_threshold: float = 0.8) -> None:
        self.acceptance_threshold = acceptance_threshold

    def decide(
        self,
        state: XorDemoState,
        candidate: BinaryPair,
        critique: BackpropResult,
        category: ART1Result,
    ) -> PolicyDecision:
        score = critique.outputs[0]
        accepted = (
            score >= self.acceptance_threshold
            and category.matched
            and not category.delta_vigilance
        )
        return PolicyDecision(
            accepted=accepted,
            label="accept" if accepted else "reject",
        )


class XorDemoTransition:
    def advance(
        self,
        state: XorDemoState,
        candidate: BinaryPair,
        critique: BackpropResult,
        category: ART1Result,
        decision: PolicyDecision,
    ) -> XorDemoState:
        if decision.accepted:
            return XorDemoState(
                accepted=state.accepted + (candidate,),
                attempts=state.attempts + 1,
            )
        return XorDemoState(
            accepted=state.accepted,
            attempts=state.attempts + 1,
        )


def build_xor_novelty_demo(
    *,
    rng_seed: int = 11,
    acceptance_threshold: float = 0.8,
    art_vigilance: float = 0.9,
) -> CooperativeSystem[XorDemoState, BinaryPair, BackpropResult, ART1Result]:
    network = BackpropNetwork.random(
        input_size=2,
        hidden_layers=(4, 4),
        output_size=1,
        learning_rate=0.8,
        momentum=0.2,
        rng=random.Random(rng_seed),
    )
    samples = (
        ((0.0, 0.0), (0.0,)),
        ((0.0, 1.0), (1.0,)),
        ((1.0, 0.0), (1.0,)),
        ((1.0, 1.0), (0.0,)),
    )
    for _ in range(6000):
        for inputs, targets in samples:
            network.train_step(inputs, targets)

    art = ART1Network(
        params=ART1Params(
            max_categories=4,
            input_length=2,
            vigilance=art_vigilance,
        )
    )

    return CooperativeSystem(
        generator=CyclicBinaryGenerator(),
        critic=BackpropXorCritic(network),
        categorizer=BinaryPairCategorizer(art),
        policy=XorNoveltyPolicy(acceptance_threshold=acceptance_threshold),
        transition=XorDemoTransition(),
    )


@dataclass(frozen=True)
class ParityPressureState:
    accepted: tuple[BinaryTriple, ...] = ()
    attempts: int = 0


class CyclicParityGenerator:
    _patterns: tuple[BinaryTriple, ...] = (
        (0, 0, 0),
        (0, 0, 1),
        (0, 1, 0),
        (0, 1, 1),
        (1, 0, 0),
        (1, 0, 1),
        (1, 1, 0),
        (1, 1, 1),
    )

    def generate(self, state: ParityPressureState) -> BinaryTriple:
        return self._patterns[state.attempts % len(self._patterns)]


class BackpropParityCritic:
    def __init__(self, network: BackpropNetwork) -> None:
        self.network = network

    def critique(self, state: ParityPressureState, candidate: BinaryTriple) -> BackpropResult:
        return self.network.predict(tuple(float(value) for value in candidate))


class BinaryTripleCategorizer:
    def __init__(self, network: ART1Network) -> None:
        self.network = network

    def categorize(self, state: ParityPressureState, candidate: BinaryTriple) -> ART1Result:
        encoded = [0] * 8
        encoded[(candidate[0] * 4) + (candidate[1] * 2) + candidate[2]] = 1
        return self.network.categorize(encoded)


class ParityPressurePolicy:
    def __init__(self, *, acceptance_threshold: float = 0.8) -> None:
        self.acceptance_threshold = acceptance_threshold

    def decide(
        self,
        state: ParityPressureState,
        candidate: BinaryTriple,
        critique: BackpropResult,
        category: ART1Result,
    ) -> PolicyDecision:
        score = critique.outputs[0]
        accepted = (
            score >= self.acceptance_threshold
            and category.matched
            and not category.delta_vigilance
        )
        return PolicyDecision(
            accepted=accepted,
            label="accept" if accepted else "reject",
        )


class ParityPressureTransition:
    def advance(
        self,
        state: ParityPressureState,
        candidate: BinaryTriple,
        critique: BackpropResult,
        category: ART1Result,
        decision: PolicyDecision,
    ) -> ParityPressureState:
        if decision.accepted:
            return ParityPressureState(
                accepted=state.accepted + (candidate,),
                attempts=state.attempts + 1,
            )
        return ParityPressureState(
            accepted=state.accepted,
            attempts=state.attempts + 1,
        )


def build_parity_pressure_demo(
    *,
    rng_seed: int = 21,
    acceptance_threshold: float = 0.8,
    art_vigilance: float = 0.95,
    art_max_categories: int = 2,
) -> CooperativeSystem[ParityPressureState, BinaryTriple, BackpropResult, ART1Result]:
    network = BackpropNetwork.random(
        input_size=3,
        hidden_layers=(6, 4),
        output_size=1,
        learning_rate=0.8,
        momentum=0.2,
        rng=random.Random(rng_seed),
    )
    samples = tuple(
        (
            tuple(float(bit) for bit in bits),
            (float(sum(bits) % 2),),
        )
        for bits in CyclicParityGenerator._patterns
    )
    for _ in range(8000):
        for inputs, targets in samples:
            network.train_step(inputs, targets)

    art = ART1Network(
        params=ART1Params(
            max_categories=art_max_categories,
            input_length=8,
            vigilance=art_vigilance,
        )
    )

    return CooperativeSystem(
        generator=CyclicParityGenerator(),
        critic=BackpropParityCritic(network),
        categorizer=BinaryTripleCategorizer(art),
        policy=ParityPressurePolicy(acceptance_threshold=acceptance_threshold),
        transition=ParityPressureTransition(),
    )
