#!/usr/bin/env python3
from __future__ import annotations
"""
Phase 11: Auto-Experimenter
---------------------------
Runs batched experiments over:
    - gate policies
    - helix_program sequences
    - regrow depths
    - lattice_D settings

For each combination:
    - Runs the task with BloomOSClient
    - Computes Se, fitness
    - Records pheromone entries
    - Bundles receipts into an experiment pack

Outputs:
    - ExperimentResult objects
    - ExperimentSummary
    - JSONL receipts
    - Compatible with EvolutionEngine + ProgramController

Usage:
    PYTHONPATH=. python3 -m helix_cdc.bloomos.auto_experimenter \\
      /tmp/experiment_results.json \\
      --goal "integrate sin(x) from 0 to pi" \\
      --pheromone /tmp/pheromones.jsonl \\
      --gate balanced --gate math_bias \\
      --program ABCDEFGH --program BGGDDGEA \\
      --depth 64 --depth 128 \\
      --D 48 --D 72
"""

import json
import time
from dataclasses import dataclass, field, asdict
from pathlib import Path
from typing import Any, Dict, List, Optional, Iterator
from itertools import product

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


# ------------------------------------------------------
# DATA MODELS
# ------------------------------------------------------

@dataclass
class ExperimentConfig:
    """Configuration for a single experiment run."""
    gate: str
    program: str
    regrow_depth: int
    lattice_D: int


@dataclass
class ExperimentResult:
    """Result of a single experiment run."""
    idx: int
    goal: str
    task_type: str
    gate: str
    program: str
    regrow_depth: int
    lattice_D: int
    status: str
    logits_mean: Optional[float]
    logits_std: Optional[float]
    se_H: Optional[float]
    se_C: Optional[float]
    se_D: Optional[float]
    se_Se: Optional[float]
    fitness_score: Optional[float]
    elapsed_ms: float
    text_output: Optional[str] = None


@dataclass
class ExperimentSummary:
    """Summary of all experiment runs."""
    total: int
    successful: int
    failed: int
    avg_fitness: Optional[float]
    avg_Se: Optional[float]
    avg_logits: Optional[float]
    best_fitness: Optional[float]
    best_config: Optional[Dict[str, Any]]
    worst_fitness: Optional[float]
    worst_config: Optional[Dict[str, Any]]


# ------------------------------------------------------
# HELPER FUNCTIONS
# ------------------------------------------------------

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)


def _compute_fitness_from_result(result: ExperimentResult) -> Optional[float]:
    """Compute fitness from experiment result."""
    class FakeEntry:
        def __init__(self, r: ExperimentResult):
            self.se_Se = r.se_Se
            self.logits_mean = r.logits_mean
            self.seed_entropy = 5.0  # Default mid-range

    try:
        return fitness(FakeEntry(result))
    except Exception:
        return None


def _generate_configs(
    gates: List[str],
    programs: List[str],
    regrow_depths: List[int],
    lattice_Ds: List[int],
) -> Iterator[ExperimentConfig]:
    """Generate all experiment configurations."""
    for gate, prog, depth, D in product(gates, programs, regrow_depths, lattice_Ds):
        yield ExperimentConfig(
            gate=gate,
            program=prog,
            regrow_depth=depth,
            lattice_D=D,
        )


# ------------------------------------------------------
# AUTO EXPERIMENTER
# ------------------------------------------------------

class AutoExperimenter:
    """
    Auto-Experimenter for systematic parameter sweeps.

    Runs experiments across combinations of gates, programs,
    regrow_depth, and lattice_D, recording results and pheromones.
    """

    def __init__(
        self,
        base_url: str = "http://localhost:7070",
        pheromone_path: Optional[str] = None,
        out_path: Optional[str] = None,
        timeout: float = 60.0,
    ) -> None:
        """
        Initialize the auto-experimenter.

        Args:
            base_url: BloomOS runtime URL
            pheromone_path: Path to pheromone JSONL (for recording)
            out_path: Path to output experiment results JSON
            timeout: HTTP timeout
        """
        self.client = BloomOSClient(base_url=base_url, timeout=timeout)
        self.pheromones = PheromoneBoard(pheromone_path) if pheromone_path else None
        self.out_path = Path(out_path) if out_path else None
        if self.out_path:
            self.out_path.parent.mkdir(parents=True, exist_ok=True)

    def run_single(
        self,
        goal: str,
        task_type: str,
        config: ExperimentConfig,
        idx: int = 0,
    ) -> ExperimentResult:
        """
        Run a single experiment.

        Args:
            goal: Natural language task description
            task_type: Task type (math, code, json, etc.)
            config: Experiment configuration
            idx: Experiment index

        Returns:
            ExperimentResult
        """
        # Build extra payload
        extra: Dict[str, Any] = {
            "helix_program": config.program,
            "regrow_depth": config.regrow_depth,
            "lattice_D": config.lattice_D,
        }

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

            raw = result.raw
            status = result.status
            logits_mean = result.logits_mean
            logits_std = result.logits_std

            # Extract text output
            text_out = _extract_output_text(raw)

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

        except Exception as e:
            elapsed_ms = (time.perf_counter() - t0) * 1000.0
            return ExperimentResult(
                idx=idx,
                goal=goal,
                task_type=task_type,
                gate=config.gate,
                program=config.program,
                regrow_depth=config.regrow_depth,
                lattice_D=config.lattice_D,
                status=f"error:{type(e).__name__}",
                logits_mean=None,
                logits_std=None,
                se_H=None,
                se_C=None,
                se_D=None,
                se_Se=None,
                fitness_score=None,
                elapsed_ms=elapsed_ms,
                text_output=None,
            )

        exp_result = ExperimentResult(
            idx=idx,
            goal=goal,
            task_type=task_type,
            gate=config.gate,
            program=config.program,
            regrow_depth=config.regrow_depth,
            lattice_D=config.lattice_D,
            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,
            fitness_score=None,  # Computed below
            elapsed_ms=elapsed_ms,
            text_output=text_out[:500] if text_out else None,  # Truncate for storage
        )

        # Compute fitness
        exp_result.fitness_score = _compute_fitness_from_result(exp_result)

        # Record pheromone
        if self.pheromones:
            try:
                entry = make_entry_from_step(
                    task_type=task_type,
                    gate=config.gate,
                    helix_program=config.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": goal,
                        "elapsed_ms": elapsed_ms,
                        "experiment_idx": idx,
                    },
                    regrow_depth=config.regrow_depth,
                    lattice_D=config.lattice_D,
                )
                self.pheromones.record(entry)
            except Exception:
                pass

        return exp_result

    def run_sweep(
        self,
        goal: str,
        task_type: str = "experiment",
        gates: Optional[List[str]] = None,
        programs: Optional[List[str]] = None,
        regrow_depths: Optional[List[int]] = None,
        lattice_Ds: Optional[List[int]] = None,
        limit: Optional[int] = None,
    ) -> tuple[List[ExperimentResult], ExperimentSummary]:
        """
        Run a parameter sweep experiment.

        Args:
            goal: Natural language task description
            task_type: Task type
            gates: List of gate policies to test
            programs: List of helix_programs to test
            regrow_depths: List of regrow depths to test
            lattice_Ds: List of lattice_D values to test
            limit: Maximum number of experiments to run

        Returns:
            Tuple of (results list, summary)
        """
        # Defaults
        gates = gates or ["balanced"]
        programs = programs or ["ABCDEFGH"]
        regrow_depths = regrow_depths or [64]
        lattice_Ds = lattice_Ds or [72]

        results: List[ExperimentResult] = []
        idx = 0

        for config in _generate_configs(gates, programs, regrow_depths, lattice_Ds):
            if limit and idx >= limit:
                break

            result = self.run_single(goal, task_type, config, idx)
            results.append(result)
            idx += 1

        # Compute summary
        summary = self._compute_summary(results)

        # Save results if out_path is set
        if self.out_path:
            self._save_results(results, summary)

        return results, summary

    def _compute_summary(self, results: List[ExperimentResult]) -> ExperimentSummary:
        """Compute summary statistics from results."""
        if not results:
            return ExperimentSummary(
                total=0,
                successful=0,
                failed=0,
                avg_fitness=None,
                avg_Se=None,
                avg_logits=None,
                best_fitness=None,
                best_config=None,
                worst_fitness=None,
                worst_config=None,
            )

        successful = [r for r in results if r.status == "ok"]
        failed = [r for r in results if r.status != "ok"]

        fitness_vals = [r.fitness_score for r in successful if r.fitness_score is not None]
        se_vals = [r.se_Se for r in successful if r.se_Se is not None]
        logits_vals = [r.logits_mean for r in successful if r.logits_mean is not None]

        best_result = max(successful, key=lambda r: r.fitness_score or 0) if successful else None
        worst_result = min(successful, key=lambda r: r.fitness_score or float('inf')) if successful else None

        return ExperimentSummary(
            total=len(results),
            successful=len(successful),
            failed=len(failed),
            avg_fitness=sum(fitness_vals) / len(fitness_vals) if fitness_vals else None,
            avg_Se=sum(se_vals) / len(se_vals) if se_vals else None,
            avg_logits=sum(logits_vals) / len(logits_vals) if logits_vals else None,
            best_fitness=best_result.fitness_score if best_result else None,
            best_config={
                "gate": best_result.gate,
                "program": best_result.program,
                "regrow_depth": best_result.regrow_depth,
                "lattice_D": best_result.lattice_D,
            } if best_result else None,
            worst_fitness=worst_result.fitness_score if worst_result else None,
            worst_config={
                "gate": worst_result.gate,
                "program": worst_result.program,
                "regrow_depth": worst_result.regrow_depth,
                "lattice_D": worst_result.lattice_D,
            } if worst_result else None,
        )

    def _save_results(
        self, results: List[ExperimentResult], summary: ExperimentSummary
    ) -> None:
        """Save results and summary to JSON file."""
        output = {
            "summary": asdict(summary),
            "results": [asdict(r) for r in results],
        }
        with open(self.out_path, "w", encoding="utf-8") as f:
            json.dump(output, f, indent=2)


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

def _cli() -> None:
    import argparse

    parser = argparse.ArgumentParser(
        prog="auto_experimenter",
        description="Phase 11: Auto-Experimenter for systematic parameter sweeps.",
    )
    parser.add_argument(
        "out",
        help="Output experiment results JSON path.",
    )
    parser.add_argument(
        "--goal",
        required=True,
        help="Natural language goal/task to run.",
    )
    parser.add_argument(
        "--task-type",
        default="experiment",
        help="Task type (default: experiment).",
    )
    parser.add_argument(
        "--pheromone",
        help="Path to pheromone JSONL file for recording.",
    )
    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 test (can be repeated).",
    )
    parser.add_argument(
        "--program",
        action="append",
        dest="programs",
        help="Helix program to test (can be repeated).",
    )
    parser.add_argument(
        "--depth",
        action="append",
        type=int,
        dest="depths",
        help="Regrow depth to test (can be repeated).",
    )
    parser.add_argument(
        "--D",
        action="append",
        type=int,
        dest="lattice_Ds",
        help="Lattice D to test (can be repeated).",
    )
    parser.add_argument(
        "--limit",
        type=int,
        help="Maximum number of experiments to run.",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Print experiment plan without running.",
    )
    parser.add_argument(
        "--json",
        action="store_true",
        help="Output summary as JSON.",
    )

    args = parser.parse_args()

    # Defaults
    gates = args.gates or ["balanced"]
    programs = args.programs or ["ABCDEFGH"]
    depths = args.depths or [64]
    lattice_Ds = args.lattice_Ds or [72]

    # Compute total experiments
    total = len(gates) * len(programs) * len(depths) * len(lattice_Ds)
    if args.limit:
        total = min(total, args.limit)

    if args.dry_run:
        print("=" * 60)
        print("AUTO-EXPERIMENTER DRY RUN")
        print("=" * 60)
        print(f"Goal: {args.goal}")
        print(f"Task type: {args.task_type}")
        print(f"Gates: {gates}")
        print(f"Programs: {programs}")
        print(f"Depths: {depths}")
        print(f"Lattice Ds: {lattice_Ds}")
        print(f"Total experiments: {total}")
        print(f"Output: {args.out}")
        print(f"Pheromone: {args.pheromone}")
        return

    experimenter = AutoExperimenter(
        base_url=args.base_url,
        pheromone_path=args.pheromone,
        out_path=args.out,
        timeout=args.timeout,
    )

    print(f"Running {total} experiments...")
    results, summary = experimenter.run_sweep(
        goal=args.goal,
        task_type=args.task_type,
        gates=gates,
        programs=programs,
        regrow_depths=depths,
        lattice_Ds=lattice_Ds,
        limit=args.limit,
    )

    if args.json:
        print(json.dumps(asdict(summary), indent=2))
    else:
        print()
        print("=" * 60)
        print("EXPERIMENT SUMMARY")
        print("=" * 60)
        print(f"Total:           {summary.total}")
        print(f"Successful:      {summary.successful}")
        print(f"Failed:          {summary.failed}")
        if summary.avg_fitness is not None:
            print(f"Avg Fitness:     {summary.avg_fitness:.4f}")
        if summary.avg_Se is not None:
            print(f"Avg Se:          {summary.avg_Se:.2f}")
        if summary.avg_logits is not None:
            print(f"Avg Logits:      {summary.avg_logits:.4f}")
        print()
        if summary.best_config:
            print("Best Config:")
            print(f"  Gate:        {summary.best_config['gate']}")
            print(f"  Program:     {summary.best_config['program']}")
            print(f"  Depth:       {summary.best_config['regrow_depth']}")
            print(f"  Lattice D:   {summary.best_config['lattice_D']}")
            print(f"  Fitness:     {summary.best_fitness:.4f}")
        print()
        print(f"Results saved to: {args.out}")


if __name__ == "__main__":
    _cli()
