import json
import os
import torch
import torch.nn as nn
import torch.nn.functional as F

CONFIG = None


def _config():
    global CONFIG
    if CONFIG is None:
        here = os.path.dirname(os.path.abspath(__file__))
        with open(os.path.join(here, "config.json")) as f:
            CONFIG = json.load(f)
    return CONFIG


class RevIN(nn.Module):
    def __init__(self, C):
        super().__init__()
        self.g = nn.Parameter(torch.ones(C))
        self.b = nn.Parameter(torch.zeros(C))

    def norm(self, x):
        self.m = x.mean(1, keepdim=True)
        self.s = x.std(1, keepdim=True) + 1e-05
        return (x - self.m) / self.s * self.g + self.b


class FNO1D(nn.Module):
    def __init__(self, D, modes):
        super().__init__()
        self.modes = modes
        s = 1 / (D * D)
        self.w = nn.Parameter(s * torch.rand(modes, D, D, dtype=torch.cfloat))

    def forward(self, x):
        P = x.shape[1]
        xf = torch.fft.rfft(x, dim=1)
        m = min(self.modes, xf.shape[1])
        o = torch.zeros_like(xf)
        o[:, :m] = torch.einsum("bpd,pde->bpe", xf[:, :m], self.w[:m])
        return torch.fft.irfft(o, n=P, dim=1)


class Block(nn.Module):
    def __init__(self, D, modes, ff=2, drop=0.0):
        super().__init__()
        self.n1 = nn.LayerNorm(D)
        self.fno = FNO1D(D, modes)
        self.d1 = nn.Dropout(drop)
        self.n2 = nn.LayerNorm(D)
        self.ff = nn.Sequential(
            nn.Linear(D, D * ff), nn.GELU(), nn.Dropout(drop), nn.Linear(D * ff, D)
        )

    def forward(self, x):
        x = x + self.d1(self.fno(self.n1(x)))
        return x + self.ff(self.n2(x))


class FELA_Grid(nn.Module):
    def __init__(self, Fin, L, D=96, modes=6, nblk=3, nq=99, arch="dual"):
        super().__init__()
        self.L = L
        self.arch = arch
        self.center = L // 2
        self.nq = nq
        self.revin = RevIN(Fin)
        self.embed = nn.Linear(Fin, D)
        self.pos = nn.Parameter(0.02 * torch.randn(1, L, D))
        self.blocks = nn.ModuleList([Block(D, modes) for _ in range(nblk)])
        self.norm = nn.LayerNorm(D)
        if arch == "dual":
            self.direct = nn.Sequential(
                nn.Linear(Fin, D), nn.GELU(), nn.Linear(D, D), nn.GELU()
            )
            fuse_in = 2 * D
        else:
            fuse_in = D
        self.med = nn.Linear(fuse_in, 1)
        self.spread = nn.Linear(fuse_in, nq)
        self.register_buffer("qidx", torch.arange(nq))

    def forward(self, x):
        xc = x[:, self.center]
        xn = self.revin.norm(x)
        h = self.embed(xn) + self.pos
        for b in self.blocks:
            h = b(h)
        ctx = self.norm(h)[:, self.center]
        if self.arch == "dual":
            z = torch.cat([ctx, self.direct(xc)], dim=1)
        else:
            z = ctx
        med = torch.sigmoid(self.med(z))
        w = F.softplus(self.spread(z))
        half = self.nq // 2
        below = -torch.flip(torch.cumsum(torch.flip(w[:, :half], [1]), 1), [1])
        above = torch.cumsum(w[:, half + 1 :], 1)
        offs = (
            torch.cat([below, torch.zeros_like(w[:, half : half + 1]), above], dim=1)
            * 0.05
        )
        return torch.clamp(med + offs, 0, 1)


def expected_shape(track):
    t = _config()["tracks"][track]
    return (t["input_steps"], t["input_features"])


def validate_input(x, track):
    steps, feats = expected_shape(track)
    if not isinstance(x, torch.Tensor):
        raise TypeError(f"Expected a torch.Tensor, got {type(x)}")
    if x.dim() != 3:
        raise ValueError(
            f"{track}: expected a 3-D tensor (batch, steps, features), got shape {tuple(x.shape)}"
        )
    if x.shape[1] != steps or x.shape[2] != feats:
        raise ValueError(
            f"{track}: expected window (batch, {steps}, {feats}), got {tuple(x.shape)}. Solar windows are (.,6,20); wind windows are (.,12,15)."
        )
    return x


def preprocess_nwp(raw_window, track, mean=None, std=None):
    x = torch.as_tensor(raw_window, dtype=torch.float32)
    if x.dim() == 2:
        x = x.unsqueeze(0)
    if mean is not None and std is not None:
        mean = torch.as_tensor(mean, dtype=torch.float32)
        std = torch.as_tensor(std, dtype=torch.float32)
        x = (x - mean) / torch.clamp(std, min=1e-06)
    else:
        m = x.mean(dim=(0, 1), keepdim=True)
        s = x.std(dim=(0, 1), keepdim=True)
        x = (x - m) / torch.clamp(s, min=1e-06)
    return validate_input(x, track)


def _build_from_state(state, track):
    dims = _config()["tracks"][track]["dims"]
    model = FELA_Grid(
        dims["Fin"],
        dims["L"],
        D=dims["D"],
        modes=dims["modes"],
        nblk=dims["nblk"],
        nq=dims["nq"],
        arch=dims.get("arch", "dual"),
    )
    ref = model.state_dict()
    fixed = {}
    for k, v in state.items():
        t = ref.get(k)
        if t is not None and t.is_complex() and (not v.is_complex()):
            v = torch.view_as_complex(v.contiguous())
        fixed[k] = v
    fixed["qidx"] = model.qidx.clone()
    model.load_state_dict(fixed, strict=True)
    return model


def load_model(path_or_repo, track="solar", filename=None):
    path = path_or_repo
    fname = filename or _config()["tracks"][track]["weights_safetensors"]
    if os.path.isdir(path):
        path = os.path.join(path, fname)
    elif not os.path.exists(path):
        from huggingface_hub import hf_hub_download

        path = hf_hub_download(path_or_repo, fname)
    from safetensors.torch import load_file

    state = load_file(path)
    model = _build_from_state(state, track)
    model.eval()
    return model


def from_pretrained(repo_id, track="solar"):
    return load_model(repo_id, track=track)