#!/usr/bin/env python3
from __future__ import annotations
"""
Phase 13: Multi-Task Co-Evolver
-------------------------------

Given a pheromone JSONL file, this module:

- Looks across *all* task_types (math, code, json, reason, image, audio, …)
- Aggregates per-task stats:
    - best gate
    - best helix_program
    - typical regrow_depth
    - typical lattice_D
    - average fitness, Se, logits_mean
- Proposes a "co-evolution plan":
    - one recommended configuration per task_type
    - global view of which programs/gates are carrying which tasks

This is *analysis-only* (no runtime calls) and sits on top of:
    - PheromoneBoard (pheromone.py)
    - fitness() (fitness.py)
    - Se metrics baked into PheromoneEntry

CLI:
    PYTHONPATH=. python3 -m helix_cdc.bloomos.multitask_coevolver \
        /path/to/pheromones.jsonl \
        --min-samples 3 \
        --pretty

    PYTHONPATH=. python3 -m helix_cdc.bloomos.multitask_coevolver \
        /path/to/pheromones.jsonl \
        --json > /tmp/multitask_plan.json
"""

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

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


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


@dataclass
class TaskConfig:
    """Recommended configuration for a single task_type."""
    task_type: str
    gate: Optional[str]
    helix_program: Optional[str]
    regrow_depth: Optional[int]
    lattice_D: Optional[int]
    avg_fitness: Optional[float]
    avg_Se: Optional[float]
    avg_logits_mean: Optional[float]
    num_samples: int
    notes: List[str] = field(default_factory=list)


@dataclass
class TaskStats:
    """Raw stats per (task_type, gate, helix_program) bucket."""
    task_type: str
    gate: Optional[str]
    helix_program: Optional[str]
    regrow_depths: List[int] = field(default_factory=list)
    lattice_Ds: List[int] = field(default_factory=list)
    fitness_values: List[float] = field(default_factory=list)
    Se_values: List[float] = field(default_factory=list)
    logits_means: List[float] = field(default_factory=list)
    count: int = 0

    def add_entry(self, entry: PheromoneEntry) -> None:
        self.count += 1
        # Fitness
        try:
            f = fitness(entry)
        except Exception:
            f = 0.0
        self.fitness_values.append(f)

        # Se metrics
        if entry.se_Se is not None:
            self.Se_values.append(entry.se_Se)
        # logits_mean
        if entry.logits_mean is not None:
            self.logits_means.append(entry.logits_mean)
        # regrow_depth / lattice_D (Phase 5+)
        if entry.regrow_depth is not None:
            self.regrow_depths.append(entry.regrow_depth)
        if entry.lattice_D is not None:
            self.lattice_Ds.append(entry.lattice_D)

    def _avg(self, xs: List[float]) -> Optional[float]:
        if not xs:
            return None
        return sum(xs) / len(xs)

    def summarize(self) -> TaskConfig:
        avg_f = self._avg(self.fitness_values)
        avg_Se = self._avg(self.Se_values)
        avg_log = self._avg(self.logits_means)
        avg_depth = int(round(self._avg(self.regrow_depths))) if self.regrow_depths else None
        avg_D = int(round(self._avg(self.lattice_Ds))) if self.lattice_Ds else None

        notes: List[str] = []
        if avg_f is None:
            notes.append("no_fitness")
        if self.count < 3:
            notes.append("low_samples")
        if avg_Se is None:
            notes.append("no_Se")

        return TaskConfig(
            task_type=self.task_type,
            gate=self.gate,
            helix_program=self.helix_program,
            regrow_depth=avg_depth,
            lattice_D=avg_D,
            avg_fitness=avg_f,
            avg_Se=avg_Se,
            avg_logits_mean=avg_log,
            num_samples=self.count,
            notes=notes,
        )


@dataclass
class MultiTaskPlan:
    """Global co-evolution plan across all task_types."""
    path: str
    total_entries: int
    task_configs: List[TaskConfig] = field(default_factory=list)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "path": self.path,
            "total_entries": self.total_entries,
            "task_configs": [asdict(cfg) for cfg in self.task_configs],
        }


# ---------------------------------------------------------------------
# Core analysis
# ---------------------------------------------------------------------


def _bucket_key(entry: PheromoneEntry) -> Tuple[str, Optional[str], Optional[str]]:
    return (
        entry.task_type or "unknown",
        entry.gate,
        entry.helix_program,
    )


def analyze_multitask(
    path: str,
    min_samples: int = 3,
) -> MultiTaskPlan:
    """
    Analyze pheromones across all task_types and produce
    a co-evolution plan (best config per task_type).

    Strategy:
        1. Load all pheromones.
        2. Bucket by (task_type, gate, helix_program).
        3. Compute TaskStats per bucket.
        4. For each task_type, pick bucket with highest avg fitness
           (with at least min_samples, if possible).
        5. Return MultiTaskPlan.

    This does *not* mutate or write anything; it's read-only.
    """
    p = Path(path)
    board = PheromoneBoard(path)
    entries = list(board.iter_entries())
    total = len(entries)

    # 1) Build buckets
    buckets: Dict[Tuple[str, Optional[str], Optional[str]], TaskStats] = {}
    for e in entries:
        key = _bucket_key(e)
        if key not in buckets:
            buckets[key] = TaskStats(
                task_type=key[0],
                gate=key[1],
                helix_program=key[2],
            )
        buckets[key].add_entry(e)

    # 2) Group by task_type
    by_task: Dict[str, List[TaskStats]] = {}
    for stats in buckets.values():
        by_task.setdefault(stats.task_type, []).append(stats)

    # 3) Choose best config per task_type
    task_configs: List[TaskConfig] = []

    for task_type, stats_list in sorted(by_task.items()):
        # Summarize each bucket
        cfgs = [s.summarize() for s in stats_list]

        # Prefer configs with >= min_samples; fall back if needed
        eligible = [c for c in cfgs if c.num_samples >= min_samples and c.avg_fitness is not None]
        if not eligible:
            eligible = [c for c in cfgs if c.avg_fitness is not None]

        if not eligible:
            # no fitness at all – pick arbitrary, mark as such
            if cfgs:
                c = cfgs[0]
                c.notes.append("no_fitness_for_task")
                task_configs.append(c)
            continue

        # Pick highest fitness
        best = max(eligible, key=lambda c: c.avg_fitness or 0.0)

        # Annotate with a couple of qualitative notes
        if best.avg_Se is not None and best.avg_Se > 100:
            best.notes.append("high_symbolic_entropy")
        if best.avg_logits_mean is not None and best.avg_logits_mean > 0.5:
            best.notes.append("high_logits_mean")

        task_configs.append(best)

    return MultiTaskPlan(
        path=str(p),
        total_entries=total,
        task_configs=task_configs,
    )


# ---------------------------------------------------------------------
# Pretty-print
# ---------------------------------------------------------------------


def _print_pretty(plan: MultiTaskPlan) -> None:
    print("=" * 60)
    print(f"MULTI-TASK CO-EVOLVER PLAN - {plan.path}")
    print("=" * 60)
    print(f"Total pheromone entries: {plan.total_entries}")
    print(f"Task types covered:      {len(plan.task_configs)}")
    print()

    if not plan.task_configs:
        print("No task configurations found.")
        return

    # per-task table
    header = (
        "task_type",
        "gate",
        "program",
        "depth",
        "D",
        "fitness",
        "Se",
        "logits",
        "n",
        "notes",
    )
    print(f"{header[0]:<10} {header[1]:<10} {header[2]:<10} {header[3]:>5} {header[4]:>4} "
          f"{header[5]:>8} {header[6]:>8} {header[7]:>8} {header[8]:>4} {header[9]}")
    print("-" * 100)

    for cfg in sorted(plan.task_configs, key=lambda c: c.task_type):
        fitness_str = f"{cfg.avg_fitness:.3f}" if cfg.avg_fitness is not None else "  -   "
        Se_str = f"{cfg.avg_Se:.1f}" if cfg.avg_Se is not None else "  -  "
        log_str = f"{cfg.avg_logits_mean:.3f}" if cfg.avg_logits_mean is not None else "  -   "
        depth_str = str(cfg.regrow_depth) if cfg.regrow_depth is not None else "-"
        D_str = str(cfg.lattice_D) if cfg.lattice_D is not None else "-"
        notes_str = ",".join(cfg.notes) if cfg.notes else ""

        print(
            f"{cfg.task_type:<10} "
            f"{(cfg.gate or '-'): <10} "
            f"{(cfg.helix_program or '-'): <10} "
            f"{depth_str:>5} {D_str:>4} "
            f"{fitness_str:>8} {Se_str:>8} {log_str:>8} "
            f"{cfg.num_samples:>4} {notes_str}"
        )


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


def _cli() -> None:
    import argparse

    parser = argparse.ArgumentParser(
        description="Phase 13: Multi-Task Co-Evolver (analyze pheromones across all tasks)."
    )
    parser.add_argument(
        "path",
        help="Pheromone JSONL path (e.g., /tmp/pheromones.jsonl)",
    )
    parser.add_argument(
        "--min-samples",
        type=int,
        default=3,
        help="Minimum samples per (task,gate,program) bucket to be considered (default: 3).",
    )
    parser.add_argument(
        "--json",
        action="store_true",
        help="Emit machine-readable JSON instead of pretty table.",
    )
    parser.add_argument(
        "--pretty",
        action="store_true",
        help="Force pretty-print table output.",
    )

    args = parser.parse_args()

    plan = analyze_multitask(args.path, min_samples=args.min_samples)

    if args.json and not args.pretty:
        print(json.dumps(plan.to_dict(), indent=2))
    else:
        _print_pretty(plan)


if __name__ == "__main__":
    _cli()
