#!/usr/bin/env python3
from __future__ import annotations

"""
Phase 6: Autonomous Gate Controller

This module provides adaptive gate selection based on pheromone history.
It reads recent pheromone entries and computes optimal:
    - gate_policy
    - regrow_depth
    - lattice_D

The scoring algorithm:
    score = Se * 0.4 + logits_mean * 0.4 + seed_entropy * 0.2

The controller selects the gate with the highest rolling score for a given task_type.

Usage:
    from helix_cdc.bloomos.gate_controller import GateController

    controller = GateController("/path/to/pheromones.jsonl")
    gate = controller.choose_gate_policy("reason")
    depth = controller.choose_regrow_depth("reason")
    d = controller.choose_lattice_D("reason")

CLI:
    PYTHONPATH=. python3 helix_cdc/bloomos/gate_controller.py --inspect /path/to/pheromones.jsonl
"""

import argparse
import json
from collections import defaultdict
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

from .pheromone import PheromoneBoard, PheromoneEntry
from .fitness import fitness


# Default gate policies in preference order (fallback if no history)
DEFAULT_GATE_POLICIES = ["balanced", "math_bias", "code_bias", "json_bias"]
DEFAULT_REGROW_DEPTH = 64
DEFAULT_LATTICE_D = 72


@dataclass
class GatePreference:
    """Computed preference for a gate policy."""
    gate: str
    score: float
    sample_count: int
    avg_se: float
    avg_logits_mean: float
    avg_seed_entropy: float


@dataclass
class ControllerRecommendation:
    """Full recommendation from the controller."""
    task_type: str
    preferred_gate: str
    preferred_regrow_depth: int
    preferred_lattice_D: int
    gate_scores: List[GatePreference]
    total_samples: int


class GateController:
    """
    Adaptive gate controller based on pheromone feedback.

    Reads historical pheromone entries and computes the best gate policy
    for each task_type based on Se, logits_mean, and seed_entropy.
    """

    def __init__(
        self,
        pheromone_path: Optional[str] = None,
        history_window: int = 100,
        min_samples: int = 3,
    ) -> None:
        """
        Initialize the gate controller.

        Args:
            pheromone_path: Path to pheromone JSONL file
            history_window: Number of most recent entries to consider
            min_samples: Minimum samples required for a gate to be considered
        """
        self.pheromone_path = Path(pheromone_path) if pheromone_path else None
        self.history_window = history_window
        self.min_samples = min_samples
        self._cache: Dict[str, ControllerRecommendation] = {}

    def _load_entries(self, task_type: Optional[str] = None) -> List[PheromoneEntry]:
        """Load recent pheromone entries, optionally filtered by task_type."""
        if self.pheromone_path is None or not self.pheromone_path.exists():
            return []

        board = PheromoneBoard(self.pheromone_path)
        entries = list(board.iter_entries())

        # Filter by task_type if specified
        if task_type:
            entries = [e for e in entries if e.task_type == task_type]

        # Take most recent entries within window
        entries = entries[-self.history_window:]
        return entries

    def _compute_gate_scores(
        self, entries: List[PheromoneEntry]
    ) -> Dict[str, GatePreference]:
        """Compute scores for each gate based on entries."""
        # Group by gate
        gate_entries: Dict[str, List[PheromoneEntry]] = defaultdict(list)
        for e in entries:
            gate = e.gate or "balanced"
            gate_entries[gate].append(e)

        scores: Dict[str, GatePreference] = {}
        for gate, gate_list in gate_entries.items():
            # Compute averages
            se_vals = [e.se_Se for e in gate_list if e.se_Se is not None]
            logits_vals = [e.logits_mean for e in gate_list if e.logits_mean is not None]
            entropy_vals = [e.seed_entropy for e in gate_list if e.seed_entropy is not None]

            avg_se = sum(se_vals) / len(se_vals) if se_vals else 0.0
            avg_logits = sum(logits_vals) / len(logits_vals) if logits_vals else 0.0
            avg_entropy = sum(entropy_vals) / len(entropy_vals) if entropy_vals else 4.0  # Default mid-range

            # Score formula: Se * 0.4 + logits_mean * 0.4 + seed_entropy * 0.2
            # Normalize Se (typically 0-100) and logits_mean (typically 0-1)
            # seed_entropy is 0-8 bits/byte
            score = (avg_se / 100.0) * 0.4 + avg_logits * 0.4 + (avg_entropy / 8.0) * 0.2

            scores[gate] = GatePreference(
                gate=gate,
                score=score,
                sample_count=len(gate_list),
                avg_se=avg_se,
                avg_logits_mean=avg_logits,
                avg_seed_entropy=avg_entropy,
            )

        return scores

    def _compute_regrow_depth(self, entries: List[PheromoneEntry]) -> int:
        """Compute preferred regrow depth from successful entries."""
        # Filter to successful entries with regrow_depth
        successful = [
            e for e in entries
            if e.status == "ok" and e.regrow_depth is not None
        ]
        if not successful:
            return DEFAULT_REGROW_DEPTH

        # Weight by fitness
        weighted_sum = 0.0
        weight_total = 0.0
        for e in successful:
            fit = fitness(e)
            weighted_sum += e.regrow_depth * fit
            weight_total += fit

        if weight_total == 0:
            return DEFAULT_REGROW_DEPTH

        # Round to nearest power of 2 (typical regrow depths: 32, 64, 128, 256)
        avg = weighted_sum / weight_total
        powers = [32, 64, 128, 256, 512, 1024]
        return min(powers, key=lambda x: abs(x - avg))

    def _compute_lattice_D(self, entries: List[PheromoneEntry]) -> int:
        """Compute preferred lattice D from successful entries."""
        successful = [
            e for e in entries
            if e.status == "ok" and e.lattice_D is not None
        ]
        if not successful:
            return DEFAULT_LATTICE_D

        # Weight by fitness
        weighted_sum = 0.0
        weight_total = 0.0
        for e in successful:
            fit = fitness(e)
            weighted_sum += e.lattice_D * fit
            weight_total += fit

        if weight_total == 0:
            return DEFAULT_LATTICE_D

        # Clamp to valid range [1, 72]
        return max(1, min(72, int(weighted_sum / weight_total)))

    def get_recommendation(self, task_type: str) -> ControllerRecommendation:
        """
        Get full controller recommendation for a task type.

        Args:
            task_type: The task type to get recommendations for

        Returns:
            ControllerRecommendation with preferred gate, depth, D, and scores
        """
        entries = self._load_entries(task_type)

        if not entries:
            return ControllerRecommendation(
                task_type=task_type,
                preferred_gate="balanced",
                preferred_regrow_depth=DEFAULT_REGROW_DEPTH,
                preferred_lattice_D=DEFAULT_LATTICE_D,
                gate_scores=[],
                total_samples=0,
            )

        gate_scores = self._compute_gate_scores(entries)

        # Filter gates with sufficient samples
        qualified_gates = {
            g: s for g, s in gate_scores.items()
            if s.sample_count >= self.min_samples
        }

        # Select best gate
        if qualified_gates:
            best_gate = max(qualified_gates.values(), key=lambda x: x.score).gate
        else:
            # Fall back to gate with most samples
            best_gate = max(gate_scores.values(), key=lambda x: x.sample_count).gate if gate_scores else "balanced"

        # Sort scores for output
        sorted_scores = sorted(gate_scores.values(), key=lambda x: x.score, reverse=True)

        return ControllerRecommendation(
            task_type=task_type,
            preferred_gate=best_gate,
            preferred_regrow_depth=self._compute_regrow_depth(entries),
            preferred_lattice_D=self._compute_lattice_D(entries),
            gate_scores=sorted_scores,
            total_samples=len(entries),
        )

    def choose_gate_policy(self, task_type: str) -> str:
        """
        Choose the best gate policy for a task type.

        Args:
            task_type: The task type (e.g., "reason", "code", "math")

        Returns:
            Gate policy name (e.g., "math_bias", "balanced")
        """
        rec = self.get_recommendation(task_type)
        return rec.preferred_gate

    def choose_regrow_depth(self, task_type: str) -> int:
        """Choose optimal regrow depth for a task type."""
        rec = self.get_recommendation(task_type)
        return rec.preferred_regrow_depth

    def choose_lattice_D(self, task_type: str) -> int:
        """Choose optimal lattice D for a task type."""
        rec = self.get_recommendation(task_type)
        return rec.preferred_lattice_D

    def inspect(self) -> Dict[str, Any]:
        """
        Inspect the pheromone history and return analysis.

        Returns:
            Dict with task_types found and their recommendations
        """
        if self.pheromone_path is None or not self.pheromone_path.exists():
            return {"error": "No pheromone file", "path": str(self.pheromone_path)}

        board = PheromoneBoard(self.pheromone_path)
        entries = list(board.iter_entries())

        # Group by task_type
        task_types = set(e.task_type for e in entries)

        result = {
            "path": str(self.pheromone_path),
            "total_entries": len(entries),
            "task_types": {},
        }

        for tt in task_types:
            rec = self.get_recommendation(tt)
            result["task_types"][tt] = {
                "preferred_gate": rec.preferred_gate,
                "preferred_regrow_depth": rec.preferred_regrow_depth,
                "preferred_lattice_D": rec.preferred_lattice_D,
                "total_samples": rec.total_samples,
                "gate_scores": [
                    {
                        "gate": g.gate,
                        "score": round(g.score, 4),
                        "samples": g.sample_count,
                        "avg_se": round(g.avg_se, 2),
                        "avg_logits_mean": round(g.avg_logits_mean, 4),
                        "avg_seed_entropy": round(g.avg_seed_entropy, 3),
                    }
                    for g in rec.gate_scores
                ],
            }

        return result


def _main() -> None:
    parser = argparse.ArgumentParser(
        prog="gate_controller",
        description="Phase 6: Autonomous Gate Controller",
    )
    parser.add_argument(
        "--inspect",
        metavar="PATH",
        help="Inspect pheromone JSONL file and print analysis",
    )
    parser.add_argument(
        "--task-type",
        help="Get recommendation for specific task type",
    )
    parser.add_argument(
        "--history-window",
        type=int,
        default=100,
        help="Number of recent entries to consider (default: 100)",
    )
    parser.add_argument(
        "--min-samples",
        type=int,
        default=3,
        help="Minimum samples required for gate selection (default: 3)",
    )
    args = parser.parse_args()

    if args.inspect:
        controller = GateController(
            pheromone_path=args.inspect,
            history_window=args.history_window,
            min_samples=args.min_samples,
        )
        result = controller.inspect()
        print(json.dumps(result, indent=2))

        if args.task_type:
            print(f"\n=== Recommendation for '{args.task_type}' ===")
            rec = controller.get_recommendation(args.task_type)
            print(f"Gate:        {rec.preferred_gate}")
            print(f"Regrow:      {rec.preferred_regrow_depth}")
            print(f"Lattice D:   {rec.preferred_lattice_D}")
            print(f"Samples:     {rec.total_samples}")
    else:
        parser.print_help()


if __name__ == "__main__":
    _main()
