"""Temporal alignment check for robot training datasets.

Detects state-action misalignment (lag), FPS jitter, missing frames,
and constant state dimensions that shouldn't be constant. Misaligned
data causes policies to learn the wrong state-action mapping.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field

import numpy as np

logger = logging.getLogger(__name__)


@dataclass
class TemporalAlignmentResult:
    """Result of temporal alignment analysis."""

    aligned: bool  # True if data appears well-aligned
    estimated_lag: int  # frames of lag between state and action (0 = good)
    fps_jitter: float  # std of frame intervals (lower = more stable)
    missing_frames: int  # estimated number of missing/dropped frames
    constant_dimensions: list[int]  # state dims that are unexpectedly constant
    severity: str  # "ok", "mild", "moderate", "severe"
    recommendation: str
    details: dict = field(default_factory=dict)


def compute_temporal_alignment(
    action_episodes: list[np.ndarray],
    state_episodes: list[np.ndarray] | None = None,
    timestamps: list[np.ndarray] | None = None,
    expected_fps: float | None = None,
) -> TemporalAlignmentResult:
    """Check temporal alignment between states and actions.

    Args:
        action_episodes: List of per-episode action arrays (T_i, action_dim).
        state_episodes: Optional list of per-episode state arrays.
        timestamps: Optional list of per-episode timestamp arrays.
        expected_fps: Expected frames per second (for jitter calculation).

    Returns:
        TemporalAlignmentResult with alignment status and diagnostics.
    """
    if not action_episodes:
        return TemporalAlignmentResult(
            aligned=True, estimated_lag=0, fps_jitter=0.0,
            missing_frames=0, constant_dimensions=[],
            severity="ok", recommendation="No data to check.",
        )

    # Step 1: Check state-action lag via cross-correlation
    estimated_lag = 0
    lag_details = {}
    if state_episodes and len(state_episodes) == len(action_episodes):
        estimated_lag, lag_details = _estimate_lag(action_episodes, state_episodes)

    # Step 2: Check FPS jitter and missing frames
    fps_jitter = 0.0
    missing_frames = 0
    fps_details = {}
    if timestamps:
        fps_jitter, missing_frames, fps_details = _check_fps(timestamps, expected_fps)

    # Step 3: Check for constant state dimensions
    constant_dims = []
    if state_episodes:
        constant_dims = _find_constant_dimensions(state_episodes)

    # Step 4: Determine severity — normalize lag by episode length
    mean_ep_length = float(np.mean([len(ep) for ep in action_episodes])) if action_episodes else 100.0
    lag_ratio = abs(estimated_lag) / max(mean_ep_length, 1.0)

    issues = []
    if abs(estimated_lag) >= 2:
        issues.append(f"State-action lag of {estimated_lag} frames ({lag_ratio:.1%} of episode)")
    elif abs(estimated_lag) == 1:
        issues.append(f"State-action lag of {estimated_lag} frame")

    if fps_jitter > 0.15:
        issues.append(f"FPS jitter: {fps_jitter:.3f}")
    if missing_frames > 0:
        issues.append(f"{missing_frames} missing frames detected")
    if constant_dims:
        issues.append(f"{len(constant_dims)} constant state dimension(s)")

    # Severity based on lag as fraction of episode, not raw frames
    if lag_ratio >= 0.05 or fps_jitter > 0.3 or missing_frames > 10:
        severity = "severe"
    elif lag_ratio >= 0.02 or fps_jitter > 0.15 or missing_frames > 0 or len(constant_dims) > 2:
        severity = "moderate"
    elif lag_ratio > 0 and lag_ratio < 0.02:
        severity = "negligible"
    elif abs(estimated_lag) > 0 or fps_jitter > 0.05 or constant_dims:
        severity = "mild"
    else:
        severity = "ok"

    aligned = severity in ("ok", "mild", "negligible")

    # Recommendation
    if severity == "severe":
        rec = (
            "Severe temporal misalignment detected. "
            "Check data collection pipeline for synchronization issues. "
            f"Issues: {'; '.join(issues)}"
        )
    elif severity == "moderate":
        rec = (
            "Moderate alignment issues detected. "
            f"Issues: {'; '.join(issues)}. "
            "Consider re-checking timestamps and data recording."
        )
    elif severity == "mild":
        rec = (
            "Minor alignment issues detected but should not significantly "
            f"impact training. Issues: {'; '.join(issues)}"
        )
    else:
        rec = "State-action alignment looks good. No timing issues detected."

    details = {**lag_details, **fps_details}
    if constant_dims:
        details["constant_state_dims"] = constant_dims

    return TemporalAlignmentResult(
        aligned=aligned,
        estimated_lag=estimated_lag,
        fps_jitter=round(fps_jitter, 4),
        missing_frames=missing_frames,
        constant_dimensions=constant_dims,
        severity=severity,
        recommendation=rec,
        details=details,
    )


def _estimate_lag(
    action_episodes: list[np.ndarray],
    state_episodes: list[np.ndarray],
    max_lag: int = 5,
) -> tuple[int, dict]:
    """Estimate lag between state velocity and action via cross-correlation.

    For each episode, computes the cross-correlation between state velocity
    magnitude and action magnitude at different lags. The lag with peak
    correlation across episodes is the estimated alignment offset.

    Returns:
        (estimated_lag, details_dict)
    """
    lag_votes = []

    for actions, states in zip(action_episodes, state_episodes):
        actions = np.asarray(actions, dtype=np.float64)
        states = np.asarray(states, dtype=np.float64)

        if actions.ndim == 1:
            actions = actions.reshape(-1, 1)
        if states.ndim == 1:
            states = states.reshape(-1, 1)

        T = min(len(actions), len(states))
        if T < max_lag * 3:
            continue

        actions = actions[:T]
        states = states[:T]

        # State velocity magnitude
        state_vel = np.diff(states, axis=0)
        state_vel_mag = np.linalg.norm(state_vel, axis=1)

        # Action magnitude
        action_mag = np.linalg.norm(actions[:-1], axis=1)  # match length with diff

        if len(state_vel_mag) < max_lag * 2 + 1:
            continue

        # Normalize
        sv = state_vel_mag - np.mean(state_vel_mag)
        am = action_mag - np.mean(action_mag)
        sv_std = np.std(sv)
        am_std = np.std(am)
        if sv_std < 1e-8 or am_std < 1e-8:
            continue

        sv = sv / sv_std
        am = am / am_std

        # Cross-correlation at different lags
        best_lag = 0
        best_corr = -1.0
        for lag in range(-max_lag, max_lag + 1):
            if lag >= 0:
                s = sv[lag:]
                a = am[:len(s)]
            else:
                a = am[-lag:]
                s = sv[:len(a)]
            n = min(len(s), len(a))
            if n < 5:
                continue
            corr = float(np.mean(s[:n] * a[:n]))
            if corr > best_corr:
                best_corr = corr
                best_lag = lag

        lag_votes.append(best_lag)

    if not lag_votes:
        return 0, {"lag_confidence": "no_data"}

    # Majority vote
    from collections import Counter
    vote_counts = Counter(lag_votes)
    estimated_lag = vote_counts.most_common(1)[0][0]
    agreement = vote_counts[estimated_lag] / len(lag_votes)

    return estimated_lag, {
        "lag_agreement": round(agreement, 3),
        "lag_votes": dict(vote_counts),
        "lag_confidence": "high" if agreement > 0.7 else ("medium" if agreement > 0.4 else "low"),
    }


def _check_fps(
    timestamps: list[np.ndarray],
    expected_fps: float | None,
) -> tuple[float, int, dict]:
    """Check FPS consistency and detect missing frames.

    Returns:
        (fps_jitter, missing_frame_count, details_dict)
    """
    all_intervals = []
    missing = 0

    for ts in timestamps:
        ts = np.asarray(ts, dtype=np.float64)
        if len(ts) < 2:
            continue

        intervals = np.diff(ts)
        intervals = intervals[intervals > 0]  # filter zero/negative
        if len(intervals) == 0:
            continue

        all_intervals.extend(intervals.tolist())

        # Detect missing frames: intervals > 1.8x median
        median_interval = float(np.median(intervals))
        if median_interval > 0:
            gaps = intervals > 1.8 * median_interval
            missing += int(np.sum(gaps))

    if not all_intervals:
        return 0.0, 0, {"fps": "unknown"}

    intervals_arr = np.array(all_intervals)
    mean_interval = float(np.mean(intervals_arr))
    std_interval = float(np.std(intervals_arr))
    jitter = std_interval / mean_interval if mean_interval > 0 else 0.0

    measured_fps = 1.0 / mean_interval if mean_interval > 0 else 0.0

    details = {
        "measured_fps": round(measured_fps, 1),
        "fps_std": round(std_interval, 4),
    }
    if expected_fps:
        details["expected_fps"] = expected_fps
        details["fps_deviation"] = round(abs(measured_fps - expected_fps), 1)

    return jitter, missing, details


def _find_constant_dimensions(
    state_episodes: list[np.ndarray],
    threshold: float = 1e-6,
) -> list[int]:
    """Find state dimensions that are constant across all episodes.

    These may indicate broken sensors or unused state channels.
    """
    all_states = []
    for s in state_episodes:
        arr = np.asarray(s, dtype=np.float64)
        if arr.ndim == 1:
            arr = arr.reshape(-1, 1)
        all_states.append(arr)

    if not all_states:
        return []

    # Check dimension consistency
    dims = [s.shape[1] for s in all_states if s.ndim == 2]
    if not dims:
        return []
    state_dim = min(dims)

    # Concatenate and check variance
    concat = np.concatenate([s[:, :state_dim] for s in all_states], axis=0)
    variances = np.var(concat, axis=0)

    constant = [int(d) for d in range(state_dim) if variances[d] < threshold]
    return constant
