"""Policy fit analyzer — estimates how well data matches a specific policy architecture.

Encodes policy-specific knowledge about ACT, Diffusion Policy, SmolVLA, DP3,
and general BC variants to predict data-policy compatibility.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field

import numpy as np

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Policy name normalization
# ---------------------------------------------------------------------------

POLICY_ALIASES = {
    # ── ACT family ──────────────────────────────────────────────────────
    "act": "act",
    "action_chunking": "act",
    "action_chunking_transformer": "act",
    "act_plus": "act",
    "act_3d": "act",
    "act_dp": "act",
    "act_image": "act",
    "act_lowdim": "act",
    "act_state": "act",
    "act_vq": "act",
    "act_large": "act",
    "act_small": "act",
    "act_100": "act",
    # RT-1 is a robotics transformer (VLA-like), NOT the ACT architecture
    "rt1": "smolvla",
    "rt_1": "smolvla",
    "rt1_x": "smolvla",
    "rt_1_x": "smolvla",
    "rt1x": "smolvla",
    "rt1_bridge": "smolvla",
    "rt_1_bridge": "smolvla",
    # ── Diffusion Policy family ─────────────────────────────────────────
    "diffusion_policy": "diffusion_policy",
    "diffusion": "diffusion_policy",
    "diffusion_policy_cnn": "diffusion_policy",
    "diffusion_policy_transformer": "diffusion_policy",
    "diffusion_policy_image": "diffusion_policy",
    "diffusion_policy_lowdim": "diffusion_policy",
    "diffusion_policy_state": "diffusion_policy",
    "dp": "diffusion_policy",
    "dp_cnn": "diffusion_policy",
    "dp_transformer": "diffusion_policy",
    "dp_t": "diffusion_policy",
    "dppo": "diffusion_policy",
    "ddpm": "diffusion_policy",
    "ddim": "diffusion_policy",
    "dit_policy": "diffusion_policy",
    "consistency_policy": "diffusion_policy",
    "score_policy": "diffusion_policy",
    "octo": "diffusion_policy",
    "octo_base": "diffusion_policy",
    "octo_small": "diffusion_policy",
    "octo_finetuned": "diffusion_policy",
    "octo_large": "diffusion_policy",
    "hpt": "diffusion_policy",
    "hpt_base": "diffusion_policy",
    "hpt_small": "diffusion_policy",
    "pi0": "diffusion_policy",
    "pi_0": "diffusion_policy",
    "pi_zero": "diffusion_policy",
    "pi0_fast": "diffusion_policy",
    "rdt": "diffusion_policy",
    "rdt_1b": "diffusion_policy",
    "rdt_170m": "diffusion_policy",
    "crossformer": "diffusion_policy",
    "sus": "diffusion_policy",
    # ── BC family ───────────────────────────────────────────────────────
    "bc": "bc",
    "behavioral_cloning": "bc",
    "bet": "bc",
    "vq_bet": "bc",
    "vqbet": "bc",
    "ibc": "bc",
    "implicit_bc": "bc",
    "gcbc": "bc",
    "gc_ibc": "bc",
    "goal_bc": "bc",
    "goal_conditioned_bc": "bc",
    "vinn": "bc",
    "bc_z": "bc",
    "bcz": "bc",
    "bc_z_policy": "bc",
    "roboflamingo": "bc",
    "robo_flamingo": "bc",
    "hulc": "bc",
    "mcil": "bc",
    "lcbc": "bc",
    "lcbc_policy": "bc",
    "cql": "bc",
    "bcq": "bc",
    "td3_bc": "bc",
    "td3_plus_bc": "bc",
    "awac": "bc",
    "residual_bc": "bc",
    "r3m_bc": "bc",
    "play_lmp": "bc",
    "spil": "bc",
    "cbet": "bc",
    "c_bet": "bc",
    "flat_bc": "bc",
    "mlp_bc": "bc",
    "decision_transformer": "bc",
    "dt": "bc",
    "gato": "bc",
    "rvt": "bc",
    "rvt2": "bc",
    "perceiver_actor": "bc",
    "peract": "bc",
    "peract2": "bc",
    "c2farm": "bc",
    "arm": "bc",
    # ── BC-RNN family ───────────────────────────────────────────────────
    "bc_rnn": "bc_rnn",
    "bcrnn": "bc_rnn",
    "lstm_gmm": "bc_rnn",
    "bc_rnn_gmm": "bc_rnn",
    "lstm": "bc_rnn",
    "rnn_bc": "bc_rnn",
    "rnn_gmm": "bc_rnn",
    "gru_bc": "bc_rnn",
    "bcrnn_image": "bc_rnn",
    "bcrnn_lowdim": "bc_rnn",
    "bc_rnn_lowdim": "bc_rnn",
    "bc_rnn_image": "bc_rnn",
    # ── VLA family ──────────────────────────────────────────────────────
    "smolvla": "smolvla",
    "smol_vla": "smolvla",
    "vla": "smolvla",
    "openvla": "smolvla",
    "open_vla": "smolvla",
    "openvla_7b": "smolvla",
    "openvla_base": "smolvla",
    "rt2": "smolvla",
    "rt_2": "smolvla",
    "rt_2_x": "smolvla",
    "rt2_x": "smolvla",
    "rt_2_x_55b": "smolvla",
    "rt_2_55b": "smolvla",
    "pali_gemma": "smolvla",
    "paligemma": "smolvla",
    "octo_vla": "smolvla",
    "susie": "smolvla",
    # ── 3D Diffusion family ─────────────────────────────────────────────
    "dp3": "dp3",
    "3d_diffusion": "dp3",
    "3d_diffusion_policy": "dp3",
    "3ddp": "dp3",
    "dp3_cnn": "dp3",
    "dp3_transformer": "dp3",
    "pointcloud_dp": "dp3",
    "point_cloud_dp": "dp3",
    "gnfactor": "dp3",
}


def normalize_policy_name(raw: str) -> str:
    """Map any policy name to one of our 6 canonical profiles."""
    import re

    if not raw:
        return "diffusion_policy"
    cleaned = raw.strip().lower()
    # Remove common suffixes like "(ours)", "(es)", etc.
    cleaned = re.sub(r'\s*\(.*?\)\s*$', '', cleaned)
    cleaned = cleaned.replace(" ", "_").replace("-", "_")
    # Normalize Unicode: π → pi, subscript digits → ascii, + → _plus_
    cleaned = cleaned.replace("π", "pi").replace("₀", "0").replace("₁", "1").replace("₂", "2")
    cleaned = cleaned.replace("+", "_plus_")
    cleaned = re.sub(r'_+', '_', cleaned).strip("_")
    # Direct lookup
    if cleaned in POLICY_ALIASES:
        return POLICY_ALIASES[cleaned]
    # Try without underscores
    no_under = cleaned.replace("_", "")
    for key, val in POLICY_ALIASES.items():
        if key.replace("_", "") == no_under:
            return val
    # Default — log so we can discover unmapped policy names
    logger.warning(
        "Unknown policy name '%s' (cleaned: '%s'), defaulting to diffusion_policy",
        raw, cleaned,
    )
    return "diffusion_policy"


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


@dataclass
class PolicyFit:
    """Assessment of how well data fits a specific policy architecture."""

    fit_score: float  # 0-1 overall fit
    episode_sufficiency: float  # 0-1 enough demos?
    action_space_compatibility: float  # 0-1 action format works?
    consistency_compatibility: float  # 0-1 how well does consistency match policy needs?
    policy_type: str  # detected or specified policy
    specific_warnings: list[str] = field(default_factory=list)
    fit_breakdown: dict[str, float] = field(default_factory=dict)
    recommended_policies: list[dict] = field(default_factory=list)


# ---------------------------------------------------------------------------
# Policy knowledge base
# ---------------------------------------------------------------------------

POLICY_PROFILES = {
    "act": {
        "full_name": "Action Chunking Transformer",
        "min_episodes_simple": 25,
        "min_episodes_complex": 50,
        "sweet_spot_low": 50,
        "sweet_spot_high": 200,
        "max_useful_episodes": 500,
        "consistency_sensitivity": 0.8,  # High — ACT needs consistent demos
        "multimodality_tolerance": 0.3,  # Low — struggles with multimodal
        "action_dim_range": (6, 14),  # Works best with arm-scale
        "strengths": ["bimanual", "action_chunking", "precise_manipulation"],
        "weaknesses": ["multimodal_distributions", "noisy_demos"],
    },
    "diffusion_policy": {
        "full_name": "Diffusion Policy",
        "min_episodes_simple": 50,
        "min_episodes_complex": 100,
        "sweet_spot_low": 100,
        "sweet_spot_high": 500,
        "max_useful_episodes": 2000,
        "consistency_sensitivity": 0.4,  # Low — handles inconsistency well
        "multimodality_tolerance": 0.9,  # High — excels at multimodal
        "action_dim_range": (2, 14),
        "strengths": ["multimodal_distributions", "diverse_demos", "noise_robust"],
        "weaknesses": ["very_small_datasets", "high_latency"],
    },
    "smolvla": {
        "full_name": "SmolVLA / Vision-Language-Action",
        "min_episodes_simple": 100,
        "min_episodes_complex": 200,
        "sweet_spot_low": 200,
        "sweet_spot_high": 1000,
        "max_useful_episodes": 10000,
        "consistency_sensitivity": 0.5,
        "multimodality_tolerance": 0.7,
        "action_dim_range": (2, 14),
        "strengths": ["visual_diversity", "language_conditioning", "generalization"],
        "weaknesses": ["small_datasets", "action_precision"],
    },
    "dp3": {
        "full_name": "3D Diffusion Policy",
        "min_episodes_simple": 50,
        "min_episodes_complex": 100,
        "sweet_spot_low": 100,
        "sweet_spot_high": 500,
        "max_useful_episodes": 2000,
        "consistency_sensitivity": 0.4,
        "multimodality_tolerance": 0.8,
        "action_dim_range": (3, 14),
        "strengths": ["spatial_tasks", "point_cloud", "3d_reasoning"],
        "weaknesses": ["requires_3d_data", "computational_cost"],
    },
    "bc": {
        "full_name": "Behavioral Cloning",
        "min_episodes_simple": 50,
        "min_episodes_complex": 200,
        "sweet_spot_low": 100,
        "sweet_spot_high": 500,
        "max_useful_episodes": 5000,
        "consistency_sensitivity": 0.9,  # Very high — BC is fragile
        "multimodality_tolerance": 0.1,  # Very low — BC collapses on multimodal
        "action_dim_range": (2, 14),
        "strengths": ["simple", "fast_training", "interpretable"],
        "weaknesses": ["multimodal_catastrophic", "needs_very_clean_data"],
    },
    "bc_rnn": {
        "full_name": "BC-RNN (Recurrent Behavioral Cloning)",
        "min_episodes_simple": 50,
        "min_episodes_complex": 150,
        "sweet_spot_low": 100,
        "sweet_spot_high": 500,
        "max_useful_episodes": 5000,
        "consistency_sensitivity": 0.6,
        "multimodality_tolerance": 0.5,  # Better than BC at multimodal
        "action_dim_range": (2, 14),
        "strengths": ["temporal_reasoning", "moderate_multimodal"],
        "weaknesses": ["needs_more_data_than_bc", "still_fragile_on_hard_tasks"],
    },
}


# ---------------------------------------------------------------------------
# Scoring functions
# ---------------------------------------------------------------------------


def _episode_sufficiency(
    episode_count: int,
    task_complexity: float,
    profile: dict,
) -> float:
    """Score whether there are enough episodes for this policy + task combo."""
    # Interpolate minimum episodes based on task complexity
    min_simple = profile["min_episodes_simple"]
    min_complex = profile["min_episodes_complex"]
    min_needed = min_simple + task_complexity * (min_complex - min_simple)

    sweet_low = profile["sweet_spot_low"]
    sweet_high = profile["sweet_spot_high"]
    sweet_needed = sweet_low + task_complexity * (sweet_high - sweet_low)

    if episode_count >= sweet_needed:
        return 1.0
    elif episode_count >= min_needed:
        # Between minimum and sweet spot
        return 0.6 + 0.4 * (episode_count - min_needed) / max(1, sweet_needed - min_needed)
    elif episode_count > 0:
        # Below minimum
        return 0.6 * (episode_count / max(1, min_needed))
    else:
        return 0.0


def _action_space_compatibility(
    action_dim: int,
    profile: dict,
) -> float:
    """Score how well the action space matches the policy's expected range."""
    low, high = profile["action_dim_range"]

    if low <= action_dim <= high:
        return 1.0
    elif action_dim < low:
        # Below range — might still work but not optimal
        return max(0.5, 1.0 - (low - action_dim) * 0.15)
    else:
        # Above range — higher dims need more capacity
        return max(0.3, 1.0 - (action_dim - high) * 0.05)


def _consistency_compatibility(
    consistency_score: float,
    multimodality_score: float,
    profile: dict,
) -> float:
    """Score how well data consistency matches policy expectations.

    ACT needs high consistency. Diffusion Policy handles multimodality.
    BC collapses on multimodal data.
    """
    sensitivity = profile["consistency_sensitivity"]
    multimodal_tolerance = profile["multimodality_tolerance"]

    # Consistency match: if policy is sensitive, low consistency hurts
    consistency_fit = 1.0 - sensitivity * (1.0 - consistency_score)

    # Multimodality: if data is multimodal and policy can't handle it, penalize
    multimodal_penalty = max(0.0, multimodality_score - multimodal_tolerance) * 0.5

    return float(np.clip(consistency_fit - multimodal_penalty, 0.0, 1.0))


def _detect_best_policy(
    action_dim: int,
    episode_count: int,
    consistency: float,
    multimodality: float,
    is_bimanual: bool,
    task_complexity: float,
) -> list[dict]:
    """Rank policies by expected fit for this data, return top 3."""
    candidates = []

    for name, profile in POLICY_PROFILES.items():
        ep_suff = _episode_sufficiency(episode_count, task_complexity, profile)
        action_compat = _action_space_compatibility(action_dim, profile)
        consist_compat = _consistency_compatibility(consistency, multimodality, profile)

        # Bonus for bimanual + ACT
        bimanual_bonus = 0.0
        if is_bimanual and "bimanual" in profile.get("strengths", []):
            bimanual_bonus = 0.10

        overall = ep_suff * 0.30 + action_compat * 0.20 + consist_compat * 0.50 + bimanual_bonus
        overall = float(np.clip(overall, 0.0, 1.0))

        candidates.append(
            {
                "policy": name,
                "full_name": profile["full_name"],
                "fit_score": round(overall, 3),
                "episode_sufficiency": round(ep_suff, 3),
                "action_compatibility": round(action_compat, 3),
                "consistency_compatibility": round(consist_compat, 3),
            }
        )

    candidates.sort(key=lambda x: x["fit_score"], reverse=True)
    return candidates[:3]


# ---------------------------------------------------------------------------
# Main analyzer
# ---------------------------------------------------------------------------


class PolicyFitAnalyzer:
    """Estimates how well data fits a specific policy architecture."""

    def analyze(
        self,
        action_episodes: list[np.ndarray],
        policy_type: str = "auto",
        episode_count: int | None = None,
        metadata: dict | None = None,
        task_complexity: float = 0.5,
        consistency_score: float = 0.7,
        multimodality_score: float = 0.3,
    ) -> PolicyFit:
        """Analyze policy-data fit.

        Args:
            action_episodes: List of per-episode action arrays.
            policy_type: Policy name or "auto" to detect best fit.
            episode_count: Override episode count (default: len(action_episodes)).
            metadata: Optional dict with 'is_bimanual', etc.
            task_complexity: Task complexity score from TaskContextAnalyzer (0-1).
            consistency_score: Episode consistency score (0-1).
            multimodality_score: Trajectory multimodality score (0-1).

        Returns:
            PolicyFit with scores and recommendations.
        """
        metadata = metadata or {}

        if not action_episodes:
            return PolicyFit(
                fit_score=0.5,
                episode_sufficiency=0.5,
                action_space_compatibility=0.5,
                consistency_compatibility=0.5,
                policy_type=policy_type,
            )

        # Determine action dim
        first = action_episodes[0]
        action_dim = first.shape[1] if first.ndim == 2 else 1
        n_episodes = episode_count or len(action_episodes)
        is_bimanual = metadata.get("is_bimanual", action_dim >= 12)

        # Auto-detect best policy
        recommended = _detect_best_policy(
            action_dim,
            n_episodes,
            consistency_score,
            multimodality_score,
            is_bimanual,
            task_complexity,
        )

        # If auto, use the top recommended policy
        if policy_type == "auto":
            policy_type = recommended[0]["policy"] if recommended else "diffusion_policy"
        else:
            # Normalize user-provided policy name to canonical form
            policy_type = normalize_policy_name(policy_type)

        # Look up the profile
        profile = POLICY_PROFILES.get(policy_type)
        if profile is None:
            # Unknown policy — use diffusion_policy as default
            logger.warning("Unknown policy '%s', defaulting to diffusion_policy", policy_type)
            profile = POLICY_PROFILES["diffusion_policy"]
            policy_type = "diffusion_policy"

        # Compute fit scores
        ep_suff = _episode_sufficiency(n_episodes, task_complexity, profile)
        action_compat = _action_space_compatibility(action_dim, profile)
        consist_compat = _consistency_compatibility(
            consistency_score,
            multimodality_score,
            profile,
        )

        # Overall fit
        fit_score = ep_suff * 0.30 + action_compat * 0.20 + consist_compat * 0.50
        fit_score = float(np.clip(fit_score, 0.0, 1.0))

        # Generate warnings
        warnings = []
        if ep_suff < 0.5:
            min_needed = profile["min_episodes_simple"] + task_complexity * (
                profile["min_episodes_complex"] - profile["min_episodes_simple"]
            )
            warnings.append(
                f"Only {n_episodes} episodes — {policy_type} needs ~{int(min_needed)} "
                f"minimum for this task complexity"
            )

        if multimodality_score > profile["multimodality_tolerance"]:
            warnings.append(
                f"Data shows multimodal strategies (score={multimodality_score:.2f}) — "
                f"{policy_type} has tolerance of {profile['multimodality_tolerance']:.2f}. "
                f"Consider Diffusion Policy instead."
            )

        if consistency_score < 0.4 and profile["consistency_sensitivity"] > 0.6:
            warnings.append(
                f"Low demonstration consistency ({consistency_score:.2f}) — "
                f"{policy_type} is sensitive to inconsistency"
            )

        if action_compat < 0.7:
            warnings.append(
                f"Action dim {action_dim} is outside {policy_type}'s optimal range "
                f"{profile['action_dim_range']}"
            )

        breakdown = {
            "episode_sufficiency": ep_suff,
            "action_compatibility": action_compat,
            "consistency_compatibility": consist_compat,
        }

        return PolicyFit(
            fit_score=fit_score,
            episode_sufficiency=ep_suff,
            action_space_compatibility=action_compat,
            consistency_compatibility=consist_compat,
            policy_type=policy_type,
            specific_warnings=warnings,
            fit_breakdown=breakdown,
            recommended_policies=recommended,
        )
