"""Training Loop -- Algebraic Routing with I-Logic State Composition

The routing table is NOT learned. It IS the I-Logic multiplication table
applied to the 19 agent definitions. Pure computation, deterministic,
computed once in __init__.

For each condition state C and agent state A:
    multiply(C, A) == DIRECT  →  that agent resolves that condition

The answer is known the instant you define the agents.
No episodes. No iteration. No for-loops.

ROUTING TABLE (algebraic, computed once):
    DIRECT      → DIRECT agents       (1 × 1 = 1)
    INVERSE     → INVERSE agents      (-1 × -1 = 1)
    SPIRAL      → ANTI_SPIRAL agents   (i × -i = 1)
    ANTI_SPIRAL → SPIRAL agents        (-i × i = 1)

USAGE:
    from cubesquare.training import TrainingLoop
    from cubesquare.swarm import Swarm

    swarm = Swarm()
    loop = TrainingLoop(swarm)

    # Routing table is already computed — no training step needed
    table = loop.get_routing_table()
    print(table)  # {INVERSE: [Healer, Validator, ...], SPIRAL: [...]}

    # Batch dispatch — one call, all results
    episode = loop.run_episode([
        "fix authentication bug",
        "create database schema",
        "find all test files",
    ])
    print(episode.resolution_rate, episode.total_reward)

The coordinate IS the routing. The algebra IS the policy.
"""

from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple
import math
import time

from .i_logic.states import ILogicState
from .i_logic.algebra import multiply, compose
from .swarm import Swarm, AgentRole, DispatchResult


# =========================================================================
# DATA STRUCTURES
# =========================================================================

@dataclass
class Transition:
    """A single state-action-reward tuple."""
    task: str
    task_state: complex          # I-Logic state of the task
    agent_id: int                # Which agent handled it
    agent_state: complex         # Agent's I-Logic state
    outcome: complex             # Final outcome (DIRECT if chain resolved)
    reward: float                # +1 DIRECT, -1 ABSORBED (costly path, energy conserved), 0 other
    resolved: bool               # True if resolved (single or chained)
    chained: bool = False        # True if resolved via multi-agent chain

    def to_dict(self) -> dict:
        d = {
            "task": self.task,
            "task_state": ILogicState.name(self.task_state),
            "agent_id": self.agent_id,
            "agent_state": ILogicState.name(self.agent_state),
            "outcome": ILogicState.name(self.outcome),
            "reward": self.reward,
            "resolved": self.resolved,
        }
        if self.chained:
            d["chained"] = True
        return d


@dataclass
class Episode:
    """A batch of routing decisions."""
    episode_id: int
    transitions: List[Transition] = field(default_factory=list)
    total_reward: float = 0.0
    resolution_rate: float = 0.0
    coherence: float = 0.0
    duration_ms: float = 0.0
    timestamp: float = field(default_factory=time.time)

    def to_dict(self) -> dict:
        return {
            "episode_id": self.episode_id,
            "transitions": len(self.transitions),
            "total_reward": round(self.total_reward, 2),
            "resolution_rate": round(self.resolution_rate, 3),
            "coherence": round(self.coherence, 3),
            "duration_ms": round(self.duration_ms, 2),
        }


@dataclass
class RoutingEntry:
    """Learned routing preference for a condition state."""
    condition_state: complex
    preferred_agents: List[Tuple[int, float]]  # (agent_id, success_rate)
    total_attempts: int = 0
    total_successes: int = 0

    @property
    def success_rate(self) -> float:
        return self.total_successes / max(self.total_attempts, 1)

    def to_dict(self) -> dict:
        return {
            "condition": ILogicState.name(self.condition_state),
            "preferred_agents": [
                {"agent_id": aid, "success_rate": round(sr, 3)}
                for aid, sr in self.preferred_agents
            ],
            "total_attempts": self.total_attempts,
            "success_rate": round(self.success_rate, 3),
        }


# =========================================================================
# REWARD FUNCTION
# =========================================================================

def compute_reward(outcome: complex) -> float:
    """Algebraic reward from I-Logic outcome.

    DIRECT (+1):       Task resolved. Maximum reward.
    INVERSE (+0.3):    Reflection state -- partial progress.
    SPIRAL (+0.1):     Expansion -- doing something, not resolved.
    ANTI_SPIRAL (0):   Contraction -- neutral.
    ABSORBED (-1):     Costly path — energy transforms at C=0 origin.
                       NOT a dead end. Enters Fernandez sphere (full identity),
                       hits L9 Membrane (K11 QRS): transmits to L10-L12 or
                       reflects back into L1-L9. H=T+V conserved.
                       Re-emergence carries stored potential (Hamiltonian bounce).

    The reward landscape mirrors the I-Logic algebra:
    Resolution (DIRECT) is the target.
    ABSORBED is the origin well — maximum undirected potential.
    """
    if outcome == ILogicState.DIRECT:
        return 1.0
    elif outcome == ILogicState.INVERSE:
        return 0.3
    elif outcome == ILogicState.SPIRAL:
        return 0.1
    elif outcome == ILogicState.ANTI_SPIRAL:
        return 0.0
    elif outcome == ILogicState.ABSORBED:
        return -1.0
    return 0.0


def compute_coherence(transitions: List[Transition]) -> float:
    """Measure how coherent a sequence of transitions is.

    Coherence = fraction of transitions that produced non-ABSORBED outcomes.
    ABSORBED transitions are origin-collapses (energy conserved via L9 membrane),
    not dead ends. They represent costly paths where potential energy stores
    at C=0 for re-emergence.
    """
    if not transitions:
        return 0.0
    alive = sum(1 for t in transitions if t.outcome != ILogicState.ABSORBED)
    return alive / len(transitions)


# =========================================================================
# TRAINING LOOP
# =========================================================================

class TrainingLoop:
    """Algebraic routing with I-Logic state composition.

    The routing table is COMPUTED, not learned.
    For each condition state, multiply by each agent's state.
    If the product is DIRECT, that agent resolves the condition.
    This is 4 × 19 = 76 multiplications. Done once. Forever correct.

    Usage:
        loop = TrainingLoop(Swarm())

        # Table is already computed — the algebra IS the answer
        table = loop.get_routing_table()
        # {INVERSE: [Healer, Validator, ...], SPIRAL: [...]}

        # Batch dispatch — one call
        episode = loop.run_episode(["fix auth", "create schema", "find tests"])
    """

    def __init__(self, swarm: Swarm):
        self._swarm = swarm
        self._episodes: List[Episode] = []
        self._routing_table: Dict[complex, RoutingEntry] = {}
        self._agent_stats: Dict[int, Dict[str, int]] = {}
        self._episode_counter = 0

        # ── ALGEBRAIC ROUTING TABLE ──────────────────────────────
        # Computed once from the I-Logic multiplication table.
        # For each condition C, find which agents A give C × A = DIRECT.
        # This is pure algebra — no iteration over experience needed.
        conditions = [ILogicState.DIRECT, ILogicState.INVERSE,
                      ILogicState.SPIRAL, ILogicState.ANTI_SPIRAL]
        agents = swarm.list_agents()

        self._routing_table = {
            cond: RoutingEntry(
                condition_state=cond,
                preferred_agents=[
                    (a.agent_id, 1.0)
                    for a in agents
                    if multiply(cond, a.i_logic_state) == ILogicState.DIRECT
                ],
                total_attempts=len(agents),
                total_successes=sum(
                    1 for a in agents
                    if multiply(cond, a.i_logic_state) == ILogicState.DIRECT
                ),
            )
            for cond in conditions
        }

    def run_episode(self, tasks: List[str]) -> Episode:
        """Batch-dispatch tasks and record transitions.

        One call → all tasks dispatched → all transitions recorded.
        The routing table is already computed algebraically.
        This method validates what the algebra predicted.

        Args:
            tasks: List of task description strings.

        Returns:
            Episode with transitions, reward, resolution rate, coherence.
        """
        self._episode_counter += 1
        start = time.perf_counter()

        # Batch dispatch — swarm handles the batch
        dispatch_results = self._swarm.dispatch_batch(tasks)

        # Map results to transitions (one comprehension, no explicit loop)
        transitions = [
            self._record_transition(task, result)
            for task, result in zip(tasks, dispatch_results)
        ]

        elapsed_ms = (time.perf_counter() - start) * 1000.0

        total_reward = sum(t.reward for t in transitions)
        resolved_count = sum(1 for t in transitions if t.resolved)
        resolution_rate = resolved_count / max(len(transitions), 1)
        coherence = compute_coherence(transitions)

        episode = Episode(
            episode_id=self._episode_counter,
            transitions=transitions,
            total_reward=total_reward,
            resolution_rate=resolution_rate,
            coherence=coherence,
            duration_ms=elapsed_ms,
        )
        self._episodes.append(episode)
        return episode

    def _record_transition(self, task: str, result: DispatchResult) -> Transition:
        """Record a single dispatch result as a transition. O(1)."""
        outcome = result.outcome
        chained = False
        if result.chain and result.chain.resolved:
            outcome = result.chain.final_state
            chained = True

        reward = compute_reward(outcome)

        transition = Transition(
            task=task,
            task_state=result.task_state,
            agent_id=result.agent.agent_id,
            agent_state=result.agent.i_logic_state,
            outcome=outcome,
            reward=reward,
            resolved=result.resolved,
            chained=chained,
        )

        # Update agent stats
        self._update_agent_stats(transition)
        return transition

        elapsed_ms = (time.perf_counter() - start) * 1000.0

        total_reward = sum(t.reward for t in transitions)
        resolved_count = sum(1 for t in transitions if t.resolved)
        resolution_rate = resolved_count / max(len(transitions), 1)
        coherence = compute_coherence(transitions)

        episode = Episode(
            episode_id=self._episode_counter,
            transitions=transitions,
            total_reward=total_reward,
            resolution_rate=resolution_rate,
            coherence=coherence,
            duration_ms=elapsed_ms,
        )
        self._episodes.append(episode)
        return episode

    def _update_agent_stats(self, transition: Transition):
        """Track per-agent performance stats."""
        aid = transition.agent_id
        if aid not in self._agent_stats:
            self._agent_stats[aid] = {"dispatched": 0, "resolved": 0, "absorbed": 0}
        stats = self._agent_stats[aid]
        stats["dispatched"] += 1
        if transition.resolved:
            stats["resolved"] += 1
        if transition.outcome == ILogicState.ABSORBED:
            stats["absorbed"] += 1

    # -----------------------------------------------------------------
    # QUERY API
    # -----------------------------------------------------------------

    def get_routing_table(self) -> Dict[str, dict]:
        """Get the learned routing table.

        Returns a dict mapping condition state names to routing entries
        with preferred agents and success rates.
        """
        return {
            ILogicState.name(state): entry.to_dict()
            for state, entry in self._routing_table.items()
        }

    def get_agent_stats(self) -> Dict[int, dict]:
        """Get per-agent performance statistics."""
        result = {}
        for aid, stats in self._agent_stats.items():
            agent = self._swarm.get_agent(aid)
            name = agent.name if agent else f"Agent-{aid}"
            dispatched = stats["dispatched"]
            resolved = stats["resolved"]
            result[aid] = {
                "name": name,
                "dispatched": dispatched,
                "resolved": resolved,
                "absorbed": stats["absorbed"],
                "resolution_rate": round(resolved / max(dispatched, 1), 3),
            }
        return result

    def get_episode_history(self) -> List[dict]:
        """Get summary of all episodes."""
        return [ep.to_dict() for ep in self._episodes]

    def get_best_resolver(self, condition_state: complex) -> Optional[int]:
        """Get the best agent ID for a condition based on learned experience.

        Falls back to algebraic prediction if no experience exists.
        """
        entry = self._routing_table.get(condition_state)
        if entry and entry.preferred_agents:
            # Return the agent with highest success rate
            return entry.preferred_agents[0][0]

        # Fallback: algebraic prediction
        resolver = self._swarm.find_resolver(condition_state)
        if resolver:
            return resolver.agent_id
        return None

    @property
    def total_episodes(self) -> int:
        return len(self._episodes)

    @property
    def total_transitions(self) -> int:
        return sum(len(ep.transitions) for ep in self._episodes)

    @property
    def overall_resolution_rate(self) -> float:
        total = self.total_transitions
        if total == 0:
            return 0.0
        resolved = sum(
            sum(1 for t in ep.transitions if t.resolved) for ep in self._episodes
        )
        return resolved / total

    @property
    def overall_coherence(self) -> float:
        if not self._episodes:
            return 0.0
        return sum(ep.coherence for ep in self._episodes) / len(self._episodes)

    def summary(self) -> dict:
        """Full training summary."""
        chained = sum(
            1 for ep in self._episodes
            for t in ep.transitions if t.chained
        )
        direct = sum(
            1 for ep in self._episodes
            for t in ep.transitions if t.resolved and not t.chained
        )
        return {
            "total_episodes": self.total_episodes,
            "total_transitions": self.total_transitions,
            "overall_resolution_rate": round(self.overall_resolution_rate, 3),
            "overall_coherence": round(self.overall_coherence, 3),
            "direct_resolutions": direct,
            "chained_resolutions": chained,
            "routing_table": self.get_routing_table(),
            "agent_stats": self.get_agent_stats(),
        }
