#!/usr/bin/env python3
from __future__ import annotations
"""
Phase 10: Substrate Evolution Engine
------------------------------------
This module continuously analyzes pheromone history, drift scores, program/gate
fitness, and substrate inspector outputs to automatically propose:

    - Mutated gate policies
    - Mutated helix_program sequences
    - Adjusted regrow_depth
    - Adjusted lattice_D preference
    - Evolution "patches" for ProgramController & GateController

It outputs EvolutionResult objects containing recommended updates.

Usage:
    PYTHONPATH=. python3 -m helix_cdc.bloomos.evolution_engine \\
      /tmp/pheromones.jsonl --task-type math

    PYTHONPATH=. python3 -m helix_cdc.bloomos.evolution_engine \\
      /tmp/pheromones.jsonl --task-type code --json
"""

import json
import random
from dataclasses import dataclass, field, asdict
from pathlib import Path
from typing import Dict, List, Optional, Any, Tuple

from .pheromone import PheromoneBoard, PheromoneEntry
from .fitness import fitness
from .substrate_inspector import (
    analyze_substrate,
    SubstrateSummary,
    TaskSummary,
    GateStats,
    ProgramStats,
)


# ------------------------------
# Constants
# ------------------------------

ALPHABET = "ABCDEFGH"
KNOWN_GATES = ["balanced", "math_bias", "code_bias", "json_bias"]
DEFAULT_REGROW_DEPTH = 64
DEFAULT_LATTICE_D = 72


# ------------------------------
# EvolutionResult
# ------------------------------

@dataclass
class EvolutionResult:
    """Result of an evolution cycle for a task_type."""
    generation: int
    task_type: str
    recommended_gate: Optional[str]
    recommended_program: Optional[str]
    recommended_regrow_depth: Optional[int]
    recommended_lattice_D: Optional[int]
    current_fitness: Optional[float]
    previous_fitness: Optional[float]
    drift: Optional[float]
    mutation_applied: bool
    mutation_trace: Dict[str, Any] = field(default_factory=dict)


# ------------------------------
# Utility functions
# ------------------------------

def mutate_program(seq: str, rate: float = 0.12) -> str:
    """
    Randomly mutate a helix_program sequence.

    Args:
        seq: 8-codon sequence (e.g., "ABCDEFGH")
        rate: Probability of mutating each codon

    Returns:
        Mutated sequence
    """
    out = []
    for c in seq:
        if random.random() < rate:
            # Mutate to a different codon
            choices = [x for x in ALPHABET if x != c]
            out.append(random.choice(choices))
        else:
            out.append(c)
    return "".join(out)


def crossover_programs(seq1: str, seq2: str) -> str:
    """
    Single-point crossover between two helix_program sequences.

    Args:
        seq1, seq2: 8-codon sequences

    Returns:
        Child sequence
    """
    point = random.randint(1, 7)
    return seq1[:point] + seq2[point:]


def _get_recent_entries(
    path: Path, window: int, task_type: Optional[str] = None
) -> List[PheromoneEntry]:
    """Get the most recent N entries, optionally filtered by task_type."""
    board = PheromoneBoard(str(path))
    entries = list(board.iter_entries())

    if task_type:
        entries = [e for e in entries if e.task_type == task_type]

    if len(entries) > window:
        entries = entries[-window:]

    return entries


def _compute_avg_fitness(entries: List[PheromoneEntry]) -> Optional[float]:
    """Compute average fitness from entries."""
    if not entries:
        return None
    fitness_vals = []
    for e in entries:
        try:
            fitness_vals.append(fitness(e))
        except Exception:
            pass
    if not fitness_vals:
        return None
    return sum(fitness_vals) / len(fitness_vals)


def _find_task_summary(
    summary: SubstrateSummary, task_type: str
) -> Optional[TaskSummary]:
    """Find TaskSummary for a specific task_type."""
    for ts in summary.task_summaries:
        if ts.task_type == task_type:
            return ts
    return None


def _get_best_gate_stats(ts: TaskSummary) -> Optional[GateStats]:
    """Get the GateStats for the best gate."""
    if not ts.gates:
        return None
    # gates are already sorted by fitness descending
    return ts.gates[0]


def _get_best_program_stats(ts: TaskSummary) -> Optional[ProgramStats]:
    """Get the ProgramStats for the best program."""
    if not ts.programs:
        return None
    # programs are already sorted by fitness descending
    return ts.programs[0]


# ------------------------------
# Evolution Engine
# ------------------------------

class EvolutionEngine:
    """
    Substrate Evolution Engine.

    Analyzes pheromone history and substrate state to propose
    evolved configurations (gates, programs, regrow_depth, lattice_D).
    """

    def __init__(
        self,
        pheromone_path: str,
        window: int = 200,
        prev_window: int = 100,
        mutation_rate: float = 0.12,
        min_samples: int = 5,
        drift_threshold: float = -0.05,
    ) -> None:
        """
        Initialize the evolution engine.

        Args:
            pheromone_path: Path to pheromone JSONL file
            window: Number of recent entries to consider for current fitness
            prev_window: Number of entries before current window for drift comparison
            mutation_rate: Probability of mutating each codon
            min_samples: Minimum samples required for evolution
            drift_threshold: Fitness drift threshold to trigger mutation (negative = regression)
        """
        self.path = Path(pheromone_path)
        self.window = window
        self.prev_window = prev_window
        self.mutation_rate = mutation_rate
        self.min_samples = min_samples
        self.drift_threshold = drift_threshold
        self.generation = 0

    def _compute_drift(
        self, task_type: str
    ) -> Tuple[Optional[float], Optional[float], Optional[float]]:
        """
        Compute fitness drift by comparing current window vs previous window.

        Returns:
            Tuple of (current_fitness, previous_fitness, drift)
        """
        board = PheromoneBoard(str(self.path))
        entries = [e for e in board.iter_entries() if e.task_type == task_type]

        if len(entries) < self.window + self.prev_window:
            # Not enough data for drift calculation
            current_entries = entries[-self.window:] if len(entries) > self.window else entries
            current_fitness = _compute_avg_fitness(current_entries)
            return current_fitness, None, None

        # Split into current and previous windows
        current_entries = entries[-self.window:]
        prev_entries = entries[-(self.window + self.prev_window):-self.window]

        current_fitness = _compute_avg_fitness(current_entries)
        prev_fitness = _compute_avg_fitness(prev_entries)

        drift = None
        if current_fitness is not None and prev_fitness is not None:
            drift = current_fitness - prev_fitness

        return current_fitness, prev_fitness, drift

    def evolve_for_task(self, task_type: str) -> EvolutionResult:
        """
        Run one evolution cycle for a task_type.

        Args:
            task_type: The task type to evolve (e.g., "math", "code", "json")

        Returns:
            EvolutionResult with recommended configuration
        """
        self.generation += 1

        # Get recent entries
        entries = _get_recent_entries(self.path, self.window, task_type)

        # Fallback for insufficient data
        if len(entries) < self.min_samples:
            return EvolutionResult(
                generation=self.generation,
                task_type=task_type,
                recommended_gate="balanced",
                recommended_program="ABCDEFGH",
                recommended_regrow_depth=DEFAULT_REGROW_DEPTH,
                recommended_lattice_D=DEFAULT_LATTICE_D,
                current_fitness=None,
                previous_fitness=None,
                drift=None,
                mutation_applied=False,
                mutation_trace={"reason": "insufficient_samples", "count": len(entries)},
            )

        # Analyze substrate
        summary = analyze_substrate(str(self.path))
        task_summary = _find_task_summary(summary, task_type)

        if task_summary is None:
            return EvolutionResult(
                generation=self.generation,
                task_type=task_type,
                recommended_gate="balanced",
                recommended_program="ABCDEFGH",
                recommended_regrow_depth=DEFAULT_REGROW_DEPTH,
                recommended_lattice_D=DEFAULT_LATTICE_D,
                current_fitness=None,
                previous_fitness=None,
                drift=None,
                mutation_applied=False,
                mutation_trace={"reason": "task_not_found"},
            )

        # Get best gate and program from substrate analysis
        best_gate = task_summary.best_gate or "balanced"
        best_program = task_summary.best_program or "ABCDEFGH"

        # Compute drift
        current_fitness, prev_fitness, drift = self._compute_drift(task_type)

        # Decide whether to mutate
        mutate_flag = False
        if drift is not None and drift < self.drift_threshold:
            mutate_flag = True

        mutation_trace: Dict[str, Any] = {
            "best_gate_from_substrate": best_gate,
            "best_program_from_substrate": best_program,
            "drift_threshold": self.drift_threshold,
        }

        # Apply mutation if regression detected
        recommended_gate = best_gate
        recommended_program = best_program

        if mutate_flag:
            # Mutate program
            original_program = best_program
            recommended_program = mutate_program(best_program, self.mutation_rate)
            mutation_trace["original_program"] = original_program
            mutation_trace["mutated_program"] = recommended_program

            # Potentially switch gate
            if random.random() < 0.3:  # 30% chance to also mutate gate
                other_gates = [g for g in KNOWN_GATES if g != best_gate]
                recommended_gate = random.choice(other_gates)
                mutation_trace["original_gate"] = best_gate
                mutation_trace["mutated_gate"] = recommended_gate

        # Compute recommended regrow_depth and lattice_D from entries
        depths = [e.regrow_depth for e in entries if e.regrow_depth is not None]
        Ds = [e.lattice_D for e in entries if e.lattice_D is not None]

        recommended_depth = int(sum(depths) / len(depths)) if depths else DEFAULT_REGROW_DEPTH
        recommended_D = int(sum(Ds) / len(Ds)) if Ds else DEFAULT_LATTICE_D

        # Round regrow_depth to nearest power of 2
        powers = [32, 64, 128, 256, 512]
        recommended_depth = min(powers, key=lambda x: abs(x - recommended_depth))

        # Clamp lattice_D to valid range
        recommended_D = max(1, min(72, recommended_D))

        return EvolutionResult(
            generation=self.generation,
            task_type=task_type,
            recommended_gate=recommended_gate,
            recommended_program=recommended_program,
            recommended_regrow_depth=recommended_depth,
            recommended_lattice_D=recommended_D,
            current_fitness=round(current_fitness, 4) if current_fitness else None,
            previous_fitness=round(prev_fitness, 4) if prev_fitness else None,
            drift=round(drift, 4) if drift else None,
            mutation_applied=mutate_flag,
            mutation_trace=mutation_trace,
        )

    def evolve_all_tasks(self) -> List[EvolutionResult]:
        """
        Run evolution for all task_types found in pheromone history.

        Returns:
            List of EvolutionResult for each task_type
        """
        summary = analyze_substrate(str(self.path))
        results = []
        for ts in summary.task_summaries:
            result = self.evolve_for_task(ts.task_type)
            results.append(result)
        return results


# ------------------------------
# CLI
# ------------------------------

def _cli() -> None:
    import argparse

    parser = argparse.ArgumentParser(
        prog="evolution_engine",
        description="Phase 10: Substrate Evolution Engine - propose evolved configurations.",
    )
    parser.add_argument(
        "path",
        help="Path to pheromone JSONL file.",
    )
    parser.add_argument(
        "--task-type",
        help="Evolve for specific task_type. If omitted, evolves all tasks.",
    )
    parser.add_argument(
        "--window",
        type=int,
        default=200,
        help="Current window size for fitness calculation (default: 200).",
    )
    parser.add_argument(
        "--prev-window",
        type=int,
        default=100,
        help="Previous window size for drift comparison (default: 100).",
    )
    parser.add_argument(
        "--mutation-rate",
        type=float,
        default=0.12,
        help="Mutation rate for program codons (default: 0.12).",
    )
    parser.add_argument(
        "--min-samples",
        type=int,
        default=5,
        help="Minimum samples required (default: 5).",
    )
    parser.add_argument(
        "--drift-threshold",
        type=float,
        default=-0.05,
        help="Drift threshold to trigger mutation (default: -0.05).",
    )
    parser.add_argument(
        "--json",
        action="store_true",
        help="Output as JSON.",
    )
    parser.add_argument(
        "--seed",
        type=int,
        help="Random seed for reproducibility.",
    )

    args = parser.parse_args()

    if args.seed is not None:
        random.seed(args.seed)

    engine = EvolutionEngine(
        pheromone_path=args.path,
        window=args.window,
        prev_window=args.prev_window,
        mutation_rate=args.mutation_rate,
        min_samples=args.min_samples,
        drift_threshold=args.drift_threshold,
    )

    if args.task_type:
        results = [engine.evolve_for_task(args.task_type)]
    else:
        results = engine.evolve_all_tasks()

    if args.json:
        print(json.dumps([asdict(r) for r in results], indent=2))
    else:
        print("=" * 70)
        print("SUBSTRATE EVOLUTION ENGINE - Phase 10")
        print("=" * 70)
        print()

        for r in results:
            print(f"=== Task: {r.task_type} (Generation {r.generation}) ===")
            print(f"Recommended Gate:      {r.recommended_gate}")
            print(f"Recommended Program:   {r.recommended_program}")
            print(f"Recommended Depth:     {r.recommended_regrow_depth}")
            print(f"Recommended Lattice D: {r.recommended_lattice_D}")
            print()
            print(f"Current Fitness:       {r.current_fitness}")
            print(f"Previous Fitness:      {r.previous_fitness}")
            print(f"Drift:                 {r.drift}")
            print(f"Mutation Applied:      {r.mutation_applied}")
            print()
            if r.mutation_trace:
                print("Mutation Trace:")
                for k, v in r.mutation_trace.items():
                    print(f"  {k}: {v}")
            print()


if __name__ == "__main__":
    _cli()
