"""Proxy training — tiny behavior cloning model for data quality assessment.

Trains a 2-layer MLP on the user's data for 100-500 steps (~2-3 minutes
on CPU) and checks if the validation loss drops. This gives a real,
non-heuristic training signal:

- Val loss drops fast and stays low → data is learnable, go train full model
- Val loss doesn't drop → data is too noisy or contradictory, don't waste GPU
- Val loss drops then explodes → overfitting, need more/different data
- Train loss low but val loss high → multimodal data, consider Diffusion Policy

Requires: torch (optional dependency).
"""

from __future__ import annotations

import logging
import time
from dataclasses import dataclass, field

import numpy as np

logger = logging.getLogger(__name__)


@dataclass
class ProxyTrainingResult:
    """Result of proxy training for quality estimation."""

    validation_loss_final: float  # final MSE on held-out episodes
    training_loss_final: float  # final MSE on training episodes
    training_loss_curve: list[float]  # loss at each checkpoint
    validation_loss_curve: list[float]  # val loss at each checkpoint
    convergence_step: int | None  # step where loss reached 50% of final
    overfitting_ratio: float  # val_loss / train_loss (>2.0 = bad)
    go_no_go: str  # "GO", "CAUTION", or "NO-GO"
    reason: str  # human-readable explanation
    training_time_seconds: float
    steps_run: int
    state_dim: int
    action_dim: int


def _prepare_data(
    action_episodes: list[np.ndarray],
    state_episodes: list[np.ndarray] | None,
    validation_split: float,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Split episodes into train/val and flatten to (state, action) pairs.

    If no state data, uses lagged actions as proxy states:
    state_t = action_{t-1}, target = action_t.
    """
    n_episodes = len(action_episodes)
    rng = np.random.RandomState(42)
    indices = rng.permutation(n_episodes)
    n_val = max(1, int(n_episodes * validation_split))
    val_idx = set(indices[:n_val])
    train_idx = set(indices[n_val:])

    has_states = (
        state_episodes is not None
        and len(state_episodes) == n_episodes
        and all(len(s) > 0 for s in state_episodes)
    )

    train_states, train_actions = [], []
    val_states, val_actions = [], []

    for i in range(n_episodes):
        actions = np.asarray(action_episodes[i], dtype=np.float32)
        if actions.ndim == 1:
            actions = actions.reshape(-1, 1)

        if has_states:
            states = np.asarray(state_episodes[i], dtype=np.float32)
            if states.ndim == 1:
                states = states.reshape(-1, 1)
            min_len = min(len(states), len(actions))
            states = states[:min_len]
            actions = actions[:min_len]
        else:
            # Use lagged actions as proxy states
            if len(actions) < 2:
                continue
            states = actions[:-1]
            actions = actions[1:]

        if len(actions) < 2:
            continue

        if i in val_idx:
            val_states.append(states)
            val_actions.append(actions)
        else:
            train_states.append(states)
            train_actions.append(actions)

    if not train_states or not val_states:
        raise ValueError("Not enough episodes for train/val split")

    return (
        np.concatenate(train_states, axis=0),
        np.concatenate(train_actions, axis=0),
        np.concatenate(val_states, axis=0),
        np.concatenate(val_actions, axis=0),
    )


def run_proxy_training(
    action_episodes: list[np.ndarray],
    state_episodes: list[np.ndarray] | None = None,
    validation_split: float = 0.2,
    max_steps: int = 500,
    hidden_dim: int = 128,
    batch_size: int = 256,
    lr: float = 1e-3,
    device: str | None = None,
    checkpoint_every: int = 25,
) -> ProxyTrainingResult:
    """Train a tiny BC model and assess data quality from the training signal.

    Args:
        action_episodes: List of per-episode action arrays.
        state_episodes: Optional state arrays. If None, uses lagged actions.
        validation_split: Fraction of episodes held out for validation.
        max_steps: Maximum training steps (100-500 recommended).
        hidden_dim: Hidden layer size for the MLP.
        batch_size: Mini-batch size.
        lr: Learning rate.
        device: "cpu", "cuda", or "mps". Auto-detected if None.
        checkpoint_every: Record loss every N steps.

    Returns:
        ProxyTrainingResult with go/no-go signal and training curves.
    """
    try:
        import torch
        import torch.nn as nn
    except ImportError:
        raise ImportError(
            "Proxy training requires PyTorch. "
            "Install with: pip install orbit-robot[proxy]"
        )

    start_time = time.time()

    # Prepare data
    train_s, train_a, val_s, val_a = _prepare_data(
        action_episodes, state_episodes, validation_split
    )

    state_dim = train_s.shape[1]
    action_dim = train_a.shape[1]

    # Auto-detect device
    if device is None:
        if torch.cuda.is_available():
            device = "cuda"
        elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            device = "mps"
        else:
            device = "cpu"

    dev = torch.device(device)

    # Convert to tensors
    train_s_t = torch.from_numpy(train_s).to(dev)
    train_a_t = torch.from_numpy(train_a).to(dev)
    val_s_t = torch.from_numpy(val_s).to(dev)
    val_a_t = torch.from_numpy(val_a).to(dev)

    # Normalize inputs (important for stable training)
    s_mean = train_s_t.mean(dim=0)
    s_std = train_s_t.std(dim=0).clamp(min=1e-6)
    a_mean = train_a_t.mean(dim=0)
    a_std = train_a_t.std(dim=0).clamp(min=1e-6)

    train_s_t = (train_s_t - s_mean) / s_std
    val_s_t = (val_s_t - s_mean) / s_std
    train_a_t = (train_a_t - a_mean) / a_std
    val_a_t = (val_a_t - a_mean) / a_std

    # Model
    model = nn.Sequential(
        nn.Linear(state_dim, hidden_dim),
        nn.ReLU(),
        nn.Linear(hidden_dim, hidden_dim),
        nn.ReLU(),
        nn.Linear(hidden_dim, action_dim),
    ).to(dev)

    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    loss_fn = nn.MSELoss()

    n_train = len(train_s_t)
    train_losses = []
    val_losses = []

    # Training loop
    model.train()
    for step in range(max_steps):
        # Mini-batch
        idx = torch.randint(0, n_train, (min(batch_size, n_train),), device=dev)
        s_batch = train_s_t[idx]
        a_batch = train_a_t[idx]

        pred = model(s_batch)
        loss = loss_fn(pred, a_batch)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if (step + 1) % checkpoint_every == 0 or step == 0:
            model.eval()
            with torch.no_grad():
                train_loss = loss_fn(model(train_s_t), train_a_t).item()
                val_loss = loss_fn(model(val_s_t), val_a_t).item()
            train_losses.append(round(train_loss, 6))
            val_losses.append(round(val_loss, 6))
            model.train()

    # Final evaluation
    model.eval()
    with torch.no_grad():
        final_train_loss = loss_fn(model(train_s_t), train_a_t).item()
        final_val_loss = loss_fn(model(val_s_t), val_a_t).item()

    training_time = time.time() - start_time

    # Compute diagnostics
    overfitting_ratio = final_val_loss / max(final_train_loss, 1e-8)

    # Find convergence step (where loss first reached 50% of initial drop)
    convergence_step = None
    if len(val_losses) >= 2:
        initial_loss = val_losses[0]
        final_loss = val_losses[-1]
        halfway = initial_loss - (initial_loss - final_loss) * 0.5
        for i, vl in enumerate(val_losses):
            if vl <= halfway:
                convergence_step = (i + 1) * checkpoint_every
                break

    # Determine go/no-go
    loss_dropped = len(val_losses) >= 2 and val_losses[-1] < val_losses[0] * 0.8
    loss_exploded = len(val_losses) >= 3 and val_losses[-1] > val_losses[-2] * 1.5

    if loss_dropped and overfitting_ratio < 2.0 and not loss_exploded:
        go_no_go = "GO"
        reason = (
            f"Validation loss dropped to {final_val_loss:.4f} "
            f"(from {val_losses[0]:.4f}) — data is learnable"
        )
    elif loss_dropped and overfitting_ratio >= 2.0:
        go_no_go = "CAUTION"
        reason = (
            f"Model is overfitting (val/train ratio: {overfitting_ratio:.1f}x). "
            f"Consider more data or a policy that handles multimodality "
            f"(e.g., Diffusion Policy)"
        )
    elif loss_exploded:
        go_no_go = "CAUTION"
        reason = (
            f"Validation loss is rising — possible overfitting or "
            f"contradictory demonstrations. Consider orbit curate to remove outliers"
        )
    elif not loss_dropped:
        go_no_go = "NO-GO"
        reason = (
            f"Validation loss did not decrease ({val_losses[0]:.4f} → "
            f"{final_val_loss:.4f}). Data may be too noisy or contradictory "
            f"to learn from"
        )
    else:
        go_no_go = "CAUTION"
        reason = f"Inconclusive — val loss: {final_val_loss:.4f}"

    return ProxyTrainingResult(
        validation_loss_final=round(final_val_loss, 6),
        training_loss_final=round(final_train_loss, 6),
        training_loss_curve=train_losses,
        validation_loss_curve=val_losses,
        convergence_step=convergence_step,
        overfitting_ratio=round(overfitting_ratio, 2),
        go_no_go=go_no_go,
        reason=reason,
        training_time_seconds=round(training_time, 1),
        steps_run=max_steps,
        state_dim=state_dim,
        action_dim=action_dim,
    )
