import argparse, glob, os, time
import numpy as np, torch, torch.nn as nn, torch.nn.functional as F
import scipy.io as sio
from modeling import FELAPDM, PDMConfig

cwru_classes = [
    "Normal",
    "IR007",
    "B007",
    "OR007",
    "IR014",
    "B014",
    "OR014",
    "IR021",
    "B021",
    "OR021",
]
cwru_idx = {c: i for i, c in enumerate(cwru_classes)}
sensor_cols = [6, 7, 8, 11, 12, 13, 15, 16, 17, 18, 19, 21, 24, 25]
rul_cap = 125


def cwru_label(f):
    b = os.path.basename(f).split("_")[0]
    if b in cwru_idx:
        return cwru_idx[b]
    raise ValueError(f)


def load_cwru(d, win=2048, stride=1024, loads=("0", "1", "2", "3"), seed=0):
    rng = np.random.default_rng(seed)
    xs, ys = ([], [])
    for f in sorted(glob.glob(os.path.join(d, "*.mat"))):
        load = os.path.basename(f).split("_")[1].replace(".mat", "")
        if load not in loads:
            continue
        m = sio.loadmat(f)
        de = [k for k in m if k.endswith("DE_time")]
        if not de:
            continue
        sig = m[de[0]].ravel().astype(np.float32)
        sig = (sig - sig.mean()) / (sig.std() + 1e-08)
        lab = cwru_label(f)
        for s in range(0, len(sig) - win + 1, stride):
            xs.append(sig[s : s + win])
            ys.append(lab)
    x = np.stack(xs).astype(np.float32)[..., None]
    y = np.array(ys, np.int64)
    idx = rng.permutation(len(x))
    return (x[idx], y[idx])


def cwru_split(x, y, test_frac=0.25, seed=0):
    rng = np.random.default_rng(seed)
    n = len(x)
    perm = rng.permutation(n)
    nte = int(n * test_frac)
    te, tr = (perm[:nte], perm[nte:])
    return (x[tr], y[tr], x[te], y[te])


def read_cmapss(p):
    return np.loadtxt(p)


def load_cmapss(d, subset="FD001", win=30, cap=rul_cap):
    train = read_cmapss(os.path.join(d, f"train_{subset}.txt"))
    test = read_cmapss(os.path.join(d, f"test_{subset}.txt"))
    truth = read_cmapss(os.path.join(d, f"RUL_{subset}.txt")).ravel()
    s = sensor_cols
    tr_s = train[:, s]
    smin, smax = (tr_s.min(0), tr_s.max(0))
    rng = smax - smin
    rng[rng == 0] = 1.0

    def norm(a):
        return (a[:, s] - smin) / rng

    xtr, ytr = ([], [])
    for u in np.unique(train[:, 0]):
        eng = train[train[:, 0] == u]
        feats = norm(eng).astype(np.float32)
        L = len(feats)
        mc = eng[:, 1].max()
        for e in range(win, L + 1):
            xtr.append(feats[e - win : e])
            ytr.append(min(mc - eng[e - 1, 1], cap))
        if L < win:
            pad = np.zeros((win - L, len(s)), np.float32)
            xtr.append(np.concatenate([pad, feats], 0))
            ytr.append(min(mc - eng[-1, 1], cap))
    xtr = np.stack(xtr).astype(np.float32)
    ytr = np.array(ytr, np.float32)
    xte, yte = ([], [])
    for i, u in enumerate(np.unique(test[:, 0])):
        eng = test[test[:, 0] == u]
        feats = norm(eng).astype(np.float32)
        L = len(feats)
        if L >= win:
            w = feats[L - win :]
        else:
            pad = np.zeros((win - L, len(s)), np.float32)
            w = np.concatenate([pad, feats], 0)
        xte.append(w)
        yte.append(min(truth[i], cap))
    return (xtr, ytr, np.stack(xte).astype(np.float32), np.array(yte, np.float32))


def cmapss_score(yt, yp):
    d = yp - yt
    s = np.where(d < 0, np.exp(-d / 13.0) - 1.0, np.exp(d / 10.0) - 1.0)
    return (float(np.sqrt(np.mean(d**2))), float(np.sum(s)))


def batches(x, y, bs, seed=0):
    idx = np.arange(len(x))
    np.random.default_rng(seed).shuffle(idx)
    for i in range(0, len(x), bs):
        j = idx[i : i + bs]
        yield (x[j], y[j])


def cfg_cwru(**kw):
    c = PDMConfig(
        in_channels=1,
        patch=4,
        n_embd=64,
        n_layer=4,
        n_head=4,
        fno_modes=64,
        ffn_hidden=128,
        n_classes=10,
        seq_len=2048,
    )
    for k, v in kw.items():
        setattr(c, k, v)
    return c


def cfg_cmapss(**kw):
    c = PDMConfig(
        in_channels=14,
        patch=1,
        n_embd=64,
        n_layer=4,
        n_head=4,
        fno_modes=32,
        ffn_hidden=128,
        rul_head=True,
        seq_len=64,
    )
    for k, v in kw.items():
        setattr(c, k, v)
    return c


def train_cwru(a):
    dev = torch.device(a.device)
    x, y = load_cwru(a.data, win=2048, stride=a.stride, seed=a.seed)
    xtr, ytr, xte, yte = cwru_split(x, y, test_frac=0.25, seed=a.seed)
    assert len(x) == 5886
    if a.smoke:
        print(
            f"Cwru total {len(x)} train {len(xtr)} test {len(xte)} classes {len(set(y.tolist()))}"
        )
        return
    cfg = cfg_cwru(dropout=a.dropout)
    m = FELAPDM(cfg).to(dev)
    opt = torch.optim.AdamW(m.parameters(), lr=a.lr, weight_decay=0.0001)
    sch = torch.optim.lr_scheduler.CosineAnnealingLR(opt, a.epochs)
    xte_t = torch.from_numpy(xte).to(dev)
    yte_t = torch.from_numpy(yte).to(dev)
    best = 0.0
    for ep in range(a.epochs):
        m.train()
        for xb, yb in batches(xtr, ytr, a.bs, seed=a.seed + ep):
            xb = torch.from_numpy(xb).to(dev)
            yb = torch.from_numpy(yb).to(dev)
            loss = F.cross_entropy(m(xb, task="cls"), yb)
            opt.zero_grad()
            loss.backward()
            nn.utils.clip_grad_norm_(m.parameters(), 1.0)
            opt.step()
        sch.step()
        m.eval()
        with torch.no_grad():
            pr = torch.cat(
                [
                    m(xte_t[i : i + 256], task="cls").argmax(-1)
                    for i in range(0, len(xte_t), 256)
                ]
            )
            acc = (pr == yte_t).float().mean().item()
        if acc > best:
            best = acc
            torch.save(
                {"cfg": cfg.__dict__, "model": m.state_dict(), "classes": cwru_classes},
                a.out,
            )
        print(f"Ep {ep:02d} acc {acc * 100:.2f} best {best * 100:.2f}")


def train_cmapss(a):
    dev = torch.device(a.device)
    xtr, ytr, xte, yte = load_cmapss(a.data, subset=a.subset, win=a.win)
    if a.subset == "FD001":
        assert len(xtr) == 17731
    if a.smoke:
        print(f"Cmapss {a.subset} train {xtr.shape} test {xte.shape}")
        return
    cfg = cfg_cmapss(dropout=a.dropout, seq_len=a.win)
    m = FELAPDM(cfg).to(dev)
    opt = torch.optim.AdamW(m.parameters(), lr=a.lr, weight_decay=0.0001)
    sch = torch.optim.lr_scheduler.CosineAnnealingLR(opt, a.epochs)
    cap = 125.0
    xte_t = torch.from_numpy(xte).to(dev)
    best = 1000000000.0
    for ep in range(a.epochs):
        m.train()
        for xb, yb in batches(xtr, ytr / cap, a.bs, seed=a.seed + ep):
            xb = torch.from_numpy(xb).to(dev)
            yb = torch.from_numpy(yb).to(dev)
            loss = F.mse_loss(m(xb, task="rul"), yb)
            opt.zero_grad()
            loss.backward()
            nn.utils.clip_grad_norm_(m.parameters(), 1.0)
            opt.step()
        sch.step()
        m.eval()
        with torch.no_grad():
            pr = np.clip(
                torch.cat(
                    [
                        m(xte_t[i : i + 256], task="rul")
                        for i in range(0, len(xte_t), 256)
                    ]
                )
                .cpu()
                .numpy()
                * cap,
                0,
                cap,
            )
        rmse, score = cmapss_score(yte, pr)
        if rmse < best:
            best = rmse
            torch.save(
                {"cfg": cfg.__dict__, "model": m.state_dict(), "rul_cap": cap}, a.out
            )
        print(f"Ep {ep:02d} RMSE {rmse:.2f} score {score:.0f} best {best:.2f}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--task", choices=["cwru", "cmapss"], required=True)
    ap.add_argument("--data", default=None)
    ap.add_argument("--subset", default="FD001")
    ap.add_argument("--out", default=None)
    ap.add_argument("--epochs", type=int, default=40)
    ap.add_argument("--bs", type=int, default=64)
    ap.add_argument("--lr", type=float, default=0.002)
    ap.add_argument("--win", type=int, default=30)
    ap.add_argument("--stride", type=int, default=1024)
    ap.add_argument("--dropout", type=float, default=0.0)
    ap.add_argument("--device", default="cpu")
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--smoke", action="store_true")
    a = ap.parse_args()
    if a.task == "cwru":
        a.data = a.data or "../data/cwru"
        a.out = a.out or "./pdm_cwru_fno.pt"
        train_cwru(a)
    else:
        a.data = a.data or "../data/cmapss"
        a.out = a.out or f"./pdm_cmapss_{a.subset}_fno.pt"
        train_cmapss(a)


if __name__ == "__main__":
    main()