#!/usr/bin/env python3
from __future__ import annotations
"""
Phase 12: Genetic Program Search (Helix64)
-----------------------------------------

Genetic search over helix_program sequences (8 codons A..H) plus:
    - gate policy
    - regrow_depth
    - lattice_D

It:
    - Evolves a population of (program, gate, depth, D)
    - Evaluates each via BloomOS runtime
    - Computes Se + fitness
    - Records pheromones for each eval
    - Produces a JSON summary

Usage:
    PYTHONPATH=. python3 -m helix_cdc.bloomos.genetic_program_search \\
      --pheromone-path /tmp/ga_pheromones.jsonl \\
      --out /tmp/ga_results.json \\
      --task-type math \\
      --goal "integrate sin(x) from 0 to pi" \\
      --generations 5 \\
      --population-size 16
"""

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

from .bloomos_client import BloomOSClient, TaskResult
from .pheromone import PheromoneBoard, make_entry_from_step
from .fitness import fitness
from helix_cdc.metrics import measure_text, SeMetrics


CODONS = list("ABCDEFGH")
PROGRAM_LENGTH = 8


# ---------------------------------------------------------------------
# Data models
# ---------------------------------------------------------------------

@dataclass
class Individual:
    """A single individual in the genetic population."""
    program: str
    gate: str
    regrow_depth: int
    lattice_D: int
    fitness: Optional[float] = None
    Se: Optional[float] = None
    logits_mean: Optional[float] = None
    status: Optional[str] = None
    eval_count: int = 0
    meta: Dict[str, Any] = field(default_factory=dict)


@dataclass
class GenerationSummary:
    """Summary of a single generation."""
    generation: int
    best_fitness: float
    best_program: str
    best_gate: str
    best_regrow_depth: int
    best_lattice_D: int
    avg_fitness: float
    diversity: float  # Number of unique programs
    population_size: int


@dataclass
class GeneticSearchResult:
    """Result of the genetic search."""
    task_type: str
    goal: str
    total_evaluations: int
    generations: List[GenerationSummary]
    best_overall: Individual
    final_population: List[Individual]


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

def _is_program_valid(seq: str) -> bool:
    """Check if a program sequence is valid."""
    if not isinstance(seq, str) or len(seq) != PROGRAM_LENGTH:
        return False
    return all(c in CODONS for c in seq)


def random_program(base: Optional[str] = None) -> str:
    """
    Generate a random program.
    If base is provided, randomly mutate some positions from it.
    """
    if base and len(base) == PROGRAM_LENGTH and _is_program_valid(base):
        chars = list(base)
        # Mutate 1-3 positions
        num_mutations = random.randint(1, 3)
        positions = random.sample(range(PROGRAM_LENGTH), k=num_mutations)
        for i in positions:
            choices = [c for c in CODONS if c != chars[i]]
            chars[i] = random.choice(choices)
        return "".join(chars)
    else:
        return "".join(random.choice(CODONS) for _ in range(PROGRAM_LENGTH))


def mutate_program(seq: str, mutation_rate: float) -> str:
    """
    Per-position mutation with probability mutation_rate.
    Ensures at least one position mutates if mutation is triggered.
    """
    chars = list(seq)
    mutated = False

    for i in range(PROGRAM_LENGTH):
        if random.random() < mutation_rate:
            choices = [c for c in CODONS if c != chars[i]]
            chars[i] = random.choice(choices)
            mutated = True

    # Ensure at least one mutation if rate > 0 and we got unlucky
    if not mutated and random.random() < mutation_rate * PROGRAM_LENGTH:
        i = random.randrange(PROGRAM_LENGTH)
        choices = [c for c in CODONS if c != chars[i]]
        chars[i] = random.choice(choices)

    candidate = "".join(chars)
    if not _is_program_valid(candidate):
        return seq  # Keep parent if invalid
    return candidate


def crossover_programs(a: str, b: str) -> Tuple[str, str]:
    """Single-point crossover between two programs."""
    if len(a) != PROGRAM_LENGTH or len(b) != PROGRAM_LENGTH:
        return a, b

    point = random.randint(1, PROGRAM_LENGTH - 1)
    c1 = a[:point] + b[point:]
    c2 = b[:point] + a[point:]

    # Validate children
    if not _is_program_valid(c1):
        c1 = a
    if not _is_program_valid(c2):
        c2 = b

    return c1, c2


def _extract_output_text(raw: Dict[str, Any]) -> Optional[str]:
    """Extract text output from TaskResult raw dict."""
    out = raw.get("output")
    if out is None:
        return None
    if isinstance(out, str):
        return out
    if isinstance(out, dict):
        for key in ("text", "content", "response", "output", "result", "reasoning"):
            if key in out and isinstance(out[key], str):
                return out[key]
        return json.dumps(out)
    return str(out)


# ---------------------------------------------------------------------
# Core engine
# ---------------------------------------------------------------------

class GeneticProgramSearch:
    """
    Genetic algorithm for searching optimal helix_program configurations.

    Evolves a population of (program, gate, regrow_depth, lattice_D) tuples
    using tournament selection, crossover, and mutation.
    """

    def __init__(
        self,
        base_url: str = "http://localhost:7070",
        pheromone_path: Optional[str] = None,
        task_type: str = "experiment",
        goal: str = "",
        gates: Optional[List[str]] = None,
        regrow_depths: Optional[List[int]] = None,
        lattice_Ds: Optional[List[int]] = None,
        timeout: float = 60.0,
    ) -> None:
        """
        Initialize the genetic search engine.

        Args:
            base_url: BloomOS runtime URL
            pheromone_path: Path to pheromone JSONL for recording
            task_type: Task type for BloomOS
            goal: Natural language goal/prompt
            gates: List of allowed gate policies
            regrow_depths: List of allowed regrow depths
            lattice_Ds: List of allowed lattice_D values
            timeout: HTTP timeout
        """
        self.client = BloomOSClient(base_url=base_url, timeout=timeout)
        self.pheromones = PheromoneBoard(pheromone_path) if pheromone_path else None
        self.task_type = task_type
        self.goal = goal
        self.gates = gates or ["balanced"]
        self.regrow_depths = regrow_depths or [64]
        self.lattice_Ds = lattice_Ds or [72]
        self.total_evaluations = 0

    def _eval_individual(self, ind: Individual) -> Individual:
        """
        Evaluate an individual via BloomOS runtime.

        Computes Se + fitness and records a pheromone entry.
        """
        extra: Dict[str, Any] = {
            "helix_program": ind.program,
            "regrow_depth": ind.regrow_depth,
            "lattice_D": ind.lattice_D,
        }

        t0 = time.perf_counter()
        try:
            result: TaskResult = self.client.run_task(
                task_type=self.task_type,
                description=self.goal,
                anchors=None,
                gate_policy=ind.gate,
                extra=extra,
            )
            elapsed_ms = (time.perf_counter() - t0) * 1000.0

            raw = result.raw
            output = _extract_output_text(raw) or ""
            logits_mean = result.logits_mean
            logits_std = result.logits_std
            status = result.status

            # Compute Se metrics
            se_H = se_C = se_D = se_Se = None
            try:
                m: SeMetrics = measure_text(output, auto_detect_d=True)
                se_H, se_C, se_D, se_Se = m.H, m.C, m.D, m.Se
            except Exception:
                se_Se = 0.0

            # Compute fitness
            class FakeEntry:
                pass
            fake = FakeEntry()
            fake.se_Se = se_Se
            fake.logits_mean = logits_mean
            fake.seed_entropy = 5.0  # Default mid-range

            try:
                f_val = fitness(fake)
            except Exception:
                f_val = 0.0

            ind.fitness = f_val
            ind.Se = se_Se
            ind.logits_mean = logits_mean
            ind.status = status
            ind.eval_count += 1
            ind.meta["elapsed_ms"] = elapsed_ms
            self.total_evaluations += 1

            # Record pheromone
            if self.pheromones:
                try:
                    entry = make_entry_from_step(
                        task_type=self.task_type,
                        gate=ind.gate,
                        helix_program=ind.program,
                        status=status,
                        logits_mean=logits_mean,
                        logits_std=logits_std,
                        se_H=se_H,
                        se_C=se_C,
                        se_D=se_D,
                        se_Se=se_Se,
                        meta={
                            "goal": self.goal,
                            "elapsed_ms": elapsed_ms,
                            "ga_generation": ind.meta.get("generation", 0),
                        },
                        regrow_depth=ind.regrow_depth,
                        lattice_D=ind.lattice_D,
                    )
                    self.pheromones.record(entry)
                except Exception:
                    pass

        except Exception as e:
            elapsed_ms = (time.perf_counter() - t0) * 1000.0
            ind.fitness = 0.0
            ind.Se = 0.0
            ind.logits_mean = 0.0
            ind.status = f"error:{type(e).__name__}"
            ind.eval_count += 1
            ind.meta["elapsed_ms"] = elapsed_ms
            ind.meta["error"] = str(e)
            self.total_evaluations += 1

        return ind

    def _init_population(
        self,
        population_size: int,
        seed_programs: Optional[List[str]] = None,
    ) -> List[Individual]:
        """Initialize the population with seeds and random individuals."""
        pop: List[Individual] = []
        seeds = (seed_programs or [])[:population_size]

        # Seed individuals
        for sp in seeds:
            if _is_program_valid(sp):
                pop.append(
                    Individual(
                        program=sp,
                        gate=random.choice(self.gates),
                        regrow_depth=random.choice(self.regrow_depths),
                        lattice_D=random.choice(self.lattice_Ds),
                    )
                )

        # Random fill
        while len(pop) < population_size:
            base = seeds[0] if seeds else "ABCDEFGH"
            pop.append(
                Individual(
                    program=random_program(base=base),
                    gate=random.choice(self.gates),
                    regrow_depth=random.choice(self.regrow_depths),
                    lattice_D=random.choice(self.lattice_Ds),
                )
            )

        return pop

    @staticmethod
    def _select_parents(population: List[Individual]) -> Tuple[Individual, Individual]:
        """Tournament selection (size 3)."""
        def tournament() -> Individual:
            candidates = random.sample(population, k=min(3, len(population)))
            return max(candidates, key=lambda i: i.fitness or 0.0)

        return tournament(), tournament()

    def _compute_diversity(self, population: List[Individual]) -> float:
        """Compute population diversity as ratio of unique programs."""
        unique = len(set(ind.program for ind in population))
        return unique / len(population) if population else 0.0

    def run(
        self,
        generations: int = 10,
        population_size: int = 16,
        mutation_rate: float = 0.15,
        crossover_rate: float = 0.6,
        elitism: int = 2,
        seed_programs: Optional[List[str]] = None,
        out_path: Optional[str] = None,
        seed: Optional[int] = None,
    ) -> GeneticSearchResult:
        """
        Run the genetic search.

        Args:
            generations: Number of generations to evolve
            population_size: Size of population
            mutation_rate: Probability of mutating each codon
            crossover_rate: Probability of crossover
            elitism: Number of top individuals to preserve
            seed_programs: Initial seed programs
            out_path: Path to save results JSON
            seed: Random seed for reproducibility

        Returns:
            GeneticSearchResult with best individual and history
        """
        if seed is not None:
            random.seed(seed)

        population = self._init_population(population_size, seed_programs)
        gen_summaries: List[GenerationSummary] = []
        best_ever: Optional[Individual] = None

        for gen in range(1, generations + 1):
            # Mark generation
            for ind in population:
                ind.meta["generation"] = gen

            # Evaluate population
            for ind in population:
                if ind.fitness is None:
                    self._eval_individual(ind)

            # Update best ever
            gen_best = max(population, key=lambda i: i.fitness or 0.0)
            if best_ever is None or (gen_best.fitness or 0) > (best_ever.fitness or 0):
                best_ever = Individual(
                    program=gen_best.program,
                    gate=gen_best.gate,
                    regrow_depth=gen_best.regrow_depth,
                    lattice_D=gen_best.lattice_D,
                    fitness=gen_best.fitness,
                    Se=gen_best.Se,
                    logits_mean=gen_best.logits_mean,
                    status=gen_best.status,
                    eval_count=gen_best.eval_count,
                    meta=dict(gen_best.meta),
                )

            # Stats
            avg_fit = sum((i.fitness or 0.0) for i in population) / len(population)
            diversity = self._compute_diversity(population)

            gen_summaries.append(
                GenerationSummary(
                    generation=gen,
                    best_fitness=gen_best.fitness or 0.0,
                    best_program=gen_best.program,
                    best_gate=gen_best.gate,
                    best_regrow_depth=gen_best.regrow_depth,
                    best_lattice_D=gen_best.lattice_D,
                    avg_fitness=avg_fit,
                    diversity=diversity,
                    population_size=len(population),
                )
            )

            # Skip reproduction on last generation
            if gen >= generations:
                break

            # Prepare next generation
            population.sort(key=lambda i: i.fitness or 0.0, reverse=True)
            next_pop: List[Individual] = []

            # Elitism - preserve top individuals
            for i in range(min(elitism, len(population))):
                elite = population[i]
                next_pop.append(
                    Individual(
                        program=elite.program,
                        gate=elite.gate,
                        regrow_depth=elite.regrow_depth,
                        lattice_D=elite.lattice_D,
                        # Reset fitness so it gets re-evaluated
                        fitness=None,
                    )
                )

            # Reproduction
            while len(next_pop) < population_size:
                parent1, parent2 = self._select_parents(population)

                # Crossover
                if random.random() < crossover_rate:
                    c1_prog, c2_prog = crossover_programs(parent1.program, parent2.program)
                else:
                    c1_prog, c2_prog = parent1.program, parent2.program

                # Mutate programs
                c1_prog = mutate_program(c1_prog, mutation_rate)
                c2_prog = mutate_program(c2_prog, mutation_rate)

                def make_child(parent: Individual, prog: str) -> Individual:
                    """Create child with possibly mutated hyperparameters."""
                    gate = parent.gate
                    if random.random() < 0.1:
                        gate = random.choice(self.gates)

                    depth = parent.regrow_depth
                    if random.random() < 0.1:
                        depth = random.choice(self.regrow_depths)

                    Dv = parent.lattice_D
                    if random.random() < 0.1:
                        Dv = random.choice(self.lattice_Ds)

                    return Individual(
                        program=prog,
                        gate=gate,
                        regrow_depth=depth,
                        lattice_D=Dv,
                    )

                next_pop.append(make_child(parent1, c1_prog))
                if len(next_pop) < population_size:
                    next_pop.append(make_child(parent2, c2_prog))

            population = next_pop

        # Build result
        result = GeneticSearchResult(
            task_type=self.task_type,
            goal=self.goal,
            total_evaluations=self.total_evaluations,
            generations=gen_summaries,
            best_overall=best_ever or population[0],
            final_population=population,
        )

        if out_path:
            self._save_result(result, out_path)

        return result

    def _save_result(self, res: GeneticSearchResult, path: str) -> None:
        """Save results to JSON file."""
        def ind_to_dict(ind: Individual) -> Dict[str, Any]:
            return {
                "program": ind.program,
                "gate": ind.gate,
                "regrow_depth": ind.regrow_depth,
                "lattice_D": ind.lattice_D,
                "fitness": round(ind.fitness, 4) if ind.fitness else None,
                "Se": round(ind.Se, 2) if ind.Se else None,
                "logits_mean": round(ind.logits_mean, 4) if ind.logits_mean else None,
                "status": ind.status,
                "eval_count": ind.eval_count,
            }

        out = {
            "task_type": res.task_type,
            "goal": res.goal,
            "total_evaluations": res.total_evaluations,
            "best_overall": ind_to_dict(res.best_overall),
            "generations": [
                {
                    "generation": g.generation,
                    "best_fitness": round(g.best_fitness, 4),
                    "best_program": g.best_program,
                    "best_gate": g.best_gate,
                    "best_regrow_depth": g.best_regrow_depth,
                    "best_lattice_D": g.best_lattice_D,
                    "avg_fitness": round(g.avg_fitness, 4),
                    "diversity": round(g.diversity, 3),
                    "population_size": g.population_size,
                }
                for g in res.generations
            ],
            "final_population": [ind_to_dict(i) for i in res.final_population],
        }

        p = Path(path)
        p.parent.mkdir(parents=True, exist_ok=True)
        p.write_text(json.dumps(out, indent=2), encoding="utf-8")


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

def _cli() -> None:
    import argparse

    parser = argparse.ArgumentParser(
        prog="genetic_program_search",
        description="Phase 12: Genetic Program Search for Helix64 configurations.",
    )
    parser.add_argument(
        "--pheromone-path",
        help="Pheromone JSONL path for recording evaluations.",
    )
    parser.add_argument(
        "--out",
        required=True,
        help="Output JSON summary path.",
    )
    parser.add_argument(
        "--task-type",
        default="experiment",
        help="BloomOS task_type (default: experiment).",
    )
    parser.add_argument(
        "--goal",
        required=True,
        help="Natural language goal/prompt.",
    )
    parser.add_argument(
        "--base-url",
        default="http://localhost:7070",
        help="BloomOS runtime URL (default: http://localhost:7070).",
    )
    parser.add_argument(
        "--timeout",
        type=float,
        default=60.0,
        help="HTTP timeout (default: 60.0).",
    )
    parser.add_argument(
        "--gate",
        action="append",
        dest="gates",
        help="Gate policy to include (can repeat).",
    )
    parser.add_argument(
        "--regrow-depth",
        action="append",
        type=int,
        dest="depths",
        help="Regrow depth to include (can repeat).",
    )
    parser.add_argument(
        "--lattice-D",
        action="append",
        type=int,
        dest="Ds",
        help="Lattice D to include (can repeat).",
    )
    parser.add_argument(
        "--generations",
        type=int,
        default=5,
        help="Number of generations (default: 5).",
    )
    parser.add_argument(
        "--population-size",
        type=int,
        default=16,
        help="Population size (default: 16).",
    )
    parser.add_argument(
        "--mutation-rate",
        type=float,
        default=0.15,
        help="Mutation rate per codon (default: 0.15).",
    )
    parser.add_argument(
        "--crossover-rate",
        type=float,
        default=0.6,
        help="Crossover probability (default: 0.6).",
    )
    parser.add_argument(
        "--elitism",
        type=int,
        default=2,
        help="Number of elites to preserve (default: 2).",
    )
    parser.add_argument(
        "--seed",
        type=int,
        help="Random seed for reproducibility.",
    )
    parser.add_argument(
        "--seed-program",
        action="append",
        dest="seed_programs",
        help="Seed program sequence (can repeat).",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Print configuration without running.",
    )

    args = parser.parse_args()

    # Resolve named programs if possible
    seed_programs: List[str] = []
    if args.seed_programs:
        try:
            from .program_loader import resolve_program
            for sp in args.seed_programs:
                resolved, _name = resolve_program(sp)
                seed_programs.append(resolved)
        except Exception:
            seed_programs = args.seed_programs

    gates = args.gates or ["balanced"]
    depths = args.depths or [64]
    Ds = args.Ds or [72]

    if args.dry_run:
        print("=" * 60)
        print("GENETIC PROGRAM SEARCH - DRY RUN")
        print("=" * 60)
        print(f"Goal: {args.goal}")
        print(f"Task type: {args.task_type}")
        print(f"Gates: {gates}")
        print(f"Depths: {depths}")
        print(f"Lattice Ds: {Ds}")
        print(f"Generations: {args.generations}")
        print(f"Population size: {args.population_size}")
        print(f"Mutation rate: {args.mutation_rate}")
        print(f"Crossover rate: {args.crossover_rate}")
        print(f"Elitism: {args.elitism}")
        print(f"Seed programs: {seed_programs or ['ABCDEFGH']}")
        print(f"Output: {args.out}")
        return

    engine = GeneticProgramSearch(
        base_url=args.base_url,
        pheromone_path=args.pheromone_path,
        task_type=args.task_type,
        goal=args.goal,
        gates=gates,
        regrow_depths=depths,
        lattice_Ds=Ds,
        timeout=args.timeout,
    )

    print(f"Running genetic search for {args.generations} generations...")
    res = engine.run(
        generations=args.generations,
        population_size=args.population_size,
        mutation_rate=args.mutation_rate,
        crossover_rate=args.crossover_rate,
        elitism=args.elitism,
        seed_programs=seed_programs or ["ABCDEFGH"],
        out_path=args.out,
        seed=args.seed,
    )

    print()
    print("=" * 60)
    print("GENETIC PROGRAM SEARCH COMPLETE")
    print("=" * 60)
    print(f"Task type:         {res.task_type}")
    print(f"Goal:              {res.goal}")
    print(f"Total evaluations: {res.total_evaluations}")
    print()
    print("Generation History:")
    for g in res.generations:
        print(
            f"  Gen {g.generation}: best={g.best_fitness:.2f} "
            f"avg={g.avg_fitness:.2f} diversity={g.diversity:.2f} "
            f"program={g.best_program}"
        )
    print()
    print("Best Overall:")
    print(f"  Program:      {res.best_overall.program}")
    print(f"  Gate:         {res.best_overall.gate}")
    print(f"  Regrow depth: {res.best_overall.regrow_depth}")
    print(f"  Lattice D:    {res.best_overall.lattice_D}")
    if res.best_overall.fitness is not None:
        print(f"  Fitness:      {res.best_overall.fitness:.4f}")
    if res.best_overall.Se is not None:
        print(f"  Se:           {res.best_overall.Se:.2f}")
    print()
    print(f"Results saved to: {args.out}")


if __name__ == "__main__":
    _cli()
