"""
bert/trt_session.py
TensorRT ORT session with optimization profiles.

Profiles define min/opt/max tensor shapes so TRT compiles the most efficient
CUDA kernels for the expected (batch_size, sequence_length) distribution.
First load: ~5 min to compile the .plan engine.
Subsequent loads: instant from cache.
"""

from __future__ import annotations

import argparse
import numpy as np
from pathlib import Path
from typing import Dict, List, Tuple

import numpy as np


def build_trt_session(
    model_path: str | Path,
    cache_dir: str | Path = "./trt_cache",
    device_id: int = 0,
    # Optimization profile bounds: (min, opt, max) for (batch, seq_len)
    min_shape:  Tuple[int, int] = (1, 16),
    opt_shape:  Tuple[int, int] = (32, 128),
    max_shape:  Tuple[int, int] = (128, 512),
) -> "ort.InferenceSession":
    """
    Build an ORT InferenceSession backed by TensorrtExecutionProvider.

    The engine is cached to `cache_dir` after first compilation so
    subsequent cold-starts load in milliseconds.
    """
    try:
        import onnxruntime as ort
    except ImportError:
        raise RuntimeError("pip install onnxruntime-gpu")

    Path(cache_dir).mkdir(parents=True, exist_ok=True)

    trt_providers = [
        (
            "TensorrtExecutionProvider",
            {
                "device_id":                 device_id,
                "trt_fp16_enable":           True,
                "trt_engine_cache_enable":   True,
                "trt_engine_cache_path":     str(cache_dir),
                # Profile: input_ids and attention_mask share the same shape bounds
                "trt_profile_min_shapes":    f"input_ids:{min_shape[0]}x{min_shape[1]},attention_mask:{min_shape[0]}x{min_shape[1]},token_type_ids:{min_shape[0]}x{min_shape[1]}",
                "trt_profile_opt_shapes":    f"input_ids:{opt_shape[0]}x{opt_shape[1]},attention_mask:{opt_shape[0]}x{opt_shape[1]},token_type_ids:{opt_shape[0]}x{opt_shape[1]}",
                "trt_profile_max_shapes":    f"input_ids:{max_shape[0]}x{max_shape[1]},attention_mask:{max_shape[0]}x{max_shape[1]},token_type_ids:{max_shape[0]}x{max_shape[1]}",
                "trt_int8_enable":           False,   # FP16 already halves memory
            },
        ),
        "CUDAExecutionProvider",    # fallback if TRT fails a subgraph
        "CPUExecutionProvider",
    ]

    opts = ort.SessionOptions()
    opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
    opts.intra_op_num_threads = 4

    print(f"[trt_session] loading {model_path} ...")
    print(f"[trt_session] cache_dir={cache_dir}  (first load may take ~5 min)")
    session = ort.InferenceSession(
        str(model_path),
        sess_options=opts,
        providers=trt_providers,
    )
    actual_providers = session.get_providers()
    print(f"[trt_session] active providers: {actual_providers}")
    return session


def run_batch(
    session: "ort.InferenceSession",
    input_ids: np.ndarray,       # (B, L) int64
    attention_mask: np.ndarray,  # (B, L) int64
    token_type_ids: np.ndarray | None = None,
) -> np.ndarray:
    """
    Run one batch through the TRT session.
    Returns logits: (B, 3) float32.
    """
    if token_type_ids is None:
        token_type_ids = np.zeros_like(input_ids)

    feeds = {
        "input_ids":      input_ids.astype(np.int64),
        "attention_mask": attention_mask.astype(np.int64),
        "token_type_ids": token_type_ids.astype(np.int64),
    }
    outputs = session.run(["logits"], feeds)
    return outputs[0]   # (B, 3)


def softmax(x: np.ndarray) -> np.ndarray:
    e = np.exp(x - x.max(axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)


def entailment_scores(logits: np.ndarray) -> np.ndarray:
    """Softmax probability of Entailment class (index 2), shape (B,)."""
    return softmax(logits)[:, 2]


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", type=str, required=True, help="path to FP16 ONNX model")
    parser.add_argument("--cache", type=str, default="./trt_cache")
    args = parser.parse_args()

    session = build_trt_session(args.model, args.cache)

    # Smoke test: random batch of 4 sequences length 128
    ids   = np.random.randint(0, 30000, (4, 128), dtype=np.int64)
    mask  = np.ones((4, 128), dtype=np.int64)
    logits = run_batch(session, ids, mask)
    scores = entailment_scores(logits)
    print(f"[smoke test] entailment scores: {scores}")