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 SpectralConv2d(nn.Module):
    def __init__(self, ci, co, m1, m2):
        super().__init__()
        self.m1, self.m2 = (m1, m2)
        s = 1 / (ci * co)
        self.w1 = nn.Parameter(s * torch.rand(ci, co, m1, m2, dtype=torch.cfloat))
        self.w2 = nn.Parameter(s * torch.rand(ci, co, m1, m2, dtype=torch.cfloat))

    def forward(self, x):
        B, C, Hh, Ww = x.shape
        xf = torch.fft.rfft2(x)
        o = torch.zeros(
            B, self.w1.shape[1], Hh, Ww // 2 + 1, dtype=torch.cfloat, device=x.device
        )
        o[:, :, : self.m1, : self.m2] = torch.einsum(
            "bixy,ioxy->boxy", xf[:, :, : self.m1, : self.m2], self.w1
        )
        o[:, :, -self.m1 :, : self.m2] = torch.einsum(
            "bixy,ioxy->boxy", xf[:, :, -self.m1 :, : self.m2], self.w2
        )
        return torch.fft.irfft2(o, s=(Hh, Ww))


class FNO2d(nn.Module):
    def __init__(self, in_ch=8, modes=12, width=32, L=3, proj_hidden=128):
        super().__init__()
        self.lift = nn.Conv2d(in_ch, width, 1)
        self.sp = nn.ModuleList(
            [SpectralConv2d(width, width, modes, modes) for _ in range(L)]
        )
        self.w = nn.ModuleList([nn.Conv2d(width, width, 1) for _ in range(L)])
        self.proj = nn.Sequential(
            nn.Conv2d(width, proj_hidden, 1), nn.GELU(), nn.Conv2d(proj_hidden, 1, 1)
        )

    def forward(self, x):
        h = self.lift(x)
        for sp, w in zip(self.sp, self.w):
            h = h + F.gelu(sp(h) + w(h))
        return self.proj(h)


def validate_input(x):
    a = _config()["arch"]
    if not isinstance(x, torch.Tensor):
        raise TypeError(f"Expected a torch.Tensor, got {type(x)}")
    if (
        x.dim() != 4
        or x.shape[1] != a["in_ch"]
        or x.shape[2] != a["grid_h"]
        or (x.shape[3] != a["grid_w"])
    ):
        raise ValueError(
            f"Expected an input of shape (batch, {a['in_ch']}, {a['grid_h']}, {a['grid_w']}), got {tuple(x.shape)}."
        )
    return x


def preprocess(raw_field, mean=None, std=None):
    x = torch.as_tensor(raw_field, dtype=torch.float32)
    if x.dim() == 3:
        x = x.unsqueeze(0)
    norm = _config()["norm"]
    m = torch.as_tensor(
        mean if mean is not None else norm["x_mean"], dtype=torch.float32
    ).reshape(1, -1, 1, 1)
    s = torch.as_tensor(
        std if std is not None else norm["x_std"], dtype=torch.float32
    ).reshape(1, -1, 1, 1)
    x = (x - m) / torch.clamp(s, min=1e-06)
    return validate_input(x)


def denormalize(y_norm):
    norm = _config()["norm"]
    return torch.as_tensor(y_norm, dtype=torch.float32) * norm["y_std"] + norm["y_mean"]


def _build():
    a = _config()["arch"]
    return FNO2d(
        in_ch=a["in_ch"],
        modes=a["modes"],
        width=a["width"],
        L=a["layers"],
        proj_hidden=a["proj_hidden"],
    )


def load_model(path_or_repo, filename=None):
    path = path_or_repo
    fname = filename or _config()["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)
    cplx = set(_config().get("complex_keys", []))
    state = {
        k: (torch.view_as_complex(v.contiguous()) if k in cplx else v)
        for k, v in state.items()
    }
    model = _build()
    model.load_state_dict(state, strict=True)
    model.eval()
    return model


def from_pretrained(repo_id):
    return load_model(repo_id)