Import('env', 'envCython', 'arch', 'msgq_python', 'common_python', 'np_version', 'acados')

gen = "c_generated_code"

casadi_model = [
  f'{gen}/long_model/long_expl_ode_fun.c',
  f'{gen}/long_model/long_expl_vde_forw.c',
]

casadi_cost_y = [
  f'{gen}/long_cost/long_cost_y_fun.c',
  f'{gen}/long_cost/long_cost_y_fun_jac_ut_xt.c',
  f'{gen}/long_cost/long_cost_y_hess.c',
]

casadi_cost_e = [
  f'{gen}/long_cost/long_cost_y_e_fun.c',
  f'{gen}/long_cost/long_cost_y_e_fun_jac_ut_xt.c',
  f'{gen}/long_cost/long_cost_y_e_hess.c',
]

casadi_cost_0 = [
  f'{gen}/long_cost/long_cost_y_0_fun.c',
  f'{gen}/long_cost/long_cost_y_0_fun_jac_ut_xt.c',
  f'{gen}/long_cost/long_cost_y_0_hess.c',
]

casadi_constraints = [
  f'{gen}/long_constraints/long_constr_h_fun.c',
  f'{gen}/long_constraints/long_constr_h_fun_jac_uxt_zt.c',
]

build_files = [f'{gen}/acados_solver_long.c'] + casadi_model + casadi_cost_y + casadi_cost_e +  \
              casadi_cost_0 + casadi_constraints

# extra generated files used to trigger a rebuild
generated_files = [
  f'{gen}/Makefile',

  f'{gen}/main_long.c',
  f'{gen}/main_sim_long.c',
  f'{gen}/acados_solver_long.h',
  f'{gen}/acados_sim_solver_long.h',
  f'{gen}/acados_sim_solver_long.c',
  f'{gen}/acados_solver.pxd',

  f'{gen}/long_model/long_expl_vde_adj.c',

  f'{gen}/long_model/long_model.h',
  f'{gen}/long_constraints/long_constraints.h',
  f'{gen}/long_cost/long_cost.h',
] + build_files

acados_include_dir = Dir(acados.INCLUDE_DIR)
acados_template_dir = Dir(acados.TEMPLATE_DIR)

source_list = ['long_mpc.py',
  '#openpilot/selfdrive/modeld/constants.py',
  acados_include_dir.File('acados_c/ocp_nlp_interface.h'),
  acados_template_dir.File('c_templates_tera/acados_solver.in.c'),
]

lenv = env.Clone()
copied_acados_libs = []
if arch != "Darwin":
  for lib in ["libacados.so", "libblasfeo.so", "libhpipm.so", "libqpOASES_e.so.3.1"]:
    copied_acados_libs += lenv.Command(f"{gen}/{lib}", Dir(acados.LIB_DIR).File(lib), [Mkdir(Dir(gen)), Copy("$TARGET", "$SOURCE")])
  lenv["RPATH"] += [lenv.Literal('\\$$ORIGIN')]
else:
  acados_rel_path = Dir(gen).rel_path(Dir(acados.LIB_DIR))
  lenv["RPATH"] += [lenv.Literal(f'\\$$ORIGIN/{acados_rel_path}')]
lenv.Clean(generated_files, Dir(gen))
generated_long = lenv.Command(generated_files,
                              source_list,
                              f"cd {Dir('.').abspath} && python3 long_mpc.py")
lenv.Depends(generated_long, [msgq_python, common_python])

lenv["CFLAGS"].append("-DACADOS_WITH_QPOASES")
lenv["CXXFLAGS"].append("-DACADOS_WITH_QPOASES")
lenv["CCFLAGS"].append("-Wno-unused")
if arch != "Darwin":
  lenv["LINKFLAGS"].append("-Wl,--disable-new-dtags")
else:
  lenv["LINKFLAGS"].append("-Wl,-install_name,@loader_path/libacados_ocp_solver_long.dylib")
  lenv["LINKFLAGS"].append(f"-Wl,-rpath,@loader_path/{acados_rel_path}")
lib_solver = lenv.SharedLibrary(f"{gen}/acados_ocp_solver_long",
                                build_files,
                                LIBS=['m', 'acados', 'hpipm', 'blasfeo', 'qpOASES_e'])

# generate cython stuff
acados_ocp_solver_pyx = acados_template_dir.File('acados_ocp_solver_pyx.pyx')
acados_ocp_solver_common = acados_template_dir.File('acados_solver_common.pxd')
libacados_ocp_solver_pxd = File(f'{gen}/acados_solver.pxd')
libacados_ocp_solver_c = File(f'{gen}/acados_ocp_solver_pyx.c')

lenv2 = envCython.Clone()
lenv2["LIBPATH"] += [lib_solver[0].dir.abspath]
lenv2["RPATH"] += [lenv2.Literal('\\$$ORIGIN')]
lenv2.Command(libacados_ocp_solver_c,
  [acados_ocp_solver_pyx, acados_ocp_solver_common, libacados_ocp_solver_pxd],
  f'cython' + \
  f' -o {libacados_ocp_solver_c.get_labspath()}' + \
  f' -I {libacados_ocp_solver_pxd.get_dir().get_labspath()}' + \
  f' -I {acados_ocp_solver_common.get_dir().get_labspath()}' + \
  f' {acados_ocp_solver_pyx.get_labspath()}')
lib_cython = lenv2.Program(f'{gen}/acados_ocp_solver_pyx.so', [libacados_ocp_solver_c], LIBS=['acados_ocp_solver_long'])
lenv2.Depends(lib_cython, lib_solver)
lenv2.Depends(lib_cython, copied_acados_libs)
lenv2.Depends(libacados_ocp_solver_c, np_version)
