"""
Training utilities for ScalarGrad.
"""

import time
from dataclasses import dataclass
from typing import List, Optional

from rich import box
from rich.live import Live
from rich.table import Table
from rich.progress import (
    Progress,
    BarColumn,
    TextColumn,
    TimeElapsedColumn,
    TimeRemainingColumn,
    MofNCompleteColumn,
)

from .core import Scalar
from .nn import Module
from .optim import Optimizer
from .loss import Loss
from .config import logger
from .theme import get_console


def _make_metrics_table() -> Table:
    """Build an empty training-metrics table."""
    tbl = Table(
        title="[table_header]Training Progress[/table_header]",
        box=box.ROUNDED,
        border_style="table_border",
        header_style="table_header",
        show_lines=False,
        expand=False,
        min_width=56,
    )
    tbl.add_column("Epoch",    style="label",   justify="right", width=8)
    tbl.add_column("Loss",     style="value",   justify="right", width=12)
    tbl.add_column("Accuracy", style="success", justify="right", width=12)
    tbl.add_column("Time (s)", style="info",    justify="right", width=10)
    return tbl


@dataclass
class TrainingMetrics:
    """Container for training metrics."""
    epoch: int
    loss: float
    accuracy: float = 0.0
    val_loss: Optional[float] = None
    val_accuracy: Optional[float] = None


class Trainer:
    """
    Trainer class for managing the training loop.
    
    Provides utilities for training, validation, and metric tracking.
    """
    
    def __init__(
        self,
        model: Module,
        optimizer: Optimizer,
        loss_fn: Loss
    ):
        """
        Initialize trainer.
        
        Args:
            model: Neural network model
            optimizer: Optimizer instance
            loss_fn: Loss function
        """
        self.model = model
        self.optimizer = optimizer
        self.loss_fn = loss_fn
        self.history: List[TrainingMetrics] = []
    
    def train_epoch(
        self,
        X: List[List[float]],
        y: List[float],
        batch_size: Optional[int] = None
    ) -> TrainingMetrics:
        """
        Train for one epoch.
        
        Args:
            X: Input data
            y: Target values
            batch_size: Batch size (None for full batch)
            
        Returns:
            TrainingMetrics for this epoch
        """
        self.model.train()
        
        # Convert to batches
        if batch_size is None:
            batch_size = len(X)
        
        total_loss = 0.0
        correct = 0
        
        for i in range(0, len(X), batch_size):
            batch_X = X[i:i + batch_size]
            batch_y = y[i:i + batch_size]
            
            # Forward pass
            inputs = [[Scalar(x) for x in row] for row in batch_X]
            predictions = [self.model(x) for x in inputs]
            
            # Ensure predictions is a list of Scalars
            if not isinstance(predictions[0], list):
                predictions = predictions
            else:
                # Handle multi-output case
                predictions = [p[0] if isinstance(p, list) else p for p in predictions]
            
            # Compute loss
            loss = self.loss_fn(predictions, batch_y)
            
            # Backward pass
            self.optimizer.zero_grad()
            loss.backward()
            
            # Update weights
            self.optimizer.step()
            
            # Track metrics
            total_loss += loss.data
            for pred, target in zip(predictions, batch_y):
                if (pred.data > 0) == (target > 0):
                    correct += 1
        
        avg_loss = total_loss / (len(X) // batch_size)
        accuracy = correct / len(X)
        
        return TrainingMetrics(epoch=len(self.history), loss=avg_loss, accuracy=accuracy)
    
    def fit(
        self,
        X: List[List[float]],
        y: List[float],
        epochs: int,
        batch_size: Optional[int] = None,
        verbose: bool = True
    ) -> List[TrainingMetrics]:
        """
        Train the model.

        Args:
            X: Input data
            y: Target values
            epochs: Number of training epochs
            batch_size: Batch size
            verbose: Whether to show a live progress table

        Returns:
            List of training metrics
        """
        if not verbose:
            return self._fit_silent(X, y, epochs, batch_size)

        # Read display preference from global config (import lazily)
        try:
            from .config import config as _cfg
            display = _cfg.training_display
        except Exception:
            display = "progress"

        if display == "table":
            return self._fit_table(X, y, epochs, batch_size)
        return self._fit_progress(X, y, epochs, batch_size)

    # ------------------------------------------------------------------
    # Progress-bar display (works reliably on every terminal)
    # ------------------------------------------------------------------
    def _fit_progress(
        self,
        X: List[List[float]],
        y: List[float],
        epochs: int,
        batch_size: Optional[int] = None,
    ) -> List[TrainingMetrics]:
        """Train with a single-line updating progress bar."""
        con = get_console()
        progress = Progress(
            TextColumn("[table_header]Training[/table_header]"),
            BarColumn(bar_width=30, style="progress_bar", complete_style="progress_bar"),
            MofNCompleteColumn(),
            TextColumn("[label]loss[/label] [value]{task.fields[loss]}[/value]"),
            TextColumn("[label]acc[/label]  [success]{task.fields[acc]}[/success]"),
            TimeElapsedColumn(),
            TimeRemainingColumn(),
            console=con,
        )
        with progress:
            task = progress.add_task(
                "training", total=epochs, loss="-.----", acc="--.--%"
            )
            for _epoch in range(epochs):
                metrics = self.train_epoch(X, y, batch_size)
                self.history.append(metrics)
                progress.update(
                    task,
                    advance=1,
                    loss=f"{metrics.loss:.4f}",
                    acc=f"{metrics.accuracy:.2%}",
                )

        self._print_summary(con)
        return self.history

    # ------------------------------------------------------------------
    # Table display (rich Live with auto_refresh=False)
    # auto_refresh=False means Live ONLY renders when we explicitly call
    # live.refresh() — no timer-driven repaints, no repeated headers.
    # ------------------------------------------------------------------
    def _fit_table(
        self,
        X: List[List[float]],
        y: List[float],
        epochs: int,
        batch_size: Optional[int] = None,
    ) -> List[TrainingMetrics]:
        """Train with a live-updating table (epoch / loss / accuracy / time)."""
        con = get_console()
        tbl = _make_metrics_table()
        log_every = max(1, epochs // 10)

        with Live(tbl, console=con, auto_refresh=False) as live:
            for epoch in range(epochs):
                t0 = time.monotonic()
                metrics = self.train_epoch(X, y, batch_size)
                elapsed = time.monotonic() - t0
                self.history.append(metrics)

                if epoch % log_every == 0 or epoch == epochs - 1:
                    tbl.add_row(
                        str(epoch),
                        f"{metrics.loss:.4f}",
                        f"{metrics.accuracy:.2%}",
                        f"{elapsed:.3f}",
                    )
                    live.refresh()  # redraw only here — no periodic repaints

        self._print_summary(con)
        return self.history

    def _print_summary(self, con) -> None:  # type: ignore[no-untyped-def]
        """Print the final training summary line."""
        if self.history:
            final = self.history[-1]
            con.print(
                f"[success]>> Training complete[/success] -- "
                f"[label]Final loss:[/label] [value]{final.loss:.4f}[/value]  "
                f"[label]Accuracy:[/label] [success]{final.accuracy:.2%}[/success]"
            )

    def _fit_silent(
        self,
        X: List[List[float]],
        y: List[float],
        epochs: int,
        batch_size: Optional[int] = None,
    ) -> List[TrainingMetrics]:
        """Train silently (no output) and return history."""
        for epoch in range(epochs):
            metrics = self.train_epoch(X, y, batch_size)
            self.history.append(metrics)
        return self.history