"""Offline inference with the full DeepSeek-V4.1-Flash on one trn2.48xlarge (Neuron basic mode).

Same flow as examples/offline_inference.py (launch_dp -> one engine process per DP rank, each owning one
NeuronCore), with the settings basic mode needs:
  * text-only prompts (VL kernels are stubbed in basic mode);
  * no DSpark/MTP drafting (DEEPINFER_USE_NEXTN_DECODE=0) and no EPLB;
  * a fixed KV-block pool instead of the engine's memory-profiling prefill (a max_num_batched_tokens dummy
    prefill does not fit next to the eager pure-torch kernels);
  * small batches and few output tokens: eager execution compiles every new op/shape on the first run.

    export DEEPINFER_ACCELERATOR=neuron WORLD_SIZE=1 RANK=0
    python examples/neuron/offline_inference_neuron.py --model_name /mnt/nvme/models/DeepSeek-V4.1-Flash-1017 --ep 64
"""

from __future__ import annotations

import argparse
import math
import os
import sys
from pathlib import Path

EXAMPLES = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(EXAMPLES))

from _config import DEFAULT_MODEL_NAME, create_llm, setup_common_env, setup_cuda_visible_devices  # noqa: E402
from _launch import launch_dp  # noqa: E402

PROMPTS = [
    "问题：法国的首都是哪里？\n答案：巴黎\n\n问题：一年有几个月？\n答案：12 个月\n\n问题：中国的首都是哪里？\n答案：",
    "Question: What is the largest planet in the solar system?\nAnswer: Jupiter\n\n"
    "Question: How many legs does a spider have?\nAnswer: Eight\n\n"
    "Question: Who wrote Romeo and Juliet?\nAnswer:",
]

parser = argparse.ArgumentParser(description="V4.1 offline inference on Neuron (basic mode)")
parser.add_argument("--model_name", default=DEFAULT_MODEL_NAME, help="HF model id or local HF checkpoint directory")
parser.add_argument("--ep", default=64, type=int)
parser.add_argument("--max_tokens", default=16, type=int)
parser.add_argument("--num_prompts", default=len(PROMPTS), type=int, help="run only the first N prompts (eager prefill compiles per prompt length)")
parser.add_argument("--num_kv_blocks", default=400, type=int, help="fixed KV pool per rank; the scheduler admits a request only if max_model_len tokens fit, per running sequence")
parser.add_argument("--max_model_len", default=16384, type=int, help="the candidate-block filter needs >= candidate_topk_blocks * candidate_block_size (2048 * 8)")
parser.add_argument("--max_num_batched_tokens", default=1024, type=int)
parser.add_argument("--max_num_seqs", default=1, type=int,
                    help="decode batch; every captured batch size adds one set of region NEFFs per layer (HBM: model code)")
args = parser.parse_args()
dp_size = args.ep


def _fixed_kv_pool(num_blocks: int) -> None:
    """Skip WorkerBase.profile_num_available_blocks (a max_num_batched_tokens dummy prefill)."""
    from deepinfer.worker.base import WorkerBase

    def profile_num_available_blocks(self):  # noqa: ANN001, ANN202
        self.model_runner.set_block_size(self.cache_config.block_size)
        self.cache_config.num_gpu_blocks = num_blocks
        return num_blocks

    WorkerBase.profile_num_available_blocks = profile_num_available_blocks


def main(dp_size: int, dp_rank: int) -> None:
    os.environ["DEEPINFER_USE_NEXTN_DECODE"] = "0"
    os.environ.setdefault("DEEPINFER_ACCELERATOR", "neuron")
    # sampler_force_topk=0: torch.topk with k=4096 over the vocabulary does not compile on the device (issue #22).
    setup_common_env(model_name=args.model_name, sampler_force_topk=0, set_override_experts=False, enable_eplb=False)
    # Pick this rank's NeuronCore before anything initializes deepinfer.PLATFORM (importing deepinfer.worker does):
    # PlatformNeuron claims the core from CUDA_VISIBLE_DEVICES once, at init.
    from deep_cluster.hardware import get_actual_num_gpus

    setup_cuda_visible_devices(dp_rank, 1, get_actual_num_gpus())
    _fixed_kv_pool(args.num_kv_blocks)

    llm = create_llm(
        args.model_name,
        dp_rank,
        ep_size=args.ep,
        cp_size=1,
        env={"DEEPINFER_LOGGING_INTERVAL_STEP": "1"},
        data_parallel_size=dp_size,
        max_num_batched_tokens=args.max_num_batched_tokens,
        max_model_len=args.max_model_len,
        max_num_seqs=args.max_num_seqs,
        gpu_memory_utilization=0.85,
        lookahead_tokens=0,
    )

    from deepinfer.sampling_params import SamplingParams

    prompts = PROMPTS[: args.num_prompts]
    params = [SamplingParams(temperature=0.0, max_tokens=args.max_tokens) for _ in prompts]
    outputs = llm.generate(prompts=prompts, sampling_params=params, use_tqdm=False)
    for i, output in enumerate(outputs):
        text = llm.llm_engine.tokenizer.decode([x.token_id for x in output])
        logps = [x.token_logprobs for x in output]
        ppl = math.exp(-sum(logps) / len(logps)) if logps else float("nan")
        if dp_rank == 0:
            print(f"DP {dp_rank}; INPUT: {PROMPTS[i]!r}; OUTPUT: {text!r}; output_tokens={len(output)}; decode_ppl={ppl}", flush=True)
    llm.close()


if __name__ == "__main__":
    launch_dp(main, dp_size)