"""Proxy training — policy-matched models for data quality assessment.

Trains a small model matched to the target policy architecture on the user's
data for 500-1000 steps 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

Policy-matched architectures (from proxy_architectures.py):
- proxy_bc:        2-layer MLP, ~50K params   — basic BC learnability
- proxy_act:       Tiny Transformer, ~200-500K — action chunking learnability
- proxy_diffusion: 1D UNet + DDPM, ~300K      — multimodal distributions
- proxy_rnn:       1-layer LSTM, ~30K          — temporal sequence learning

References:
- DataMIL (2025): BC val loss is a good proxy for downstream rollout performance
- DataDecide (2025): small-model rankings transfer to large with 80% accuracy

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 loss on held-out episodes
    training_loss_final: float  # final loss 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
    # Per-episode prediction error (episode_index → mean loss)
    per_episode_errors: dict[int, float] = field(default_factory=dict)
    worst_episodes: list[int] = field(default_factory=list)
    # Architecture used
    architecture: str = "mlp"


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, list[int], list[int]]:
    """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.

    Returns:
        (train_states, train_actions, val_states, val_actions,
         train_episode_indices, val_episode_indices)
    """
    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 = [], []
    train_ep_indices: list[int] = []
    val_ep_indices: list[int] = []

    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)
            val_ep_indices.append(i)
        else:
            train_states.append(states)
            train_actions.append(actions)
            train_ep_indices.append(i)

    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),
        train_ep_indices,
        val_ep_indices,
    )


def _prepare_sequential_data(
    action_episodes: list[np.ndarray],
    state_episodes: list[np.ndarray] | None,
    validation_split: float,
    seq_len: int,
) -> tuple:
    """Prepare sequential (windowed) data for RNN/ACT/Diffusion training.

    Returns state/action windows of shape (N, seq_len, dim) plus episode indices.
    """
    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])

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

    train_s_windows, train_a_windows = [], []
    val_s_windows, val_a_windows = [], []
    train_ep_indices: list[int] = []
    val_ep_indices: list[int] = []

    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:
            if len(actions) < 2:
                continue
            states = actions[:-1]
            actions = actions[1:]

        T = len(actions)
        if T < seq_len:
            # Pad short episodes
            pad_s = np.zeros((seq_len - T, states.shape[1]), dtype=np.float32)
            pad_a = np.zeros((seq_len - T, actions.shape[1]), dtype=np.float32)
            states = np.concatenate([states, pad_s], axis=0)
            actions = np.concatenate([actions, pad_a], axis=0)
            if i in val_idx:
                val_s_windows.append(states[:seq_len])
                val_a_windows.append(actions[:seq_len])
                val_ep_indices.append(i)
            else:
                train_s_windows.append(states[:seq_len])
                train_a_windows.append(actions[:seq_len])
                train_ep_indices.append(i)
        else:
            # Sliding window
            for start in range(0, T - seq_len + 1, max(1, seq_len // 2)):
                s_win = states[start : start + seq_len]
                a_win = actions[start : start + seq_len]
                if i in val_idx:
                    val_s_windows.append(s_win)
                    val_a_windows.append(a_win)
                    val_ep_indices.append(i)
                else:
                    train_s_windows.append(s_win)
                    train_a_windows.append(a_win)
                    train_ep_indices.append(i)

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

    return (
        np.array(train_s_windows),
        np.array(train_a_windows),
        np.array(val_s_windows),
        np.array(val_a_windows),
        train_ep_indices,
        val_ep_indices,
    )


def _compute_per_episode_errors(
    model,
    action_episodes: list[np.ndarray],
    state_episodes: list[np.ndarray] | None,
    s_mean,
    s_std,
    a_mean,
    a_std,
    dev,
    architecture: str,
) -> dict[int, float]:
    """Compute per-episode prediction error for identifying problem episodes."""
    import torch

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

    errors: dict[int, float] = {}
    model.eval()

    with torch.no_grad():
        for i, ep_actions in enumerate(action_episodes):
            actions = np.asarray(ep_actions, 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:
                if len(actions) < 2:
                    continue
                states = actions[:-1]
                actions = actions[1:]

            if len(actions) < 2:
                continue

            s_t = torch.from_numpy(states).to(dev)
            a_t = torch.from_numpy(actions).to(dev)

            # Normalize
            s_t = (s_t - s_mean) / s_std
            a_t = (a_t - a_mean) / a_std

            if architecture == "mlp":
                pred = model(s_t)
                loss = torch.mean((pred - a_t) ** 2).item()
            elif architecture == "rnn":
                pred = model(s_t.unsqueeze(0))
                loss = torch.mean((pred.squeeze(0) - a_t) ** 2).item()
            else:
                # For ACT/diffusion, use simple MLP-style error as proxy
                # (the full architecture-specific error is captured in training)
                pred = model(s_t) if hasattr(model, 'net') else None
                if pred is not None:
                    loss = torch.mean((pred - a_t) ** 2).item()
                else:
                    continue

            errors[i] = round(loss, 6)

    return errors


# ---------------------------------------------------------------------------
# Training loops per architecture
# ---------------------------------------------------------------------------


def _train_mlp(
    train_s, train_a, val_s, val_a, state_dim, action_dim,
    dev, max_steps, hidden_dim, batch_size, lr, checkpoint_every,
):
    """Train MLP proxy (basic BC)."""
    import torch
    import torch.nn as nn

    from orbit.analyzer.proxy_architectures import ProxyMLP

    model = ProxyMLP(state_dim, action_dim, hidden_dim).to(dev)
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    loss_fn = nn.MSELoss()

    n_train = len(train_s)
    train_losses, val_losses = [], []

    model.train()
    for step in range(max_steps):
        idx = torch.randint(0, n_train, (min(batch_size, n_train),), device=dev)
        pred = model(train_s[idx])
        loss = loss_fn(pred, train_a[idx])

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

        if (step + 1) % checkpoint_every == 0 or step == 0:
            model.eval()
            with torch.no_grad():
                tl = loss_fn(model(train_s), train_a).item()
                vl = loss_fn(model(val_s), val_a).item()
            train_losses.append(round(tl, 6))
            val_losses.append(round(vl, 6))
            model.train()

    model.eval()
    with torch.no_grad():
        final_tl = loss_fn(model(train_s), train_a).item()
        final_vl = loss_fn(model(val_s), val_a).item()

    return model, train_losses, val_losses, final_tl, final_vl


def _train_rnn(
    action_episodes, state_episodes, validation_split,
    state_dim, action_dim, dev, max_steps, batch_size, lr, checkpoint_every,
    s_mean, s_std, a_mean, a_std,
):
    """Train LSTM proxy (temporal learning)."""
    import torch
    import torch.nn as nn

    from orbit.analyzer.proxy_architectures import ProxyRNN

    seq_len = 32
    train_s, train_a, val_s, val_a, _, _ = _prepare_sequential_data(
        action_episodes, state_episodes, validation_split, seq_len
    )

    # Convert and normalize
    train_s_t = (torch.from_numpy(train_s).to(dev) - s_mean) / s_std
    train_a_t = (torch.from_numpy(train_a).to(dev) - a_mean) / a_std
    val_s_t = (torch.from_numpy(val_s).to(dev) - s_mean) / s_std
    val_a_t = (torch.from_numpy(val_a).to(dev) - a_mean) / a_std

    model = ProxyRNN(state_dim, action_dim, hidden_dim=64).to(dev)
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    loss_fn = nn.MSELoss()

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

    model.train()
    for step in range(max_steps):
        idx = torch.randint(0, n_train, (min(batch_size, n_train),), device=dev)
        pred = model(train_s_t[idx])
        loss = loss_fn(pred, train_a_t[idx])

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

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

    model.eval()
    with torch.no_grad():
        final_tl = loss_fn(model(train_s_t), train_a_t).item()
        final_vl = loss_fn(model(val_s_t), val_a_t).item()

    return model, train_losses, val_losses, final_tl, final_vl


def _train_act(
    action_episodes, state_episodes, validation_split,
    state_dim, action_dim, dev, max_steps, batch_size, lr, checkpoint_every,
    s_mean, s_std, a_mean, a_std,
):
    """Train tiny Transformer proxy (ACT-style action chunking)."""
    import torch
    import torch.nn as nn

    from orbit.analyzer.proxy_architectures import ProxyACT

    chunk_size = 10
    train_s, train_a, val_s, val_a, _, _ = _prepare_sequential_data(
        action_episodes, state_episodes, validation_split, chunk_size
    )

    # For ACT: state is the first timestep, target is the full action chunk
    # train_s shape: (N, chunk_size, state_dim) → use first state
    # train_a shape: (N, chunk_size, action_dim) → predict all
    train_s_t = (torch.from_numpy(train_s[:, 0, :]).to(dev) - s_mean) / s_std
    train_a_t = (torch.from_numpy(train_a).to(dev) - a_mean) / a_std
    val_s_t = (torch.from_numpy(val_s[:, 0, :]).to(dev) - s_mean) / s_std
    val_a_t = (torch.from_numpy(val_a).to(dev) - a_mean) / a_std

    model = ProxyACT(state_dim, action_dim, chunk_size=chunk_size).to(dev)
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    loss_fn = nn.MSELoss()

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

    model.train()
    for step in range(max_steps):
        idx = torch.randint(0, n_train, (min(batch_size, n_train),), device=dev)
        pred = model(train_s_t[idx])
        loss = loss_fn(pred, train_a_t[idx])

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

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

    model.eval()
    with torch.no_grad():
        final_tl = loss_fn(model(train_s_t), train_a_t).item()
        final_vl = loss_fn(model(val_s_t), val_a_t).item()

    return model, train_losses, val_losses, final_tl, final_vl


def _train_diffusion(
    action_episodes, state_episodes, validation_split,
    state_dim, action_dim, dev, max_steps, batch_size, lr, checkpoint_every,
    s_mean, s_std, a_mean, a_std,
):
    """Train 1D UNet diffusion proxy (multimodal distributions)."""
    import torch

    from orbit.analyzer.proxy_architectures import ProxyDiffusion

    horizon = 16
    train_s, train_a, val_s, val_a, _, _ = _prepare_sequential_data(
        action_episodes, state_episodes, validation_split, horizon
    )

    # State: first timestep. Actions: (N, horizon, action_dim) → (N, action_dim, horizon)
    train_s_t = (torch.from_numpy(train_s[:, 0, :]).to(dev) - s_mean) / s_std
    train_a_t = (torch.from_numpy(train_a).to(dev) - a_mean) / a_std
    train_a_t = train_a_t.permute(0, 2, 1)  # (N, action_dim, horizon)

    val_s_t = (torch.from_numpy(val_s[:, 0, :]).to(dev) - s_mean) / s_std
    val_a_t = (torch.from_numpy(val_a).to(dev) - a_mean) / a_std
    val_a_t = val_a_t.permute(0, 2, 1)

    model = ProxyDiffusion(state_dim, action_dim, horizon=horizon).to(dev)
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)

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

    model.train()
    for step in range(max_steps):
        idx = torch.randint(0, n_train, (min(batch_size, n_train),), device=dev)
        loss = model.compute_loss(train_a_t[idx], train_s_t[idx])

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

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

    model.eval()
    with torch.no_grad():
        final_tl = model.compute_loss(train_a_t, train_s_t).item()
        final_vl = model.compute_loss(val_a_t, val_s_t).item()

    return model, train_losses, val_losses, final_tl, final_vl


# ---------------------------------------------------------------------------
# Main entry point
# ---------------------------------------------------------------------------


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,
    policy_type: str | None = None,
) -> ProxyTrainingResult:
    """Train a policy-matched proxy model and assess data quality.

    The proxy architecture is automatically selected based on the target
    policy type:
    - act → tiny Transformer with action chunking (1000 steps)
    - diffusion_policy/dp3 → 1D UNet with DDPM (1000 steps)
    - bc_rnn → single-layer LSTM (500 steps)
    - bc/smolvla/other → 2-layer MLP (500 steps)

    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.
        hidden_dim: Hidden layer size (MLP only).
        batch_size: Mini-batch size.
        lr: Learning rate.
        device: "cpu", "cuda", or "mps". Auto-detected if None.
        checkpoint_every: Record loss every N steps.
        policy_type: Target policy (act, diffusion_policy, bc, etc.).
            Auto-selects proxy architecture. None defaults to MLP.

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

    from orbit.analyzer.proxy_architectures import get_proxy_architecture

    start_time = time.time()

    # Select architecture
    architecture = get_proxy_architecture(policy_type or "bc")

    # For complex architectures, use more training steps
    if architecture in ("act", "diffusion") and max_steps <= 500:
        max_steps = 1000

    # Prepare flat data for normalization stats and MLP
    train_s, train_a, val_s, val_a, train_ep_idx, val_ep_idx = _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)

    # Compute normalization stats from flat training data
    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)

    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)

    # Normalize flat data (used by MLP and per-episode error computation)
    train_s_norm = (train_s_t - s_mean) / s_std
    train_a_norm = (train_a_t - a_mean) / a_std
    val_s_norm = (val_s_t - s_mean) / s_std
    val_a_norm = (val_a_t - a_mean) / a_std

    # Train the selected architecture
    logger.info("Proxy training: architecture=%s, steps=%d", architecture, max_steps)

    if architecture == "mlp":
        model, train_losses, val_losses, final_tl, final_vl = _train_mlp(
            train_s_norm, train_a_norm, val_s_norm, val_a_norm,
            state_dim, action_dim, dev, max_steps, hidden_dim,
            batch_size, lr, checkpoint_every,
        )
    elif architecture == "rnn":
        model, train_losses, val_losses, final_tl, final_vl = _train_rnn(
            action_episodes, state_episodes, validation_split,
            state_dim, action_dim, dev, max_steps, batch_size, lr,
            checkpoint_every, s_mean, s_std, a_mean, a_std,
        )
    elif architecture == "act":
        model, train_losses, val_losses, final_tl, final_vl = _train_act(
            action_episodes, state_episodes, validation_split,
            state_dim, action_dim, dev, max_steps, batch_size, lr,
            checkpoint_every, s_mean, s_std, a_mean, a_std,
        )
    elif architecture == "diffusion":
        model, train_losses, val_losses, final_tl, final_vl = _train_diffusion(
            action_episodes, state_episodes, validation_split,
            state_dim, action_dim, dev, max_steps, batch_size, lr,
            checkpoint_every, s_mean, s_std, a_mean, a_std,
        )
    else:
        raise ValueError(f"Unknown proxy architecture: {architecture}")

    training_time = time.time() - start_time

    # Per-episode error tracking — identifies which episodes the model
    # struggles with most.  High-error episodes are likely problematic.
    per_episode_errors: dict[int, float] = {}
    if architecture == "mlp":
        per_episode_errors = _compute_per_episode_errors(
            model, action_episodes, state_episodes,
            s_mean, s_std, a_mean, a_std, dev, architecture,
        )

    # For non-MLP architectures, compute per-episode errors using an
    # MLP-compatible forward pass on the flat data.  The episode-level
    # errors are approximate but still useful for identifying outliers.
    if architecture in ("rnn", "act", "diffusion") and not per_episode_errors:
        # Train a quick MLP on the same data for per-episode errors
        try:
            mlp_model, _, _, _, _ = _train_mlp(
                train_s_norm, train_a_norm, val_s_norm, val_a_norm,
                state_dim, action_dim, dev, 200, hidden_dim,
                batch_size, lr, checkpoint_every,
            )
            per_episode_errors = _compute_per_episode_errors(
                mlp_model, action_episodes, state_episodes,
                s_mean, s_std, a_mean, a_std, dev, "mlp",
            )
        except Exception:
            logger.debug("Per-episode error computation failed, skipping")

    # Identify worst episodes
    worst_episodes = []
    if per_episode_errors:
        ranked = sorted(per_episode_errors.items(), key=lambda x: x[1], reverse=True)
        worst_episodes = [idx for idx, _ in ranked[:5]]

    # Compute diagnostics
    overfitting_ratio = final_vl / max(final_tl, 1e-8)

    # Find convergence step
    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

    arch_label = {
        "mlp": "BC-MLP",
        "act": "ACT-mini",
        "diffusion": "DP-mini",
        "rnn": "LSTM",
    }.get(architecture, architecture.upper())

    if loss_dropped and overfitting_ratio < 2.0 and not loss_exploded:
        go_no_go = "GO"
        reason = (
            f"Validation loss dropped to {final_vl:.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_vl:.4f}). Data may be too noisy or contradictory "
            f"to learn from"
        )
    else:
        go_no_go = "CAUTION"
        reason = f"Inconclusive — val loss: {final_vl:.4f}"

    return ProxyTrainingResult(
        validation_loss_final=round(final_vl, 6),
        training_loss_final=round(final_tl, 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,
        per_episode_errors=per_episode_errors,
        worst_episodes=worst_episodes,
        architecture=architecture,
    )
