import sys, numpy as np, pandas as pd, torch, torch.nn as nn, torch.nn.functional as F

dev = "cuda" if torch.cuda.is_available() else "cpu"
torch.manual_seed(0)
np.random.seed(0)
csv, L, H = ("/workspace/data/electricity.csv", 512, 96)
smoke = "--smoke" in sys.argv
save = (
    sys.argv[sys.argv.index("--save") + 1]
    if "--save" in sys.argv
    else "/workspace/ts_demo/fela_ts_electricity.pt"
)
epochs = int(sys.argv[sys.argv.index("--epochs") + 1]) if "--epochs" in sys.argv else 30


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

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

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


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

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


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

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


class FELA_TS(nn.Module):
    def __init__(s, C, L, H, patch=16, stride=8, D=128, modes=16, nblk=3):
        super().__init__()
        s.C, s.L, s.H, s.patch, s.stride = (C, L, H, patch, stride)
        s.revin = RevIN(C)
        s.np_ = (L - patch) // stride + 1
        s.embed = nn.Linear(patch, D)
        s.blocks = nn.ModuleList([Block(D, modes) for _ in range(nblk)])
        s.head = nn.Linear(s.np_ * D, H)

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


df = pd.read_csv(csv)
cols = [c for c in df.columns if c != "date"]
data = df[cols].values.astype(np.float32)
n = len(data)
ntr, nva = (int(n * 0.7), int(n * 0.1))
mu = data[:ntr].mean(0)
sd = data[:ntr].std(0) + 1e-08
data = (data - mu) / sd


def win(a):
    xs, ys = ([], [])
    for i in range(0, len(a) - L - H + 1, 1):
        xs.append(a[i : i + L])
        ys.append(a[i + L : i + L + H])
    return (torch.tensor(np.array(xs)), torch.tensor(np.array(ys)))


Xtr, Ytr = win(data[:ntr])
Xte, Yte = win(data[ntr + nva :])
C = data.shape[1]
assert len(Xtr) == 17805
if smoke:
    print(
        f"Electricity C={C} n={n} ntr={ntr} nva={nva} train {len(Xtr)} test {len(Xte)}"
    )
    sys.exit()
m = FELA_TS(C, L, H).to(dev)
opt = torch.optim.Adam(m.parameters(), lr=0.001)
sch = torch.optim.lr_scheduler.CosineAnnealingLR(opt, epochs)
bs = 64
print(f"[Ts] electricity C={C} train {len(Xtr)} test {len(Xte)}")
for ep in range(epochs):
    m.train()
    p = torch.randperm(len(Xtr))
    for i in range(0, len(Xtr) - bs, bs):
        idx = p[i : i + bs]
        loss = F.l1_loss(m(Xtr[idx].to(dev)), Ytr[idx].to(dev))
        opt.zero_grad()
        loss.backward()
        opt.step()
    sch.step()
m.eval()
se = ae = cnt = 0
with torch.no_grad():
    for i in range(0, len(Xte), 256):
        pr = m(Xte[i : i + 256].to(dev))
        y = Yte[i : i + 256].to(dev)
        se += F.mse_loss(pr, y, reduction="sum").item()
        ae += (pr - y).abs().sum().item()
        cnt += y.numel()
mse, mae = (se / cnt, ae / cnt)
print(f"[Ts] electricity/96 TEST MSE {mse:.4f} MAE {mae:.4f}")
torch.save(m.state_dict(), save)
print(f"SAVED {save}")