"""MaskDiT-style denoiser for native packed radar wedges."""

from __future__ import annotations

import math
import json
from pathlib import Path
import numpy as np
import yaml

import torch
from torch import nn
from torch.nn import functional as F
from torch.utils.data import Dataset


def load_config(path):
    with open(path, encoding="utf-8") as handle:
        return yaml.safe_load(handle)


def geometry_mask(config):
    data = config["data"]
    mask = np.zeros((data["frames"], 4, 366, 180), dtype=bool)
    for elevation, (azimuths, ranges) in enumerate(zip(data["azimuth_counts"], data["range_bins"])):
        mask[:, elevation, :azimuths, :ranges] = True
    return mask


def make_sample(config, seed):
    rng = np.random.default_rng(seed); data = config["data"]
    volume = np.zeros(tuple(data["storage_shape"]), dtype=np.float32)
    phase = rng.uniform(0, 2*np.pi)
    for t in range(data["frames"]):
        for e, (na, nr) in enumerate(zip(data["azimuth_counts"], data["range_bins"])):
            az=np.arange(na,dtype=np.float32)[:,None]; radius=np.arange(nr,dtype=np.float32)[None]
            center=(80+12*t+8*e+7*np.sin(phase+t/3))%na; center_r=min(nr-15,42+4*t-5*e)
            da=np.minimum(abs(az-center),na-abs(az-center)); core=np.exp(-.5*((da/(14+2*e))**2+((radius-center_r)/18)**2))
            volume[t,e,:na,:nr]=np.clip(58*np.exp(-.22*e)*core+2*np.sin(radius/8+az/30+phase+t/2),0,70)
    valid=geometry_mask(config); observed=valid.copy()
    for t in range(data["history"]):
        for e,(na,nr) in enumerate(zip(data["azimuth_counts"],data["range_bins"])):
            width=max(1,int(round(na*data["missing_ratio"]))); start=int(rng.integers(0,na-width+1)); observed[t,e,start:start+width,:nr]=False
    packed=volume.copy();packed[~observed]=0
    return {"values":packed,"truth":volume,"validity":valid,"observed":observed}


class PackedRadarDataset(Dataset):
    def __init__(self,directory,split):
        self.files=sorted(Path(directory).glob(f"{split}_*.npz"))
        if not self.files: raise FileNotFoundError(f"no {split} samples; run fake_data.py")
    def __len__(self): return len(self.files)
    def __getitem__(self,index):
        with np.load(self.files[index]) as sample:
            return {key:torch.from_numpy(sample[key].copy()) for key in ("values","truth","validity","observed")}


class DiffusionSchedule:
    def __init__(self,steps,beta_start,beta_end,device):
        self.steps=steps;self.betas=torch.linspace(beta_start,beta_end,steps,device=device);self.alphas=1-self.betas;self.alpha_bars=torch.cumprod(self.alphas,0)
    def add_noise(self,clean,timestep):
        noise=torch.randn_like(clean);alpha=self.alpha_bars[timestep].view(-1,1,1,1,1);return alpha.sqrt()*clean+(1-alpha).sqrt()*noise,noise
    @torch.no_grad()
    def sample(self,model,history,observed):
        target=torch.randn(history.shape[0],5,*history.shape[2:],device=history.device)
        for step in reversed(range(self.steps)):
            t=torch.full((history.shape[0],),step,device=history.device,dtype=torch.long);clean=model(history,target,observed,t,None)
            target=clean if step==0 else self.alpha_bars[step-1].sqrt()*clean+(1-self.alpha_bars[step-1]).sqrt()*torch.randn_like(clean)
        return target


def evaluate_ensemble(ensemble,truth,valid,thresholds):
    mean=ensemble.mean(0);error=mean-truth;result={"mae":float(abs(error)[valid].mean()),"rmse":float(np.sqrt((error**2)[valid].mean()))}
    result["crps"]=float(abs(ensemble-truth[None])[:,valid].mean()-.5*abs(ensemble[:,None]-ensemble[None,:])[:,:,valid].mean())
    result["categorical"]={}
    for threshold in thresholds:
        forecast,observed=mean>=threshold,truth>=threshold;hits=int((forecast&observed&valid).sum());miss=int((~forecast&observed&valid).sum());false=int((forecast&~observed&valid).sum())
        result["categorical"][str(threshold)]={"csi":hits/max(hits+miss+false,1),"far":false/max(hits+false,1),"pod":hits/max(hits+miss,1)}
    result["coverage_ratio"]=[[float(np.isfinite(ensemble).all(0)[lead,e][valid[lead,e]].mean()) for e in range(4)] for lead in range(5)]
    return result


class DiTBlock(nn.Module):
    """Linear-cost latent attention block suitable for 24k native wedge tokens."""

    def __init__(self, hidden: int, heads: int, latent_tokens: int):
        super().__init__()
        self.latents = nn.Parameter(torch.randn(1, latent_tokens, hidden) / math.sqrt(hidden))
        self.read = nn.MultiheadAttention(hidden, heads, batch_first=True)
        self.write = nn.MultiheadAttention(hidden, heads, batch_first=True)
        self.norm1 = nn.LayerNorm(hidden)
        self.norm2 = nn.LayerNorm(hidden)
        self.mlp = nn.Sequential(nn.Linear(hidden, hidden * 4), nn.GELU(), nn.Linear(hidden * 4, hidden))

    def forward(self, tokens: torch.Tensor, conditioning: torch.Tensor) -> torch.Tensor:
        latent = self.latents.expand(tokens.shape[0], -1, -1) + conditioning[:, None]
        latent = latent + self.read(latent, self.norm1(tokens), self.norm1(tokens), need_weights=False)[0]
        tokens = tokens + self.write(self.norm1(tokens), latent, latent, need_weights=False)[0]
        return tokens + self.mlp(self.norm2(tokens + conditioning[:, None]))


class EchoCast3D(nn.Module):
    def __init__(self, config: dict):
        super().__init__()
        self.config = config
        model = config["model"]
        self.patch = tuple(model["patch_size"])
        patch_area = self.patch[0] * self.patch[1]
        self.embed = nn.Linear(config["data"]["frames"] * patch_area, model["hidden_dim"])
        self.time_embed = nn.Embedding(config["diffusion"]["steps"], model["hidden_dim"])
        self.elevation_embed = nn.Embedding(4, model["hidden_dim"])
        self.mask_token = nn.Parameter(torch.zeros(1, 1, model["hidden_dim"]))
        self.encoder = nn.ModuleList(
            [DiTBlock(model["hidden_dim"], model["heads"], model["latent_tokens"]) for _ in range(model["depth"])]
        )
        self.decoder = nn.ModuleList(
            [DiTBlock(model["hidden_dim"], model["heads"], model["latent_tokens"]) for _ in range(model["depth"])]
        )
        self.head = nn.Linear(model["hidden_dim"], config["data"]["target"] * patch_area)
        self.layouts = []
        offset = 0
        for e, (na, nr) in enumerate(zip(config["data"]["azimuth_counts"], config["data"]["range_bins"])):
            rows, cols = na // self.patch[0], nr // self.patch[1]
            count = rows * cols
            self.layouts.append((e, na, nr, rows, cols, offset, offset + count))
            offset += count
        self.token_count = offset

    def patchify(self, volume: torch.Tensor) -> torch.Tensor:
        parts = []
        for e, na, nr, _, _, _, _ in self.layouts:
            wedge = volume[:, :, e, :na, :nr]
            patches = F.unfold(wedge, kernel_size=self.patch, stride=self.patch).transpose(1, 2)
            parts.append(patches)
        return torch.cat(parts, dim=1)

    def token_validity(self, validity: torch.Tensor) -> torch.Tensor:
        parts = []
        for e, na, nr, _, _, _, _ in self.layouts:
            wedge = validity[:, :, e, :na, :nr].float()
            patch = F.unfold(wedge, kernel_size=self.patch, stride=self.patch)
            parts.append((patch > 0.5).all(dim=1))
        return torch.cat(parts, dim=1)

    def unpatchify(self, patches: torch.Tensor) -> torch.Tensor:
        batch = patches.shape[0]
        output = patches.new_zeros((batch, 5, 4, 366, 180))
        for e, na, nr, rows, cols, start, end in self.layouts:
            wedge = F.fold(
                patches[:, start:end].transpose(1, 2),
                output_size=(na, nr),
                kernel_size=self.patch,
                stride=self.patch,
            )
            output[:, :, e, :na, :nr] = wedge
        return output

    def forward(
        self,
        history: torch.Tensor,
        noisy_target: torch.Tensor,
        observed: torch.Tensor,
        timestep: torch.Tensor,
        token_mask: torch.Tensor | None,
    ) -> torch.Tensor:
        full = torch.cat((history, noisy_target), dim=1)
        tokens = self.embed(self.patchify(full))
        elevation_ids = torch.cat(
            [torch.full((end - start,), e, device=full.device) for e, _, _, _, _, start, end in self.layouts]
        )
        tokens = tokens + self.elevation_embed(elevation_ids)[None]
        if token_mask is not None:
            tokens = torch.where(token_mask[..., None], self.mask_token.expand_as(tokens), tokens)
        condition = self.time_embed(timestep)
        for block in self.encoder:
            tokens = block(tokens, condition)
        for block in self.decoder:
            tokens = block(tokens, condition)
        return self.unpatchify(self.head(tokens))


def masked_objective(
    model: EchoCast3D,
    prediction: torch.Tensor,
    truth: torch.Tensor,
    validity: torch.Tensor,
    token_mask: torch.Tensor,
    reconstruction_weight: float,
) -> tuple[torch.Tensor, dict[str, float]]:
    target_valid = validity[:, 3:]
    token_pixels = token_mask[:, :, None].expand(-1, -1, 5 * 9)
    masked_pixels = model.unpatchify(token_pixels.float()).bool()
    squared = (prediction - truth) ** 2
    unmasked = target_valid & ~masked_pixels
    masked = target_valid & masked_pixels
    score = squared[unmasked].mean()
    reconstruction = squared[masked].mean()
    loss = score + reconstruction_weight * reconstruction
    return loss, {"score": score.item(), "reconstruction": reconstruction.item()}