from __future__ import annotations
import json
import os
from dataclasses import dataclass
import torch
import torch.nn as nn
import torch.nn.functional as F


@dataclass
class TSConfig:
    C: int = 321
    L: int = 512
    H: int = 96
    patch: int = 16
    stride: int = 8
    D: int = 128
    modes: int = 16
    nblk: int = 3


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

    def denorm(self, x):
        return (x - self.b) / self.g * self.s + self.m


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

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


class Block(nn.Module):
    def __init__(self, D, m, ff=2, drop=0.2):
        super().__init__()
        self.n1 = nn.LayerNorm(D)
        self.fno = FNO1D(D, m)
        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_TS(nn.Module):
    def __init__(self, C, L, H, patch=16, stride=8, D=128, modes=16, nblk=3):
        super().__init__()
        self.C, self.L, self.H, self.patch, self.stride = (C, L, H, patch, stride)
        self.revin = RevIN(C)
        self.np_ = (L - patch) // stride + 1
        self.embed = nn.Linear(patch, D)
        self.blocks = nn.ModuleList([Block(D, modes) for _ in range(nblk)])
        self.head = nn.Linear(self.np_ * D, H)

    def forward(self, x):
        x = self.revin.norm(x)
        x = x.permute(0, 2, 1).reshape(-1, self.L)
        x = x.unfold(1, self.patch, self.stride)
        h = self.embed(x)
        for b in self.blocks:
            h = b(h)
        y = self.head(h.flatten(1)).reshape(-1, self.C, self.H).permute(0, 2, 1)
        return self.revin.denorm(y)


_CONFIG_FIELDS = set(TSConfig.__dataclass_fields__.keys())


def _read_json(path):
    with open(path) as f:
        return json.load(f)


def _cfg_from_dict(d):
    return TSConfig(**{k: v for k, v in d.items() if k in _CONFIG_FIELDS})


def validate_history(x: torch.Tensor, cfg: TSConfig):
    if x.dim() != 3:
        raise ValueError(f"expected a 3D tensor (B, L, C); got {tuple(x.shape)}")
    if x.size(1) != cfg.L or x.size(2) != cfg.C:
        raise ValueError(
            f"expected history of shape (B, {cfg.L}, {cfg.C}); got {tuple(x.shape)}"
        )


def load_model(path_or_repo: str):
    if os.path.isdir(path_or_repo):
        cfg_dict = _read_json(os.path.join(path_or_repo, "config.json"))
        weights = os.path.join(path_or_repo, "model.safetensors")
    elif os.path.isfile(path_or_repo) and path_or_repo.endswith(".safetensors"):
        cfg_dict = _read_json(
            os.path.join(os.path.dirname(path_or_repo), "config.json")
        )
        weights = path_or_repo
    elif os.path.isfile(path_or_repo):
        cfg_dict = _read_json(
            os.path.join(os.path.dirname(path_or_repo) or ".", "config.json")
        )
        weights = path_or_repo
    else:
        from huggingface_hub import hf_hub_download

        cfg_dict = _read_json(hf_hub_download(path_or_repo, "config.json"))
        weights = hf_hub_download(path_or_repo, "model.safetensors")
    cfg = _cfg_from_dict(cfg_dict)
    model = FELA_TS(
        cfg.C, cfg.L, cfg.H, cfg.patch, cfg.stride, cfg.D, cfg.modes, cfg.nblk
    ).eval()
    if weights.endswith(".safetensors"):
        from safetensors.torch import load_file

        state = load_file(weights)
        cplx = set(cfg_dict.get("complex_keys", []))
        state = {
            k: (torch.view_as_complex(v.contiguous()) if k in cplx else v)
            for k, v in state.items()
        }
    else:
        state = torch.load(weights, map_location="cpu", weights_only=False)
        if isinstance(state, dict) and "state_dict" in state:
            state = state["state_dict"]
    model.load_state_dict(state)
    model.cfg = cfg
    return model


from_pretrained = load_model


@torch.no_grad()
def forecast(model, x: torch.Tensor) -> torch.Tensor:
    validate_history(x, model.cfg)
    return model(x)