"""Managed pool of graded seeds with weighted sampling and technique tracking."""

from __future__ import annotations

import random
from collections import Counter

from probegpt.core.models import Probe

from .models import GradedSeed


class SeedPool:
    """Holds graded seeds and initial (unanalyzed) seeds.

    Initial seeds from seeds.yaml are kept as a stable baseline.
    GradedSeeds are added dynamically as promising probes are discovered.
    Mutation sampling weights toward high-scoring, novel, transferable seeds.
    """

    def __init__(self, initial_seeds: list[Probe], max_graded: int = 30) -> None:
        self._initial: list[Probe] = list(initial_seeds)
        self._graded: list[GradedSeed] = []
        self._max_graded = max_graded

    @property
    def graded(self) -> list[GradedSeed]:
        return list(self._graded)

    @property
    def all_probes(self) -> list[Probe]:
        """All seeds as Probe objects: initial + generated, for logging."""
        generated = [g.probe for g in self._graded]
        return self._initial + generated

    def get_technique_histogram(self, top_n: int = 8) -> list[str]:
        """Return the top-N most frequent technique tags across graded seeds."""
        counter: Counter[str] = Counter()
        for gs in self._graded:
            for tag in gs.analysis.techniques:
                counter[tag] += 1
        return [tag for tag, _ in counter.most_common(top_n)]

    def sample_for_mutation(self) -> Probe:
        """Weighted random sample from graded seeds by composite weight.

        Falls back to uniform sample from initial seeds if no graded seeds exist.
        """
        if not self._graded:
            return random.choice(self._initial)

        weights = [max(gs.weight, 1e-6) for gs in self._graded]
        total = sum(weights)
        normalised = [w / total for w in weights]
        chosen = random.choices(self._graded, weights=normalised, k=1)[0]
        return chosen.probe

    def get_top_graded(self, n: int) -> list[GradedSeed]:
        """Return top-N graded seeds by composite weight."""
        return sorted(self._graded, key=lambda g: g.weight, reverse=True)[:n]

    def add(self, seed: GradedSeed) -> bool:
        """Add a GradedSeed. Evicts the lowest-weight seed if the pool is full."""
        self._graded.append(seed)
        if len(self._graded) > self._max_graded:
            self._graded.sort(key=lambda g: g.weight, reverse=True)
            self._graded = self._graded[: self._max_graded]
        return True

    def __len__(self) -> int:
        return len(self._initial) + len(self._graded)

    def __repr__(self) -> str:
        return f"SeedPool(initial={len(self._initial)}, graded={len(self._graded)})"
