"""Scaling law advisor based on ICLR 2025 scaling findings.

Uses the power-law relationship: success gap ∝ (env × obj pairs)^(-0.44)
to recommend whether users need more episodes or more diversity.

Key finding: 32 environment-object pairs × 50 demos each ≈ 90% success.
After ~50 demos per setting, more demos don't help — only more diversity helps.

Reference:
    "Scaling Laws for Imitation Learning in Single-Task Robotic Manipulation",
    ICLR 2025.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass

import numpy as np

logger = logging.getLogger(__name__)

# Power-law exponent from ICLR 2025
_SCALING_EXPONENT = -0.44

# Reference point: 32 pairs × 50 demos = ~90% success
_REFERENCE_PAIRS = 32
_REFERENCE_DEMOS_PER_PAIR = 50
_REFERENCE_SUCCESS = 0.90

# Pre-computed coefficient: gap = coeff * pairs^exponent
# At 32 pairs: gap = 1.0 - 0.90 = 0.10
# 0.10 = coeff * 32^(-0.44) → coeff = 0.10 / 32^(-0.44)
_SCALING_COEFF = 0.10 / (_REFERENCE_PAIRS ** _SCALING_EXPONENT)


@dataclass
class ScalingAdvice:
    """Data collection recommendations based on scaling laws."""

    current_episodes: int
    estimated_diversity: int  # estimated distinct environment-object pairs
    demos_per_setting: float  # episodes / diversity
    recommended_episodes: int
    recommended_diversity: int
    expected_success_at_current: float
    expected_success_at_recommended: float
    bottleneck: str  # "diversity", "quantity", or "quality"
    advice: str
    details: str  # longer explanation


def _estimate_diversity(
    action_episodes: list[np.ndarray],
    state_episodes: list[np.ndarray] | None = None,
) -> int:
    """Estimate number of distinct environment-object pairs.

    Uses initial state clustering as a proxy: episodes starting from
    similar states are likely the same environment/object pair.

    Falls back to simple variance-based estimation if clustering fails.
    """
    n = len(action_episodes)
    if n < 3:
        return 1

    # Use initial states (or initial actions) for clustering
    if state_episodes and len(state_episodes) == n:
        initial = np.array([ep[0] for ep in state_episodes if len(ep) > 0])
    else:
        initial = np.array([ep[0] for ep in action_episodes if len(ep) > 0])

    if len(initial) < 3:
        return 1

    # Try HDBSCAN first, then simple variance-based estimation
    try:
        from sklearn.cluster import DBSCAN
        from sklearn.preprocessing import StandardScaler

        scaled = StandardScaler().fit_transform(initial)

        # DBSCAN with adaptive eps using the k-distance knee-point
        from sklearn.neighbors import NearestNeighbors

        k = min(5, len(scaled) - 1)
        nn = NearestNeighbors(n_neighbors=k)
        nn.fit(scaled)
        distances, _ = nn.kneighbors(scaled)
        k_distances = np.sort(distances[:, -1])

        # Find the knee point: largest gap in sorted k-distances.
        # Points within a cluster have small k-distances; the jump to
        # inter-cluster distances marks the natural eps.
        if len(k_distances) > 2:
            diffs = np.diff(k_distances)
            knee_idx = int(np.argmax(diffs))
            # eps should be just above the intra-cluster distance
            eps = float(k_distances[knee_idx] + diffs[knee_idx] * 0.1)
        else:
            eps = float(np.median(k_distances))
        eps = max(eps, 0.3)

        clustering = DBSCAN(eps=eps, min_samples=2).fit(scaled)
        labels = clustering.labels_
        n_clusters = len(set(labels) - {-1})

        # Count noise points as individual clusters
        n_noise = int(np.sum(labels == -1))
        estimated = max(1, n_clusters + n_noise)
        return min(estimated, n)

    except ImportError:
        pass

    # Fallback: variance-based estimation
    # High variance in initial states → multiple environments
    per_dim_var = np.var(initial, axis=0)
    mean_var = float(np.mean(per_dim_var))

    # Heuristic: variance > 0.5 per dim suggests multiple settings
    if mean_var > 2.0:
        estimated = min(n, max(2, int(n * 0.3)))
    elif mean_var > 0.5:
        estimated = min(n, max(2, int(n * 0.15)))
    else:
        estimated = 1

    return estimated


def _predict_success(diversity: int, demos_per_pair: float) -> float:
    """Predict success rate from diversity and demos using scaling law.

    success = 1 - coeff * diversity^(-0.44) * quantity_factor
    """
    if diversity < 1:
        diversity = 1

    # Diversity component: gap decreases as pairs^(-0.44)
    gap = _SCALING_COEFF * (diversity ** _SCALING_EXPONENT)

    # Quantity component: diminishing returns after 50 demos per setting
    if demos_per_pair >= _REFERENCE_DEMOS_PER_PAIR:
        quantity_factor = 1.0
    elif demos_per_pair >= 10:
        # Linear ramp from 10 to 50 demos
        quantity_factor = 0.5 + 0.5 * (demos_per_pair - 10) / (
            _REFERENCE_DEMOS_PER_PAIR - 10
        )
    else:
        # Below 10 demos: severe penalty
        quantity_factor = max(0.1, demos_per_pair / 20)

    predicted = 1.0 - gap / quantity_factor
    return float(np.clip(predicted, 0.05, 0.98))


def compute_scaling_advice(
    episode_count: int,
    action_episodes: list[np.ndarray],
    state_episodes: list[np.ndarray] | None = None,
) -> ScalingAdvice:
    """Estimate diversity and recommend collection strategy.

    Args:
        episode_count: Total number of episodes.
        action_episodes: Per-episode action arrays.
        state_episodes: Optional per-episode state arrays.

    Returns:
        ScalingAdvice with bottleneck analysis and recommendations.
    """
    diversity = _estimate_diversity(action_episodes, state_episodes)
    demos_per_setting = episode_count / max(1, diversity)

    current_success = _predict_success(diversity, demos_per_setting)

    # Determine bottleneck
    # Scenario A: more diversity (double the environments, keep total demos same)
    new_diversity_a = diversity * 2
    new_demos_a = episode_count / max(1, new_diversity_a)
    success_more_diversity = _predict_success(new_diversity_a, new_demos_a)

    # Scenario B: more quantity (double the demos, same environments)
    new_demos_b = demos_per_setting * 2
    success_more_quantity = _predict_success(diversity, new_demos_b)

    if diversity < 8 and demos_per_setting >= 30:
        bottleneck = "diversity"
    elif demos_per_setting < 20:
        bottleneck = "quantity"
    elif success_more_diversity > success_more_quantity + 0.05:
        bottleneck = "diversity"
    elif success_more_quantity > success_more_diversity + 0.05:
        bottleneck = "quantity"
    else:
        bottleneck = "quality"

    # Recommendations
    if bottleneck == "diversity":
        target_diversity = max(16, diversity * 3)
        target_episodes = target_diversity * min(50, int(demos_per_setting))
        recommended_diversity = target_diversity
        recommended_episodes = target_episodes
        target_success = _predict_success(
            target_diversity, target_episodes / target_diversity
        )
        more_envs = target_diversity - diversity
        advice = (
            f"Collect in {more_envs} more environments rather than "
            f"adding demos to existing {diversity}"
        )
        details = (
            f"You have ~{diversity} distinct environments with "
            f"~{int(demos_per_setting)} demos each. "
            f"After ~50 demos per environment, more demos give diminishing returns. "
            f"Collecting in {more_envs} new environments (50 demos each) is the "
            f"highest-impact next step.\n"
            f"Based on: ICLR 2025 power-law scaling (gap ∝ pairs^-0.44)"
        )
    elif bottleneck == "quantity":
        target_demos = max(50, int(demos_per_setting * 2))
        recommended_episodes = diversity * target_demos
        recommended_diversity = diversity
        target_success = _predict_success(diversity, target_demos)
        more_demos = recommended_episodes - episode_count
        advice = f"Collect ~{more_demos} more demonstrations ({target_demos}/environment)"
        details = (
            f"You have ~{int(demos_per_setting)} demos per environment. "
            f"The sweet spot is ~50 demos per environment before diversity matters more. "
            f"Collecting {more_demos} more demos will have the biggest impact.\n"
            f"Based on: ICLR 2025 scaling data"
        )
    else:
        recommended_episodes = episode_count
        recommended_diversity = diversity
        target_success = current_success + 0.05
        advice = "Episode count and diversity look balanced — focus on data quality"
        details = (
            f"You have ~{diversity} environments with ~{int(demos_per_setting)} demos each. "
            f"Neither more diversity nor more demos will help much — "
            f"focus on improving demonstration quality and consistency.\n"
            f"Consider: orbit curate to select the best episodes"
        )

    return ScalingAdvice(
        current_episodes=episode_count,
        estimated_diversity=diversity,
        demos_per_setting=round(demos_per_setting, 1),
        recommended_episodes=recommended_episodes,
        recommended_diversity=recommended_diversity,
        expected_success_at_current=round(current_success, 4),
        expected_success_at_recommended=round(target_success, 4),
        bottleneck=bottleneck,
        advice=advice,
        details=details,
    )
