"""Compact consistency model for full-grid precipitation downscaling."""
import json
from pathlib import Path

import numpy as np
import torch
from torch import nn
import torch.nn.functional as F
import yaml


def load_config(root):
    return yaml.safe_load((Path(root) / "conf/config.yaml").read_text())


class TimeBlock(nn.Module):
    def __init__(self, cin, cout, time_dim):
        super().__init__()
        groups = min(4, cout)
        self.conv = nn.Conv2d(cin, cout, 3, padding=1)
        self.norm = nn.GroupNorm(groups, cout)
        self.time = nn.Linear(time_dim, cout)

    def forward(self, x, embedding):
        return F.silu(self.norm(self.conv(x)) + self.time(embedding)[:, :, None, None])


class ScaleAdaptiveCM(nn.Module):
    def __init__(self, channels=(4, 8, 16), time_dim=16, sigma_data=0.5):
        super().__init__()
        self.sigma_data = float(sigma_data)
        self.time_dim = int(time_dim)
        self.time_mlp = nn.Sequential(nn.Linear(time_dim, time_dim), nn.SiLU(), nn.Linear(time_dim, time_dim))
        self.enc1 = TimeBlock(1, channels[0], time_dim)
        self.enc2 = TimeBlock(channels[0], channels[1], time_dim)
        self.mid = TimeBlock(channels[1], channels[2], time_dim)
        self.dec2 = TimeBlock(channels[2] + channels[1], channels[1], time_dim)
        self.dec1 = TimeBlock(channels[1] + channels[0], channels[0], time_dim)
        self.out = nn.Conv2d(channels[0], 1, 1)
        self.model_config = {"channels": list(channels), "time_dim": time_dim, "sigma_data": sigma_data}

    def embed_time(self, t):
        half = self.time_dim // 2
        freq = torch.exp(torch.linspace(0, -7, half, device=t.device))
        emb = torch.cat((torch.sin(t[:, None] * freq), torch.cos(t[:, None] * freq)), 1)
        return self.time_mlp(emb)

    def forward(self, noisy, t):
        if noisy.ndim != 4 or noisy.shape[1] != 1:
            raise ValueError("expected [B,1,H,W]")
        emb = self.embed_time(t.float())
        e1 = self.enc1(noisy, emb)
        e2 = self.enc2(F.avg_pool2d(e1, 2), emb)
        mid = self.mid(F.avg_pool2d(e2, 2), emb)
        d2 = self.dec2(torch.cat((F.interpolate(mid, e2.shape[-2:], mode="bilinear", align_corners=False), e2), 1), emb)
        d1 = self.dec1(torch.cat((F.interpolate(d2, e1.shape[-2:], mode="bilinear", align_corners=False), e1), 1), emb)
        raw = self.out(d1)
        sigma2 = self.sigma_data ** 2
        cskip = sigma2 / ((t[:, None, None, None] - 0.002).square() + sigma2)
        cout = self.sigma_data * t[:, None, None, None] / torch.sqrt(t[:, None, None, None].square() + sigma2)
        return cskip * noisy + cout * raw


def structured_fields(samples, high_h, high_w, seed):
    rng = np.random.default_rng(seed)
    yy, xx = np.mgrid[-1:1:complex(high_h), -1:1:complex(high_w)]
    fields = []
    for i in range(samples):
        phase = 2 * np.pi * i / max(samples, 4)
        itcz = 9 * np.exp(-((yy - .12 * np.sin(phase)) / .16) ** 2)
        storms = 18 * np.exp(-((xx - .45 * np.cos(phase)) ** 2 + (yy - .3 * np.sin(phase)) ** 2) / .025)
        texture = 2 * np.maximum(0, np.sin(18 * xx + phase) * np.cos(13 * yy - phase))
        fields.append(np.maximum(0, itcz + storms + texture + rng.normal(0, .15, yy.shape)))
    return np.asarray(fields, np.float32)[:, None]


def radial_spectrum(field):
    power = np.abs(np.fft.fftshift(np.fft.fft2(field))) ** 2
    y, x = np.indices(field.shape); r = np.sqrt((y-field.shape[0]/2)**2 + (x-field.shape[1]/2)**2).astype(int)
    return np.bincount(r.ravel(), power.ravel()) / np.maximum(np.bincount(r.ravel()), 1)


def write_json(path, value):
    path = Path(path); path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(value, indent=2) + "\n")