"""Frame-level phase detection from visual embeddings or action data.

Clusters frames WITHIN episodes temporally to find where the robot transitions
between phases (approach, grasp, transport, place) without any labels.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field

import numpy as np

from orbit.analyzer.task_inference import TaskProfile

logger = logging.getLogger(__name__)


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


@dataclass
class DetectedPhase:
    """A single detected phase within the episode timeline."""

    name: str
    start_frac: float  # 0.0 to 1.0
    end_frac: float
    episodes_with_data: int = 0
    embedding_variance: float = 0.0
    coverage_score: float = 0.0
    issues: list[str] = field(default_factory=list)


@dataclass
class PhaseDetectionResult:
    """Full result from phase detection."""

    detected_phases: list[DetectedPhase] = field(default_factory=list)
    truncated_episodes: list[int] = field(default_factory=list)
    phase_coverage_scores: dict[str, float] = field(default_factory=dict)
    confidence: str = "low"  # "high", "medium", "low"
    visual_change_signal: np.ndarray | None = None

    def to_dict(self) -> dict:
        """Serialize to a JSON-safe dictionary."""
        return {
            "detected_phases": [
                {
                    "name": p.name,
                    "start_pct": round(p.start_frac * 100),
                    "end_pct": round(p.end_frac * 100),
                    "episodes_with_data": p.episodes_with_data,
                    "coverage_score": round(p.coverage_score, 2),
                    "issues": p.issues,
                }
                for p in self.detected_phases
            ],
            "truncated_episodes": self.truncated_episodes,
            "phase_coverage_scores": {
                k: round(v, 2) for k, v in self.phase_coverage_scores.items()
            },
            "confidence": self.confidence,
        }


# ---------------------------------------------------------------------------
# Core algorithm
# ---------------------------------------------------------------------------


def _compute_change_rate(embeddings: np.ndarray) -> np.ndarray:
    """Compute cosine distance between consecutive frames.

    Parameters
    ----------
    embeddings:
        Shape (N_frames, D) — embeddings for one episode.

    Returns
    -------
    change_rate:
        Shape (N_frames - 1,) — cosine distance between consecutive frames.
    """
    if len(embeddings) < 2:
        return np.zeros(1)

    # Normalize
    norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
    norms = np.where(norms == 0, 1.0, norms)
    normed = embeddings / norms

    # Cosine distance = 1 - cosine_similarity
    similarities = np.sum(normed[:-1] * normed[1:], axis=1)
    distances = 1.0 - np.clip(similarities, -1.0, 1.0)

    return distances


def _normalize_to_timesteps(signal: np.ndarray, n_steps: int = 100) -> np.ndarray:
    """Interpolate a 1-D signal to exactly n_steps points."""
    if len(signal) == 0:
        return np.zeros(n_steps)
    if len(signal) == 1:
        return np.full(n_steps, signal[0])

    x_orig = np.linspace(0, 1, len(signal))
    x_new = np.linspace(0, 1, n_steps)
    return np.interp(x_new, x_orig, signal)


def detect_phases(
    frame_embeddings: np.ndarray,
    task_profile: TaskProfile,
    episode_ids: list[int] | None = None,
) -> PhaseDetectionResult:
    """Detect task phases from per-episode frame embeddings.

    Parameters
    ----------
    frame_embeddings:
        Shape (N_episodes, N_frames, D) — dense frame embeddings.
    task_profile:
        Task profile with expected phases.
    episode_ids:
        Optional list of episode IDs for truncation reporting.

    Returns
    -------
    PhaseDetectionResult with detected phases and coverage scores.
    """
    n_episodes = frame_embeddings.shape[0]
    n_frames = frame_embeddings.shape[1]

    if n_episodes == 0 or n_frames < 2:
        return PhaseDetectionResult(confidence="low")

    n_steps = 100

    # (a) Compute visual change rate per episode
    all_change_rates: list[np.ndarray] = []
    episode_lengths: list[int] = []
    for ep_idx in range(n_episodes):
        ep_embs = frame_embeddings[ep_idx]
        # Skip episodes that are all zeros (padding)
        if np.all(ep_embs == 0):
            continue
        change = _compute_change_rate(ep_embs)
        normalized = _normalize_to_timesteps(change, n_steps)
        all_change_rates.append(normalized)
        episode_lengths.append(n_frames)

    if len(all_change_rates) < 2:
        return PhaseDetectionResult(confidence="low")

    # (b) Average change rate across episodes
    avg_change = np.mean(all_change_rates, axis=0)

    # (c) Detect boundaries using peaks
    try:
        from scipy.signal import find_peaks

        n_phases = len(task_profile.phases)
        n_boundaries = n_phases - 1

        if n_boundaries < 1:
            # Single phase — no boundaries needed
            boundaries = []
        else:
            # Find peaks with minimum distance
            min_dist = max(5, n_steps // (n_phases + 1))
            peaks, properties = find_peaks(avg_change, distance=min_dist)

            if len(peaks) >= n_boundaries:
                # Select top N-1 peaks by prominence
                prominences = avg_change[peaks]
                top_indices = np.argsort(prominences)[-n_boundaries:]
                boundaries = sorted(peaks[top_indices].tolist())
            elif len(peaks) > 0:
                # Use all available peaks
                boundaries = sorted(peaks.tolist())
            else:
                # No peaks — use uniform division
                boundaries = [int(n_steps * i / n_phases) for i in range(1, n_phases)]

    except ImportError:
        # scipy not available — use uniform division
        n_phases = len(task_profile.phases)
        boundaries = [int(n_steps * i / n_phases) for i in range(1, n_phases)]

    # Build phase segments
    phase_starts = [0] + boundaries
    phase_ends = boundaries + [n_steps]

    detected_phases: list[DetectedPhase] = []
    phase_coverage_scores: dict[str, float] = {}

    for i, phase_def in enumerate(task_profile.phases):
        if i < len(phase_starts):
            start = phase_starts[i]
            end = phase_ends[i] if i < len(phase_ends) else n_steps
        else:
            # More phases than boundaries — assign remaining to last segment
            start = phase_starts[-1] if phase_starts else 0
            end = n_steps

        start_frac = start / n_steps
        end_frac = end / n_steps

        # (d) Score each phase
        # Check how many episodes have meaningful data in this segment
        eps_with_data = 0
        segment_variances: list[float] = []

        for ep_idx in range(min(n_episodes, len(all_change_rates))):
            ep_embs = frame_embeddings[ep_idx]
            seg_start = int(start_frac * n_frames)
            seg_end = int(end_frac * n_frames)
            seg_end = max(seg_end, seg_start + 1)

            segment = ep_embs[seg_start:seg_end]
            if len(segment) > 0 and not np.all(segment == 0):
                eps_with_data += 1
                segment_variances.append(float(np.var(segment)))

        # Coverage score based on: episode coverage * variance
        coverage_frac = eps_with_data / max(n_episodes, 1)
        mean_var = float(np.mean(segment_variances)) if segment_variances else 0.0
        # Normalize variance to 0-1 range (sigmoid)
        var_score = float(2.0 * mean_var / (mean_var + 0.01)) if mean_var > 0 else 0.0
        coverage_score = float(np.clip(0.6 * coverage_frac + 0.4 * var_score, 0, 1))

        issues: list[str] = []
        if coverage_frac < 0.3:
            issues.append(f"Only {eps_with_data}/{n_episodes} episodes have data for this phase")
        if coverage_frac < 0.5:
            issues.append("MISSING — many episodes lack this phase")

        detected_phases.append(
            DetectedPhase(
                name=phase_def.name,
                start_frac=start_frac,
                end_frac=end_frac,
                episodes_with_data=eps_with_data,
                embedding_variance=mean_var,
                coverage_score=coverage_score,
                issues=issues,
            )
        )
        phase_coverage_scores[phase_def.name] = coverage_score

    # (e) Detect truncation
    truncated: list[int] = []
    if len(task_profile.phases) >= 3:
        # Check if last phase has significantly fewer episodes
        last_phase = detected_phases[-1] if detected_phases else None
        if last_phase and last_phase.episodes_with_data < n_episodes * 0.5:
            # Many episodes end before the last phase
            for ep_idx in range(n_episodes):
                ep_embs = frame_embeddings[ep_idx]
                # Check if the last 20% of frames are zero/padding
                cutoff = int(0.8 * n_frames)
                if np.all(ep_embs[cutoff:] == 0):
                    ep_id = (
                        episode_ids[ep_idx] if episode_ids and ep_idx < len(episode_ids) else ep_idx
                    )
                    truncated.append(ep_id)

    # Determine confidence
    valid_episodes = len(all_change_rates)
    peak_count = len(boundaries)
    expected_boundaries = len(task_profile.phases) - 1

    if valid_episodes >= 10 and peak_count >= expected_boundaries:
        confidence = "high"
    elif valid_episodes >= 5 and peak_count >= max(1, expected_boundaries - 1):
        confidence = "medium"
    else:
        confidence = "low"

    return PhaseDetectionResult(
        detected_phases=detected_phases,
        truncated_episodes=truncated,
        phase_coverage_scores=phase_coverage_scores,
        confidence=confidence,
        visual_change_signal=avg_change,
    )


# ---------------------------------------------------------------------------
# Action-based phase detection (no embeddings needed)
# ---------------------------------------------------------------------------


def detect_phases_from_actions(
    action_data: np.ndarray,
    task_profile: TaskProfile,
    episode_ids: list[int] | None = None,
) -> PhaseDetectionResult:
    """Detect phases from action data when embeddings aren't available.

    Parameters
    ----------
    action_data:
        Shape (N_episodes, N_frames, action_dim).
    task_profile:
        Task profile with expected phases.
    """
    n_episodes = action_data.shape[0]
    n_frames = action_data.shape[1]

    if n_episodes == 0 or n_frames < 4:
        return PhaseDetectionResult(confidence="low")

    n_steps = 100

    # Compute action magnitude change rate per episode
    all_change_rates: list[np.ndarray] = []
    for ep_idx in range(n_episodes):
        actions = action_data[ep_idx]  # (N_frames, action_dim)
        # Compute L2 norm of action differences
        if len(actions) < 2:
            continue
        diffs = np.linalg.norm(np.diff(actions, axis=0), axis=1)
        normalized = _normalize_to_timesteps(diffs, n_steps)
        all_change_rates.append(normalized)

    if len(all_change_rates) < 2:
        return PhaseDetectionResult(confidence="low")

    # Construct pseudo-embeddings from action data for the main algorithm
    # Shape: (N_episodes, N_frames, action_dim)
    return detect_phases(action_data, task_profile, episode_ids=episode_ids)
