"""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=100)

    # 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))

    # Step 5b: Detect centroid opposition — clusters in similar states
    # whose mean actions point in opposite directions.  Within-cluster CV
    # misses this when opposing strategies have the same shape (e.g.
    # sinusoids flipped 180°), because individual cluster variance stays
    # low even though the demonstrator did contradictory things.
    opposition_boost = _centroid_opposition_score(
        states_norm, actions, labels, n_clusters,
    )
    if opposition_boost > 0:
        # Blend: keep the higher of CV-based score and opposition score,
        # but allow opposition to lift the final score when CV is low.
        score = max(score, opposition_boost)

    # 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 _centroid_opposition_score(
    states_norm: np.ndarray,
    actions: np.ndarray,
    labels: np.ndarray,
    n_clusters: int,
) -> float:
    """Detect clusters in similar states whose action centroids oppose each other.

    For each pair of clusters that are close in state space, compute the
    cosine similarity of their mean action vectors.  Negative cosine
    similarity means the demonstrator chose opposite actions in similar
    states — a strong signal of conflicting strategies that within-cluster
    CV alone misses (e.g. opposing sinusoids have low per-cluster variance).

    Returns a score in [0, 1] representing the worst opposition found,
    weighted by cluster sizes.  Returns 0 if no meaningful opposition.
    """
    # Compute per-cluster state centroid and action centroid
    cluster_ids = []
    state_centroids = []
    action_centroids = []
    cluster_sizes = []
    for k in range(n_clusters):
        mask = labels == k
        n_k = int(np.sum(mask))
        if n_k < 2:
            continue
        cluster_ids.append(k)
        state_centroids.append(np.mean(states_norm[mask], axis=0))
        action_centroids.append(np.mean(actions[mask], axis=0))
        cluster_sizes.append(n_k)

    if len(cluster_ids) < 2:
        return 0.0

    state_centroids = np.array(state_centroids)
    action_centroids = np.array(action_centroids)
    cluster_sizes = np.array(cluster_sizes, dtype=float)

    # Compute pairwise state distances between cluster centroids
    n_c = len(cluster_ids)
    state_dists = np.zeros((n_c, n_c))
    for i in range(n_c):
        for j in range(i + 1, n_c):
            d = float(np.linalg.norm(state_centroids[i] - state_centroids[j]))
            state_dists[i, j] = d
            state_dists[j, i] = d

    # "Nearby" threshold: clusters within the 30th percentile of all
    # pairwise state distances are considered to be in similar states.
    upper_tri = state_dists[np.triu_indices(n_c, k=1)]
    if len(upper_tri) == 0:
        return 0.0
    proximity_threshold = float(np.percentile(upper_tri, 30))
    if proximity_threshold < 1e-8:
        proximity_threshold = float(np.median(upper_tri))
    if proximity_threshold < 1e-8:
        return 0.0

    # For nearby cluster pairs, compute cosine similarity of action centroids
    opposition_scores = []
    pair_weights = []
    for i in range(n_c):
        for j in range(i + 1, n_c):
            if state_dists[i, j] > proximity_threshold:
                continue
            a_i = action_centroids[i]
            a_j = action_centroids[j]
            norm_i = float(np.linalg.norm(a_i))
            norm_j = float(np.linalg.norm(a_j))
            if norm_i < 1e-8 or norm_j < 1e-8:
                continue
            cos_sim = float(np.dot(a_i, a_j) / (norm_i * norm_j))
            # Negative cosine = opposite directions.  Map to [0, 1]:
            # cos=-1 → opposition=1, cos=0 → opposition=0.5, cos>0 → <0.5
            opposition = float(np.clip((1.0 - cos_sim) / 2.0, 0.0, 1.0))
            # Only count meaningful opposition (cos < -0.2 → opposition > 0.6)
            if opposition > 0.6:
                opposition_scores.append(opposition)
                pair_weights.append(float(cluster_sizes[i] + cluster_sizes[j]))

    if not opposition_scores:
        return 0.0

    pair_weights = np.array(pair_weights)
    pair_weights /= pair_weights.sum()
    raw = float(np.average(opposition_scores, weights=pair_weights))

    # Scale: opposition of 0.6 → 0.15, opposition of 1.0 → 0.50
    # This puts centroid-opposed datasets solidly into "moderate" to "high"
    # divergence territory where the grader applies meaningful penalties.
    scaled = float(np.clip((raw - 0.6) * 1.25 + 0.15, 0.0, 0.50))
    return scaled


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 _detect_action_type(actions_concat: np.ndarray) -> str:
    """Detect if actions are velocity/delta commands or absolute positions.

    Absolute position actions have high lag-1 autocorrelation (the robot
    doesn't teleport) and a non-zero mean relative to their spread.
    Delta/velocity actions are centered near zero with low autocorrelation.

    Note: array_intelligence.py has a more sophisticated version that also
    detects normalized, delta, and torque types with confidence scores.
    This version is intentionally simpler — it only needs the binary
    answer for the state-proxy builder in _build_state_from_actions().
    """
    if actions_concat.shape[0] < 3:
        return "delta_velocity"

    autocorrs = []
    for d in range(actions_concat.shape[1]):
        col = actions_concat[:, d]
        if np.std(col) < 1e-10:
            continue
        corr = np.corrcoef(col[:-1], col[1:])[0, 1]
        if np.isfinite(corr):
            autocorrs.append(corr)

    if not autocorrs:
        return "delta_velocity"

    autocorr = float(np.mean(autocorrs))
    mean_ratio = float(np.abs(np.mean(actions_concat)) / (np.std(actions_concat) + 1e-8))

    if autocorr > 0.95 and mean_ratio > 0.5:
        return "absolute_position"
    return "delta_velocity"


def _build_state_from_actions(action_episodes: list[np.ndarray]) -> np.ndarray | None:
    """Build state proxy from actions.

    For delta/velocity actions: cumulative sum integrates to position.
    For absolute position actions: use raw values directly (cumsum
    would produce nonsensical trajectories).
    """
    if not action_episodes:
        return None

    # Concatenate to detect action type
    all_arr = []
    for ep in action_episodes:
        arr = np.asarray(ep, dtype=np.float64)
        if arr.ndim == 1:
            arr = arr.reshape(-1, 1)
        all_arr.append(arr)

    action_type = _detect_action_type(np.concatenate(all_arr, axis=0))

    positions = []
    for arr in all_arr:
        if action_type == "absolute_position":
            positions.append(arr)
        else:
            positions.append(np.cumsum(arr, axis=0))

    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
