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


@dataclass
class FelaTabConfig:
    max_features: int = 100
    max_classes: int = 10
    dim: int = 1024
    n_layers: int = 28
    n_heads: int = 16
    head_dim: int = 64
    chunk: int = 64
    ffn_mult: int = 3
    dropout: float = 0.0
    use_landmark: bool = True
    n_landmarks: int = 48
    ln_eps: float = 1e-05

    def __post_init__(self):
        assert self.dim % self.n_heads == 0 or self.n_heads * self.head_dim > 0


def delta_rule_chunk(q, k, v, beta, chunk):
    B, H, T, D = q.shape
    dev, dt = (q.device, q.dtype)
    pad = (chunk - T % chunk) % chunk
    if pad:
        q = F.pad(q, (0, 0, 0, pad))
        k = F.pad(k, (0, 0, 0, pad))
        v = F.pad(v, (0, 0, 0, pad))
        beta = F.pad(beta, (0, pad))
    Tp = T + pad
    nC = Tp // chunk
    q = q.view(B, H, nC, chunk, D)
    k = k.view(B, H, nC, chunk, D)
    v = v.view(B, H, nC, chunk, D)
    beta = beta.view(B, H, nC, chunk)
    tri = torch.tril(torch.ones(chunk, chunk, device=dev, dtype=dt), -1)
    eye = torch.eye(chunk, device=dev, dtype=dt)
    S = torch.zeros(B, H, D, D, device=dev, dtype=dt)
    outs = []
    for c in range(nC):
        kc = k[:, :, c]
        vc = v[:, :, c]
        qc = q[:, :, c]
        bc = beta[:, :, c]
        KK = torch.einsum("bhtd,bhsd->bhts", kc, kc)
        A = bc.unsqueeze(-1) * KK * tri
        kS0 = torch.einsum("bhtd,bhde->bhte", kc, S)
        RHS = bc.unsqueeze(-1) * (vc - kS0)
        M = eye + A
        U = torch.linalg.solve_triangular(M, RHS, upper=False, unitriangular=True)
        qS0 = torch.einsum("bhtd,bhde->bhte", qc, S)
        QK = torch.einsum("bhtd,bhsd->bhts", qc, kc) * tri
        o = qS0 + torch.einsum("bhts,bhse->bhte", QK, U)
        outs.append(o)
        S = S + torch.einsum("bhtd,bhte->bhde", kc, U)
    o = torch.cat(outs, dim=2).view(B, H, Tp, D)[:, :, :T]
    return (o, S)


def pool_landmarks(x, wm, G):
    B, T, C = x.shape
    ns = wm.sum(1).clamp(min=1)
    gs = torch.ceil(ns / G).clamp(min=1)
    pos = torch.arange(T, device=x.device)[None, :].expand(B, T)
    gid = torch.minimum((pos / gs[:, None]).floor().long(), torch.full_like(pos, G - 1))
    valid = wm > 0
    flat = (torch.arange(B, device=x.device)[:, None] * G + gid)[valid]
    src = x[valid]
    land = torch.zeros(B * G, C, device=x.device, dtype=x.dtype)
    cnt = torch.zeros(B * G, device=x.device, dtype=x.dtype)
    land.index_add_(0, flat, src)
    cnt.index_add_(0, flat, torch.ones(src.shape[0], device=x.device, dtype=x.dtype))
    land = land / cnt.clamp(min=1)[:, None]
    return (land.view(B, G, C), (cnt > 0).view(B, G))


class DeltaAttn(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.H, self.D, self.chunk = (cfg.n_heads, cfg.head_dim, cfg.chunk)
        inner = cfg.n_heads * cfg.head_dim
        self.q = nn.Linear(cfg.dim, inner, bias=False)
        self.k = nn.Linear(cfg.dim, inner, bias=False)
        self.v = nn.Linear(cfg.dim, inner, bias=False)
        self.b = nn.Linear(cfg.dim, cfg.n_heads)
        self.o = nn.Linear(inner, cfg.dim, bias=False)

    def forward(self, x, write_mask):
        B, T, _ = x.shape
        H, D = (self.H, self.D)
        q = self.q(x).view(B, T, H, D).transpose(1, 2)
        k = self.k(x).view(B, T, H, D).transpose(1, 2)
        v = self.v(x).view(B, T, H, D).transpose(1, 2)
        q = F.normalize(q, dim=-1)
        k = F.normalize(k, dim=-1)
        beta = torch.sigmoid(self.b(x)).transpose(1, 2)
        beta = beta * write_mask.unsqueeze(1)
        o, _ = delta_rule_chunk(q, k, v, beta, self.chunk)
        o = o.transpose(1, 2).reshape(B, T, H * D)
        return self.o(o)


class LandmarkAttn(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.H, self.D, self.G = (cfg.n_heads, cfg.head_dim, cfg.n_landmarks)
        inner = cfg.n_heads * cfg.head_dim
        self.q = nn.Linear(cfg.dim, inner, bias=False)
        self.k = nn.Linear(cfg.dim, inner, bias=False)
        self.v = nn.Linear(cfg.dim, inner, bias=False)
        self.o = nn.Linear(inner, cfg.dim, bias=False)
        self.gate = nn.Parameter(torch.zeros(1))

    def forward(self, x, wm):
        B, T, _ = x.shape
        H, D, G = (self.H, self.D, self.G)
        land, lmask = pool_landmarks(x, wm, G)
        q = self.q(x).view(B, T, H, D).transpose(1, 2)
        k = self.k(land).view(B, G, H, D).transpose(1, 2)
        v = self.v(land).view(B, G, H, D).transpose(1, 2)
        att = torch.einsum("bhtd,bhgd->bhtg", q, k) / math.sqrt(D)
        att = att.masked_fill(~lmask[:, None, None, :], -1000000000.0)
        att = att.softmax(-1)
        o = torch.einsum("bhtg,bhgd->bhtd", att, v).transpose(1, 2).reshape(B, T, H * D)
        return self.o(o) * torch.tanh(self.gate)


class Block(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.n1 = nn.LayerNorm(cfg.dim, eps=cfg.ln_eps)
        self.attn = DeltaAttn(cfg)
        self.landmark = LandmarkAttn(cfg) if cfg.use_landmark else None
        self.n2 = nn.LayerNorm(cfg.dim, eps=cfg.ln_eps)
        hid = cfg.dim * cfg.ffn_mult
        self.ff = nn.Sequential(
            nn.Linear(cfg.dim, hid), nn.GELU(), nn.Linear(hid, cfg.dim)
        )

    def forward(self, x, wm):
        h = self.n1(x)
        a = self.attn(h, wm)
        if self.landmark is not None:
            a = a + self.landmark(h, wm)
        x = x + a
        x = x + self.ff(self.n2(x))
        return x


class FelaTab(nn.Module):
    def __init__(self, cfg: FelaTabConfig):
        super().__init__()
        self.cfg = cfg
        self.feat_enc = nn.Linear(cfg.max_features, cfg.dim)
        self.cls_label_emb = nn.Embedding(cfg.max_classes + 1, cfg.dim)
        self.reg_label = nn.Linear(1, cfg.dim)
        self.reg_label_gate = nn.Parameter(torch.ones(1))
        self.type_emb = nn.Embedding(2, cfg.dim)
        self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layers)])
        self.norm = nn.LayerNorm(cfg.dim, eps=cfg.ln_eps)
        self.cls_head = nn.Linear(cfg.dim, cfg.max_classes)
        self.reg_head = nn.Linear(cfg.dim, 2)

    def encode_tokens(self, X, y, n_support, task_type_id):
        B, T, _ = X.shape
        tok = self.feat_enc(X)
        tok = tok + self.type_emb(task_type_id).view(1, 1, -1)
        idx = torch.arange(T, device=X.device).unsqueeze(0)
        wm = (idx < n_support.unsqueeze(1)).float()
        if int(task_type_id) == 0:
            yl = y.clamp(min=0).long()
            lab = self.cls_label_emb(yl) * wm.unsqueeze(-1)
            tok = tok + lab
        else:
            lab = (
                self.reg_label(y.unsqueeze(-1).float())
                * wm.unsqueeze(-1)
                * self.reg_label_gate
            )
            tok = tok + lab
        return (tok, wm)

    def forward(self, X, y, n_support, task_type_id):
        tok, wm = self.encode_tokens(X, y, n_support, task_type_id)
        h = tok
        for blk in self.blocks:
            h = blk(h, wm)
        h = self.norm(h)
        if int(task_type_id) == 0:
            return self.cls_head(h)
        return self.reg_head(h)


_CONFIG_FIELDS = set(FelaTabConfig.__dataclass_fields__.keys())
TIERS = {
    "big": dict(dim=1024, n_layers=28, n_heads=16, head_dim=64),
    "small": dict(dim=512, n_layers=14, n_heads=8, head_dim=64),
}


def _to_config(cfg_dict: dict) -> FelaTabConfig:
    return FelaTabConfig(**{k: v for k, v in cfg_dict.items() if k in _CONFIG_FIELDS})


def _read_json(path: str) -> dict:
    with open(path) as f:
        return json.load(f)


def _dequant_state_dict(tensors: dict) -> dict:
    sd = {}
    for k, v in tensors.items():
        if k.endswith(".scale"):
            continue
        if v.dtype == torch.int8:
            scale = tensors[k + ".scale"]
            sd[k] = v.float() * scale[:, None]
        else:
            sd[k] = v.float()
    return sd


def load_model(path_or_repo: str = ".", tier: str = "big", config: dict = None):
    from safetensors.torch import load_file

    def _pick(dir_path):
        for base in (
            f"model_{tier}_int8.safetensors",
            f"model_{tier}.safetensors",
            "model_int8.safetensors",
            "model.safetensors",
        ):
            p = os.path.join(dir_path, base)
            if os.path.isfile(p):
                return p
        raise FileNotFoundError(f"no model safetensors for tier '{tier}' in {dir_path}")

    def _cfg(dir_path):
        for base in (f"config_{tier}.json", "config.json"):
            p = os.path.join(dir_path, base)
            if os.path.isfile(p):
                return _read_json(p)
        return dict(TIERS[tier])

    if os.path.isfile(path_or_repo) and path_or_repo.endswith(".safetensors"):
        weights_path = path_or_repo
        cfg_dict = config or _cfg(os.path.dirname(path_or_repo) or ".")
    elif os.path.isdir(path_or_repo):
        weights_path = _pick(path_or_repo)
        cfg_dict = config or _cfg(path_or_repo)
    else:
        from huggingface_hub import hf_hub_download

        cfg_dict = config
        if cfg_dict is None:
            for base in (f"config_{tier}.json", "config.json"):
                try:
                    cfg_dict = _read_json(hf_hub_download(path_or_repo, base))
                    break
                except Exception:
                    continue
            cfg_dict = cfg_dict or dict(TIERS[tier])
        weights_path = None
        for base in (f"model_{tier}_int8.safetensors", f"model_{tier}.safetensors"):
            try:
                weights_path = hf_hub_download(path_or_repo, base)
                break
            except Exception:
                continue
        if weights_path is None:
            raise FileNotFoundError(
                f"no model_{tier}[_int8].safetensors in repo {path_or_repo}"
            )
    for k, val in TIERS[tier].items():
        cfg_dict.setdefault(k, val)
    tensors = load_file(weights_path)
    sd = _dequant_state_dict(tensors)
    model = FelaTab(_to_config(cfg_dict))
    model.load_state_dict(sd)
    model.eval()
    return model


from_pretrained = load_model


def _standardize_fit(Xtr: np.ndarray):
    mean = Xtr.mean(axis=0)
    scale = np.sqrt(((Xtr - mean) ** 2).mean(axis=0))
    scale[scale == 0] = 1.0
    return (mean.astype(np.float32), scale.astype(np.float32))


def _prep(model, Xtr, Xte):
    cfg = model.cfg
    Xtr = np.asarray(Xtr, dtype=np.float32)
    Xte = np.asarray(Xte, dtype=np.float32)
    mean, scale = _standardize_fit(Xtr)
    Xtr2 = (Xtr - mean) / scale
    Xte2 = (Xte - mean) / scale
    F_ = cfg.max_features
    if Xtr2.shape[1] > F_:
        try:
            from sklearn.decomposition import PCA

            p = PCA(n_components=F_, random_state=0).fit(Xtr2)
            Xtr2, Xte2 = (p.transform(Xtr2), p.transform(Xte2))
        except Exception:
            Xtr2, Xte2 = (Xtr2[:, :F_], Xte2[:, :F_])

    def pad(A):
        out = np.zeros((A.shape[0], F_), np.float32)
        out[:, : min(A.shape[1], F_)] = A[:, :F_]
        return out

    return (np.nan_to_num(pad(Xtr2)), np.nan_to_num(pad(Xte2)))


@torch.no_grad()
def predict(
    model, Xtr, ytr, Xte, task="classification", n_classes=None, support_cap=3000
):
    is_cls = task.startswith("cls") or task == "classification"
    Xtr = np.asarray(Xtr, dtype=np.float32)
    ytr = np.asarray(ytr)
    Xte = np.asarray(Xte, dtype=np.float32)
    if len(Xtr) > support_cap:
        idx = np.random.RandomState(0).permutation(len(Xtr))[:support_cap]
        Xtr, ytr = (Xtr[idx], ytr[idx])
    Xtr2, Xte2 = _prep(model, Xtr, Xte)
    ns = len(Xtr2)
    X = np.concatenate([Xtr2, Xte2], 0)[None]
    if is_cls:
        ncls = int(n_classes if n_classes is not None else ytr.max() + 1)
        y = np.concatenate([ytr, np.zeros(len(Xte2))]).astype(np.int64)[None]
        ttype = 0
    else:
        ymu, ysd = (float(ytr.mean()), float(ytr.std() + 1e-06))
        y = np.concatenate([(ytr - ymu) / ysd, np.zeros(len(Xte2))]).astype(np.float32)[
            None
        ]
        ttype = 1
    dev = next(model.parameters()).device
    Xt = torch.from_numpy(X).to(dev)
    yt = torch.from_numpy(y).to(dev)
    nst = torch.tensor([ns]).to(dev)
    out = model(Xt, yt, nst, torch.tensor(ttype, device=dev)).float()
    q = out[0, ns:]
    if is_cls:
        return F.softmax(q[:, :ncls], dim=-1).cpu().numpy()
    mean = q[:, 0].cpu().numpy() * ysd + ymu
    std = np.sqrt(np.exp(np.clip(q[:, 1].cpu().numpy(), -8, 8))) * ysd
    return (mean, std)


def predict_bagged(
    model,
    Xtr,
    ytr,
    Xte,
    task="classification",
    n_classes=None,
    support_cap=3000,
    n_bag=16,
    base_seed=1000,
):
    is_cls = task.startswith("cls") or task == "classification"
    if not is_cls or n_bag <= 1:
        return predict(model, Xtr, ytr, Xte, task, n_classes, support_cap)
    Xtr = np.asarray(Xtr, dtype=np.float32)
    ytr = np.asarray(ytr)
    Xte = np.asarray(Xte, dtype=np.float32)
    ncls = int(n_classes if n_classes is not None else ytr.max() + 1)
    acc = []
    for b in range(n_bag):
        rng = np.random.RandomState(base_seed + b)
        rp = rng.permutation(len(Xtr))
        cperm = rng.permutation(Xtr.shape[1])
        acc.append(
            predict(
                model,
                Xtr[rp][:, cperm],
                ytr[rp],
                Xte[:, cperm],
                "classification",
                ncls,
                support_cap,
            )
        )
    return np.mean(acc, axis=0)


def count_params(m):
    return sum((p.numel() for p in m.parameters()))