import torch
import torch.nn as nn
import math
import warnings
from pathlib import Path
from transformers import AutoTokenizer, AutoModelForCausalLM

try:
    import numpy as np
    from scipy.optimize import minimize
except ImportError:  # pragma: no cover
    np = None  # type: ignore
    minimize = None  # type: ignore

try:
    from gliner import GLiNER as GLiNERType
except ImportError:
    GLiNERType = None


def preferred_device() -> str:
    if torch.cuda.is_available():
        return "cuda"
    if torch.backends.mps.is_available():
        return "mps"
    return "cpu"


DEVICE = preferred_device()

# Half precision on accelerator backends; CPU keeps float32 for speed and compatibility.
DTYPE = torch.float16 if DEVICE in ("cuda", "mps") else torch.float32

MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct"

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_NAME,
    clean_up_tokenization_spaces=False,
)
# Llama has no pad token by default; align padding with EOS for generation.
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
    device_map="auto",
)
# generate() reads pad_token_id from generation_config; set both to silence warnings.
pid = tokenizer.pad_token_id
model.config.pad_token_id = pid
if model.generation_config is not None:
    model.generation_config.pad_token_id = pid

# ============================================================
# PHASE 1 — LATENT ENTITY PROBE (NuNER-style head slot)
# ============================================================
PROBE_HIDDEN_LAYER = -4
PROBE_LABEL_NAMES = (
    "glass_broken",
    "glass_intact",
    "john_alive",
    "john_dead",
    "temp_hot",
    "temp_cold",
)


class EntityHead(nn.Module):
    def __init__(self, hidden_size: int, num_labels: int):
        super().__init__()
        self.proj = nn.Linear(hidden_size, num_labels, bias=True)

    def forward(self, hidden: torch.Tensor) -> torch.Tensor:
        return self.proj(hidden)


_param = next(model.parameters())
_entity_hidden = model.config.hidden_size
entity_head = EntityHead(_entity_hidden, len(PROBE_LABEL_NAMES)).to(
    device=_param.device, dtype=_param.dtype
)


def probe_logits_from_text(text: str) -> torch.Tensor:
    """Last-token hidden at PROBE_HIDDEN_LAYER → linear probe logits [num_labels]."""
    batch = tokenizer(text, return_tensors="pt").to(_param.device)
    out = model(**batch, output_hidden_states=True)
    hidden_stack = out.hidden_states
    layer_h = hidden_stack[PROBE_HIDDEN_LAYER]
    last_h = layer_h[0, -1]
    return entity_head(last_h)


def probe_alignment_score(logits: torch.Tensor) -> float:
    """Scalar in the spirit of causal_score: prefer broken-over-intact glass, etc."""
    p = torch.sigmoid(logits.float())
    s = (p[0] - p[1]) * 3.0
    s = s + (p[2] - p[3]) * 0.5
    s = s + (p[4] - p[5]) * 0.25
    return float(s.item())

# ============================================================
# CAUSAL STATE TRACKER
# ============================================================


class WorldState:
    def __init__(self):
        self.entities = {}
        self.events = []

    def clone(self):
        new = WorldState()
        new.entities = {
            k: v.copy() for k, v in self.entities.items()
        }
        new.events = self.events[:]
        return new


CONTRADICTIONS = [
    ("broken", "intact"),
    ("alive", "dead"),
    ("hot", "cold"),
]

# ============================================================
# NuNER (GLiNER) — unified extraction (regex-free)
# ============================================================
# https://huggingface.co/numind/NuNER_Zero

NUNER_MODEL_ID = "numind/NuNER_Zero"
NUNER_THRESHOLD = 0.22

NUNER_LABEL_METADATA: dict[str, dict] = {
    "broken glass": {"fact": ("glass", "broken"), "hawkes": "glass_break"},
    "shattered cracked or broken drinking glass": {"fact": ("glass", "broken"), "hawkes": "glass_break"},
    "intact glass": {"fact": ("glass", "intact"), "hawkes": "glass_intact"},
    "living person": {"fact": ("john", "alive")},
    "dead person": {"fact": ("john", "dead")},
    "hot temperature or fire": {"fact": ("temperature", "hot"), "hawkes": "fire_event"},
    "cold or freezing conditions": {"fact": ("temperature", "cold"), "hawkes": "freeze_event"},
    "glass object dropped or falling": {"rule": ("glass", "broken"), "hawkes": "glass_drop"},
    "person injury bleeding or hurt": {"hawkes": "john_injury"},
    "panic screaming or terrified reaction": {"hawkes": "panic"},
    "cleaning up glass shards or sweeping broken glass": {"hawkes": "cleanup"},
}

ALL_NUNER_LABELS = sorted(NUNER_LABEL_METADATA.keys())

HAWKES_KIND_ORDER = [
    "glass_drop",
    "glass_break",
    "glass_intact",
    "john_injury",
    "panic",
    "cleanup",
    "fire_event",
    "freeze_event",
]
HAWKES_K = len(HAWKES_KIND_ORDER)
HAWKES_KIND_TO_I = {k: i for i, k in enumerate(HAWKES_KIND_ORDER)}

HAWKES_BETA = 2.5
HAWKES_MU: dict[str, float] = {
    "glass_drop": 0.08,
    "glass_break": 0.04,
    "glass_intact": 0.03,
    "john_injury": 0.02,
    "panic": 0.02,
    "cleanup": 0.03,
    "fire_event": 0.03,
    "freeze_event": 0.03,
}
HAWKES_BRANCH: dict[str, list[tuple[str, float]]] = {
    "glass_drop": [
        ("glass_break", 2.2),
        ("john_injury", 0.6),
        ("panic", 0.35),
        ("cleanup", 0.45),
    ],
    "glass_break": [
        ("cleanup", 1.1),
        ("panic", 0.5),
        ("john_injury", 0.9),
    ],
    "fire_event": [("panic", 0.4), ("john_injury", 0.25)],
    "freeze_event": [("john_injury", 0.15), ("panic", 0.1)],
}

_hawkes_mu_hat = None
_hawkes_alpha_hat = None
_hawkes_beta_hat = None
_hawkes_mle_buffer: list = []
HAWKES_MLE_MIN_SEQUENCES = 4
HAWKES_MLE_BUFFER_CAP = 256
HAWKES_REFIT_EVERY_STEPS = 6

_gliner_nuner = None
_nuner_load_failed = False


def _hawkes_branch_matrix_initial():
    import numpy as np_local

    a = np_local.zeros((HAWKES_K, HAWKES_K), dtype=np_local.float64)
    for src, targets in HAWKES_BRANCH.items():
        si = HAWKES_KIND_TO_I.get(src)
        if si is None:
            continue
        for tgt, val in targets:
            ti = HAWKES_KIND_TO_I.get(tgt)
            if ti is not None:
                a[si, ti] = val
    return a


def _hawkes_mu_vector_fallback():
    import numpy as np_local

    return np_local.array(
        [HAWKES_MU.get(k, 0.02) for k in HAWKES_KIND_ORDER],
        dtype=np_local.float64,
    )


def _hawkes_mu_vector():
    if np is not None and _hawkes_mu_hat is not None:
        return _hawkes_mu_hat
    return _hawkes_mu_vector_fallback()


def _hawkes_alpha_matrix():
    if np is not None and _hawkes_alpha_hat is not None:
        return _hawkes_alpha_hat
    return _hawkes_branch_matrix_initial()


def _hawkes_beta_scalar() -> float:
    if _hawkes_beta_hat is not None:
        return float(_hawkes_beta_hat)
    return float(HAWKES_BETA)


def _get_gliner_nuner():
    global _gliner_nuner, _nuner_load_failed
    if _nuner_load_failed:
        return None
    if _gliner_nuner is not None:
        return _gliner_nuner
    if GLiNERType is None:
        warnings.warn(
            "NuNER disabled: install gliner (pip install gliner).",
            stacklevel=2,
        )
        _nuner_load_failed = True
        return None
    try:
        _gliner_nuner = GLiNERType.from_pretrained(NUNER_MODEL_ID)
    except Exception as exc:  # noqa: BLE001
        _nuner_load_failed = True
        warnings.warn(f"NuNER load failed ({exc}); continuing without it.", stacklevel=2)
        return None
    return _gliner_nuner


def nuner_predict_all(text: str) -> list:
    ner = _get_gliner_nuner()
    if ner is None or not str(text).strip():
        return []
    try:
        return ner.predict_entities(
            text,
            ALL_NUNER_LABELS,
            threshold=NUNER_THRESHOLD,
        )
    except Exception as exc:  # noqa: BLE001
        warnings.warn(f"NuNER predict_entities failed: {exc}", stacklevel=2)
        return []


def apply_nuner_entity_list(entities: list, state: WorldState) -> None:
    for ent in entities:
        label = ent.get("label")
        meta = NUNER_LABEL_METADATA.get(label)
        if not meta:
            continue
        if "fact" in meta:
            entity, attr = meta["fact"]
            if entity not in state.entities:
                state.entities[entity] = {}
            state.entities[entity][attr] = True
        if "rule" in meta:
            entity, attr = meta["rule"]
            if entity not in state.entities:
                state.entities[entity] = {}
            state.entities[entity][attr] = True


def event_history_from_entities(text: str, entities: list) -> list[tuple[float, str]]:
    denom = max(len(text), 1)
    out = []
    seen = set()
    for ent in entities:
        meta = NUNER_LABEL_METADATA.get(ent.get("label"))
        if not meta or "hawkes" not in meta:
            continue
        t = float(ent.get("start", 0)) / denom
        kind = meta["hawkes"]
        key = (round(t, 6), kind)
        if key in seen:
            continue
        seen.add(key)
        out.append((t, kind))
    out.sort(key=lambda x: x[0])
    return out


def nuner_hawkes_kind_set_entities(entities: list) -> set[str]:
    kinds: set[str] = set()
    for ent in entities:
        meta = NUNER_LABEL_METADATA.get(ent.get("label"))
        if meta and "hawkes" in meta:
            kinds.add(meta["hawkes"])
    return kinds


def novel_hawkes_events_for_scoring(
    generated: str,
    future: str,
    entities_gen: list,
    entities_future: list,
) -> list[tuple[float, str]]:
    """Event kinds that appear when parsing ``future`` but not when parsing the
    committed prefix alone.

    **Why not raw char offset ≥ len(generated)?** NuNER times sit on the prompt
    at the start of ``future``; that filter made *every* event look "old" once
    ``generated`` contained the prompt, so Hawkes stayed 0 forever.
    """
    L = max(len(future), 1)
    if future.startswith(generated):
        gsplit_char = len(generated)
    else:
        gsplit_char = 0

    before_kinds = nuner_hawkes_kind_set_entities(entities_gen)
    novel_kinds = nuner_hawkes_kind_set_entities(entities_future) - before_kinds
    if not novel_kinds:
        return []

    out: list[tuple[float, str]] = []
    for kind in sorted(novel_kinds):
        best_t = None
        best_start = None
        for ent in entities_future:
            meta = NUNER_LABEL_METADATA.get(ent.get("label"))
            if not meta or meta.get("hawkes") != kind:
                continue
            start = int(ent.get("start", 0))
            t = start / L
            if start >= gsplit_char:
                if best_start is None or start < best_start:
                    best_start = start
                    best_t = t
        if best_t is None:
            best_t = 0.999
        out.append((best_t, kind))
    out.sort(key=lambda x: x[0])
    return out


def glass_observation_broken_from_entities(entities: list) -> bool:
    for ent in entities:
        meta = NUNER_LABEL_METADATA.get(ent.get("label"))
        if meta and meta.get("fact") == ("glass", "broken"):
            return True
    return False


def glass_observation_intact_from_entities(entities: list) -> bool:
    for ent in entities:
        meta = NUNER_LABEL_METADATA.get(ent.get("label"))
        if meta and meta.get("fact") == ("glass", "intact"):
            return True
    return False


def hawkes_log_likelihood(times, types, T: float, mu, alpha, beta: float) -> float:
    """Multivariate Hawkes with exponential kernel."""
    K = int(mu.shape[0])
    M = int(len(times))
    ll = 0.0
    beta = max(float(beta), 1e-6)
    for i in range(M):
        lam = float(mu[types[i]])
        for j in range(i):
            lam += float(alpha[types[j], types[i]]) * math.exp(
                -beta * float(times[i] - times[j])
            )
        ll += math.log(max(lam, 1e-15))
    compensator = float(mu.sum()) * T
    for k in range(K):
        for i in range(M):
            dt = max(0.0, T - float(times[i]))
            compensator += (
                float(alpha[types[i], k]) * (1.0 - math.exp(-beta * dt)) / beta
            )
    return float(ll - compensator)


def total_hawkes_log_likelihood_buffer(buffer: list, mu, alpha, beta: float) -> float:
    tot = 0.0
    for times, types, T in buffer:
        if len(times) == 0:
            continue
        tot += hawkes_log_likelihood(times, types, T, mu, alpha, beta)
    return tot


def record_hawkes_training_segment(text: str, entities: list) -> None:
    global _hawkes_mle_buffer
    if np is None:
        return
    hist = event_history_from_entities(text, entities)
    if not hist:
        return
    times = np.array([t for t, _ in hist], dtype=np.float64)
    types = np.array(
        [HAWKES_KIND_TO_I[k] for _, k in hist],
        dtype=np.int32,
    )
    _hawkes_mle_buffer.append((times, types, 1.0))
    if len(_hawkes_mle_buffer) > HAWKES_MLE_BUFFER_CAP:
        del _hawkes_mle_buffer[: len(_hawkes_mle_buffer) - HAWKES_MLE_BUFFER_CAP]


def refit_hawkes_mle_if_ready() -> bool:
    global _hawkes_mu_hat, _hawkes_alpha_hat, _hawkes_beta_hat
    if np is None or minimize is None:
        return False
    if len(_hawkes_mle_buffer) < HAWKES_MLE_MIN_SEQUENCES:
        return False
    K = HAWKES_K
    mu0 = _hawkes_mu_vector_fallback()
    a0 = _hawkes_branch_matrix_initial()
    b0 = float(HAWKES_BETA)
    x0 = np.concatenate([mu0, a0.ravel(), np.array([b0])])
    buffer_copy = list(_hawkes_mle_buffer)

    def nll(x: np.ndarray) -> float:
        mu = np.maximum(x[:K], 1e-6)
        alpha = np.maximum(x[K : K + K * K].reshape(K, K), 0.0)
        beta = max(float(x[-1]), 1e-3)
        return -total_hawkes_log_likelihood_buffer(buffer_copy, mu, alpha, beta)

    bounds = [(1e-6, 8.0)] * K + [(0.0, 6.0)] * (K * K) + [(0.05, 24.0)]
    res = minimize(
        nll,
        x0,
        method="L-BFGS-B",
        bounds=bounds,
        options={"maxiter": 280},
    )
    if not res.success:
        return False
    x = res.x
    _hawkes_mu_hat = np.maximum(x[:K], 1e-6)
    _hawkes_alpha_hat = np.maximum(x[K : K + K * K].reshape(K, K), 0.0)
    _hawkes_beta_hat = max(float(x[-1]), 1e-3)
    return True


def merge_observations_into(target: WorldState, delta: WorldState):
    for ent, attrs in delta.entities.items():
        if ent not in target.entities:
            target.entities[ent] = {}
        for attr, val in attrs.items():
            if val:
                target.entities[ent][attr] = True
    return target


# ============================================================
# WORLD PARSER (NuNER-only)
# ============================================================


def extract_observations(text, state, use_nuner: bool = True, entities=None):
    if not use_nuner:
        warnings.warn(
            "extract_observations requires NuNER in regex-free mode.",
            stacklevel=2,
        )
        return state
    elist = entities if entities is not None else nuner_predict_all(text)
    apply_nuner_entity_list(elist, state)
    return state


def extract_facts(text, state, use_nuner: bool = True, entities=None):
    return extract_observations(text, state, use_nuner=use_nuner, entities=entities)


# ============================================================
# HAWKES — NuNER events + optional MLE parameters
# ============================================================


class HawkesField:
    __slots__ = ("history",)

    def __init__(self):
        self.history: list[tuple[float, str]] = []

    def sync_from_entities(self, text: str, entities: list) -> None:
        self.history = event_history_from_entities(text, entities)

    def sync_from_text(self, text: str) -> None:
        self.sync_from_entities(text, nuner_predict_all(text))

    def intensity(self, kind: str, t: float) -> float:
        mu = _hawkes_mu_vector()
        alpha = _hawkes_alpha_matrix()
        beta = _hawkes_beta_scalar()
        ki = HAWKES_KIND_TO_I[kind]
        lam = float(mu[ki])
        for ti, src in self.history:
            if ti >= t:
                continue
            sj = HAWKES_KIND_TO_I[src]
            lam += float(alpha[sj, ki]) * math.exp(-beta * (t - ti))
        return lam

    def clone(self) -> "HawkesField":
        h = HawkesField()
        h.history = self.history.copy()
        return h


def hawkes_compensator_interval(prefix: HawkesField, T_lo: float, T_hi: float) -> float:
    if T_hi <= T_lo:
        return 0.0
    mu = _hawkes_mu_vector()
    alpha = _hawkes_alpha_matrix()
    beta = _hawkes_beta_scalar()
    tot = 0.0
    b = max(beta, 1e-9)
    for k in range(HAWKES_K):
        tot += float(mu[k]) * (T_hi - T_lo)
        for ti, src in prefix.history:
            sj = HAWKES_KIND_TO_I[src]
            a = float(alpha[sj, k])
            if a <= 0 or T_hi <= ti:
                continue
            lo = max(T_lo, ti)
            tot += a * (math.exp(-b * (lo - ti)) - math.exp(-b * (T_hi - ti))) / b
    return tot


def hawkes_branch_score_partial_ll(
    prefix: HawkesField,
    generated: str,
    future: str,
    entities_future: list,
    entities_gen: list,
) -> float:
    novel = novel_hawkes_events_for_scoring(
        generated, future, entities_gen, entities_future
    )
    if not novel:
        return 0.0
    score = 0.0
    for t, kind in novel:
        score += math.log(max(prefix.intensity(kind, t), 1e-12))
    score -= hawkes_compensator_interval(prefix, novel[0][0], novel[-1][0])
    return score


# ============================================================
# PEARL FILTER
# ============================================================


def causal_score(
    state,
    prior_obs=None,
    post_obs=None,
    entities_prior=None,
    entities_post=None,
):
    score = 0.0

    for entity, attrs in state.entities.items():
        attr_keys = list(attrs.keys())
        for a, b in CONTRADICTIONS:
            if a in attr_keys and b in attr_keys:
                score -= 10.0

    # Glass transition from *NuNER fact* labels only (not "dropped glass" rule),
    # otherwise rule sets broken on both prefix and future → +3 never fires.
    if entities_prior is not None and entities_post is not None:
        if glass_observation_broken_from_entities(entities_post) and not glass_observation_broken_from_entities(
            entities_prior
        ):
            score += 3.0
        if glass_observation_intact_from_entities(entities_post) and not glass_observation_intact_from_entities(
            entities_prior
        ):
            score -= 2.0
    elif prior_obs is not None and post_obs is not None:
        g0 = prior_obs.entities.get("glass", {})
        g1 = post_obs.entities.get("glass", {})
        if g1.get("broken") and not g0.get("broken"):
            score += 3.0
        if g1.get("intact") and not g0.get("intact"):
            score -= 2.0

    return score

# ============================================================
# SHORT ROLLOUT
# ============================================================


@torch.no_grad()
def rollout(prompt, depth=20):

    inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)

    out = model.generate(
        **inputs,
        max_new_tokens=depth,
        do_sample=True,
        temperature=0.8,
        top_p=0.95,
        pad_token_id=tokenizer.pad_token_id,
    )

    text = tokenizer.decode(out[0], skip_special_tokens=True)

    return text

# ============================================================
# PEARL DECODER
# ============================================================


@torch.no_grad()
def pearl_decode(
    prompt,
    steps=80,
    top_k=8,
    rollout_depth=16,
    entity_probe_weight: float = 0.0,
    use_nuner: bool = True,
    hawkes_weight: float = 1.0,
    hawkes_mle_refit: bool = True,
):

    generated = prompt
    persistent_state = WorldState()

    for step in range(steps):

        entities_gen = nuner_predict_all(generated)

        prior_obs = WorldState()
        extract_observations(generated, prior_obs, use_nuner=use_nuner, entities=entities_gen)
        merge_observations_into(persistent_state, prior_obs)

        hawkes_prefix = HawkesField()
        hawkes_prefix.sync_from_entities(generated, entities_gen)

        record_hawkes_training_segment(generated, entities_gen)
        if hawkes_mle_refit and step > 0 and step % HAWKES_REFIT_EVERY_STEPS == 0:
            if refit_hawkes_mle_if_ready():
                print(f"[HawMLE] refit ok at step {step} (n={len(_hawkes_mle_buffer)})")

        inputs = tokenizer(generated, return_tensors="pt").to(DEVICE)

        outputs = model(**inputs)

        logits = outputs.logits[:, -1, :]

        probs = torch.softmax(logits, dim=-1)

        top_probs, top_indices = torch.topk(probs, top_k)

        candidates = []

        for prob, idx in zip(top_probs[0], top_indices[0]):

            token = tokenizer.decode([idx])

            candidate_text = generated + token

            future = rollout(
                candidate_text,
                depth=rollout_depth
            )

            entities_future = nuner_predict_all(future)

            post_obs = WorldState()
            extract_observations(future, post_obs, use_nuner=use_nuner, entities=entities_future)

            state = persistent_state.clone()
            merge_observations_into(state, post_obs)

            causal = causal_score(
                state,
                prior_obs=prior_obs,
                post_obs=post_obs,
                entities_prior=entities_gen,
                entities_post=entities_future,
            )

            probe_bonus = 0.0
            if entity_probe_weight != 0.0:
                pl = probe_logits_from_text(future)
                probe_bonus = entity_probe_weight * probe_alignment_score(pl)

            hawkes_bonus = 0.0
            if hawkes_weight != 0.0:
                hawkes_bonus = hawkes_weight * hawkes_branch_score_partial_ll(
                    hawkes_prefix,
                    generated,
                    future,
                    entities_future,
                    entities_gen,
                )

            combined = (
                prob.item() * 4.0
                + causal
                + probe_bonus
                + hawkes_bonus
            )

            candidates.append(
                (
                    combined,
                    token,
                    future,
                    causal,
                    prob.item(),
                    hawkes_bonus,
                )
            )

        candidates.sort(key=lambda x: x[0], reverse=True)

        best = candidates[0]

        generated += best[1]

        print("=" * 60)
        print(f"STEP {step}")
        print(f"TOKEN: {repr(best[1])}")
        print(f"PROB: {best[4]:.4f}")
        print(f"CAUSAL: {best[3]:.4f}")
        print(f"HAWKES: {best[5]:.4f}")
        print("-" * 60)
        print(best[2][-300:])
        print("=" * 60)

        if best[1] in ["</s>", "\n\n"]:
            break

    return generated

# ============================================================
# MAIN
# ============================================================

if __name__ == "__main__":

    _self = Path(__file__).resolve().read_text(encoding="utf-8")
    print("=" * 60)
    print("SOURCE:", Path(__file__).resolve())
    print("=" * 60)
    print(_self)
    print("=" * 60)
    print("END SOURCE")
    print("=" * 60 + "\n")

    prompt = """
John dropped the glass onto the floor.
"""

    output = pearl_decode(
        prompt,
        steps=40,
        top_k=10,
        rollout_depth=20
    )

    print("\nFINAL OUTPUT\n")
    print(output)