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

"""
Phase 8: Substrate Inspector

Read pheromone history + grammar/program metadata and produce a
human- and machine-readable summary of the current substrate state:

- Which gates are winning per task_type?
- Which helix_programs are being used and how often?
- What are the average Se / fitness values per (task_type, gate, program)?
- Are there any anomalous configurations (low fitness, high error rate)?

This is read-only analysis: no mutation, no writes.

Usage:
    PYTHONPATH=. python3 -m helix_cdc.bloomos.substrate_inspector /tmp/pheromones.jsonl --pretty
    PYTHONPATH=. python3 -m helix_cdc.bloomos.substrate_inspector /tmp/pheromones.jsonl --task-type math
"""

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

from .pheromone import PheromoneBoard, PheromoneEntry
from .fitness import fitness
from .program_loader import list_programs
from .programs import GATE_SEMANTICS

# Phase 16: Superglyph integration
from helix_cdc.superglyph_bridge import describe_superglyph_substrate


@dataclass
class GateStats:
    """Statistics for a single gate policy."""
    gate: str
    num_samples: int
    success_rate: float
    avg_fitness: Optional[float]
    avg_Se: Optional[float]
    avg_logits: Optional[float]
    avg_seed_entropy: Optional[float]


@dataclass
class ProgramStats:
    """Statistics for a single helix_program."""
    sequence: str
    name: Optional[str]
    num_samples: int
    success_rate: float
    avg_fitness: Optional[float]
    avg_Se: Optional[float]
    avg_logits: Optional[float]
    avg_seed_entropy: Optional[float]


@dataclass
class TaskSummary:
    """Summary for a single task_type."""
    task_type: str
    total_samples: int
    success_rate: float
    gates: List[GateStats]
    programs: List[ProgramStats]
    best_gate: Optional[str]
    best_program: Optional[str]


@dataclass
class SubstrateSummary:
    """Complete substrate analysis summary."""
    path: str
    total_entries: int
    unique_task_types: int
    unique_gates: int
    unique_programs: int
    task_summaries: List[TaskSummary]
    anomalies: List[Dict[str, Any]]
    # Phase 16: Physical substrate description
    superglyph: Dict[str, Any] = field(default_factory=dict)


# --------------------------
# Internal helpers
# --------------------------

def _iter_entries(path: Path) -> Iterable[PheromoneEntry]:
    """Iterate over all entries in pheromone file."""
    board = PheromoneBoard(str(path))
    return board.iter_entries()


def _group_by_task(entries: Iterable[PheromoneEntry]) -> Dict[str, List[PheromoneEntry]]:
    """Group entries by task_type."""
    buckets: Dict[str, List[PheromoneEntry]] = {}
    for e in entries:
        tt = e.task_type or "unknown"
        buckets.setdefault(tt, []).append(e)
    return buckets


def _safe_avg(vals: List[float]) -> Optional[float]:
    """Compute average, returning None for empty list."""
    if not vals:
        return None
    return sum(vals) / len(vals)


def _compute_success_rate(entries: List[PheromoneEntry]) -> float:
    """Compute success rate for a list of entries."""
    if not entries:
        return 0.0
    ok_count = sum(1 for e in entries if e.status == "ok")
    return ok_count / len(entries)


def _build_gate_stats(entries: List[PheromoneEntry]) -> List[GateStats]:
    """Build statistics for each gate from entries."""
    by_gate: Dict[str, List[PheromoneEntry]] = {}
    for e in entries:
        gate = e.gate or "unknown"
        by_gate.setdefault(gate, []).append(e)

    stats: List[GateStats] = []
    for gate, es in by_gate.items():
        fitness_vals: List[float] = []
        se_vals: List[float] = []
        logits_vals: List[float] = []
        entropy_vals: List[float] = []

        for e in es:
            try:
                fitness_vals.append(fitness(e))
            except Exception:
                pass
            if e.se_Se is not None:
                se_vals.append(e.se_Se)
            if e.logits_mean is not None:
                logits_vals.append(e.logits_mean)
            if e.seed_entropy is not None:
                entropy_vals.append(e.seed_entropy)

        stats.append(
            GateStats(
                gate=gate,
                num_samples=len(es),
                success_rate=_compute_success_rate(es),
                avg_fitness=_safe_avg(fitness_vals),
                avg_Se=_safe_avg(se_vals),
                avg_logits=_safe_avg(logits_vals),
                avg_seed_entropy=_safe_avg(entropy_vals),
            )
        )

    # Sort by avg_fitness descending
    stats.sort(key=lambda s: (s.avg_fitness or -math.inf), reverse=True)
    return stats


def _build_program_stats(entries: List[PheromoneEntry]) -> List[ProgramStats]:
    """Build statistics for each helix_program from entries."""
    by_prog: Dict[str, List[PheromoneEntry]] = {}
    for e in entries:
        seq = e.helix_program or "UNKNOWN"
        by_prog.setdefault(seq, []).append(e)

    # Map sequences to names (from helix_programs.json if available)
    name_map: Dict[str, str] = {}
    try:
        named = list_programs()
        for name, seq in named.items():
            name_map[seq] = name
    except Exception:
        pass

    stats: List[ProgramStats] = []
    for seq, es in by_prog.items():
        fitness_vals: List[float] = []
        se_vals: List[float] = []
        logits_vals: List[float] = []
        entropy_vals: List[float] = []

        for e in es:
            try:
                fitness_vals.append(fitness(e))
            except Exception:
                pass
            if e.se_Se is not None:
                se_vals.append(e.se_Se)
            if e.logits_mean is not None:
                logits_vals.append(e.logits_mean)
            if e.seed_entropy is not None:
                entropy_vals.append(e.seed_entropy)

        stats.append(
            ProgramStats(
                sequence=seq,
                name=name_map.get(seq),
                num_samples=len(es),
                success_rate=_compute_success_rate(es),
                avg_fitness=_safe_avg(fitness_vals),
                avg_Se=_safe_avg(se_vals),
                avg_logits=_safe_avg(logits_vals),
                avg_seed_entropy=_safe_avg(entropy_vals),
            )
        )

    stats.sort(key=lambda s: (s.avg_fitness or -math.inf), reverse=True)
    return stats


def _detect_anomalies(
    task_summaries: List[TaskSummary],
    fitness_threshold: float = 10.0,
    success_threshold: float = 0.5,
) -> List[Dict[str, Any]]:
    """
    Detect anomalous configurations:
    - Gates/programs with low fitness
    - Gates/programs with low success rate
    """
    anomalies: List[Dict[str, Any]] = []

    for ts in task_summaries:
        # Check gates
        for g in ts.gates:
            if g.avg_fitness is not None and g.avg_fitness < fitness_threshold:
                anomalies.append({
                    "type": "low_fitness_gate",
                    "task_type": ts.task_type,
                    "gate": g.gate,
                    "avg_fitness": round(g.avg_fitness, 3),
                    "samples": g.num_samples,
                })
            if g.success_rate < success_threshold and g.num_samples >= 3:
                anomalies.append({
                    "type": "low_success_gate",
                    "task_type": ts.task_type,
                    "gate": g.gate,
                    "success_rate": round(g.success_rate, 3),
                    "samples": g.num_samples,
                })

        # Check programs
        for p in ts.programs:
            if p.avg_fitness is not None and p.avg_fitness < fitness_threshold:
                anomalies.append({
                    "type": "low_fitness_program",
                    "task_type": ts.task_type,
                    "program": p.sequence,
                    "name": p.name,
                    "avg_fitness": round(p.avg_fitness, 3),
                    "samples": p.num_samples,
                })
            if p.success_rate < success_threshold and p.num_samples >= 3:
                anomalies.append({
                    "type": "low_success_program",
                    "task_type": ts.task_type,
                    "program": p.sequence,
                    "name": p.name,
                    "success_rate": round(p.success_rate, 3),
                    "samples": p.num_samples,
                })

    return anomalies


def analyze_substrate(path: str) -> SubstrateSummary:
    """
    Analyze substrate state from pheromone file.

    Args:
        path: Path to pheromone JSONL file

    Returns:
        SubstrateSummary with complete analysis
    """
    file_path = Path(path)
    entries = list(_iter_entries(file_path))
    by_task = _group_by_task(entries)

    # Collect unique values
    all_gates: set = set()
    all_programs: set = set()

    task_summaries: List[TaskSummary] = []
    for task_type, es in by_task.items():
        gates = _build_gate_stats(es)
        progs = _build_program_stats(es)

        # Track unique values
        for g in gates:
            all_gates.add(g.gate)
        for prog in progs:
            all_programs.add(prog.sequence)

        # Determine best gate and program
        best_gate = gates[0].gate if gates else None
        best_program = progs[0].sequence if progs else None

        task_summaries.append(
            TaskSummary(
                task_type=task_type,
                total_samples=len(es),
                success_rate=_compute_success_rate(es),
                gates=gates,
                programs=progs,
                best_gate=best_gate,
                best_program=best_program,
            )
        )

    # Detect anomalies
    anomalies = _detect_anomalies(task_summaries)

    # Phase 16: Get Superglyph physical substrate descriptor
    superglyph_desc = describe_superglyph_substrate()

    return SubstrateSummary(
        path=str(file_path),
        total_entries=len(entries),
        unique_task_types=len(by_task),
        unique_gates=len(all_gates),
        unique_programs=len(all_programs),
        task_summaries=task_summaries,
        anomalies=anomalies,
        superglyph=superglyph_desc.to_dict(),
    )


def get_known_gates() -> List[str]:
    """Get list of known gate policies from GATE_SEMANTICS."""
    return list(GATE_SEMANTICS.keys())


def get_known_programs() -> Dict[str, str]:
    """Get list of known programs from helix_programs.json."""
    try:
        return list_programs()
    except Exception:
        return {}


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

def _cli() -> None:
    import argparse

    parser = argparse.ArgumentParser(
        prog="substrate_inspector",
        description="Phase 8: Inspect substrate state from pheromone JSONL.",
    )
    parser.add_argument(
        "path",
        help="Path to pheromone JSONL file.",
    )
    parser.add_argument(
        "--task-type",
        help="Optional: print only summary for a single task_type.",
    )
    parser.add_argument(
        "--pretty",
        action="store_true",
        help="Pretty-print summary instead of raw JSON.",
    )
    parser.add_argument(
        "--show-anomalies",
        action="store_true",
        help="Highlight anomalies in pretty output.",
    )
    args = parser.parse_args()

    summary = analyze_substrate(args.path)
    data = asdict(summary)

    if args.task_type:
        filtered_tasks = [
            ts for ts in data["task_summaries"] if ts["task_type"] == args.task_type
        ]
        data["task_summaries"] = filtered_tasks

    if args.pretty:
        # Human-readable view
        print(f"{'=' * 60}")
        print(f"SUBSTRATE INSPECTOR - {data['path']}")
        print(f"{'=' * 60}")
        print(f"Total entries:      {data['total_entries']}")
        print(f"Unique task_types:  {data['unique_task_types']}")
        print(f"Unique gates:       {data['unique_gates']}")
        print(f"Unique programs:    {data['unique_programs']}")
        print()

        for ts in data["task_summaries"]:
            print(f"=== Task: {ts['task_type']} ===")
            print(f"Samples: {ts['total_samples']}  |  Success rate: {ts['success_rate']:.1%}")
            print(f"Best gate: {ts['best_gate']}  |  Best program: {ts['best_program']}")
            print()

            print("  Gates (ranked by fitness):")
            for g in ts["gates"]:
                fitness_str = f"{g['avg_fitness']:.2f}" if g['avg_fitness'] is not None else "N/A"
                se_str = f"{g['avg_Se']:.1f}" if g['avg_Se'] is not None else "N/A"
                logits_str = f"{g['avg_logits']:.3f}" if g['avg_logits'] is not None else "N/A"
                print(
                    f"    {g['gate']:<15} n={g['num_samples']:<4} "
                    f"success={g['success_rate']:.0%}  "
                    f"fitness={fitness_str:<8} Se={se_str:<8} logits={logits_str}"
                )
            print()

            print("  Programs (ranked by fitness):")
            for p in ts["programs"]:
                name = p['name'] or "custom"
                fitness_str = f"{p['avg_fitness']:.2f}" if p['avg_fitness'] is not None else "N/A"
                se_str = f"{p['avg_Se']:.1f}" if p['avg_Se'] is not None else "N/A"
                logits_str = f"{p['avg_logits']:.3f}" if p['avg_logits'] is not None else "N/A"
                print(
                    f"    {name:<15} ({p['sequence']}) "
                    f"n={p['num_samples']:<4} success={p['success_rate']:.0%}  "
                    f"fitness={fitness_str:<8} Se={se_str:<8} logits={logits_str}"
                )
            print()

        if args.show_anomalies and data["anomalies"]:
            print("=" * 60)
            print("ANOMALIES DETECTED")
            print("=" * 60)
            for a in data["anomalies"]:
                print(f"  [{a['type']}] task={a.get('task_type')} ", end="")
                if "gate" in a:
                    print(f"gate={a['gate']} ", end="")
                if "program" in a:
                    print(f"program={a['program']} ({a.get('name', '?')}) ", end="")
                if "avg_fitness" in a:
                    print(f"fitness={a['avg_fitness']}", end="")
                if "success_rate" in a:
                    print(f"success={a['success_rate']:.0%}", end="")
                print()
            print()

        # Phase 16: Superglyph physical substrate
        sg = data.get("superglyph", {})
        if sg:
            print("=" * 60)
            print("PHYSICAL SUBSTRATE: Superglyph")
            print("=" * 60)
            print(f"  name:      {sg.get('name')}")
            print(f"  version:   {sg.get('version')}")
            caps = sg.get("capabilities", {})
            print(f"  genome:          {caps.get('has_genome')}")
            print(f"  lab_tools:       {caps.get('has_lab_tools')}")
            print(f"  agentic_cells:   {caps.get('has_agentic_cells')}")
            print(f"  swarm_scripts:   {caps.get('has_swarm_scripts')}")
            print(f"  regen_core:      {caps.get('has_regen_core')}")
            genome = sg.get("genome", {})
            if genome.get("path"):
                print(f"  genome_path:     {genome.get('path')}")
                if genome.get("size_bytes"):
                    print(f"  genome_size:     {genome.get('size_bytes')} bytes")
            print()
    else:
        # Round floats for cleaner JSON output
        def _round_floats(obj: Any) -> Any:
            if isinstance(obj, float):
                return round(obj, 4)
            elif isinstance(obj, dict):
                return {k: _round_floats(v) for k, v in obj.items()}
            elif isinstance(obj, list):
                return [_round_floats(v) for v in obj]
            return obj

        print(json.dumps(_round_floats(data), indent=2))


if __name__ == "__main__":
    _cli()
