"""
kiri/model.py

Base Model class — the friendly high-level API.
"""

import numpy as np
from kiri.backend.detect import BACKEND
from kiri.nn.layers import Module
from kiri.nn import cross_entropy, mse_loss
from kiri.optim import Adam

if BACKEND == "mlx":
    import mlx.core as mx


class Model(Module):
    """
    Base class for all Kiri models.
    Subclass this and implement forward().

    Example
    -------
    class MyCNN(kiri.Model):
        def __init__(self):
            self.conv = kiri.nn.Conv2d(1, 16, 3, padding=1)
            self.relu = kiri.nn.ReLU()
            self.fc   = kiri.nn.Linear(16*28*28, 10)

        def forward(self, x):
            x = self.relu(self.conv(x))
            x = kiri.nn.Flatten()(x)
            return self.fc(x)
    """

    def fit(self, data, y=None, epochs=10, lr=1e-3, batch_size=32,
            loss_fn=None, optimizer=None, verbose=True, val_data=None,
            scheduler=None):
        """
        Train the model.

        Parameters
        ----------
        data      : np.ndarray OR kiri.DataLoader
                    If ndarray, y must also be provided.
        y         : np.ndarray labels (only when data is ndarray)
        epochs    : number of full passes
        lr        : learning rate (ignored if optimizer provided)
        batch_size: mini-batch size (only when data is ndarray)
        loss_fn   : kiri.nn.cross_entropy | mse_loss | any callable
                    auto-detected from label dtype if not provided
        optimizer : kiri.optim instance — defaults to Adam
        verbose   : print epoch summaries
        val_data  : (X_val, y_val) tuple OR DataLoader
        scheduler : kiri.optim scheduler — .step() called each epoch
        """
        from kiri.data import DataLoader, TensorDataset

        # ── build loader ────────────────────────────────────────────────────
        if isinstance(data, DataLoader):
            loader = data
            # infer loss_fn from first batch dtype
            if loss_fn is None:
                sample_X, sample_y = next(iter(loader))
                loss_fn = cross_entropy if np.issubdtype(sample_y.dtype, np.integer) else mse_loss
        else:
            X, y_arr = data, y
            if loss_fn is None:
                loss_fn = cross_entropy if (
                    np.issubdtype(y_arr.dtype, np.integer) or
                    (y_arr.ndim == 1 and y_arr.dtype != np.float32)
                ) else mse_loss
            loader = DataLoader((X, y_arr), batch_size=batch_size, shuffle=True)

        params = self.parameters()
        if optimizer is None:
            optimizer = Adam(params, lr=lr)

        history = {"loss": [], "val_loss": []}
        ep_width = len(str(epochs))

        for epoch in range(epochs):
            epoch_loss = 0.0
            n_batches  = 0

            for Xb, yb in loader:
                if BACKEND == "mlx":
                    loss_val = self._train_step_mlx(Xb, yb, loss_fn, optimizer)
                else:
                    loss_val = self._train_step_cpu(Xb, yb, loss_fn, optimizer)
                epoch_loss += loss_val
                n_batches  += 1

            avg_loss = epoch_loss / max(n_batches, 1)
            history["loss"].append(avg_loss)

            if scheduler is not None:
                if hasattr(scheduler, 'step'):
                    # ReduceLROnPlateau needs the metric
                    import inspect
                    sig = inspect.signature(scheduler.step)
                    if len(sig.parameters) > 0 and 'metric' in str(sig):
                        scheduler.step(avg_loss)
                    else:
                        scheduler.step()

            metrics = {f"loss": avg_loss}

            if val_data is not None:
                if isinstance(val_data, DataLoader):
                    val_losses = []
                    for Xv, yv in val_data:
                        val_losses.append(self._eval_loss(Xv, yv, loss_fn))
                    val_loss = float(np.mean(val_losses))
                else:
                    X_val, y_val = val_data
                    val_loss = self._eval_loss(X_val, y_val, loss_fn)
                history["val_loss"].append(val_loss)
                metrics["val_loss"] = val_loss

            if verbose:
                lr_now = optimizer.lr
                parts  = [f"Epoch {epoch+1:>{ep_width}}/{epochs}"]
                for k, v in metrics.items():
                    parts.append(f"{k}: \033[1m{v:.4f}\033[0m")
                parts.append(f"lr: {lr_now:.2e}")
                print("  " + "  │  ".join(parts))

        return history

    # ── internal train steps ───────────────────────────────────────────────

    def _train_step_cpu(self, Xb, yb, loss_fn, optimizer):
        from kiri.tensor import _CPUTensor
        optimizer.zero_grad()

        x_t = _CPUTensor(Xb.astype(np.float32), requires_grad=False)
        y_t = yb if isinstance(yb, _CPUTensor) else yb

        logits = self.forward(x_t)
        loss   = loss_fn(logits, y_t)
        loss.backward()
        optimizer.step()

        return loss.item()

    def _train_step_mlx(self, Xb, yb, loss_fn, optimizer):
        import mlx.nn as mlx_nn
        x_t = mx.array(Xb.astype(np.float32))
        y_t = mx.array(yb)

        # mlx_nn.value_and_grad differentiates w.r.t. the module's
        # trainable parameters — requires self to be an mlx.nn.Module
        def loss_wrapper(model, x, y):
            return loss_fn(model.forward(x), y)

        loss_val, grads = mlx_nn.value_and_grad(self, loss_wrapper)(self, x_t, y_t)

        # Use MLX optimizer directly with the module + grad dict
        optimizer._mlx_opt.update(self, grads)
        mx.eval(self.parameters())
        return float(loss_val)

    def _eval_loss(self, X, y, loss_fn):
        from kiri.tensor import _CPUTensor
        if BACKEND == "mlx":
            x_t = mx.array(X.astype(np.float32))
            y_t = mx.array(y)
            loss = loss_fn(self.forward(x_t), y_t)
            mx.eval(loss)
            return float(loss)
        else:
            x_t = _CPUTensor(X.astype(np.float32), requires_grad=False)
            logits = self.forward(x_t)
            loss = loss_fn(logits, y)
            return loss.item()

    def predict(self, X, batch_size=256):
        """
        Run inference. Returns np.ndarray of class indices (for classification)
        or raw output array.
        """
        results = []
        for start in range(0, len(X), batch_size):
            Xb = X[start:start+batch_size].astype(np.float32)
            if BACKEND == "mlx":
                out = self.forward(mx.array(Xb))
                mx.eval(out)
                results.append(np.array(out))
            else:
                from kiri.tensor import _CPUTensor
                out = self.forward(_CPUTensor(Xb, requires_grad=False))
                results.append(out.data)

        return np.concatenate(results, axis=0)

    def predict_classes(self, X, batch_size=256):
        """Returns predicted class indices."""
        logits = self.predict(X, batch_size)
        return np.argmax(logits, axis=-1)

    def accuracy(self, X, y, batch_size=256):
        """Convenience method: returns accuracy as a float."""
        preds = self.predict_classes(X, batch_size)
        return float((preds == y).mean())

    def save(self, path):
        """Save model weights to a .npz file."""
        params = self.parameters()
        weights = {}
        for i, p in enumerate(params):
            if BACKEND == "mlx":
                mx.eval(p)
                weights[f"param_{i}"] = np.array(p)
            else:
                weights[f"param_{i}"] = p.data
        np.savez(path, **weights)
        print(f"✓ Model saved to {path}")

    def load(self, path):
        """Load weights from a .npz file."""
        data = np.load(path)
        params = self.parameters()
        for i, p in enumerate(params):
            key = f"param_{i}"
            if key not in data:
                print(f"  Warning: {key} not found in checkpoint")
                continue
            if BACKEND == "mlx":
                # in-place update
                p[...] = mx.array(data[key])
            else:
                p.data = data[key].astype(np.float32)
        print(f"✓ Model loaded from {path}")
