# Addon version
bl_info = {
    "name" : "Hi Phy",
    "description": "",
    "author" : "HiPhyEngine Team",
    "version" : (1, 0, 1),
    "blender" : (4, 5, 0),
    "location" : "Properties > Physics > Hi Phy",
    "warning" : "",
    "wiki_url" : "HiPhyEngine.github.io",
    "doc_url" : "HiPhyEngine.github.io",
    "category" : "Simulation"
}

ADDON_VERSION = ".".join(map(str, bl_info["version"]))

import os
import ctypes

from . import variables
import subprocess
import sys
import bpy

# Get the addon directory
addon_dir = os.path.dirname(__file__)
# sys.path.append(os.path.dirname(addon_dir))
libs_path = os.path.join(addon_dir, "libs")

# print (sys.executable)
def install_package(package):
    try:
        __import__(package)
    except ImportError:
        subprocess.check_call([sys.executable, "-m", "pip", "install", package])


# Trial version will not be packaged with the wheel files
install_package("mkl")
install_package("mkl_service")
install_package("mkl_random")
install_package("pybind11")


# the user sight may not be in the import path, we need to add it
import site
user_site = site.getusersitepackages()
if user_site not in sys.path:
    sys.path.append(user_site)

def add_dll_path():
    if (os.name == "nt"):
        # on windows, the mkl library path may not by in the enviroment, we manually add it here
        executable_dll_path = os.path.join(os.path.dirname(sys.executable), "../Library/bin")
        if os.path.exists(executable_dll_path):
            os.add_dll_directory(executable_dll_path)
        site_dll_path = os.path.join(user_site, "../../Library/bin")
        if os.path.exists(site_dll_path):
            os.add_dll_directory(site_dll_path)

try:
    add_dll_path()
    import mkl
    import pybind11
except:
    from pathlib import Path
    # install the wheels
    for root, dirs, files in os.walk(os.path.join(addon_dir, "wheels")):
        for file in files:
            file_path = os.path.join(root, file)
            print (file_path)
            install_package(file_path)
    add_dll_path()
    import mkl
    import pybind11

# Add your library path
# os.environ['LD_LIBRARY_PATH'] = libs_path + ":" + os.environ.get('LD_LIBRARY_PATH', '')
libs = []
if (os.name == "posix") :
    libs = [
        "libiomp5.so",
        "libtbbmalloc.so.2.12",
        "libtbb.so.12.12",
        "libmetis.so.5",
        "libembree4.so.4",
        "libSeExpr2.so.2",
        # MKL libraries
        # We can't load mkl libraries if we import mkl, it will cause symbol conflicts
        # "libmkl_intel_lp64.so.2",
        # "libmkl_core.so.2",
        # "libmkl_intel_thread.so.2",
        # ICC libraries
        "libintlc.so.5",
        "libimf.so",
        "libsvml.so",
        "libirng.so",
        "libsodium.so.23.3.0",
        "libHiPhyEngine.so"
    ]
if (os.name == "nt"):
    libs = [
        "embree4.dll",
        "svml_dispmd.dll",
        "libmmd.dll",
        "libsodium.dll",
        "HiPhyEngine.dll",
    ]

for lib in libs:
    lib_path = os.path.join(libs_path, lib)
    # Pre-load the library
    ctypes.CDLL(lib_path, mode=ctypes.RTLD_GLOBAL)

if (os.name == "nt"):
    try:
        ctypes.CDLL(os.path.join(addon_dir, "cAPI\SimulationAPI.cp311-win_amd64.pyd"))
    except OSError as e:
        print(f"Error loading .pyd file: {e}")


import mathutils

@bpy.app.handlers.persistent
def scene_update_post(scene):
    return;

from . import utils
from .utils import create_curve_parameterization_attribute
from .utils import evaluate_expression_attributes
from .utils import caches
import base64

from .cAPI import SimulationAPI

SimulationAPI.CheckLicense()

# --- License / AddonPreferences ---
class HiPhyAddonPreferences(bpy.types.AddonPreferences):
    bl_idname = __package__
    def draw(self, context):
        layout = self.layout
        if (SimulationAPI.GetLicenseTier() == 0):
            days_left = SimulationAPI.GetLicenseDaysLeft()
            layout.label(text="This is a trial version of HiPhyEngine. " + str(days_left) + " day(s) left for the trial period.")
        elif (SimulationAPI.GetLicenseTier() == 1):
            layout.label(text="This is a retail version of HiPhyEngine.")

# stop all simulations
@bpy.app.handlers.persistent
def load_pre(dummy):
    for solver in utils.caches:
        if caches[solver].IsSimulationRunning():
            caches[solver].driver.Interrupt();
            # wait until the sim has been interrupted
            while True:
                if not caches[solver].IsSimulationRunning():
                    break;

# clear cache when a new file is loaded
@bpy.app.handlers.persistent
def load_post(dummy):
    utils.caches.clear()
    # set the renderer to be locked to prevent crash
    scene = bpy.context.scene
    if scene:
        scene.render.use_lock_interface = True

# To prevent bad position been saved, have to set the scene time to the start of the simulation frame to reset the position of everything
saved_frame = 0
@bpy.app.handlers.persistent
def save_pre(filename):
    scene = bpy.context.scene
    global saved_frame
    saved_frame = scene.frame_current
    solvers = utils.GetSolvers();
    start_frames = [caches[solver.name].GetStartFrame() for solver in solvers if solver.name in caches]
    start_frames.append(saved_frame)
    start_frame = min(start_frames)
    bpy.context.scene.frame_set(int(start_frame))

@bpy.app.handlers.persistent
def save_post(filename):
    bpy.context.scene.frame_set(saved_frame)

import numpy as np
import bmesh

@bpy.app.handlers.persistent
def frame_change_post(scene, depsgraph=None):
    frame = scene.frame_current
    solvers = utils.GetSolvers();
    # Use the cache to set object position
    for solver in solvers:
        if (not solver.name in utils.caches) :
            continue
        # If the solver is exporting, we don't need to set the position
        if (solver.hi_phy.solver_is_exporting):
            continue;
        cache = utils.caches[solver.name]
        if (not cache or not cache.scene):
            continue;
        if frame > cache.scene.simulation_end_time:
            frame = cache.scene.simulation_end_time;
        for node in cache.objs:
            position_accessor = cache.scene.GetVector3DataAccessor(node, "position", frame)
            obj = bpy.data.objects.get(node.name)
            if not obj or not obj.hi_phy.is_active:
                print ("Simulation object " + node.name + " no longer exists/active")
                continue;
            if obj.hi_phy.driven_shape:
                obj = obj.hi_phy.driven_shape
            obj_data = obj.data
            position = obj_data.attributes['position']
            n_verts = obj_data.attributes.domain_size('POINT')
            if (position_accessor.GetDataArraySize() != n_verts):
                print ("simulation vertex data no longer matches the object " + node.name)
                continue;

            position_values = np.zeros((n_verts, 3), dtype=np.float32)
            position_accessor.CopyRawVector3ArrayToPython(position_values)
            position.data.foreach_set("vector", np.ravel(position_values))
            # TODO: Also set the frame information for the rods
            if (obj.type == 'MESH'):
                bm = bmesh.new()
                bm.from_mesh(obj.data)       # read mesh into BMesh
                bm.to_mesh(obj.data)         # write back to mesh
                bm.free()

    depsgraph = bpy.context.evaluated_depsgraph_get()
    for obj in bpy.data.objects:
        # update expressions
        if (obj.hi_phy.is_active and (obj.hi_phy.expression_attribute_live_update or variables.expression_attribute_live_update)) :
            evaluate_expression_attributes.evaluate_expressions(obj)
        # Update constraints
        # Mesh constraint and Root constraint
        # We should be ok doing root constraints directly in the local coordinate
        # reason: 1. length does not matter, so scaling with the local triangle size change should be ok
        #         2. the normal frame will always be perpendicular to the triangle and the basenormal will always on the triangle plane, so they will always be orthogonal
        #         3. we reorthogonalize and normalize when loading the constraint into the simulator anyway
        if (obj.hi_phy.is_active and obj.hi_phy.object_type == "CONSTRAINT" and
            (obj.hi_phy.constraint_type == SimulationAPI.ConstraintType.MeshConstraint or
             obj.hi_phy.constraint_type == SimulationAPI.ConstraintType.RootConstraint)):
            # First get the mesh
            # TODO: this is expensive, figuring out is there a way not to copy the whole mesh data each frame
            mesh = obj.hi_phy.constraint_object1.evaluated_get(depsgraph)
            faces = mesh.data.polygons
            n_faces = len(faces);
            tris = SimulationAPI.CreateIntArray(n_faces * 3)
            for f in range(0, n_faces):
                tris[f * 3 + 0] = faces[f].vertices[0];
                tris[f * 3 + 1] = faces[f].vertices[1];
                tris[f * 3 + 2] = faces[f].vertices[2];

            positions_attr = mesh.data.attributes['position']
            n_verts = mesh.data.attributes.domain_size('POINT')
            positions_data = np.zeros((n_verts, 3), dtype=np.float32)
            positions_attr.data.foreach_get('vector', np.ravel(positions_data))
            verts = SimulationAPI.CreateVector3Array(n_verts)
            SimulationAPI.CopyVector3Array(verts, positions_data)

            # Get the binding info
            n_points = obj.data.attributes.domain_size('POINT')

            bond_face_ids_np = np.zeros(n_points, dtype=np.int32)
            obj.data.attributes["hi_phy_constraint_point_indices1_0"].data.foreach_get('value', bond_face_ids_np);
            bond_face_ids = SimulationAPI.CreateIntArray(n_points)
            SimulationAPI.CopyIntArray(bond_face_ids, bond_face_ids_np)

            bond_barycentric_weights_on_mesh_np = np.zeros((n_points, 3), dtype=np.float32)
            obj.data.attributes["hi_phy_constraint_weight1"].data.foreach_get('vector', np.ravel(bond_barycentric_weights_on_mesh_np));
            bond_barycentric_weights_on_mesh = SimulationAPI.CreateVector3Array(n_points)
            SimulationAPI.CopyVector3Array(bond_barycentric_weights_on_mesh, bond_barycentric_weights_on_mesh_np)

            bond_local_coord_np = np.zeros((n_points, 3), dtype=np.float32)
            obj.data.attributes["hi_phy_constraint_local_coord"].data.foreach_get('vector', np.ravel(bond_local_coord_np));
            bond_local_coord = SimulationAPI.CreateVector3Array(n_points)
            SimulationAPI.CopyVector3Array(bond_local_coord, bond_local_coord_np)

            points = SimulationAPI.CreateVector3Array(n_points)
            mesh_transform = utils.createTransform(mesh.matrix_world)
            points_transform = utils.createTransform(obj.matrix_world)
            if not SimulationAPI.Utils.DeformPointsOnMesh(points, tris, verts, bond_face_ids, bond_barycentric_weights_on_mesh,
                                                          bond_local_coord, points_transform, mesh_transform):
                print ("Unable to deform points on mesh")
                continue

            # Now set the position
            position_values = np.zeros((n_points, 3), dtype=np.float32)
            SimulationAPI.CopyVector3ArrayToPython(position_values, points)
            obj.data.attributes['position'].data.foreach_set("vector", np.ravel(position_values))
            # use bmesh to force blender update
            bm = bmesh.new()
            bm.from_mesh(obj.data)       # read mesh into BMesh
            bm.to_mesh(obj.data)         # write back to mesh
            bm.free()

        # Binding constraint
        if (obj.hi_phy.is_active and obj.hi_phy.object_type == "CONSTRAINT" and
            obj.hi_phy.constraint_type == SimulationAPI.ConstraintType.BindingConstraint
            and obj.hi_phy.constraint_live_update):
            mesh_0 = obj.hi_phy.constraint_object0.evaluated_get(depsgraph)
            positions_attr = mesh_0.data.attributes['position']
            n_verts = mesh_0.data.attributes.domain_size('POINT')
            positions_data_0 = np.zeros((n_verts, 3), dtype=np.float32)
            positions_attr.data.foreach_get('vector', np.ravel(positions_data_0))

            mesh_1 = obj.hi_phy.constraint_object1.evaluated_get(depsgraph)
            positions_attr = mesh_1.data.attributes['position']
            n_verts = mesh_1.data.attributes.domain_size('POINT')
            positions_data_1 = np.zeros((n_verts, 3), dtype=np.float32)
            positions_attr.data.foreach_get('vector', np.ravel(positions_data_1))

            # Get the binding info
            n_points = obj.data.attributes.domain_size('EDGE')

            indices_0_0 = np.zeros(n_points, dtype=np.int32)
            obj.data.attributes["hi_phy_constraint_point_indices0_0"].data.foreach_get('value', indices_0_0);
            indices_0_1 = np.zeros(n_points, dtype=np.int32)
            obj.data.attributes["hi_phy_constraint_point_indices0_1"].data.foreach_get('value', indices_0_1);
            indices_0_2 = np.zeros(n_points, dtype=np.int32)
            obj.data.attributes["hi_phy_constraint_point_indices0_2"].data.foreach_get('value', indices_0_2);

            indices_1_0 = np.zeros(n_points, dtype=np.int32)
            obj.data.attributes["hi_phy_constraint_point_indices1_0"].data.foreach_get('value', indices_1_0);
            indices_1_1 = np.zeros(n_points, dtype=np.int32)
            obj.data.attributes["hi_phy_constraint_point_indices1_1"].data.foreach_get('value', indices_1_1);
            indices_1_2 = np.zeros(n_points, dtype=np.int32)
            obj.data.attributes["hi_phy_constraint_point_indices1_2"].data.foreach_get('value', indices_1_2);

            weights_0 = np.zeros((n_points, 3), dtype=np.float32)
            obj.data.attributes["hi_phy_constraint_weight0"].data.foreach_get('vector', np.ravel(weights_0));
            weights_1 = np.zeros((n_points, 3), dtype=np.float32)
            obj.data.attributes["hi_phy_constraint_weight1"].data.foreach_get('vector', np.ravel(weights_1));

            mesh0_transform = mesh_0.matrix_world
            mesh1_transform = mesh_1.matrix_world
            obj_transform = obj.matrix_world
            obj_transform.invert()

            mesh0_to_obj = obj_transform @ mesh0_transform
            mesh1_to_obj = obj_transform @ mesh1_transform

            points_0 = [positions_data_0[indices_0_0[p]] * weights_0[p][0] +
                        positions_data_0[indices_0_1[p]] * weights_0[p][1] +
                        positions_data_0[indices_0_2[p]] * weights_0[p][2] for p in range(n_points)];
            points_1 = [positions_data_1[indices_1_0[p]] * weights_1[p][0] +
                        positions_data_1[indices_1_1[p]] * weights_1[p][1] +
                        positions_data_1[indices_1_2[p]] * weights_1[p][2] for p in range(n_points)];

            positions_data = np.zeros((n_points * 2, 3), dtype=np.float32)
            for i in range(0, n_points):
                positions_data[i           ] = mesh0_to_obj @ mathutils.Vector((points_0[i][0], points_0[i][1], points_0[i][2]))
                positions_data[i + n_points] = mesh1_to_obj @ mathutils.Vector((points_1[i][0], points_1[i][1], points_1[i][2]))
            positions_attr = obj.data.attributes['position']
            positions_attr.data.foreach_set('vector', np.ravel(positions_data))

            # use bmesh to force blender update
            bm = bmesh.new()
            bm.from_mesh(obj.data)       # read mesh into BMesh
            bm.to_mesh(obj.data)         # write back to mesh
            bm.free()
            continue;

        # Update forcefield
        if (obj.hi_phy.is_active and obj.hi_phy.object_type == "FORCE_FIELD" and obj.hi_phy.force_field_visualization_live_update and obj.hi_phy.force_field_visualization_initialized):
            positions_attr = obj.data.attributes['position']
            n_verts = obj.data.attributes.domain_size('POINT')
            points = np.zeros((n_verts, 3), dtype=np.float32)
            positions_attr.data.foreach_get('vector', np.ravel(points))

            n_points = int(n_verts / 2)
            p = points[0:n_points, :]
            f = np.zeros((n_points, 3), dtype=np.float32)

            points_final = utils.force_field.ComputeVisualizationPoints(obj.hi_phy, p, f)
            if (points_final is None):
                print ("Failed to evaluate the force")
                continue;
            positions_attr.data.foreach_set('vector', np.ravel(points_final))
            # use bmesh to force blender update
            bm = bmesh.new()
            bm.from_mesh(obj.data)       # read mesh into BMesh
            bm.to_mesh(obj.data)         # write back to mesh
            bm.free()

from . import properties
from . import operators
from . import ui

modules = [
    properties,
    operators,
    ui
]

# TODO call upgrade script here for the future
@bpy.app.handlers.persistent
def check_addon_version(dummy):
    scene = bpy.context.scene
    if (scene.hiphy_version != ADDON_VERSION):
        print ("HiPhy versionof the scene " + scene.hiphy_version + " does not match the current hi phy version " + ADDON_VERSION)
    scene.hiphy_version = ADDON_VERSION

def register():
    bpy.utils.register_class(HiPhyAddonPreferences)
    for module in modules:
        module.register()
    bpy.types.Scene.hiphy_version = bpy.props.StringProperty(
        name="HiPhy Version",
        description="The version of the hiphy used to create the scene",
        default=ADDON_VERSION)
    bpy.app.handlers.depsgraph_update_post.append(scene_update_post)
    bpy.app.handlers.frame_change_post.append(frame_change_post)
    bpy.app.handlers.save_pre.append(save_pre)
    bpy.app.handlers.save_post.append(save_post)
    bpy.app.handlers.load_post.append(load_post)
    bpy.app.handlers.load_pre.append(load_pre)
    bpy.app.handlers.load_post.append(check_addon_version)

def unregister():
    for module in reversed(modules):
        module.unregister()
    bpy.utils.unregister_class(HiPhyAddonPreferences)
    bpy.app.handlers.depsgraph_update_post.remove(scene_update_post)
    bpy.app.handlers.frame_change_post.remove(frame_change_post)
    bpy.app.handlers.save_pre.remove(save_pre)
    bpy.app.handlers.save_post.remove(save_post)
    bpy.app.handlers.load_pre.remove(load_pre)
    bpy.app.handlers.load_post.remove(check_addon_version)
    del bpy.types.Scene.hiphy_version
