#!/usr/bin/env python3
"""
build_scene_trt.py — Build TensorRT FP16 engines for saltybot_scene.

Builds two engines:
  1. yolov8n_scene.engine  — YOLOv8-nano for scene object detection
  2. room_classifier.engine — MobileNetV2 for room classification (--room flag)

Prerequisites:
  pip install ultralytics onnx
  # TensorRT + pycuda available on Jetson (JetPack 6)

Usage:
  # Build YOLOv8n engine (downloads weights, exports to ONNX, converts to TRT FP16)
  ros2 run saltybot_scene build_scene_trt

  # Also build room classifier engine
  ros2 run saltybot_scene build_scene_trt -- --room

  # Custom output directory
  ros2 run saltybot_scene build_scene_trt -- --output /config

Engine performance targets (Jetson Orin Nano Super 8GB):
  yolov8n_scene:    ~20-25 FPS at FP16 (640×640 input, 8400 anchors)
  room_classifier:  ~200 FPS at FP16 (224×224 input, 8 classes)
"""

import argparse
import os
import sys
from pathlib import Path


def build_yolov8_engine(output_dir: Path) -> None:
    print('[build_scene_trt] Building YOLOv8-nano scene engine...')

    onnx_path   = output_dir / 'yolov8n_scene.onnx'
    engine_path = output_dir / 'yolov8n_scene.engine'

    if not onnx_path.exists():
        print(f'  Exporting YOLOv8n to ONNX → {onnx_path}')
        from ultralytics import YOLO
        model = YOLO('yolov8n.pt')
        model.export(format='onnx', imgsz=640, half=False, dynamic=False,
                     simplify=True)
        # ultralytics exports to same directory as weights; move to output
        exported = Path('yolov8n.onnx')
        if exported.exists():
            exported.rename(onnx_path)
        else:
            print(f'  [ERROR] ONNX export not found at {exported}')
            return

    print(f'  Converting ONNX → TensorRT FP16 → {engine_path}')
    import tensorrt as trt
    logger = trt.Logger(trt.Logger.WARNING)
    builder = trt.Builder(logger)
    network = builder.create_network(
        1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
    )
    parser = trt.OnnxParser(network, logger)
    with open(onnx_path, 'rb') as f:
        if not parser.parse(f.read()):
            for i in range(parser.num_errors):
                print(f'  [TRT parse error] {parser.get_error(i)}')
            return

    config = builder.create_builder_config()
    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)   # 1 GB
    config.set_flag(trt.BuilderFlag.FP16)

    serialized = builder.build_serialized_network(network, config)
    with open(engine_path, 'wb') as f:
        f.write(serialized)
    print(f'  [OK] Engine saved to {engine_path}')
    _print_latency(engine_path, (1, 3, 640, 640))


def build_room_classifier_engine(output_dir: Path) -> None:
    """
    Build MobileNetV2 room classifier TRT engine.

    The ONNX model is a MobileNetV2 with the classifier head replaced by
    an 8-class linear layer (one per ROOM_* constant).  Pre-trained weights
    are loaded from a checkpoint if available, otherwise random init (for
    structure verification only — train before deploying).
    """
    print('[build_scene_trt] Building room classifier engine...')
    onnx_path   = output_dir / 'room_classifier.onnx'
    engine_path = output_dir / 'room_classifier.engine'

    if not onnx_path.exists():
        print(f'  Exporting MobileNetV2 room classifier to ONNX → {onnx_path}')
        import torch
        import torch.nn as nn
        from torchvision.models import mobilenet_v2

        N_ROOMS = 8
        model = mobilenet_v2(weights=None)
        model.classifier[1] = nn.Linear(model.last_channel, N_ROOMS)

        # Load pretrained weights if available
        ckpt = output_dir / 'room_classifier.pth'
        if ckpt.exists():
            model.load_state_dict(torch.load(ckpt, map_location='cpu'))
            print(f'  Loaded weights from {ckpt}')
        else:
            print(f'  [WARN] No checkpoint at {ckpt} — using random weights. '
                  f'Train before deploying.')

        model.eval()
        dummy = torch.zeros(1, 3, 224, 224)
        torch.onnx.export(
            model, dummy, str(onnx_path),
            input_names=['input'], output_names=['logits'],
            opset_version=12, dynamic_axes=None,
        )

    print(f'  Converting ONNX → TensorRT FP16 → {engine_path}')
    import tensorrt as trt
    logger = trt.Logger(trt.Logger.WARNING)
    builder = trt.Builder(logger)
    network = builder.create_network(
        1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
    )
    parser = trt.OnnxParser(network, logger)
    with open(onnx_path, 'rb') as f:
        if not parser.parse(f.read()):
            for i in range(parser.num_errors):
                print(f'  [TRT parse error] {parser.get_error(i)}')
            return

    config = builder.create_builder_config()
    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 512 << 20)
    config.set_flag(trt.BuilderFlag.FP16)

    serialized = builder.build_serialized_network(network, config)
    with open(engine_path, 'wb') as f:
        f.write(serialized)
    print(f'  [OK] Engine saved to {engine_path}')


def _print_latency(engine_path: Path, input_shape: tuple) -> None:
    """Quick warmup + timing of the built engine."""
    try:
        import tensorrt as trt
        import pycuda.driver as cuda
        import pycuda.autoinit  # noqa: F401
        import numpy as np
        import time

        logger = trt.Logger(trt.Logger.WARNING)
        with open(engine_path, 'rb') as f, trt.Runtime(logger) as rt:
            engine = rt.deserialize_cuda_engine(f.read())
        context = engine.create_execution_context()

        # Allocate buffers
        inputs, outputs, bindings = [], [], []
        stream = cuda.Stream()
        for i in range(engine.num_io_tensors):
            name   = engine.get_tensor_name(i)
            dtype  = trt.nptype(engine.get_tensor_dtype(name))
            shape  = tuple(engine.get_tensor_shape(name))
            nbytes = int(np.prod(shape)) * np.dtype(dtype).itemsize
            host   = cuda.pagelocked_empty(shape, dtype)
            device = cuda.mem_alloc(nbytes)
            bindings.append(int(device))
            if engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
                inputs.append({'host': host, 'device': device})
            else:
                outputs.append({'host': host, 'device': device})

        # Warmup
        for _ in range(5):
            cuda.memcpy_htod_async(inputs[0]['device'], inputs[0]['host'], stream)
            context.execute_async_v2(bindings, stream.handle)
            stream.synchronize()

        # Benchmark
        N = 20
        t0 = time.monotonic()
        for _ in range(N):
            cuda.memcpy_htod_async(inputs[0]['device'], inputs[0]['host'], stream)
            context.execute_async_v2(bindings, stream.handle)
            stream.synchronize()
        ms = (time.monotonic() - t0) / N * 1000
        print(f'  [benchmark] {ms:.1f} ms/frame → {1000/ms:.1f} FPS')
    except Exception as exc:
        print(f'  [benchmark skipped] {exc}')


def main():
    parser = argparse.ArgumentParser(description='Build TRT engines for saltybot_scene')
    parser.add_argument('--output', default='/config',
                        help='Output directory for engine files')
    parser.add_argument('--room', action='store_true',
                        help='Also build room classifier engine')
    args, _ = parser.parse_known_args()

    output_dir = Path(args.output)
    output_dir.mkdir(parents=True, exist_ok=True)

    build_yolov8_engine(output_dir)
    if args.room:
        build_room_classifier_engine(output_dir)

    print('[build_scene_trt] Done.')


if __name__ == '__main__':
    main()
