"""Learning Loop -- Reinforcement Learning with I-Logic State Transitions

The state space is deterministic (basin, layer, I-Logic state).
The action space is I-Logic multiplication (5 possible transitions).
The reward is algebraic: DIRECT = +1 (resolution), ABSORBED = -1 (dead end).

This is NOT neural RL. The I-Logic multiplication table IS the policy.
What the learning loop discovers is:
    1. Which agents handle which basins best (specialization)
    2. Which agent sequences resolve problems (compose paths)
    3. Which state transitions to prefer (routing table)

The routing table is built from EXPERIENCE, not training:
    - Run episodes (batches of routing decisions)
    - Track state-action-reward tuples
    - Update routing preferences
    - Measure resolution rate and coherence

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

    swarm = Swarm()
    loop = TrainingLoop(swarm)

    # Run a training episode
    episode = loop.run_episode([
        "fix authentication bug",
        "create database schema",
        "find all test files",
    ])
    print(episode.resolution_rate, episode.total_reward)

    # Get learned routing table
    table = loop.get_routing_table()
    print(table)  # {INVERSE: [Healer, Validator], SPIRAL: [Generator, ...]}

NO PYTHON LOOPS in the I-Logic computation.
Episodes iterate tasks (bounded, not hot path).
All I-Logic ops are O(1).
"""

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, 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):     Dead end. Negative reward.

    The reward landscape mirrors the I-Logic algebra:
    Resolution (DIRECT) is the target.
    ABSORBED (the absorbing element) is the trap.
    """
    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.
    A fully coherent episode has no dead ends.
    """
    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:
    """Reinforcement learning loop with I-Logic state transitions.

    The loop runs EPISODES -- batches of task routing decisions.
    Each episode:
        1. Takes a list of task descriptions
        2. Dispatches each to the swarm
        3. Records the transition (state, action, reward)
        4. Updates the routing table
        5. Computes episode metrics

    The routing table is a learned mapping:
        condition_state -> [(agent_id, success_rate), ...]

    Over many episodes, the table converges to the optimal routing:
    - INVERSE conditions -> Healer/Validator agents
    - SPIRAL conditions -> Generator/Coordinator agents
    - etc.

    Usage:
        loop = TrainingLoop(Swarm())
        for batch in task_batches:
            episode = loop.run_episode(batch)
            print(f"Episode {episode.episode_id}: "
                  f"reward={episode.total_reward}, "
                  f"resolution={episode.resolution_rate:.1%}")
    """

    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

        # Initialize routing table for all non-trivial states
        for state in [ILogicState.DIRECT, ILogicState.INVERSE,
                      ILogicState.SPIRAL, ILogicState.ANTI_SPIRAL]:
            self._routing_table[state] = RoutingEntry(
                condition_state=state,
                preferred_agents=[],
            )

    def run_episode(self, tasks: List[str]) -> Episode:
        """Run one training episode with a batch of tasks.

        Each task is dispatched to the swarm, the outcome is recorded,
        and the routing table is updated.

        Args:
            tasks: List of task description strings.

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

        transitions = []
        for task in tasks:
            result = self._swarm.dispatch(task)
            outcome = result.outcome
            chained = False

            # If resolved via chain, use chain's final state for reward
            if result.chain and result.chain.resolved:
                outcome = result.chain.final_state  # DIRECT
                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,
            )
            transitions.append(transition)

            # Update routing table
            self._update_routing(transition)

            # Update agent stats
            self._update_agent_stats(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_routing(self, transition: Transition):
        """Update routing table from a transition."""
        state = transition.task_state
        if state not in self._routing_table:
            self._routing_table[state] = RoutingEntry(
                condition_state=state, preferred_agents=[],
            )
        entry = self._routing_table[state]
        entry.total_attempts += 1
        if transition.resolved:
            entry.total_successes += 1

        # Rebuild preferred agents list sorted by success rate
        agent_success = {}
        agent_attempts = {}
        for ep in self._episodes:
            for t in ep.transitions:
                if t.task_state == state:
                    aid = t.agent_id
                    agent_attempts[aid] = agent_attempts.get(aid, 0) + 1
                    if t.resolved:
                        agent_success[aid] = agent_success.get(aid, 0) + 1

        ranked = []
        for aid, attempts in agent_attempts.items():
            successes = agent_success.get(aid, 0)
            rate = successes / attempts
            ranked.append((aid, rate))
        ranked.sort(key=lambda x: x[1], reverse=True)
        entry.preferred_agents = ranked[:5]

    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(),
        }
