"""Loss functions and scheduling utilities for NDVI forecasting."""

from __future__ import annotations

import math
import numpy as np
import torch


def horizon_weights(T: int, gamma: float = 0.6, device=None) -> torch.Tensor:
    """Exponential horizon weighting emphasizing later timesteps."""
    t = torch.linspace(0, 1, steps=T, device=device)
    w = torch.exp(gamma * t)
    return w / w.mean().clamp_min(1e-8)


def make_time_mask(lengths: torch.Tensor, T: int, device=None) -> torch.Tensor:
    """Boolean mask (B, T) where True indicates valid timesteps."""
    device = device or lengths.device
    return (torch.arange(T, device=device).unsqueeze(0) < lengths.unsqueeze(1)).float()


def masked_mse(
    pred: torch.Tensor, tgt: torch.Tensor, lengths: torch.Tensor, w: torch.Tensor | None = None
) -> torch.Tensor:
    """MSE with masking for variable-length sequences."""
    B, T = pred.shape
    m = make_time_mask(lengths, T, device=pred.device)
    err2 = (pred - tgt) ** 2
    if w is not None:
        w = w[:T].to(pred.device)
        err2 = err2 * w.unsqueeze(0)
        denom = (m * w.unsqueeze(0)).sum().clamp_min(1)
    else:
        denom = m.sum().clamp_min(1)
    return (err2 * m).sum() / denom


def masked_slope_mse(
    pred: torch.Tensor, tgt: torch.Tensor, lengths: torch.Tensor, w: torch.Tensor | None = None
) -> torch.Tensor:
    """First-difference MSE with masking."""
    sp = pred[:, 1:] - pred[:, :-1]
    st = tgt[:, 1:] - tgt[:, :-1]
    m = make_time_mask((lengths - 1).clamp_min(0), sp.size(1), device=pred.device)
    err2 = (sp - st) ** 2
    if w is not None:
        w = w[: sp.size(1)].to(pred.device)
        err2 = err2 * w.unsqueeze(0)
        denom = (m * w.unsqueeze(0)).sum().clamp_min(1)
    else:
        denom = m.sum().clamp_min(1)
    return (err2 * m).sum() / denom


def masked_curvature_mse(
    pred: torch.Tensor, tgt: torch.Tensor, lengths: torch.Tensor, w: torch.Tensor | None = None
) -> torch.Tensor:
    """Second-difference MSE with masking."""
    cp = pred[:, 2:] - 2 * pred[:, 1:-1] + pred[:, :-2]
    ct = tgt[:, 2:] - 2 * tgt[:, 1:-1] + tgt[:, :-2]
    m = make_time_mask((lengths - 2).clamp_min(0), cp.size(1), device=pred.device)
    err2 = (cp - ct) ** 2
    if w is not None:
        w = w[: cp.size(1)].to(pred.device)
        err2 = err2 * w.unsqueeze(0)
        denom = (m * w.unsqueeze(0)).sum().clamp_min(1)
    else:
        denom = m.sum().clamp_min(1)
    return (err2 * m).sum() / denom


def simple_ema_smooth_targets(y: torch.Tensor, lengths: torch.Tensor, alpha: float = 0.1) -> torch.Tensor:
    """Training-only EMA smoothing on targets."""
    B, T = y.shape
    out = y.clone()
    for b in range(B):
        L = int(lengths[b].item())
        if L <= 1:
            continue
        prev = out[b, 0]
        for t in range(1, L):
            prev = alpha * out[b, t] + (1 - alpha) * prev
            out[b, t] = prev
    return out


def cosine_decay(init_val: float, final_val: float, progress: float) -> float:
    """Cosine annealing from init_val to final_val. progress in [0, 1]."""
    return final_val + 0.5 * (init_val - final_val) * (1 + math.cos(math.pi * progress))
