#!/usr/bin/env python3
"""
helix_cdc/training/fibpi_train_loop.py

Generic FibPi-aware training loop that wraps each optimizer step
in runtime.on_core() for Tesla valve core cycling.

Usage:
    from helix_cdc.runtime import FibPiRuntime
    from helix_cdc.training import train_with_fibpi

    runtime = FibPiRuntime.from_env()
    stats = train_with_fibpi(model, optimizer, data_loader, runtime, max_steps=40)
    print(stats["core_counts"])
"""

from __future__ import annotations

import time
from collections import Counter
from typing import Any, Callable, Dict, Iterable, List, Optional, Union

try:
    import torch
    import torch.nn as nn
    TORCH_AVAILABLE = True
except ImportError:
    TORCH_AVAILABLE = False


def _move_to_device(batch: Any, device: str) -> Any:
    """Move batch tensors to device."""
    if not TORCH_AVAILABLE:
        return batch

    if isinstance(batch, torch.Tensor):
        return batch.to(device)
    elif isinstance(batch, dict):
        return {k: _move_to_device(v, device) for k, v in batch.items()}
    elif isinstance(batch, (list, tuple)):
        return type(batch)(_move_to_device(item, device) for item in batch)
    return batch


def train_with_fibpi(
    model: "nn.Module",
    optimizer: "torch.optim.Optimizer",
    data_loader: Iterable,
    runtime: Any,  # FibPiRuntime
    *,
    device: str = "cpu",
    max_steps: Optional[int] = None,
    loss_fn: Optional[Callable] = None,
    verbose: bool = False,
) -> Dict[str, Any]:
    """
    Run a standard training loop where every optimizer step
    is wrapped in runtime.on_core() for FibPi core cycling.

    Args:
        model: PyTorch model to train.
        optimizer: Optimizer for model parameters.
        data_loader: Iterable yielding batches.
        runtime: FibPiRuntime instance (or any object with on_core() context manager).
        device: Device to move batches to ("cpu" or "cuda").
        max_steps: Maximum number of optimizer steps. None = run full epoch.
        loss_fn: Optional loss function. If None, expects model to return loss
                 (e.g., HuggingFace style with labels).
        verbose: If True, print progress.

    Returns:
        Dict with:
            steps: int - total optimizer steps taken
            loss_start: float - first step's loss
            loss_end: float - last step's loss
            losses: List[float] - all step losses
            core_counts: Dict[int, int] - steps per core
            elapsed_sec: float - total wall time
            mean_step_time: float - average seconds per step
    """
    if not TORCH_AVAILABLE:
        raise ImportError("torch is required for train_with_fibpi")

    model.train()

    steps = 0
    losses: List[float] = []
    core_counts: Counter = Counter()
    loss_start: Optional[float] = None
    loss_end: float = 0.0

    start_time = time.perf_counter()

    for batch in data_loader:
        if max_steps is not None and steps >= max_steps:
            break

        # Move batch to device
        batch = _move_to_device(batch, device)

        # Run step on next FibPi core
        with runtime.on_core() as core:
            optimizer.zero_grad()

            # Handle different batch formats
            if isinstance(batch, dict):
                # HuggingFace style: model(**batch) returns object with .loss
                outputs = model(**batch)
                if hasattr(outputs, "loss"):
                    loss = outputs.loss
                elif loss_fn is not None:
                    # Custom loss function
                    loss = loss_fn(outputs, batch)
                else:
                    raise ValueError("Model output has no .loss and no loss_fn provided")
            elif isinstance(batch, (list, tuple)) and len(batch) >= 2:
                # Standard (input, target) format
                inputs, targets = batch[0], batch[1]
                outputs = model(inputs)
                if loss_fn is not None:
                    loss = loss_fn(outputs, targets)
                else:
                    raise ValueError("No loss_fn provided for (input, target) batch format")
            else:
                raise ValueError(f"Unsupported batch format: {type(batch)}")

            loss.backward()
            optimizer.step()

        # Track stats
        loss_val = loss.item()
        losses.append(loss_val)
        core_counts[core] += 1

        if loss_start is None:
            loss_start = loss_val
        loss_end = loss_val

        steps += 1

        if verbose and steps % max(1, (max_steps or 100) // 10) == 0:
            print(f"[FIBPI_TRAIN] step {steps}: loss={loss_val:.4f} core={core}")

    elapsed = time.perf_counter() - start_time

    return {
        "steps": steps,
        "loss_start": loss_start if loss_start is not None else 0.0,
        "loss_end": loss_end,
        "losses": losses,
        "core_counts": dict(core_counts),
        "elapsed_sec": elapsed,
        "mean_step_time": elapsed / max(steps, 1),
    }


def train_with_fibpi_se(
    model: "nn.Module",
    optimizer: "torch.optim.Optimizer",
    data_loader: Iterable,
    runtime: Any,  # FibPiRuntime
    *,
    se_fn: Callable[[Any], float],
    device: str = "cpu",
    max_steps: Optional[int] = None,
    loss_fn: Optional[Callable] = None,
    verbose: bool = False,
) -> Dict[str, Any]:
    """
    Run a training loop with Se-based (Symbolic Entropy) core routing.

    Instead of round-robin, uses the provided se_fn to compute an Se score
    for each batch, then routes to the appropriate core via SeRouter bucketing.

    High Se → high core index (complex patterns)
    Low Se  → low core index (simple patterns)

    Args:
        model: PyTorch model to train.
        optimizer: Optimizer for model parameters.
        data_loader: Iterable yielding batches.
        runtime: FibPiRuntime with a SeRouter attached via set_router().
        se_fn: Function that takes a batch and returns its Se score (float).
        device: Device to move batches to ("cpu" or "cuda").
        max_steps: Maximum number of optimizer steps. None = run full epoch.
        loss_fn: Optional loss function.
        verbose: If True, print progress including Se scores.

    Returns:
        Dict with:
            steps: int
            loss_start: float
            loss_end: float
            losses: List[float]
            core_counts: Dict[int, int]
            se_scores: List[float] - Se score per step
            mean_se: float - average Se across all steps
            elapsed_sec: float
            mean_step_time: float
    """
    if not TORCH_AVAILABLE:
        raise ImportError("torch is required for train_with_fibpi_se")

    model.train()

    steps = 0
    losses: List[float] = []
    se_scores: List[float] = []
    core_counts: Counter = Counter()
    loss_start: Optional[float] = None
    loss_end: float = 0.0

    start_time = time.perf_counter()

    for batch in data_loader:
        if max_steps is not None and steps >= max_steps:
            break

        # Move batch to device
        batch = _move_to_device(batch, device)

        # Compute Se score for this batch
        se = se_fn(batch)
        se_scores.append(se)

        # Get core from Se-routing (uses SeRouter if attached)
        core = runtime.next_core_for_se(se)

        # Run step on Se-selected core
        with runtime.on_core(core=core) as chosen_core:
            optimizer.zero_grad()

            # Handle different batch formats
            if isinstance(batch, dict):
                outputs = model(**batch)
                if hasattr(outputs, "loss"):
                    loss = outputs.loss
                elif loss_fn is not None:
                    loss = loss_fn(outputs, batch)
                else:
                    raise ValueError("Model output has no .loss and no loss_fn provided")
            elif isinstance(batch, (list, tuple)) and len(batch) >= 2:
                inputs, targets = batch[0], batch[1]
                outputs = model(inputs)
                if loss_fn is not None:
                    loss = loss_fn(outputs, targets)
                else:
                    raise ValueError("No loss_fn provided for (input, target) batch format")
            else:
                raise ValueError(f"Unsupported batch format: {type(batch)}")

            loss.backward()
            optimizer.step()

        # Track stats
        loss_val = loss.item()
        losses.append(loss_val)
        core_counts[chosen_core] += 1

        if loss_start is None:
            loss_start = loss_val
        loss_end = loss_val

        steps += 1

        if verbose and steps % max(1, (max_steps or 100) // 10) == 0:
            print(f"[FIBPI_SE] step {steps}: loss={loss_val:.4f} se={se:.3f} core={chosen_core}")

    elapsed = time.perf_counter() - start_time

    return {
        "steps": steps,
        "loss_start": loss_start if loss_start is not None else 0.0,
        "loss_end": loss_end,
        "losses": losses,
        "core_counts": dict(core_counts),
        "se_scores": se_scores,
        "mean_se": sum(se_scores) / max(len(se_scores), 1),
        "elapsed_sec": elapsed,
        "mean_step_time": elapsed / max(steps, 1),
    }


def train_epochs_with_fibpi(
    model: "nn.Module",
    optimizer: "torch.optim.Optimizer",
    data_loader: Iterable,
    runtime: Any,
    *,
    epochs: int = 1,
    device: str = "cpu",
    loss_fn: Optional[Callable] = None,
    verbose: bool = False,
) -> Dict[str, Any]:
    """
    Run multiple epochs of training with FibPi core cycling.

    Returns:
        Dict with:
            epochs: int
            total_steps: int
            epoch_losses: List[float] - average loss per epoch
            loss_start: float
            loss_end: float
            core_counts: Dict[int, int]
            elapsed_sec: float
    """
    if not TORCH_AVAILABLE:
        raise ImportError("torch is required for train_epochs_with_fibpi")

    model.train()

    total_steps = 0
    all_losses: List[float] = []
    epoch_losses: List[float] = []
    core_counts: Counter = Counter()
    loss_start: Optional[float] = None

    start_time = time.perf_counter()

    for epoch in range(epochs):
        epoch_step_losses: List[float] = []

        for batch in data_loader:
            batch = _move_to_device(batch, device)

            with runtime.on_core() as core:
                optimizer.zero_grad()

                if isinstance(batch, dict):
                    outputs = model(**batch)
                    if hasattr(outputs, "loss"):
                        loss = outputs.loss
                    elif loss_fn is not None:
                        loss = loss_fn(outputs, batch)
                    else:
                        raise ValueError("Model output has no .loss and no loss_fn provided")
                elif isinstance(batch, (list, tuple)) and len(batch) >= 2:
                    inputs, targets = batch[0], batch[1]
                    outputs = model(inputs)
                    if loss_fn is not None:
                        loss = loss_fn(outputs, targets)
                    else:
                        raise ValueError("No loss_fn provided")
                else:
                    raise ValueError(f"Unsupported batch format: {type(batch)}")

                loss.backward()
                optimizer.step()

            loss_val = loss.item()
            epoch_step_losses.append(loss_val)
            all_losses.append(loss_val)
            core_counts[core] += 1

            if loss_start is None:
                loss_start = loss_val

            total_steps += 1

        epoch_avg = sum(epoch_step_losses) / max(len(epoch_step_losses), 1)
        epoch_losses.append(epoch_avg)

        if verbose:
            print(f"[FIBPI_TRAIN] epoch {epoch+1}/{epochs}: avg_loss={epoch_avg:.4f}")

    elapsed = time.perf_counter() - start_time

    return {
        "epochs": epochs,
        "total_steps": total_steps,
        "epoch_losses": epoch_losses,
        "loss_start": loss_start if loss_start is not None else 0.0,
        "loss_end": all_losses[-1] if all_losses else 0.0,
        "core_counts": dict(core_counts),
        "elapsed_sec": elapsed,
        "mean_step_time": elapsed / max(total_steps, 1),
    }
