"""Action divergence analysis — detects inconsistent demonstrator behavior.

When similar states lead to different actions, behavior cloning policies
get confused. This module clusters states and measures action variance
within each cluster to identify where the demonstrator was inconsistent.

Key metric: action_divergence_score (0-1, lower is better)
- 0.0 = perfectly consistent (same action in same state every time)
- 1.0 = maximally inconsistent (random actions regardless of state)
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field

import numpy as np

logger = logging.getLogger(__name__)


@dataclass
class DivergentCluster:
    """A state region where the demonstrator was inconsistent."""

    cluster_id: int
    num_frames: int
    divergence: float  # coefficient of variation within this cluster
    frame_indices: list[int] = field(default_factory=list)


@dataclass
class ActionDivergenceResult:
    """Result of action divergence analysis."""

    score: float | None  # 0-1, None if not enough data
    severity: str  # "low", "moderate", "high", "critical", or "unknown"
    num_clusters: int
    worst_clusters: list[DivergentCluster]
    estimated_success_penalty: float  # 0-0.20
    recommendation: str
    note: str | None = None  # extra context (e.g. "only 1 episode")


def compute_action_divergence(
    action_episodes: list[np.ndarray],
    state_episodes: list[np.ndarray] | None = None,
    max_clusters: int = 50,
) -> ActionDivergenceResult:
    """Compute action divergence score for a dataset.

    Clusters similar states and measures action variance within each cluster.
    High variance within a cluster means the demonstrator did different things
    in similar states — this confuses behavior cloning policies.

    Args:
        action_episodes: List of per-episode action arrays (T_i, action_dim).
        state_episodes: Optional list of per-episode state arrays. If None,
            uses action-integrated positions as state proxy.
        max_clusters: Maximum number of K-means clusters.

    Returns:
        ActionDivergenceResult with score, severity, and worst clusters.
    """
    if not action_episodes:
        return ActionDivergenceResult(
            score=None, severity="unknown", num_clusters=0,
            worst_clusters=[], estimated_success_penalty=0.0,
            recommendation="No data available for divergence analysis.",
            note="No action episodes provided",
        )

    # Need at least 2 episodes for meaningful divergence analysis
    if len(action_episodes) < 2:
        return ActionDivergenceResult(
            score=None, severity="unknown", num_clusters=0,
            worst_clusters=[], estimated_success_penalty=0.0,
            recommendation="Need at least 2 episodes for divergence analysis.",
            note="Only 1 episode — cannot measure cross-episode consistency",
        )

    # Step 1: Build state representations
    if state_episodes and len(state_episodes) == len(action_episodes):
        states = _build_state_matrix(state_episodes)
    else:
        # Use cumulative sum of actions as position proxy
        states = _build_state_from_actions(action_episodes)

    actions = np.concatenate(action_episodes, axis=0)
    if actions.ndim == 1:
        actions = actions.reshape(-1, 1)

    if states is None or len(states) != len(actions):
        return ActionDivergenceResult(
            score=None, severity="unknown", num_clusters=0,
            worst_clusters=[], estimated_success_penalty=0.0,
            recommendation="Could not build state representations.",
            note="State/action length mismatch",
        )

    total_frames = len(states)
    if total_frames < 10:
        return ActionDivergenceResult(
            score=None, severity="unknown", num_clusters=0,
            worst_clusters=[], estimated_success_penalty=0.0,
            recommendation="Not enough frames for divergence analysis.",
            note=f"Only {total_frames} total frames",
        )

    # Step 2: Normalize states
    state_std = np.std(states, axis=0, keepdims=True)
    state_std[state_std < 1e-8] = 1.0
    states_norm = (states - np.mean(states, axis=0, keepdims=True)) / state_std

    # Step 3: Cluster similar states via K-means
    n_clusters = min(len(action_episodes) * 2, max_clusters, total_frames // 5)
    n_clusters = max(2, n_clusters)

    labels = _simple_kmeans(states_norm, n_clusters, max_iter=20)

    # Step 4: Measure action variance within each cluster
    cluster_divergences = []
    cluster_infos = []

    for k in range(n_clusters):
        mask = labels == k
        n_in_cluster = int(np.sum(mask))
        if n_in_cluster < 2:
            continue

        cluster_actions = actions[mask]

        # Measure action spread within cluster relative to global action spread
        # Using global std as normalizer avoids division-by-near-zero when
        # cluster means are close to zero (common for centered actions).
        stds = np.std(cluster_actions, axis=0)
        global_std = np.std(actions, axis=0)
        global_std[global_std < 1e-8] = 1e-8

        cv = stds / global_std
        cluster_cv = float(np.mean(np.clip(cv, 0, 5)))  # clip extreme outliers

        cluster_divergences.append((cluster_cv, n_in_cluster))

        frame_indices = list(np.where(mask)[0][:20])  # store up to 20 frame indices
        cluster_infos.append(DivergentCluster(
            cluster_id=k,
            num_frames=n_in_cluster,
            divergence=round(cluster_cv, 4),
            frame_indices=frame_indices,
        ))

    if not cluster_divergences:
        return ActionDivergenceResult(
            score=0.0, severity="low", num_clusters=n_clusters,
            worst_clusters=[], estimated_success_penalty=0.0,
            recommendation="Clusters too small for reliable divergence measurement.",
        )

    # Step 5: Weighted average divergence (larger clusters matter more)
    cvs = np.array([c[0] for c in cluster_divergences])
    weights = np.array([c[1] for c in cluster_divergences], dtype=float)
    weights /= weights.sum()

    raw_score = float(np.average(cvs, weights=weights))

    # Normalize to 0-1 range: CV of 2.0+ maps to 1.0
    score = float(np.clip(raw_score / 2.0, 0.0, 1.0))

    # Sort clusters by divergence (worst first), take top 5
    cluster_infos.sort(key=lambda c: c.divergence, reverse=True)
    worst_5 = cluster_infos[:5]

    # Step 6: Estimate impact
    if score > 0.5:
        severity = "critical"
        penalty = 0.15 + (score - 0.5) * 0.10  # 0.15-0.20
    elif score > 0.3:
        severity = "high"
        penalty = 0.08 + (score - 0.3) * 0.35  # 0.08-0.15
    elif score > 0.15:
        severity = "moderate"
        penalty = 0.03 + (score - 0.15) * 0.33  # 0.03-0.08
    else:
        severity = "low"
        penalty = score * 0.20  # 0.0-0.03

    penalty = round(min(0.20, penalty), 3)

    # Recommendation
    if severity == "critical":
        rec = (
            "High action divergence detected. The demonstrator used very different "
            "strategies in similar states. Consider filtering to a single strategy "
            "or using Diffusion Policy which handles multimodality better."
        )
    elif severity == "high":
        rec = (
            "Significant action divergence. Review the worst clusters and consider "
            "removing inconsistent episodes or switching to a multimodal policy."
        )
    elif severity == "moderate":
        rec = (
            "Moderate action divergence. Some inconsistency in similar states. "
            "This may slightly reduce success rate."
        )
    else:
        rec = "Action divergence is low — demonstrator behavior is consistent."

    return ActionDivergenceResult(
        score=round(score, 4),
        severity=severity,
        num_clusters=n_clusters,
        worst_clusters=worst_5,
        estimated_success_penalty=penalty,
        recommendation=rec,
    )


def _build_state_matrix(state_episodes: list[np.ndarray]) -> np.ndarray | None:
    """Concatenate state episodes into a single matrix."""
    cleaned = []
    for s in state_episodes:
        arr = np.asarray(s, dtype=np.float64)
        if arr.ndim == 1:
            arr = arr.reshape(-1, 1)
        cleaned.append(arr)

    if not cleaned:
        return None

    return np.concatenate(cleaned, axis=0)


def _build_state_from_actions(action_episodes: list[np.ndarray]) -> np.ndarray | None:
    """Build state proxy from cumulative sum of actions (position integration)."""
    positions = []
    for ep in action_episodes:
        arr = np.asarray(ep, dtype=np.float64)
        if arr.ndim == 1:
            arr = arr.reshape(-1, 1)
        pos = np.cumsum(arr, axis=0)
        positions.append(pos)

    if not positions:
        return None

    return np.concatenate(positions, axis=0)


def _simple_kmeans(
    data: np.ndarray,
    k: int,
    max_iter: int = 20,
    seed: int = 42,
) -> np.ndarray:
    """Simple K-means clustering without sklearn dependency.

    Args:
        data: (N, D) array of data points.
        k: Number of clusters.
        max_iter: Maximum iterations.
        seed: Random seed for reproducibility.

    Returns:
        (N,) array of cluster labels.
    """
    rng = np.random.RandomState(seed)
    n = len(data)
    k = min(k, n)

    # Initialize centroids via random selection
    indices = rng.choice(n, size=k, replace=False)
    centroids = data[indices].copy()

    labels = np.zeros(n, dtype=int)

    for _ in range(max_iter):
        # Assign each point to nearest centroid
        # Use chunked computation to avoid huge memory allocation
        chunk_size = 5000
        for start in range(0, n, chunk_size):
            end = min(start + chunk_size, n)
            chunk = data[start:end]
            # (chunk_size, 1, D) - (1, k, D) -> (chunk_size, k)
            dists = np.sum((chunk[:, np.newaxis, :] - centroids[np.newaxis, :, :]) ** 2, axis=2)
            labels[start:end] = np.argmin(dists, axis=1)

        # Update centroids
        new_centroids = np.zeros_like(centroids)
        for j in range(k):
            mask = labels == j
            if np.any(mask):
                new_centroids[j] = np.mean(data[mask], axis=0)
            else:
                new_centroids[j] = centroids[j]

        if np.allclose(centroids, new_centroids, atol=1e-6):
            break
        centroids = new_centroids

    return labels
