import pickle
import time
import numpy as np
from pathlib import Path
from tinygrad.tensor import Tensor
from tinygrad.engine.jit import TinyJit
from tinygrad.device import Device

from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE
from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye
from openpilot.selfdrive.modeld.compile_modeld import NV12Frame, make_frame_prepare as _make_frame_prepare

CAMERA_CONFIGS = [
  (_ar_ox_fisheye.width, _ar_ox_fisheye.height),
  (_os_fisheye.width, _os_fisheye.height),
]


def make_frame_prepare(cam_w, cam_h, model_w, model_h):
  nv12 = NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h))
  return _make_frame_prepare(nv12, model_w, model_h)


def warp_pkl_path(w, h):
  from openpilot.selfdrive.modeld.helpers import MODELS_DIR
  return MODELS_DIR / f'warp_{w}x{h}_tinygrad.pkl'


def make_update_img_input(frame_prepare, model_w, model_h):
  def update_img_input_tinygrad(tensor, frame, M_inv):
    M_inv = M_inv.to(Device.DEFAULT)
    new_img = frame_prepare(frame, M_inv)
    tensor.assign(tensor[6:].cat(new_img, dim=0).contiguous())
    return Tensor.cat(tensor[:6], tensor[-6:], dim=0).contiguous().reshape(1, 12, model_h//2, model_w//2)
  return update_img_input_tinygrad


def make_update_both_imgs(frame_prepare, model_w, model_h):
  update_img = make_update_img_input(frame_prepare, model_w, model_h)
  def update_both_imgs_tinygrad(calib_img_buffer, new_img, M_inv,
                                calib_big_img_buffer, new_big_img, M_inv_big):
    calib_img_pair = update_img(calib_img_buffer, new_img, M_inv)
    calib_big_img_pair = update_img(calib_big_img_buffer, new_big_img, M_inv_big)
    return calib_img_pair, calib_big_img_pair
  return update_both_imgs_tinygrad

MODELS_DIR = Path(__file__).parent / 'models'
MODEL_W, MODEL_H = MEDMODEL_INPUT_SIZE
UPSTREAM_BUFFER_LENGTH = 5


def v2_warp_pkl_path(cam_w, cam_h, buffer_length):
  return MODELS_DIR / f'warp_{cam_w}x{cam_h}_b{buffer_length}_tinygrad.pkl'


def compile_v2_warp(cam_w, cam_h, buffer_length):
  _, _, _, yuv_size = get_nv12_info(cam_w, cam_h)
  img_buffer_shape = (buffer_length * 6, MODEL_H // 2, MODEL_W // 2)

  print(f"Compiling v2 warp for {cam_w}x{cam_h} buffer_length={buffer_length}...")

  frame_prepare = make_frame_prepare(cam_w, cam_h, MODEL_W, MODEL_H)
  update_both_imgs = make_update_both_imgs(frame_prepare, MODEL_W, MODEL_H)
  update_img_jit = TinyJit(update_both_imgs, prune=True)

  full_buffer = Tensor.zeros(img_buffer_shape, dtype='uint8').contiguous().realize()
  big_full_buffer = Tensor.zeros(img_buffer_shape, dtype='uint8').contiguous().realize()
  new_frame_np = np.random.randint(0, 256, yuv_size, dtype=np.uint8)
  new_big_frame_np = np.random.randint(0, 256, yuv_size, dtype=np.uint8)
  for i in range(10):
    img_inputs = [full_buffer,
                  Tensor.from_blob(new_frame_np.ctypes.data, (yuv_size,), dtype='uint8').realize(),
                  Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY')]
    big_img_inputs = [big_full_buffer,
                      Tensor.from_blob(new_big_frame_np.ctypes.data, (yuv_size,), dtype='uint8').realize(),
                      Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY')]
    inputs = img_inputs + big_img_inputs
    Device.default.synchronize()

    st = time.perf_counter()
    _ = update_img_jit(*inputs)
    mt = time.perf_counter()
    Device.default.synchronize()
    et = time.perf_counter()
    print(f"  [{i+1}/10] enqueue {(mt-st)*1e3:6.2f} ms -- total {(et-st)*1e3:6.2f} ms")

  pkl_path = v2_warp_pkl_path(cam_w, cam_h, buffer_length)
  with open(pkl_path, "wb") as f:
    pickle.dump(update_img_jit, f)
  print(f"  Saved to {pkl_path}")

  jit = pickle.load(open(pkl_path, "rb"))
  verify_frame = np.random.randint(0, 256, yuv_size, dtype=np.uint8)
  verify_big_frame = np.random.randint(0, 256, yuv_size, dtype=np.uint8)
  fresh_inputs = [
    Tensor.zeros(img_buffer_shape, dtype='uint8').contiguous().realize(),
    Tensor.from_blob(verify_frame.ctypes.data, (yuv_size,), dtype='uint8').realize(),
    Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY'),
    Tensor.zeros(img_buffer_shape, dtype='uint8').contiguous().realize(),
    Tensor.from_blob(verify_big_frame.ctypes.data, (yuv_size,), dtype='uint8').realize(),
    Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY'),
  ]
  jit(*fresh_inputs)


class Warp:
  def __init__(self, buffer_length=2):
    self.buffer_length = buffer_length
    self.img_buffer_shape = (buffer_length * 6, MODEL_H // 2, MODEL_W // 2)

    self.jit_cache = {}
    self.full_buffers = {k: Tensor.zeros(self.img_buffer_shape, dtype='uint8').contiguous().realize() for k in ['img', 'big_img']}
    self._blob_cache: dict[int, Tensor] = {}
    self._nv12_cache: dict[tuple[int, int], int] = {}
    self.transforms_np = {k: np.zeros((3, 3), dtype=np.float32) for k in ['img', 'big_img']}
    self.transforms = {k: Tensor(v, device='NPY').realize() for k, v in self.transforms_np.items()}

  def process(self, bufs, transforms):
    if not bufs:
      return {}
    road = next(n for n in bufs if 'big' not in n)
    wide = next(n for n in bufs if 'big' in n)
    cam_w, cam_h = bufs[road].width, bufs[road].height
    key = (cam_w, cam_h)

    if key not in self.jit_cache:
      v2_pkl = v2_warp_pkl_path(cam_w, cam_h, self.buffer_length)
      if v2_pkl.exists():
        with open(v2_pkl, 'rb') as f:
          self.jit_cache[key] = pickle.load(f)
      elif self.buffer_length == UPSTREAM_BUFFER_LENGTH:
        upstream_pkl = warp_pkl_path(cam_w, cam_h)
        if upstream_pkl.exists():
          with open(upstream_pkl, 'rb') as f:
            self.jit_cache[key] = pickle.load(f)
      if key not in self.jit_cache:
        frame_prepare = make_frame_prepare(cam_w, cam_h, MODEL_W, MODEL_H)
        update_both_imgs = make_update_both_imgs(frame_prepare, MODEL_W, MODEL_H)
        self.jit_cache[key] = TinyJit(update_both_imgs, prune=True)

    if key not in self._nv12_cache:
      self._nv12_cache[key] = get_nv12_info(cam_w, cam_h)[3]
    yuv_size = self._nv12_cache[key]

    road_ptr = bufs[road].data.ctypes.data
    wide_ptr = bufs[wide].data.ctypes.data
    if road_ptr not in self._blob_cache:
      self._blob_cache[road_ptr] = Tensor.from_blob(road_ptr, (yuv_size,), dtype='uint8')
    if wide_ptr not in self._blob_cache:
      self._blob_cache[wide_ptr] = Tensor.from_blob(wide_ptr, (yuv_size,), dtype='uint8')
    road_blob = self._blob_cache[road_ptr]
    wide_blob = self._blob_cache[wide_ptr] if wide_ptr != road_ptr else Tensor.from_blob(wide_ptr, (yuv_size,), dtype='uint8')
    np.copyto(self.transforms_np['img'], transforms[road].reshape(3, 3))
    np.copyto(self.transforms_np['big_img'], transforms[wide].reshape(3, 3))

    Device.default.synchronize()
    res = self.jit_cache[key](
      self.full_buffers['img'], road_blob, self.transforms['img'],
      self.full_buffers['big_img'], wide_blob, self.transforms['big_img'],
    )
    out_road = res[0].realize()
    out_wide = res[1].realize()

    return {road: out_road, wide: out_wide}


if __name__ == "__main__":
  for cam_w, cam_h in CAMERA_CONFIGS:
    for bl in [2, 5]:
      compile_v2_warp(cam_w, cam_h, bl)