import argparse
import random
from pathlib import Path

import torch

from model.sutra import SutraModel
from tokenizer import SutraTokenizerFast


# ============================================================
# NeemSutra-125M-Python-Instruct
# Reproducible inference script
# ============================================================

ROOT = Path(__file__).resolve().parent

CHECKPOINT = ROOT / "neemsutra_125m_python_instruct.pt"
TOKENIZER_DIR = ROOT


# Exact architecture of the released model
VOCAB_SIZE = 8000
CONTEXT_LENGTH = 256
D_MODEL = 768
N_LAYERS = 16
N_HEADS = 12

# Training used dropout, but inference must disable it.
DROPOUT = 0.0


def get_device(requested):
    if requested == "cpu":
        return torch.device("cpu")

    if requested == "cuda":
        if not torch.cuda.is_available():
            raise RuntimeError("CUDA was requested but is not available.")
        return torch.device("cuda")

    if requested == "mps":
        if not (
            hasattr(torch.backends, "mps")
            and torch.backends.mps.is_available()
        ):
            raise RuntimeError("MPS was requested but is not available.")
        return torch.device("mps")

    # auto
    if torch.cuda.is_available():
        return torch.device("cuda")

    if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
        return torch.device("mps")

    return torch.device("cpu")


def set_seed(seed):
    random.seed(seed)
    torch.manual_seed(seed)

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


def load_model(device):
    if not CHECKPOINT.exists():
        raise FileNotFoundError(
            f"Checkpoint not found:\n{CHECKPOINT}"
        )

    model = SutraModel(
        vocab_size=VOCAB_SIZE,
        context_length=CONTEXT_LENGTH,
        d_model=D_MODEL,
        n_layers=N_LAYERS,
        n_heads=N_HEADS,
        dropout=DROPOUT,
    )

    checkpoint = torch.load(
        CHECKPOINT,
        map_location="cpu",
        weights_only=False,
    )

    # --------------------------------------------------------
    # Support common checkpoint formats
    # --------------------------------------------------------

    if isinstance(checkpoint, dict):

        if "model_state_dict" in checkpoint:
            state_dict = checkpoint["model_state_dict"]

        elif "state_dict" in checkpoint:
            state_dict = checkpoint["state_dict"]

        elif "model" in checkpoint and isinstance(
            checkpoint["model"], dict
        ):
            state_dict = checkpoint["model"]

        else:
            # The checkpoint itself may be a state_dict.
            state_dict = checkpoint

    else:
        raise RuntimeError(
            f"Unsupported checkpoint type: {type(checkpoint)}"
        )

    # --------------------------------------------------------
    # Remove DDP "module." prefix if present
    # --------------------------------------------------------

    cleaned_state_dict = {}

    for key, value in state_dict.items():
        if key.startswith("module."):
            key = key[len("module."):]
        cleaned_state_dict[key] = value

    missing, unexpected = model.load_state_dict(
        cleaned_state_dict,
        strict=False,
    )

    if missing:
        raise RuntimeError(
            "Missing model parameters:\n"
            + "\n".join(missing)
        )

    if unexpected:
        print(
            "Warning: unexpected checkpoint parameters:"
        )
        for key in unexpected:
            print(f"  {key}")

    model = model.to(device)
    model.eval()

    return model


def sample_top_k(logits, top_k):
    if top_k <= 0:
        return logits

    k = min(top_k, logits.size(-1))

    values, _ = torch.topk(logits, k)

    threshold = values[..., -1, None]

    return torch.where(
        logits < threshold,
        torch.full_like(logits, float("-inf")),
        logits,
    )


def sample_top_p(logits, top_p):
    if top_p >= 1.0:
        return logits

    if top_p <= 0.0:
        return logits

    sorted_logits, sorted_indices = torch.sort(
        logits,
        descending=True,
    )

    probabilities = torch.softmax(
        sorted_logits,
        dim=-1,
    )

    cumulative_probabilities = torch.cumsum(
        probabilities,
        dim=-1,
    )

    remove = cumulative_probabilities > top_p

    # Always retain the highest-probability token.
    remove[..., 0] = False

    sorted_logits = sorted_logits.masked_fill(
        remove,
        float("-inf"),
    )

    filtered = torch.full_like(logits, float("-inf"))

    filtered.scatter_(
        -1,
        sorted_indices,
        sorted_logits,
    )

    return filtered


@torch.no_grad()
def generate(
    model,
    tokenizer,
    prompt,
    max_new_tokens=40,
    temperature=0.0,
    top_k=0,
    top_p=1.0,
    seed=42,
):
    set_seed(seed)

    input_ids = tokenizer.encode(prompt)

    if not input_ids:
        raise ValueError("Prompt produced no tokens.")

    if len(input_ids) > CONTEXT_LENGTH:
        print(
            f"Warning: prompt contains {len(input_ids)} tokens. "
            f"Truncating to the last {CONTEXT_LENGTH} tokens."
        )

        input_ids = input_ids[-CONTEXT_LENGTH:]

    tokens = torch.tensor(
        [input_ids],
        dtype=torch.long,
        device=next(model.parameters()).device,
    )

    generated = list(input_ids)

    eos_id = tokenizer.eos_token_id

    for _ in range(max_new_tokens):

        # Keep only the model's available context.
        context = tokens[:, -CONTEXT_LENGTH:]

        logits = model(context)

        # Prediction for the final input position.
        next_token_logits = logits[:, -1, :]

        # ----------------------------------------------------
        # Greedy / deterministic generation
        # ----------------------------------------------------

        if temperature <= 0.0:
            next_token = torch.argmax(
                next_token_logits,
                dim=-1,
                keepdim=True,
            )

        else:
            scaled_logits = (
                next_token_logits / temperature
            )

            scaled_logits = sample_top_k(
                scaled_logits,
                top_k,
            )

            scaled_logits = sample_top_p(
                scaled_logits,
                top_p,
            )

            probabilities = torch.softmax(
                scaled_logits,
                dim=-1,
            )

            next_token = torch.multinomial(
                probabilities,
                num_samples=1,
            )

        token_id = next_token.item()

        generated.append(token_id)

        # Append token for next iteration.
        tokens = torch.cat(
            [tokens, next_token],
            dim=1,
        )

        if eos_id is not None and token_id == eos_id:
            break

    # Decode only the generated continuation.
    continuation = generated[len(input_ids):]

    return tokenizer.decode(
        continuation,
        skip_special_tokens=True,
    )


def main():

    parser = argparse.ArgumentParser(
        description="NeemSutra-125M-Python-Instruct inference"
    )

    parser.add_argument(
        "--prompt",
        type=str,
        required=True,
        help="Input prompt.",
    )

    parser.add_argument(
        "--max-new-tokens",
        type=int,
        default=40,
        help="Maximum number of generated tokens.",
    )

    parser.add_argument(
        "--temperature",
        type=float,
        default=0.0,
        help="Sampling temperature. 0 = greedy/deterministic.",
    )

    parser.add_argument(
        "--top-k",
        type=int,
        default=0,
        help="Top-k sampling. 0 disables top-k.",
    )

    parser.add_argument(
        "--top-p",
        type=float,
        default=1.0,
        help="Nucleus sampling threshold.",
    )

    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="Random seed.",
    )

    parser.add_argument(
        "--device",
        choices=["auto", "cpu", "cuda", "mps"],
        default="auto",
        help="Inference device.",
    )

    args = parser.parse_args()

    if args.max_new_tokens < 1:
        raise ValueError(
            "--max-new-tokens must be greater than 0."
        )

    if args.temperature < 0:
        raise ValueError(
            "--temperature cannot be negative."
        )

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

    device = get_device(args.device)

    print("=" * 70)
    print(f"Checkpoint : {CHECKPOINT.name}")
    print(f"Device     : {device}")
    print(f"Parameters : {sum(p.numel() for p in SutraModel(
        VOCAB_SIZE, CONTEXT_LENGTH, D_MODEL, N_LAYERS, N_HEADS,
        DROPOUT
    ).parameters()):,}")
    print(f"Temperature: {args.temperature}")
    print(f"Seed       : {args.seed}")
    print("=" * 70)

    tokenizer = SutraTokenizerFast(
        TOKENIZER_DIR
    )

    model = load_model(device)

    parameter_count = sum(
        p.numel() for p in model.parameters()
    )

    print("=" * 70)
    print("NeemSutra-125M-Python-Instruct")
    print("=" * 70)
    print(f"Checkpoint : {CHECKPOINT.name}")
    print(f"Device     : {device}")
    print(f"Parameters : {parameter_count:,}")
    print(f"Temperature: {args.temperature}")
    print(f"Seed       : {args.seed}")
    print("=" * 70)

    print("\nPrompt:")
    print(args.prompt)

    output = generate(
        model=model,
        tokenizer=tokenizer,
        prompt=args.prompt,
        max_new_tokens=args.max_new_tokens,
        temperature=args.temperature,
        top_k=args.top_k,
        top_p=args.top_p,
        seed=args.seed,
    )

    print("\nGenerated:")
    print("-" * 70)
    print(output)
    print("-" * 70)


if __name__ == "__main__":
    main()