"""
kiri/nn/loss.py
"""

import numpy as np
from kiri.backend.detect import BACKEND
from kiri.tensor import _CPUTensor

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


def cross_entropy(logits, targets):
    """
    logits : (N, C) raw scores (pre-softmax)
    targets: (N,) integer class indices
    """
    if BACKEND == "mlx":
        return mlx_nn.losses.cross_entropy(logits, targets).mean()
    else:
        N = logits.data.shape[0]
        shifted = logits.data - logits.data.max(axis=1, keepdims=True)
        exp_x   = np.exp(shifted)
        probs   = exp_x / exp_x.sum(axis=1, keepdims=True)
        t = targets.data.astype(int).ravel() if isinstance(targets, _CPUTensor) else targets.astype(int).ravel()
        log_probs = -np.log(probs[np.arange(N), t] + 1e-9)
        loss = _CPUTensor(np.array(log_probs.mean(), dtype=np.float32), (logits,), "cross_entropy")

        def _backward():
            dl = probs.copy()
            dl[np.arange(N), t] -= 1
            dl /= N
            logits.grad += dl * loss.grad
        loss._backward = _backward
        return loss


def mse_loss(pred, target):
    """Mean squared error."""
    if BACKEND == "mlx":
        return mx.mean((pred - target) ** 2)
    else:
        if not isinstance(target, _CPUTensor):
            target = _CPUTensor(np.asarray(target, dtype=np.float32), requires_grad=False)
        diff = pred - target
        return (diff * diff).mean()


def binary_cross_entropy(pred, target):
    """
    pred  : (N,) sigmoid outputs
    target: (N,) 0/1 labels
    """
    if BACKEND == "mlx":
        return mlx_nn.losses.binary_cross_entropy(pred, target).mean()
    else:
        if not isinstance(target, _CPUTensor):
            target = _CPUTensor(np.asarray(target, dtype=np.float32), requires_grad=False)
        eps = 1e-9
        p = pred.data.clip(eps, 1 - eps)
        t = target.data
        val = -(t * np.log(p) + (1 - t) * np.log(1 - p)).mean()
        loss = _CPUTensor(np.array(val, dtype=np.float32), (pred,), "bce")

        def _backward():
            dl = -(t / (p + eps) - (1 - t) / (1 - p + eps)) / len(p)
            pred.grad += dl * loss.grad
        loss._backward = _backward
        return loss
