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

"""
Phase 6: Fitness Function for Pheromone Entries

This module defines the fitness function used by the GateController
to evaluate pheromone entries and compute reinforcement signals.

The fitness formula:
    fitness = Se + logits_mean + (seed_entropy / 2)

Where:
    Se          = Symbolic entropy (H * C * D), typically 0-100
    logits_mean = Model confidence, typically 0-1
    seed_entropy = Seed information density, 0-8 bits/byte

The function is robust against None values.

Usage:
    from helix_cdc.bloomos.fitness import fitness, normalized_fitness

    entry = PheromoneEntry(...)
    fit = fitness(entry)
    norm_fit = normalized_fitness(entry)  # 0-1 range
"""

from typing import Any, Optional, Union


def fitness(entry: Any) -> float:
    """
    Compute fitness score for a pheromone entry.

    Formula: fitness = Se + logits_mean + (seed_entropy / 2)

    This formula balances:
    - Symbolic entropy (Se): Measures semantic richness of output
    - logits_mean: Model confidence in the response
    - seed_entropy: Information density of the CDR seed

    Args:
        entry: A PheromoneEntry or dict-like object with se_Se, logits_mean, seed_entropy

    Returns:
        Fitness score (unbounded, typically 0-10 range)
    """
    # Extract values, handling both dataclass and dict
    if hasattr(entry, "se_Se"):
        se = entry.se_Se
        logits = entry.logits_mean
        seed_ent = entry.seed_entropy
    elif isinstance(entry, dict):
        se = entry.get("se_Se")
        logits = entry.get("logits_mean")
        seed_ent = entry.get("seed_entropy")
    else:
        return 0.0

    # Handle None values robustly
    se_val = float(se) if se is not None else 0.0
    logits_val = float(logits) if logits is not None else 0.0
    seed_ent_val = float(seed_ent) if seed_ent is not None else 4.0  # Default mid-range

    # Fitness formula
    # Note: Se is typically 0-100, logits_mean is 0-1, seed_entropy is 0-8
    # We weight them to roughly equal contribution
    return se_val + logits_val + (seed_ent_val / 2.0)


def normalized_fitness(entry: Any) -> float:
    """
    Compute normalized fitness score in [0, 1] range.

    This normalizes the raw fitness to a 0-1 scale for comparison.

    Expected ranges:
    - Se: 0-100 (scaled to 0-1)
    - logits_mean: 0-1 (already in range)
    - seed_entropy: 0-8 (scaled to 0-1)

    Args:
        entry: A PheromoneEntry or dict-like object

    Returns:
        Normalized fitness in [0, 1] range
    """
    # Extract values
    if hasattr(entry, "se_Se"):
        se = entry.se_Se
        logits = entry.logits_mean
        seed_ent = entry.seed_entropy
    elif isinstance(entry, dict):
        se = entry.get("se_Se")
        logits = entry.get("logits_mean")
        seed_ent = entry.get("seed_entropy")
    else:
        return 0.0

    # Handle None values
    se_val = float(se) if se is not None else 0.0
    logits_val = float(logits) if logits is not None else 0.0
    seed_ent_val = float(seed_ent) if seed_ent is not None else 4.0

    # Normalize each component to 0-1
    se_norm = min(1.0, se_val / 100.0)  # Cap at 100
    logits_norm = min(1.0, max(0.0, logits_val))  # Clamp to 0-1
    seed_ent_norm = seed_ent_val / 8.0  # 8 bits/byte max

    # Weighted average (equal weights)
    return (se_norm + logits_norm + seed_ent_norm) / 3.0


def fitness_components(entry: Any) -> dict:
    """
    Break down fitness into components for debugging.

    Args:
        entry: A PheromoneEntry or dict-like object

    Returns:
        Dict with individual components and total fitness
    """
    # Extract values
    if hasattr(entry, "se_Se"):
        se = entry.se_Se
        logits = entry.logits_mean
        seed_ent = entry.seed_entropy
    elif isinstance(entry, dict):
        se = entry.get("se_Se")
        logits = entry.get("logits_mean")
        seed_ent = entry.get("seed_entropy")
    else:
        return {"error": "Invalid entry type"}

    se_val = float(se) if se is not None else 0.0
    logits_val = float(logits) if logits is not None else 0.0
    seed_ent_val = float(seed_ent) if seed_ent is not None else 4.0

    return {
        "se_Se": se_val,
        "se_contribution": se_val,
        "logits_mean": logits_val,
        "logits_contribution": logits_val,
        "seed_entropy": seed_ent_val,
        "seed_entropy_contribution": seed_ent_val / 2.0,
        "fitness": se_val + logits_val + (seed_ent_val / 2.0),
        "normalized_fitness": (min(1.0, se_val / 100.0) + min(1.0, max(0.0, logits_val)) + seed_ent_val / 8.0) / 3.0,
    }


def rank_entries(entries: list, descending: bool = True) -> list:
    """
    Rank a list of entries by fitness.

    Args:
        entries: List of PheromoneEntry or dict-like objects
        descending: If True, highest fitness first

    Returns:
        List of (entry, fitness) tuples, sorted by fitness
    """
    scored = [(e, fitness(e)) for e in entries]
    scored.sort(key=lambda x: x[1], reverse=descending)
    return scored


# CLI for testing
if __name__ == "__main__":
    import json
    import sys

    print("=== Fitness Function Test ===\n")

    # Test with mock entries
    test_entries = [
        {"se_Se": 50.0, "logits_mean": 0.8, "seed_entropy": 5.0},
        {"se_Se": 80.0, "logits_mean": 0.6, "seed_entropy": 6.5},
        {"se_Se": None, "logits_mean": 0.9, "seed_entropy": None},
        {"se_Se": 100.0, "logits_mean": 1.0, "seed_entropy": 8.0},
    ]

    print("Test Entries:")
    for i, entry in enumerate(test_entries):
        components = fitness_components(entry)
        print(f"\nEntry {i + 1}: {entry}")
        print(f"  fitness: {components['fitness']:.3f}")
        print(f"  normalized: {components['normalized_fitness']:.3f}")

    print("\n\nRanked (descending):")
    ranked = rank_entries(test_entries)
    for entry, fit in ranked:
        print(f"  fitness={fit:.3f}  entry={entry}")
