"""Autoregressive forecasting utilities."""

import numpy as np
import torch
from torch.utils.data import DataLoader, Dataset

from ndvi_forecast._utils import to_numpy


def autoregressive_forecast(
    model,
    dataloader: DataLoader,
    input_len: int,
    target_len: int,
    step: int,
    device: str = "cpu",
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Run autoregressive forecast over a dataloader (teacher_forcing=0).

    Returns (preds, targets, chunked_preds) where chunked_preds is
    (N, num_chunks, step) with NaN padding in the last chunk.
    """
    model.eval()
    preds_all, targets_all, chunks_all = [], [], []

    with torch.no_grad():
        for xb, yb, in_lens, tgt_lens, gs in dataloader:
            xb = xb.to(device)
            yb = yb.to(device)
            in_lens = in_lens.to(device)
            gs = gs.to(device)

            B = xb.size(0)
            L_out = int(min(target_len, yb.size(1)))

            dummy = torch.zeros((B, L_out), device=device)
            preds = model(xb, in_lens, gs, target_seq=dummy, teacher_forcing_ratio=0.0)
            if preds.dim() == 3 and preds.size(-1) == 1:
                preds = preds.squeeze(-1)

            preds = preds[:, :L_out]
            targets = yb[:, :L_out]

            preds_all.append(preds.cpu())
            targets_all.append(targets.cpu())

            s = max(1, int(step))
            num_chunks = (L_out + s - 1) // s
            pc = torch.full((B, num_chunks, s), float("nan"))
            for k in range(num_chunks):
                a, b = k * s, min(L_out, (k + 1) * s)
                if b > a:
                    pc[:, k, : b - a] = preds[:, a:b].cpu()
            chunks_all.append(pc)

    preds = torch.cat(preds_all, dim=0) if preds_all else torch.empty(0, 0)
    targets = torch.cat(targets_all, dim=0) if targets_all else torch.empty(0, 0)
    chunks = torch.cat(chunks_all, dim=0) if chunks_all else torch.empty(0, 0, step)
    return preds, targets, chunks


def run_autoregressive_forecast(
    model,
    ndvi_row: np.ndarray,
    group_stats: np.ndarray,
    current_dap: int,
    input_len: int,
    step: int,
    max_dap: int = 127,
    device: str = "cpu",
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Single-row autoregressive forecast using raw DAP as the time channel.

    Returns (preds, targets, chunked_preds), each with batch dim 1.
    """
    step = int(step) if step and step > 0 else 1

    ndvi = np.clip(np.asarray(ndvi_row, dtype=np.float32), 0.0, 1.0)
    if not np.isfinite(ndvi).all():
        mask = np.isfinite(ndvi)
        if mask.any():
            idx = np.where(mask, np.arange(mask.size), 0)
            np.maximum.accumulate(idx, out=idx)
            ndvi = ndvi[idx]
        else:
            ndvi = np.zeros_like(ndvi)

    T = ndvi.shape[0]
    cap = min(max_dap, T)
    L_in = min(cap, int(current_dap) if current_dap is not None else cap,
               int(input_len) if input_len is not None else cap)
    L_in = max(0, L_in)

    if L_in >= cap:
        empty = torch.empty(1, 0, device=device)
        return empty, empty, torch.empty(1, 0, step, device=device)

    daps = np.arange(L_in, dtype=np.float32)
    inp_np = np.stack([ndvi[:L_in], daps], axis=1)
    input_seq = torch.from_numpy(inp_np).unsqueeze(0)

    gs = np.asarray(group_stats, dtype=np.float32).reshape(1, -1)
    if gs.shape[1] < 2:
        gs = np.concatenate([gs, np.zeros((1, 2 - gs.shape[1]), dtype=np.float32)], axis=1)
    elif gs.shape[1] > 2:
        gs = gs[:, :2]
    gs_tensor = torch.from_numpy(gs)

    tgt_np = ndvi[L_in:cap].astype(np.float32)
    if tgt_np.size == 0:
        empty = torch.empty(1, 0, device=device)
        return empty, empty, torch.empty(1, 0, step, device=device)
    target_seq = torch.from_numpy(tgt_np).unsqueeze(0)
    target_len = tgt_np.shape[0]

    class _OneRow(Dataset):
        def __len__(self):
            return 1
        def __getitem__(self, _):
            return (
                input_seq.squeeze(0),
                target_seq.squeeze(0),
                L_in,
                torch.tensor(0, dtype=torch.long),
                gs_tensor.squeeze(0),
            )

    loader = DataLoader(_OneRow(), batch_size=1, shuffle=False, num_workers=0)
    model.eval()
    with torch.no_grad():
        p, t, ch = autoregressive_forecast(model, loader, L_in, target_len, step, device)
    return p.to(device), t.to(device), ch.to(device)
