#!/usr/bin/env python3
from __future__ import annotations
"""
Phase 14: Capsule-Space Evolver
-------------------------------
Operate directly on SymCapsules recorded in pheromone JSONL logs
(or standalone capsule JSONL files) to discover "archetype capsules"
per task_type / gate / helix_program.

This is **standalone**: it does NOT mutate or patch any existing
runtime modules. It only:

- Reads a JSONL file with entries that include either:
    - a "capsule" field (SymCapsule.to_dict())
    - or a raw SymCapsule dict (if using a capsule-only file)
- Extracts:
    - task_type (if present)
    - gate
    - helix_program
    - Se metrics (H, C, D, Se)
    - logits_mean
    - fitness (via fitness.fitness() on pheromone-style entries)
- Clusters capsules inside each (task_type, gate, program) bucket
  using a simple similarity metric on tokenized text.
- Produces:
    - CapsuleArchetype objects representing cluster centers
    - A JSON summary file suitable for downstream tools.

Usage (pretty print):
    PYTHONPATH=. python3 -m helix_cdc.bloomos.capsule_evolver \\
        /tmp/pheromones.jsonl --pretty

Usage (JSON output only):
    PYTHONPATH=. python3 -m helix_cdc.bloomos.capsule_evolver \\
        /tmp/pheromones.jsonl --out /tmp/capsule_archetypes.json --json

The input format is intentionally forgiving; it will do its best with:
- pheromone entries created by PheromoneBoard.make_entry_from_step()
- EchoBridge cognitive trace entries (if they contain 'capsule')
- pure SymCapsule JSONL dumps

This file is designed to be Phase-14-only: no changes to other modules.
"""

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

from .fitness import fitness as fitness_fn

if TYPE_CHECKING:
    # For type hints only; avoids runtime import issues
    from helix_cdc.bloomos.sym import SymCapsule


# ---------------------------------------------------------------------------
# Dataclasses
# ---------------------------------------------------------------------------


@dataclass
class CapsuleRecord:
    """Single capsule-bearing record from a pheromone or capsule JSONL file."""

    index: int  # line index in the source file
    task_type: str
    gate: Optional[str]
    helix_program: Optional[str]

    se_H: Optional[float]
    se_C: Optional[float]
    se_D: Optional[float]
    se_Se: Optional[float]

    logits_mean: Optional[float]
    logits_std: Optional[float]

    fitness: Optional[float]

    capsule: Dict[str, Any]
    raw: Dict[str, Any] = field(default_factory=dict)


@dataclass
class CapsuleCluster:
    """Internal cluster representation used during evolution."""

    cluster_id: str
    task_type: str
    gate: Optional[str]
    helix_program: Optional[str]

    indices: List[int] = field(default_factory=list)
    avg_Se: Optional[float] = None
    avg_fitness: Optional[float] = None
    avg_logits_mean: Optional[float] = None

    representative_index: Optional[int] = None
    representative_capsule: Optional[Dict[str, Any]] = None
    keywords: List[str] = field(default_factory=list)


@dataclass
class CapsuleArchetype:
    """Public archetype structure, suitable for JSON output."""

    id: str
    task_type: str
    gate: Optional[str]
    helix_program: Optional[str]

    num_samples: int
    avg_Se: Optional[float]
    avg_fitness: Optional[float]
    avg_logits_mean: Optional[float]

    representative_index: Optional[int]
    representative_capsule: Optional[Dict[str, Any]]
    keywords: List[str]


@dataclass
class CapsuleEvolutionSummary:
    """Top-level summary for the whole file."""

    source_path: str
    total_records: int
    task_types: List[str]
    archetypes: List[CapsuleArchetype]


# ---------------------------------------------------------------------------
# Utility functions: text & similarity
# ---------------------------------------------------------------------------


_STOPWORDS: Set[str] = {
    "a", "an", "the", "and", "or", "of", "in", "to", "for", "on", "with",
    "this", "that", "these", "those", "is", "are", "was", "were", "be",
    "as", "by", "at", "from", "it", "its", "into", "over", "under", "about",
}


def _safe_lower(x: Any) -> str:
    try:
        return str(x).lower()
    except Exception:
        return ""


def _extract_text_from_capsule_dict(capsule: Dict[str, Any]) -> str:
    """
    Try to extract the most relevant text from a SymCapsule dict.
    We look into:
        - capsule["data"].get("text") / "content"
        - capsule["meaning"] (if string or dict with "summary"/"intent")
    Fallback: JSON dump of the capsule.
    """
    if not isinstance(capsule, dict):
        return str(capsule)

    data = capsule.get("data")
    if isinstance(data, dict):
        for key in ("text", "content", "response", "output"):
            if key in data and isinstance(data[key], str):
                return data[key]

    meaning = capsule.get("meaning")
    if isinstance(meaning, str):
        return meaning
    if isinstance(meaning, dict):
        for key in ("summary", "intent", "description"):
            if key in meaning and isinstance(meaning[key], str):
                return meaning[key]

    # Fallback: stringify
    try:
        return json.dumps(capsule, ensure_ascii=False)
    except Exception:
        return str(capsule)


def _tokenize(text: str) -> List[str]:
    """Very simple tokenizer for similarity; lowercases, splits on non-alnum."""
    if not text:
        return []
    out: List[str] = []
    current = []
    for ch in text.lower():
        if ch.isalnum() or ch in "_":
            current.append(ch)
        else:
            if current:
                token = "".join(current)
                if token not in _STOPWORDS:
                    out.append(token)
                current = []
    if current:
        token = "".join(current)
        if token not in _STOPWORDS:
            out.append(token)
    return out


def _jaccard_similarity(a: Set[str], b: Set[str]) -> float:
    if not a and not b:
        return 1.0
    if not a or not b:
        return 0.0
    inter = len(a & b)
    union = len(a | b)
    if union == 0:
        return 0.0
    return inter / union


# ---------------------------------------------------------------------------
# Loading records from JSONL
# ---------------------------------------------------------------------------


def _load_jsonl(path: Path) -> Iterable[Dict[str, Any]]:
    with path.open("r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                yield json.loads(line)
            except Exception:
                # Skip malformed lines
                continue


def _extract_capsule_from_entry(entry: Dict[str, Any]) -> Optional[Dict[str, Any]]:
    """
    Try multiple heuristics to pull out a capsule dict.

    Supported shapes:
    - { "capsule": { ... } }
    - { "result": { "capsule": { ... } } }
    - Direct capsule (we assume it already is a SymCapsule.to_dict())
    """
    if "capsule" in entry and isinstance(entry["capsule"], dict):
        return entry["capsule"]

    result = entry.get("result")
    if isinstance(result, dict) and "capsule" in result and isinstance(result["capsule"], dict):
        return result["capsule"]

    # If it looks like a SymCapsule dict: we expect top-level "soul" and "data"
    if "soul" in entry and "data" in entry:
        return entry

    return None


def _extract_task_type(entry: Dict[str, Any], default: str = "unknown") -> str:
    for key in ("task_type", "taskType", "type"):
        v = entry.get(key)
        if isinstance(v, str):
            return v
    meta = entry.get("meta")
    if isinstance(meta, dict):
        for key in ("task_type", "taskType"):
            v = meta.get(key)
            if isinstance(v, str):
                return v
    return default


def _extract_gate(entry: Dict[str, Any]) -> Optional[str]:
    for key in ("gate", "gate_policy", "gate_used"):
        v = entry.get(key)
        if isinstance(v, str):
            return v
    meta = entry.get("meta")
    if isinstance(meta, dict):
        for key in ("gate", "gate_policy", "gate_used"):
            v = meta.get(key)
            if isinstance(v, str):
                return v
    return None


def _extract_program(entry: Dict[str, Any]) -> Optional[str]:
    for key in ("helix_program", "helix_program_used", "program_after", "program"):
        v = entry.get(key)
        if isinstance(v, str):
            return v
    meta = entry.get("meta")
    if isinstance(meta, dict):
        for key in ("helix_program", "helix_program_used", "program_after", "program"):
            v = meta.get(key)
            if isinstance(v, str):
                return v
    return None


def _extract_se_metrics(entry: Dict[str, Any]) -> Tuple[Optional[float], Optional[float], Optional[float], Optional[float]]:
    se = entry.get("se_metrics")
    if isinstance(se, dict):
        H = se.get("H")
        C = se.get("C")
        D = se.get("D")
        Se = se.get("Se")
        return (
            float(H) if isinstance(H, (int, float)) else None,
            float(C) if isinstance(C, (int, float)) else None,
            float(D) if isinstance(D, (int, float)) else None,
            float(Se) if isinstance(Se, (int, float)) else None,
        )
    # Some logs might inline Se values at top level
    H = entry.get("se_H", None)
    C = entry.get("se_C", None)
    D = entry.get("se_D", None)
    Se = entry.get("se_Se", None)
    return (
        float(H) if isinstance(H, (int, float)) else None,
        float(C) if isinstance(C, (int, float)) else None,
        float(D) if isinstance(D, (int, float)) else None,
        float(Se) if isinstance(Se, (int, float)) else None,
    )


def _extract_logits(entry: Dict[str, Any]) -> Tuple[Optional[float], Optional[float]]:
    lm = entry.get("logits_mean")
    ls = entry.get("logits_std")
    return (
        float(lm) if isinstance(lm, (int, float)) else None,
        float(ls) if isinstance(ls, (int, float)) else None,
    )


def _compute_fitness_from_entry(entry: Dict[str, Any]) -> Optional[float]:
    try:
        return float(fitness_fn(entry))
    except Exception:
        return None


def load_capsule_records(
    path: Path,
    min_Se: Optional[float] = None,
    min_fitness: Optional[float] = None,
) -> List[CapsuleRecord]:
    """
    Load CapsuleRecord list from a JSONL file.

    Filters:
        - min_Se: drop records with se_Se < min_Se
        - min_fitness: drop records with fitness < min_fitness
    """
    records: List[CapsuleRecord] = []
    for idx, entry in enumerate(_load_jsonl(path)):
        capsule = _extract_capsule_from_entry(entry)
        if capsule is None:
            continue

        task_type = _extract_task_type(entry)
        gate = _extract_gate(entry)
        program = _extract_program(entry)
        se_H, se_C, se_D, se_Se = _extract_se_metrics(entry)
        logits_mean, logits_std = _extract_logits(entry)
        fit = _compute_fitness_from_entry(entry)

        if min_Se is not None and se_Se is not None and se_Se < min_Se:
            continue
        if min_fitness is not None and fit is not None and fit < min_fitness:
            continue

        records.append(
            CapsuleRecord(
                index=idx,
                task_type=task_type,
                gate=gate,
                helix_program=program,
                se_H=se_H,
                se_C=se_C,
                se_D=se_D,
                se_Se=se_Se,
                logits_mean=logits_mean,
                logits_std=logits_std,
                fitness=fit,
                capsule=capsule,
                raw=entry,
            )
        )

    return records


# ---------------------------------------------------------------------------
# Clustering in capsule-space
# ---------------------------------------------------------------------------


@dataclass
class _ClusterBuilder:
    """Helper to maintain cluster stats while building clusters."""

    cluster: CapsuleCluster
    se_values: List[float] = field(default_factory=list)
    fitness_values: List[float] = field(default_factory=list)
    logits_values: List[float] = field(default_factory=list)

    def add(self, record: CapsuleRecord, token_set: Set[str]) -> None:
        self.cluster.indices.append(record.index)
        if record.se_Se is not None:
            self.se_values.append(record.se_Se)
        if record.fitness is not None:
            self.fitness_values.append(record.fitness)
        if record.logits_mean is not None:
            self.logits_values.append(record.logits_mean)
        # If representative not set yet, use first record
        if self.cluster.representative_index is None:
            self.cluster.representative_index = record.index
            self.cluster.representative_capsule = record.capsule
            self.cluster.keywords = sorted(list(token_set))[:32]


def cluster_capsules_for_bucket(
    records: List[CapsuleRecord],
    sim_threshold: float = 0.35,
    min_cluster_size: int = 1,
) -> List[CapsuleCluster]:
    """
    Greedy clustering within a bucket (same task_type, gate, program).

    Algorithm:
        - For each record:
            - Tokenize text from capsule
            - Compute Jaccard similarity with existing clusters' keywords
            - If max similarity >= sim_threshold → assign to that cluster
            - Else, start a new cluster
        - After assignment, compute avg_Se, avg_fitness, avg_logits for each cluster
        - Drop clusters smaller than min_cluster_size
    """
    if not records:
        return []

    builders: List[_ClusterBuilder] = []
    cluster_counter = 0

    # Precompute token sets for all records
    token_sets: Dict[int, Set[str]] = {}
    for r in records:
        text = _extract_text_from_capsule_dict(r.capsule)
        tokens = set(_tokenize(text))
        token_sets[r.index] = tokens

    for r in records:
        tokens = token_sets[r.index]

        # Find best cluster
        best_idx = None
        best_sim = -1.0
        for i, b in enumerate(builders):
            cluster_tokens = set(b.cluster.keywords)
            sim = _jaccard_similarity(tokens, cluster_tokens)
            if sim > best_sim:
                best_sim = sim
                best_idx = i

        if best_idx is None or best_sim < sim_threshold:
            # New cluster
            cid = f"{r.task_type}:{cluster_counter}"
            cluster_counter += 1
            c = CapsuleCluster(
                cluster_id=cid,
                task_type=r.task_type,
                gate=r.gate,
                helix_program=r.helix_program,
            )
            builder = _ClusterBuilder(cluster=c)
            builder.add(r, tokens)
            builders.append(builder)
        else:
            builders[best_idx].add(r, tokens)

    # Finalize clusters
    clusters: List[CapsuleCluster] = []
    for b in builders:
        n = len(b.cluster.indices)
        if n < min_cluster_size:
            continue
        if b.se_values:
            b.cluster.avg_Se = sum(b.se_values) / len(b.se_values)
        if b.fitness_values:
            b.cluster.avg_fitness = sum(b.fitness_values) / len(b.fitness_values)
        if b.logits_values:
            b.cluster.avg_logits_mean = sum(b.logits_values) / len(b.logits_values)
        clusters.append(b.cluster)

    return clusters


def build_archetypes(
    records: List[CapsuleRecord],
    sim_threshold: float = 0.35,
    min_cluster_size: int = 1,
) -> CapsuleEvolutionSummary:
    """
    Build capsule archetypes from raw CapsuleRecord list.

    - Buckets by (task_type, gate, helix_program)
    - Clusters inside each bucket
    - Converts clusters into CapsuleArchetype list
    """
    if not records:
        return CapsuleEvolutionSummary(
            source_path="",
            total_records=0,
            task_types=[],
            archetypes=[],
        )

    # Group indices by bucket
    buckets: Dict[Tuple[str, Optional[str], Optional[str]], List[CapsuleRecord]] = {}
    for r in records:
        key = (r.task_type, r.gate, r.helix_program)
        buckets.setdefault(key, []).append(r)

    archetypes: List[CapsuleArchetype] = []

    for (task_type, gate, program), bucket_records in buckets.items():
        clusters = cluster_capsules_for_bucket(
            bucket_records,
            sim_threshold=sim_threshold,
            min_cluster_size=min_cluster_size,
        )
        for c in clusters:
            archetypes.append(
                CapsuleArchetype(
                    id=c.cluster_id,
                    task_type=c.task_type,
                    gate=c.gate,
                    helix_program=c.helix_program,
                    num_samples=len(c.indices),
                    avg_Se=c.avg_Se,
                    avg_fitness=c.avg_fitness,
                    avg_logits_mean=c.avg_logits_mean,
                    representative_index=c.representative_index,
                    representative_capsule=c.representative_capsule,
                    keywords=c.keywords,
                )
            )

    task_types = sorted({r.task_type for r in records})

    # We'll fill source_path later in CLI (we want the actual input path there)
    summary = CapsuleEvolutionSummary(
        source_path="",
        total_records=len(records),
        task_types=task_types,
        archetypes=archetypes,
    )
    return summary


# ---------------------------------------------------------------------------
# Public analysis API
# ---------------------------------------------------------------------------


def analyze_capsules(
    path: str | Path,
    min_Se: Optional[float] = None,
    min_fitness: Optional[float] = None,
    sim_threshold: float = 0.35,
    min_cluster_size: int = 1,
    task_type_filter: Optional[str] = None,
) -> CapsuleEvolutionSummary:
    """
    Convenience entrypoint for external callers.

    Args:
        path: JSONL path with pheromone or capsule entries
        min_Se: filter records with Se < min_Se
        min_fitness: filter records with fitness < min_fitness
        sim_threshold: Jaccard similarity threshold for clustering
        min_cluster_size: minimum records per cluster
        task_type_filter: if provided, only include this task_type

    Returns:
        CapsuleEvolutionSummary
    """
    p = Path(path)
    records = load_capsule_records(p, min_Se=min_Se, min_fitness=min_fitness)
    if task_type_filter:
        records = [r for r in records if r.task_type == task_type_filter]

    summary = build_archetypes(
        records,
        sim_threshold=sim_threshold,
        min_cluster_size=min_cluster_size,
    )
    summary.source_path = str(p)
    return summary


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


def _summary_to_dict(summary: CapsuleEvolutionSummary) -> Dict[str, Any]:
    return {
        "source_path": summary.source_path,
        "total_records": summary.total_records,
        "task_types": summary.task_types,
        "archetypes": [asdict(a) for a in summary.archetypes],
    }


def _print_pretty(summary: CapsuleEvolutionSummary) -> None:
    print("=" * 60)
    print(f"CAPSULE EVOLVER - {summary.source_path}")
    print("=" * 60)
    print(f"Total records: {summary.total_records}")
    print(f"Task types:   {', '.join(summary.task_types) if summary.task_types else '(none)'}")
    print()

    if not summary.archetypes:
        print("No archetypes found (perhaps filters were too strict).")
        return

    # Group by task_type
    by_task: Dict[str, List[CapsuleArchetype]] = {}
    for a in summary.archetypes:
        by_task.setdefault(a.task_type, []).append(a)

    for task, archetypes in sorted(by_task.items()):
        print("-" * 60)
        print(f"Task Type: {task}")
        print("-" * 60)
        print(
            f"{'ID':<16} {'Gate':<12} {'Program':<10} "
            f"{'N':>4} {'Se':>8} {'Fit':>8} {'Logits':>8}"
        )
        print("-" * 60)
        for a in sorted(archetypes, key=lambda x: (-(x.avg_fitness or 0.0), x.id)):
            se_str = f"{a.avg_Se:.2f}" if a.avg_Se is not None else "  -"
            fit_str = f"{a.avg_fitness:.2f}" if a.avg_fitness is not None else "  -"
            log_str = f"{a.avg_logits_mean:.3f}" if a.avg_logits_mean is not None else "  -"
            prog = a.helix_program or "-"
            gate = a.gate or "-"
            print(
                f"{a.id:<16} {gate:<12} {prog:<10} "
                f"{a.num_samples:>4d} {se_str:>8} {fit_str:>8} {log_str:>8}"
            )
        print()

        # Show top keywords for the best archetype
        best = max(archetypes, key=lambda x: x.avg_fitness or -math.inf)
        if best.keywords:
            kws = ", ".join(best.keywords[:12])
            print(f"  Top keywords for best archetype ({best.id}):")
            print(f"    {kws}")
        print()


def main(argv: Optional[List[str]] = None) -> int:
    parser = argparse.ArgumentParser(
        prog="capsule_evolver",
        description="Phase 14: Capsule-space archetype discovery for SymCapsules.",
    )
    parser.add_argument(
        "path",
        help="JSONL path: pheromones or capsule logs (must contain 'capsule' field).",
    )
    parser.add_argument(
        "--task-type",
        help="Filter to a specific task_type.",
    )
    parser.add_argument(
        "--min-Se",
        type=float,
        dest="min_Se",
        help="Minimum Se value for records to be included.",
    )
    parser.add_argument(
        "--min-fitness",
        type=float,
        help="Minimum fitness value for records to be included.",
    )
    parser.add_argument(
        "--sim-threshold",
        type=float,
        default=0.35,
        help="Jaccard similarity threshold for capsule clustering (default: 0.35).",
    )
    parser.add_argument(
        "--min-cluster-size",
        type=int,
        default=1,
        help="Minimum cluster size (default: 1).",
    )
    parser.add_argument(
        "--out",
        help="Optional JSON output path for archetype summary.",
    )
    parser.add_argument(
        "--json",
        action="store_true",
        help="Print JSON summary to stdout.",
    )
    parser.add_argument(
        "--pretty",
        action="store_true",
        help="Print human-readable table summary.",
    )

    args = parser.parse_args(argv)

    path = Path(args.path)
    if not path.exists():
        print(f"Error: path not found: {path}", file=sys.stderr)
        return 1

    summary = analyze_capsules(
        path=path,
        min_Se=args.min_Se,
        min_fitness=args.min_fitness,
        sim_threshold=args.sim_threshold,
        min_cluster_size=args.min_cluster_size,
        task_type_filter=args.task_type,
    )

    summary_dict = _summary_to_dict(summary)

    # Write to file if requested
    if args.out:
        out_path = Path(args.out)
        out_path.parent.mkdir(parents=True, exist_ok=True)
        with out_path.open("w", encoding="utf-8") as f:
            json.dump(summary_dict, f, indent=2, ensure_ascii=False)
        print(f"Archetype summary written to {out_path}")

    # Print to stdout
    if args.json:
        print(json.dumps(summary_dict, indent=2, ensure_ascii=False))
    if args.pretty or not (args.json or args.out):
        # Default to pretty output if neither json nor out provided
        _print_pretty(summary)

    return 0


if __name__ == "__main__":
    raise SystemExit(main())
