"""Standalone inference for Process-FM checkpoints (fm-v6 and scaling/sc*).

Dependency: torch + numpy only. The models use global activity/resource ids:
activities 4..776 (UNK=1), resources 2..2739 (UNK=1), log ids from vocab["log_ids"].
"""
import json
from datetime import datetime
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn

N_NUM = 6
ACT_OFFSET = 4
RES_OFFSET = 2
UNK_ACT = 1
UNK_RES = 1


class EventTransformer(nn.Module):
    def __init__(self, n_act: int, n_res: int, n_logs: int, n_classes: int,
                 d_model: int = 256, n_layers: int = 6, n_heads: int = 8,
                 dropout: float = 0.1, max_len: int = 256,
                 act_dim: int = 64, res_dim: int = 32, log_dim: int = 16):
        super().__init__()
        self.act_emb = nn.Embedding(n_act, act_dim, padding_idx=0)
        self.res_emb = nn.Embedding(n_res, res_dim, padding_idx=0)
        self.log_emb = nn.Embedding(n_logs, log_dim)
        self.proj = nn.Linear(act_dim + res_dim + log_dim + N_NUM, d_model)
        self.pos = nn.Embedding(max_len, d_model)
        layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=n_heads, dim_feedforward=4 * d_model,
            dropout=dropout, batch_first=True, norm_first=True, activation="gelu",
        )
        self.encoder = nn.TransformerEncoder(layer, num_layers=n_layers)
        self.norm = nn.LayerNorm(d_model)
        self.act_out = nn.Linear(d_model, n_classes)
        self.res_out = nn.Linear(d_model, n_res)
        self.max_len = max_len

    def forward(self, gact: torch.Tensor, gres: torch.Tensor, num: torch.Tensor,
                log_id: torch.Tensor, causal: bool = True,
                return_hidden: bool = False):
        b, t = gact.shape
        x = torch.cat([
            self.act_emb(gact),
            self.res_emb(gres),
            self.log_emb(log_id).unsqueeze(1).expand(b, t, -1),
            num,
        ], dim=-1)
        x = self.proj(x)
        idx = torch.arange(t, device=gact.device).unsqueeze(0).expand(b, t)
        x = x + self.pos(idx)
        if causal:
            x = self.encoder(x, mask=torch.triu(
                torch.ones(t, t, device=gact.device, dtype=torch.bool), diagonal=1))
        else:
            x = self.encoder(x)
        h = self.norm(x)
        if return_hidden:
            return h
        return self.act_out(h), self.res_out(h)


def build_model(config: dict) -> EventTransformer:
    return EventTransformer(
        n_act=config["n_act"], n_res=config["n_res"], n_logs=config["n_logs"],
        n_classes=config["n_classes"], d_model=config["d_model"],
        n_layers=config["n_layers"], n_heads=config["n_heads"],
        dropout=config["dropout"], max_len=config["max_len"],
        act_dim=config["act_dim"], res_dim=config["res_dim"], log_dim=config["log_dim"],
    )


def _to_ns(ts) -> int:
    if isinstance(ts, datetime):
        return int(ts.timestamp() * 1e9)
    return int(np.datetime64(ts, "ns").astype("int64"))


def featurize(events: list[dict], vocab: dict, amount: float = 0.0,
              log: str | None = None):
    max_len = min(int(vocab.get("max_prefix_len", 256)), 256)
    events = events[-max_len:]
    n = len(events)
    ts = np.array([_to_ns(e["timestamp"]) for e in events], dtype=np.int64)
    dt_prev = np.zeros(n, dtype=np.float64)
    dt_prev[1:] = (ts[1:] - ts[:-1]) / 1e9
    dt_start = (ts - ts[0]) / 1e9
    hours = ((ts // 3_600_000_000_000) % 24).astype(np.float64)
    weekday = ((ts // 86_400_000_000_000) + 3) % 7
    amount_feat = np.log1p(max(amount, 0.0)) / 14.0
    num = np.stack([
        np.log1p(dt_prev) / 20.0,
        np.log1p(dt_start) / 20.0,
        np.sin(2 * np.pi * hours / 24.0),
        np.cos(2 * np.pi * hours / 24.0),
        weekday / 6.0,
        np.full(n, amount_feat),
    ], axis=1).astype(np.float32)

    act_ids = {str(a): ACT_OFFSET + i for i, a in enumerate(vocab["activities"])}
    res_ids = {str(r): RES_OFFSET + i for i, r in enumerate(vocab["resources"])}
    gact = np.array([act_ids.get(str(e.get("activity", "")), UNK_ACT) for e in events],
                    dtype=np.int64)
    gres = np.array([res_ids.get(str(e.get("resource", "")), UNK_RES) for e in events],
                    dtype=np.int64)
    log_name = log or vocab.get("default_log")
    log_id = int(vocab["log_ids"].get(log_name, vocab.get("log_unk_id", vocab["n_logs"] - 1)))
    return (
        torch.from_numpy(gact)[None],
        torch.from_numpy(gres)[None],
        torch.from_numpy(num)[None],
        torch.tensor([log_id], dtype=torch.long),
    )


def load_model(release_dir: str | Path = "."):
    release_dir = Path(release_dir)
    config = json.loads((release_dir / "config.json").read_text())
    vocab = json.loads((release_dir / "vocab.json").read_text())
    model = build_model(config)
    state = torch.load(release_dir / "model.pt", map_location="cpu", weights_only=True)
    model.load_state_dict(state)
    model.eval()
    return model, vocab


@torch.no_grad()
def predict(model: EventTransformer, vocab: dict, events: list[dict],
            amount: float = 0.0, top_k: int = 5, log: str | None = None) -> list[tuple[str, float]]:
    gact, gres, num, log_id = featurize(events, vocab, amount=amount, log=log)
    h = model(gact, gres, num, log_id, causal=True, return_hidden=True)
    logits = model.act_out(h[0, gact.shape[1] - 1].float())
    probs = torch.softmax(logits, dim=-1)
    k = min(top_k, probs.numel() - 1)
    values, indices = torch.topk(probs, k)
    id_to_act = {ACT_OFFSET + i: str(a) for i, a in enumerate(vocab["activities"])}
    return [(id_to_act.get(int(i), "UNK"), float(v)) for v, i in zip(values, indices)]


if __name__ == "__main__":
    here = Path(__file__).resolve().parent
    model, vocab = load_model(here / "fm-v6")
    activities = [str(a) for a in vocab["activities"][:3]]
    resources = [str(r) for r in vocab["resources"][:2]]
    events = [
        {"activity": activities[0], "resource": resources[0], "timestamp": "2024-01-01 10:00:00"},
        {"activity": activities[1], "resource": resources[1], "timestamp": "2024-01-01 11:00:00"},
        {"activity": activities[2], "resource": resources[0], "timestamp": "2024-01-01 12:00:00"},
    ]
    print(predict(model, vocab, events, top_k=3, log=vocab.get("default_log")))