import mlx.core as mx
import mlx.nn as nn


def length_masked_ce_loss(model, inputs, targets, lengths):
    # Run model on inputs
    logits = model(inputs).astype(mx.float32)

    # Mask padding tokens
    length_mask = mx.arange(inputs.shape[1])[None, :] < lengths[:, None]

    # Calculate the loss
    ce = nn.losses.cross_entropy(logits, targets) * length_mask
    loss = ce.sum() / length_mask.sum()
    #  print(f"|    {inputs=}\n| ==>{logits=}\n| ~^~{ce=}\n| == {loss=}")
    return loss


# Naive intermediate layer loss, where we replace the targets with gradients and just multiply the output by the gradients to derive the loss. This is naive and may warrant some further iteration, but will do the job for now
def back_gradient_loss(model, inputs, gradients, lengths):
    out = model(inputs).astype(mx.float32)
    grad = gradients.astype(mx.float32)

    # Mask padding tokens
    length_mask = mx.repeat(
        mx.arange(inputs.shape[1])[None, :] < lengths[:, None], out.shape[-1]
    ).reshape(out.shape)

    masked_sum = (out * length_mask).sum(axis=1)
    gradient_lens = mx.abs(grad * masked_sum)
    loss = gradient_lens.sum() / length_mask.sum()
    #  print(f"|    {inputs=}\n"
    #      + f"| ==>{out=}\n"
    #      + f"| ~^~{masked_sum=}\n"
    #      + f"| <~>{gradient_lens=}\n"
    #      + f"| == {loss=}")
    return loss


loss_fns = {
    "back_gradient": back_gradient_loss,
    "length_masked_ce": length_masked_ce_loss,
}
