"""
helix_cdc/brainstem/band_mapper.py
==================================

Maps brainstem neural bands to QUBO weight configurations.

This is the key interface between brainstem state and QUBO optimization:
    BrainstemState -> BandMapper -> RouteQUBOConfig

Band -> Weight Mapping (from plan):
| Band  | Latency Weight | Fidelity Weight | Determinism Weight |
|-------|---------------|-----------------|-------------------|
| Theta | 0.1 (slow ok) | 0.3             | 0.5               |
| Alpha | 0.3           | 0.7 (reflection)| 0.6               |
| Beta  | 0.8 (fast!)   | 0.5             | 0.7               |
| Gamma | 0.9 (insight) | 0.4             | 0.3 (can vary)    |
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Dict, Optional

from .state_detector import BrainstemState

# Import RouteQUBOConfig - handle both installed and dev paths
try:
    from helix_cdc.routing.qubo_route_builder import RouteQUBOConfig
except ImportError:
    # For standalone usage
    RouteQUBOConfig = None


# Band weight profiles from the plan
BAND_WEIGHT_PROFILES: Dict[str, Dict[str, float]] = {
    "theta": {
        "lambda_latency": 0.1,     # Slow is OK, exploring
        "lambda_fidelity": 0.3,    # Not strict
        "lambda_determinism": 0.5, # Moderate
        "lambda_thermal": 1.0,     # Still care about hardware
        "lambda_se": 3.0,          # High Se exploration
        "lambda_memory": 0.5,      # Relaxed
    },
    "alpha": {
        "lambda_latency": 0.3,     # Some urgency
        "lambda_fidelity": 0.7,    # Reflection needs fidelity
        "lambda_determinism": 0.6, # Want reliability
        "lambda_thermal": 2.0,     # Care more
        "lambda_se": 2.0,          # Balanced
        "lambda_memory": 1.0,      # Normal
    },
    "beta": {
        "lambda_latency": 0.8,     # Problem-solving = fast
        "lambda_fidelity": 0.5,    # Balanced
        "lambda_determinism": 0.7, # Want determinism
        "lambda_thermal": 2.5,     # High care (intensive)
        "lambda_se": 1.5,          # Lower priority
        "lambda_memory": 1.5,      # Care about resources
    },
    "gamma": {
        "lambda_latency": 0.9,     # Insight = speed matters
        "lambda_fidelity": 0.4,    # Can tolerate some loss
        "lambda_determinism": 0.3, # Can tolerate variance
        "lambda_thermal": 3.0,     # Watch hardware closely
        "lambda_se": 4.0,          # High Se is the goal
        "lambda_memory": 1.0,      # Balanced
    },
}


@dataclass
class BandMappingResult:
    """Result of band -> QUBO config mapping."""
    band: str
    config_dict: Dict[str, float]
    banned_backends: list
    guardrail_penalty: float
    description: str


class BandMapper:
    """
    Maps BrainstemState to RouteQUBOConfig.

    Usage:
        from helix_cdc.brainstem import StateDetector, BandMapper

        detector = StateDetector()
        mapper = BandMapper()

        state = detector.detect(input)
        config = mapper.to_qubo_config(state)
        # Use config with build_route_qubo()
    """

    def __init__(self, custom_profiles: Optional[Dict[str, Dict[str, float]]] = None):
        """
        Initialize the band mapper.

        Args:
            custom_profiles: Optional custom band->weight profiles
        """
        self.profiles = custom_profiles or BAND_WEIGHT_PROFILES

    def map(self, state: BrainstemState) -> BandMappingResult:
        """
        Map brainstem state to configuration dict.

        Args:
            state: BrainstemState from detector

        Returns:
            BandMappingResult with config dict and metadata
        """
        band = state.band
        profile = self.profiles.get(band, self.profiles["alpha"])

        # Apply confidence-weighted interpolation toward alpha (default)
        # If confidence is low, blend toward alpha (safe default)
        if state.band_confidence < 0.8:
            alpha_profile = self.profiles["alpha"]
            blend_factor = state.band_confidence
            blended_profile = {}
            for key in profile:
                blended_profile[key] = (
                    profile[key] * blend_factor +
                    alpha_profile[key] * (1 - blend_factor)
                )
            profile = blended_profile

        # Adjust for lobe votes (trio dispatch influence)
        adjusted_profile = self._apply_lobe_influence(profile, state.lobe_votes)

        # Calculate guardrail penalty based on band
        guardrail = self._compute_guardrail_penalty(band)

        # Generate description
        description = self._generate_description(band, state)

        return BandMappingResult(
            band=band,
            config_dict=adjusted_profile,
            banned_backends=list(state.banned_backends),
            guardrail_penalty=guardrail,
            description=description,
        )

    def to_qubo_config(self, state: BrainstemState) -> "RouteQUBOConfig":
        """
        Convert BrainstemState directly to RouteQUBOConfig.

        Args:
            state: BrainstemState from detector

        Returns:
            RouteQUBOConfig ready for build_route_qubo()

        Raises:
            ImportError: If RouteQUBOConfig not available
        """
        if RouteQUBOConfig is None:
            raise ImportError(
                "RouteQUBOConfig not available. "
                "Install helix_cdc.routing or use map() for config dict."
            )

        mapping = self.map(state)

        return RouteQUBOConfig(
            lambda_latency=mapping.config_dict.get("lambda_latency", 1.0),
            lambda_memory=mapping.config_dict.get("lambda_memory", 1.0),
            lambda_fidelity=mapping.config_dict.get("lambda_fidelity", 5.0),
            lambda_se=mapping.config_dict.get("lambda_se", 2.0),
            lambda_determinism=mapping.config_dict.get("lambda_determinism", 1.0),
            lambda_thermal=mapping.config_dict.get("lambda_thermal", 2.0),
            guardrail_penalty=mapping.guardrail_penalty,
        )

    def to_config_kwargs(self, state: BrainstemState) -> Dict[str, float]:
        """
        Convert BrainstemState to kwargs dict for RouteQUBOConfig().

        Use this when you can't import RouteQUBOConfig directly.

        Args:
            state: BrainstemState from detector

        Returns:
            Dict suitable for RouteQUBOConfig(**kwargs)
        """
        mapping = self.map(state)
        return {
            "lambda_latency": mapping.config_dict.get("lambda_latency", 1.0),
            "lambda_memory": mapping.config_dict.get("lambda_memory", 1.0),
            "lambda_fidelity": mapping.config_dict.get("lambda_fidelity", 5.0),
            "lambda_se": mapping.config_dict.get("lambda_se", 2.0),
            "lambda_determinism": mapping.config_dict.get("lambda_determinism", 1.0),
            "lambda_thermal": mapping.config_dict.get("lambda_thermal", 2.0),
            "guardrail_penalty": mapping.guardrail_penalty,
        }

    def _apply_lobe_influence(
        self,
        profile: Dict[str, float],
        lobe_votes: Dict[str, float]
    ) -> Dict[str, float]:
        """
        Apply trio lobe votes to adjust profile.

        - Left lobe (logical) -> increase determinism_weight
        - Right lobe (intuitive) -> increase fidelity_weight
        - Poet lobe (creative) -> increase Se weight
        """
        adjusted = dict(profile)

        left = lobe_votes.get("left", 0.33)
        right = lobe_votes.get("right", 0.33)
        poet = lobe_votes.get("poet", 0.34)

        # Left brain boosts determinism
        if left > 0.4:
            adjusted["lambda_determinism"] = adjusted.get("lambda_determinism", 1.0) * (1 + (left - 0.4))

        # Right brain boosts fidelity
        if right > 0.4:
            adjusted["lambda_fidelity"] = adjusted.get("lambda_fidelity", 5.0) * (1 + (right - 0.4))

        # Poet brain boosts Se
        if poet > 0.4:
            adjusted["lambda_se"] = adjusted.get("lambda_se", 2.0) * (1 + (poet - 0.4))

        return adjusted

    def _compute_guardrail_penalty(self, band: str) -> float:
        """
        Compute Se guardrail penalty based on band.

        Theta/Alpha are more relaxed, Beta/Gamma need stricter guardrails.
        """
        penalties = {
            "theta": 30.0,   # Relaxed
            "alpha": 40.0,   # Moderate
            "beta": 60.0,    # Strict (problem-solving needs reliability)
            "gamma": 50.0,   # Moderate-high (speed but need some safety)
        }
        return penalties.get(band, 50.0)

    def _generate_description(self, band: str, state: BrainstemState) -> str:
        """Generate human-readable description of the mapping."""
        descriptions = {
            "theta": "Relaxed exploration mode - slow ok, creative freedom",
            "alpha": "Reflective mode - balanced, fidelity matters",
            "beta": "Problem-solving mode - speed matters, high determinism",
            "gamma": "Insight mode - maximum speed, can tolerate variance",
        }
        base = descriptions.get(band, f"Unknown band: {band}")

        # Add lobe influence note
        dominant_lobe = max(state.lobe_votes.items(), key=lambda x: x[1])
        if dominant_lobe[1] > 0.4:
            base += f" (dominant: {dominant_lobe[0]})"

        return base


def band_to_weights(band: str) -> tuple:
    """
    Quick utility to get weight triple for a band.

    Returns:
        (latency_weight, fidelity_weight, determinism_weight)
    """
    profile = BAND_WEIGHT_PROFILES.get(band, BAND_WEIGHT_PROFILES["alpha"])
    return (
        profile["lambda_latency"],
        profile["lambda_fidelity"],
        profile["lambda_determinism"],
    )


def weights_from_state(state: BrainstemState) -> Dict[str, float]:
    """
    Quick utility to get all weights from a BrainstemState.

    Convenience wrapper around BandMapper for simple cases.
    """
    mapper = BandMapper()
    return mapper.to_config_kwargs(state)
