"""
Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
"""
import glob
import os

Import("env", "envCython", "arch", "cereal", "messaging", "common", "visionipc")
lenv = env.Clone()
lenvCython = envCython.Clone()

libs = [cereal, messaging, visionipc, common, "capnp", "kj", "pthread"]
frameworks = []
core_sources = ["native/iqmodel.cc", "transforms/yuv.cc", "transforms/warp_geometry.cc"]
SMALL_MODEL_NAMES = ["supercombo", "driving_vision", "driving_off_policy", "driving_on_policy", "driving_policy"]
FUSED_TRIPLET = ["driving_vision", "driving_off_policy", "driving_on_policy"]
PC = not os.path.isfile("/TICI")


def _inject_path_define(symbol, filename):
  quoted = f'-D{symbol}_PATH=\\"{File(filename).abspath}\\"'
  for active_env in (lenv, lenvCython):
    active_env["CXXFLAGS"].append(quoted)


def _tinygrad_sources():
  root = env.Dir("#tinygrad_repo").relpath
  workspace = env.Dir("#").abspath
  return ["#" + path for path in glob.glob(root + "/**", recursive=True, root_dir=workspace) if "pycache" not in path]


def _present_models():
  return [name for name in SMALL_MODEL_NAMES if File(f"models/{name}.onnx").exists()]


def _tinygrad_flags():
  if arch == "larch64":
    return "DEV=QCOM IMAGE=2 FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 OPENPILOT_HACKS=1"
  if arch == "Darwin":
    return f'DEV=CPU HOME={os.path.expanduser("~")} IMAGE=0'
  if arch == "x86_64":
    return "DEV=CPU:LLVM IMAGE=0"
  return "DEV=CPU:LLVM IMAGE=0"


def _fused_flags():
  if arch == "larch64":
    return "DEV=QCOM IMAGE=2 FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 OPENPILOT_HACKS=1"
  if arch == "Darwin":
    return f'DEV=CPU HOME={os.path.expanduser("~")} IMAGE=0 FLOAT16=1'
  if arch == "x86_64":
    return "DEV=CPU:LLVM IMAGE=0 FLOAT16=1"
  return "DEV=CPU:LLVM IMAGE=0 FLOAT16=1"


def _queue_metadata_generation(model_names, tinygrad_files):
  if not PC:
    return
  inputs = tinygrad_files + [File(Dir("#iqpilot/selfdrive/iqmodeld/tools").File("install_models_pc.py").abspath)]
  outputs = []
  for model_name in model_names:
    inputs.extend([File(f"models/{model_name}.onnx"), File(f"models/{model_name}_tinygrad.pkl")])
    outputs.append(File(f"models/{model_name}_metadata.pkl"))
  if outputs:
    tool_dir = Dir("#iqpilot/selfdrive/iqmodeld/tools").abspath
    model_dir = Dir("models").abspath
    lenv.Command(outputs, inputs,
                 lenv.PrettyAction(f"${{PYWARN}} python3 {tool_dir}/install_models_pc.py {model_dir}", 'META'))


if arch == "Darwin":
  frameworks += ["OpenCL"]
else:
  libs += ["OpenCL"]

for symbol, filename in {"TRANSFORM": "transforms/warp_geometry.cl", "LOADYUV": "transforms/yuv.cl"}.items():
  _inject_path_define(symbol, filename)

cython_libs = envCython["LIBS"] + libs
iqmodel_lib = lenv.Library("iqmodel", core_sources)
lenvCython.Program("native/iqmodel_pyx.so", "native/iqmodel_pyx.pyx", LIBS=[iqmodel_lib, *cython_libs], FRAMEWORKS=frameworks)
tinygrad_files = _tinygrad_sources()
present_models = _present_models()
_queue_metadata_generation(present_models, tinygrad_files)


def tg_compile(flags, model_name):
  pythonpath_string = 'PYTHONPATH="${PYTHONPATH}:' + env.Dir("#tinygrad_repo").abspath + '"'
  fn = File(f"models/{model_name}").abspath
  return lenv.Command(
    fn + "_tinygrad.pkl",
    [fn + ".onnx"] + tinygrad_files,
    lenv.PrettyAction(
      f'${{PYWARN}} {pythonpath_string} {flags} python3 {Dir("#tinygrad_repo").abspath}/examples/openpilot/compile3.py {fn}.onnx {fn}_tinygrad.pkl',
      'MODEL', logfile='${TARGET}.log')
  )


for model_name in present_models:
  tg_compile(_tinygrad_flags(), model_name)

from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye
from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE

FUSED_CAMERA_CONFIGS = [(_ar_ox_fisheye.width, _ar_ox_fisheye.height), (_os_fisheye.width, _os_fisheye.height)]
FUSED_FRAME_SKIP = 4


def tg_compile_fused(file_prefix, flags):
  pythonpath_string = 'PYTHONPATH="${PYTHONPATH}:' + env.Dir("#tinygrad_repo").abspath + '"'
  model_dir = Dir("models").abspath
  model_w, model_h = MEDMODEL_INPUT_SIZE
  camera_args = " ".join(f"{cw}x{ch}" for cw, ch in FUSED_CAMERA_CONFIGS)
  out_pkl = File(f"models/{file_prefix}driving_fused_tinygrad.pkl").abspath
  onnx_deps = [File(f"models/{file_prefix}{model_name}.onnx") for model_name in FUSED_TRIPLET]
  cmd = (
    f'${{PYWARN}} {pythonpath_string} {flags} python3 {Dir("#iqpilot/selfdrive/iqmodeld/tools").abspath}/compile_daemon.py '
    f'--model-size {model_w}x{model_h} --camera-resolutions {camera_args} '
    f'--vision-onnx {model_dir}/{file_prefix}driving_vision.onnx '
    f'--off-policy-onnx {model_dir}/{file_prefix}driving_off_policy.onnx '
    f'--on-policy-onnx {model_dir}/{file_prefix}driving_on_policy.onnx '
    f'--output {out_pkl} --frame-skip {FUSED_FRAME_SKIP}'
  )
  return lenv.Command(out_pkl, onnx_deps + tinygrad_files,
                      lenv.PrettyAction(cmd, 'MODEL', logfile='${TARGET}.log'))


if all(File(f"models/{model_name}.onnx").exists() for model_name in FUSED_TRIPLET):
  tg_compile_fused("", _fused_flags())

if all(File(f"models/big_{model_name}.onnx").exists() for model_name in FUSED_TRIPLET):
  tg_compile_fused("big_", "DEV=USB+AMD:LLVM WARP_DEV=QCOM FLOAT16=1 JIT_BATCH_SIZE=0 GMMU=0")