#!/usr/bin/env python3
import argparse
import random
from typing import Dict, List, Optional, Set

import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer


IM_START = "[IM_START]"
IM_END = "[IM_END]"
NO_THINK = "/no think"


# ============================================================
# UTILS
# ============================================================

def get_dtype(name: str):
    name = str(name).lower()

    if name in {"bf16", "bfloat16"}:
        return torch.bfloat16

    if name in {"fp16", "float16", "half"}:
        return torch.float16

    if name in {"fp32", "float32", "float"}:
        return torch.float32

    raise ValueError(f"Unknown dtype: {name}")


def set_seed(seed: int):
    if seed is None:
        return

    if seed < 0:
        seed = random.randint(0, 2**31 - 1)
        print(f"[INFO] random seed: {seed}")

    random.seed(seed)
    torch.manual_seed(seed)

    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def build_prompt(args) -> str:
    if args.no_think:
        return (
            f"{IM_START}user\n"
            f"{args.prompt} {NO_THINK}"
            f"{IM_END}\n"
            f"{IM_START}assistant\n"
            "<think>\n</think>\n"
        )

    return args.prompt


def decode(tokenizer, ids: List[int]) -> str:
    return tokenizer.decode(ids, skip_special_tokens=False)


def extract_completion(full_text: str, prompt: str) -> str:
    if full_text.startswith(prompt):
        return full_text[len(prompt):]

    pos = full_text.rfind(prompt)
    if pos != -1:
        return full_text[pos + len(prompt):]

    return full_text


def strip_after_stop_text(text: str, stop_strings: List[str]) -> str:
    best = None

    for s in stop_strings:
        if not s:
            continue

        pos = text.find(s)
        if pos != -1:
            if best is None or pos < best:
                best = pos

    if best is None:
        return text

    return text[:best]


def build_stop_sequences(tokenizer, stop_strings: List[str]) -> List[List[int]]:
    out = []

    for s in stop_strings:
        ids = tokenizer.encode(s, add_special_tokens=False)
        if ids:
            out.append(ids)

    return out


def endswith_sequence(ids: List[int], suffix: List[int]) -> bool:
    if not suffix:
        return False

    if len(ids) < len(suffix):
        return False

    return ids[-len(suffix):] == suffix


# ============================================================
# SAMPLING
# ============================================================

def apply_repetition_penalty(
    logits: torch.Tensor,
    generated_ids: List[int],
    penalty: float,
) -> torch.Tensor:
    if penalty is None or penalty == 1.0:
        return logits

    if penalty <= 0:
        raise ValueError("--repetition-penalty must be > 0")

    for tid in set(generated_ids):
        if tid < 0 or tid >= logits.numel():
            continue

        if logits[tid] > 0:
            logits[tid] = logits[tid] / penalty
        else:
            logits[tid] = logits[tid] * penalty

    return logits


def apply_frequency_presence_penalty(
    logits: torch.Tensor,
    generated_ids: List[int],
    frequency_penalty: float,
    presence_penalty: float,
) -> torch.Tensor:
    if not generated_ids:
        return logits

    if frequency_penalty == 0.0 and presence_penalty == 0.0:
        return logits

    counts: Dict[int, int] = {}

    for tid in generated_ids:
        counts[tid] = counts.get(tid, 0) + 1

    for tid, count in counts.items():
        if tid < 0 or tid >= logits.numel():
            continue

        if frequency_penalty:
            logits[tid] -= frequency_penalty * count

        if presence_penalty:
            logits[tid] -= presence_penalty

    return logits


def get_banned_ngram_tokens(
    generated_ids: List[int],
    no_repeat_ngram_size: int,
) -> Set[int]:
    n = no_repeat_ngram_size
    banned = set()

    if n <= 0:
        return banned

    if len(generated_ids) + 1 < n:
        return banned

    prefix_len = n - 1
    current_prefix = tuple(generated_ids[-prefix_len:])

    ngram_map = {}

    for i in range(len(generated_ids) - n + 1):
        prefix = tuple(generated_ids[i:i + prefix_len])
        next_token = generated_ids[i + prefix_len]

        if prefix not in ngram_map:
            ngram_map[prefix] = set()

        ngram_map[prefix].add(next_token)

    banned.update(ngram_map.get(current_prefix, set()))
    return banned


def apply_no_repeat_ngram(
    logits: torch.Tensor,
    generated_ids: List[int],
    no_repeat_ngram_size: int,
) -> torch.Tensor:
    if no_repeat_ngram_size <= 0:
        return logits

    banned = get_banned_ngram_tokens(
        generated_ids=generated_ids,
        no_repeat_ngram_size=no_repeat_ngram_size,
    )

    for tid in banned:
        if 0 <= tid < logits.numel():
            logits[tid] = -float("inf")

    return logits


def apply_top_k(logits: torch.Tensor, top_k: int) -> torch.Tensor:
    if top_k is None or top_k <= 0:
        return logits

    top_k = min(top_k, logits.size(-1))
    values, _ = torch.topk(logits, top_k)
    cutoff = values[-1]

    logits[logits < cutoff] = -float("inf")
    return logits


def apply_top_p(logits: torch.Tensor, top_p: float) -> torch.Tensor:
    if top_p is None or top_p >= 1.0:
        return logits

    if top_p <= 0:
        raise ValueError("--top-p must be > 0")

    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
    sorted_probs = F.softmax(sorted_logits, dim=-1)
    cumulative = torch.cumsum(sorted_probs, dim=-1)

    remove = cumulative > top_p
    remove[1:] = remove[:-1].clone()
    remove[0] = False

    indices_to_remove = sorted_indices[remove]
    logits[indices_to_remove] = -float("inf")

    return logits


def apply_min_p(logits: torch.Tensor, min_p: float) -> torch.Tensor:
    if min_p is None or min_p <= 0:
        return logits

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

    keep = probs >= (min_p * max_prob)
    logits[~keep] = -float("inf")

    return logits


def apply_typical_p(logits: torch.Tensor, typical_p: float) -> torch.Tensor:
    if typical_p is None or typical_p >= 1.0:
        return logits

    if typical_p <= 0:
        raise ValueError("--typical-p must be > 0")

    probs = F.softmax(logits, dim=-1)
    log_probs = F.log_softmax(logits, dim=-1)

    entropy = -(probs * log_probs).sum()
    shifted_scores = torch.abs((-log_probs) - entropy)

    sorted_scores, sorted_indices = torch.sort(shifted_scores, descending=False)
    sorted_probs = probs[sorted_indices]
    cumulative_probs = torch.cumsum(sorted_probs, dim=-1)

    remove = cumulative_probs > typical_p
    remove[1:] = remove[:-1].clone()
    remove[0] = False

    indices_to_remove = sorted_indices[remove]
    logits[indices_to_remove] = -float("inf")

    return logits


def apply_bad_words(
    logits: torch.Tensor,
    tokenizer,
    bad_words: List[str],
):
    for word in bad_words:
        ids = tokenizer.encode(word, add_special_tokens=False)
        if len(ids) == 1:
            tid = ids[0]
            if 0 <= tid < logits.numel():
                logits[tid] = -float("inf")

    return logits


def sample_next_token(
    logits: torch.Tensor,
    generated_ids: List[int],
    tokenizer,
    args,
) -> int:
    logits = logits.float().clone()

    logits = apply_bad_words(
        logits=logits,
        tokenizer=tokenizer,
        bad_words=args.bad_words,
    )

    logits = apply_repetition_penalty(
        logits=logits,
        generated_ids=generated_ids,
        penalty=args.repetition_penalty,
    )

    logits = apply_frequency_presence_penalty(
        logits=logits,
        generated_ids=generated_ids,
        frequency_penalty=args.frequency_penalty,
        presence_penalty=args.presence_penalty,
    )

    logits = apply_no_repeat_ngram(
        logits=logits,
        generated_ids=generated_ids,
        no_repeat_ngram_size=args.no_repeat_ngram_size,
    )

    if args.temperature <= 0:
        return int(torch.argmax(logits).item())

    logits = logits / args.temperature

    logits = apply_top_k(logits, args.top_k)
    logits = apply_top_p(logits, args.top_p)
    logits = apply_min_p(logits, args.min_p)
    logits = apply_typical_p(logits, args.typical_p)

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

    if torch.isnan(probs).any() or torch.isinf(probs).any() or probs.sum() <= 0:
        return int(torch.argmax(logits).item())

    return int(torch.multinomial(probs, num_samples=1).item())


# ============================================================
# MODEL
# ============================================================

def model_forward_logits(model, input_ids: torch.Tensor):
    out = model(input_ids=input_ids)

    if hasattr(out, "logits"):
        return out.logits

    if isinstance(out, tuple):
        return out[0]

    raise RuntimeError("Impossible de récupérer logits depuis la sortie du modèle.")


@torch.no_grad()
def generate_manual(
    model,
    tokenizer,
    input_ids: torch.Tensor,
    args,
) -> torch.Tensor:
    idx = input_ids
    generated_after_prompt: List[int] = []

    stop_sequences = build_stop_sequences(
        tokenizer,
        stop_strings=args.stop_strings,
    )

    eos_id = tokenizer.eos_token_id

    for step in range(args.max_new_tokens):
        idx_cond = idx[:, -args.ctx_len:] if args.ctx_len > 0 else idx

        logits = model_forward_logits(model, idx_cond)
        logits = logits[:, -1, :][0]

        if step < args.min_new_tokens:
            if eos_id is not None and 0 <= eos_id < logits.numel():
                logits[eos_id] = -float("inf")

            for seq in stop_sequences:
                if len(seq) == 1:
                    tid = seq[0]
                    if 0 <= tid < logits.numel():
                        logits[tid] = -float("inf")

        next_id = sample_next_token(
            logits=logits,
            generated_ids=generated_after_prompt,
            tokenizer=tokenizer,
            args=args,
        )

        next_tensor = torch.tensor(
            [[next_id]],
            dtype=torch.long,
            device=idx.device,
        )

        idx = torch.cat([idx, next_tensor], dim=1)
        generated_after_prompt.append(next_id)

        full_ids = idx[0].tolist()

        if step >= args.min_new_tokens:
            if eos_id is not None and next_id == eos_id:
                break

            should_stop = False

            for seq in stop_sequences:
                if endswith_sequence(full_ids, seq):
                    should_stop = True
                    break

            if should_stop:
                break

    return idx


# ============================================================
# CLI
# ============================================================

def parse_args():
    p = argparse.ArgumentParser(
        description="Inference script for Arithmetic-SLM using [IM_START]/[IM_END]."
    )

    p.add_argument(
        "--model",
        default="PhysiQuanty/Arithmetic-SLM",
        help="HF model id or local path.",
    )

    p.add_argument(
        "--prompt",
        default="59 + 45 =",
        help="Raw arithmetic prompt. With --no-think, inserted in chat template.",
    )

    p.add_argument(
        "--no-think",
        action="store_true",
        help="Use production no-think template with [IM_START]/[IM_END].",
    )

    p.add_argument(
        "--device",
        default="cuda" if torch.cuda.is_available() else "cpu",
    )

    p.add_argument(
        "--dtype",
        default="bfloat16",
        choices=["bfloat16", "bf16", "float16", "fp16", "float32", "fp32"],
    )

    p.add_argument("--ctx-len", type=int, default=2048)
    p.add_argument("--max-new-tokens", type=int, default=64)
    p.add_argument("--min-new-tokens", type=int, default=1)

    p.add_argument("--temperature", type=float, default=0.7)
    p.add_argument("--top-k", type=int, default=40)
    p.add_argument("--top-p", type=float, default=0.90)
    p.add_argument("--min-p", type=float, default=0.0)
    p.add_argument("--typical-p", type=float, default=1.0)

    p.add_argument("--repetition-penalty", type=float, default=1.05)
    p.add_argument("--frequency-penalty", type=float, default=0.10)
    p.add_argument("--presence-penalty", type=float, default=0.0)
    p.add_argument("--no-repeat-ngram-size", type=int, default=4)

    p.add_argument("--seed", type=int, default=-1)

    p.add_argument(
        "--stop-string",
        action="append",
        default=None,
        help="Additional stop string. Can be passed multiple times.",
    )

    p.add_argument(
        "--bad-word",
        action="append",
        default=None,
        help="Single-token word/token to ban. Can be passed multiple times.",
    )

    p.add_argument(
        "--print-full",
        action="store_true",
        help="Print full prompt + completion.",
    )

    p.add_argument(
        "--trust-remote-code",
        action="store_true",
        default=True,
    )

    args = p.parse_args()

    if args.max_new_tokens <= 0:
        raise ValueError("--max-new-tokens must be > 0")

    if args.min_new_tokens < 0:
        raise ValueError("--min-new-tokens must be >= 0")

    if args.temperature < 0:
        raise ValueError("--temperature must be >= 0")

    if args.top_k < 0:
        raise ValueError("--top-k must be >= 0")

    if not (0 < args.top_p <= 1.0):
        raise ValueError("--top-p must be in (0, 1]")

    if args.min_p < 0:
        raise ValueError("--min-p must be >= 0")

    if not (0 < args.typical_p <= 1.0):
        raise ValueError("--typical-p must be in (0, 1]")

    if args.repetition_penalty <= 0:
        raise ValueError("--repetition-penalty must be > 0")

    if args.no_repeat_ngram_size < 0:
        raise ValueError("--no-repeat-ngram-size must be >= 0")

    stop_strings = [
        IM_END,
        IM_START,
    ]

    if args.stop_string:
        stop_strings.extend(args.stop_string)

    args.stop_strings = stop_strings

    bad_words = []

    if args.bad_word:
        bad_words.extend(args.bad_word)

    args.bad_words = bad_words

    return args


def main():
    args = parse_args()
    set_seed(args.seed)

    dtype = get_dtype(args.dtype)

    print(f"[INFO] model:       {args.model}")
    print(f"[INFO] device:      {args.device}")
    print(f"[INFO] dtype:       {args.dtype}")
    print(f"[INFO] template:    {'no_think' if args.no_think else 'raw'}")
    print(f"[INFO] IM_START:    {IM_START}")
    print(f"[INFO] IM_END:      {IM_END}")
    print(f"[INFO] NO_THINK:    {NO_THINK}")
    print(f"[INFO] temperature: {args.temperature}")
    print(f"[INFO] top_k:       {args.top_k}")
    print(f"[INFO] top_p:       {args.top_p}")
    print(f"[INFO] min_p:       {args.min_p}")
    print(f"[INFO] typical_p:   {args.typical_p}")
    print()

    tokenizer = AutoTokenizer.from_pretrained(
        args.model,
        trust_remote_code=args.trust_remote_code,
    )

    model = AutoModelForCausalLM.from_pretrained(
        args.model,
        dtype=dtype,
        trust_remote_code=args.trust_remote_code,
    ).to(args.device)

    model.eval()

    prompt = build_prompt(args)

    encoded = tokenizer(
        prompt,
        return_tensors="pt",
        add_special_tokens=False,
    )

    encoded.pop("token_type_ids", None)

    input_ids = encoded["input_ids"].to(args.device)

    output_ids = generate_manual(
        model=model,
        tokenizer=tokenizer,
        input_ids=input_ids,
        args=args,
    )

    full_text = decode(tokenizer, output_ids[0].tolist())

    if args.print_full:
        print(full_text)
        return

    completion = extract_completion(full_text, prompt)
    completion = strip_after_stop_text(completion, args.stop_strings)

    print(completion)


if __name__ == "__main__":
    main()