import glob
import os
from typing import List, Optional, Tuple

import numpy as np
import trimesh
from gym_ignition.utils import logger
from numpy.random import RandomState
from pcg_gazebo.parsers import parse_sdf
from pcg_gazebo.parsers.sdf import create_sdf_element
from scenario import gazebo as scenario_gazebo

# Note: only models with mesh geometry are supported


class ModelCollectionRandomizer:

    _class_model_paths = None
    __sdf_base_name = "model.sdf"
    __configured_sdf_base_name = "model_modified.sdf"
    __blacklisted_base_name = "BLACKLISTED"
    __collision_mesh_dir = "meshes/collision/"
    __collision_mesh_file_type = "stl"
    __original_scale_base_name = "original_scale.txt"

    def __init__(
        self,
        model_paths=None,
        owner="GoogleResearch",
        collection="Google Scanned Objects",
        server="https://fuel.ignitionrobotics.org",
        server_version="1.0",
        unique_cache=False,
        reset_collection=False,
        enable_blacklisting=True,
        np_random: Optional[RandomState] = None,
    ):

        # If enabled, the newly created objects of this class will use its own individual cache
        # for model paths and must discover/download them on its own
        self._unique_cache = unique_cache

        # Flag that determines if models that cannot be used are blacklisted
        self._enable_blacklisting = enable_blacklisting

        # If enabled, the cache of the class used to store model paths among instances will be reset
        if reset_collection and not self._unique_cache:
            self._class_model_paths = None

        # Get file path to all models from
        # a) `model_paths` arg
        # b) local cache owner (if `owner` has some models, i.e `collection` is already downloaded)
        # c) Fuel collection (if `owner` has no models in local cache)
        if model_paths is not None:
            # Use arg
            if self._unique_cache:
                self._model_paths = model_paths
            else:
                self._class_model_paths = model_paths
        else:
            # Use local cache or Fuel
            if self._unique_cache:
                self._model_paths = self.get_collection_paths(
                    owner=owner,
                    collection=collection,
                    server=server,
                    server_version=server_version,
                )
            elif self._class_model_paths is None:
                # Executed only once, unless the paths are reset with `reset_collection` arg
                self._class_model_paths = self.get_collection_paths(
                    owner=owner,
                    collection=collection,
                    server=server,
                    server_version=server_version,
                )

        # Initialise rng with (with seed is desired)
        if np_random is not None:
            self.np_random = np_random
        else:
            self.np_random = np.random.default_rng()

    @classmethod
    def get_collection_paths(
        cls,
        owner="GoogleResearch",
        collection="Google Scanned Objects",
        server="https://fuel.ignitionrobotics.org",
        server_version="1.0",
        model_name: str = "",
    ) -> List[str]:

        # First check the local cache (for performance)
        # Note: This unfortunately does not check if models belong to the specified collection
        # TODO: Make sure models belong to the collection if sampled from local cache
        model_paths = scenario_gazebo.get_local_cache_model_paths(
            owner=owner, name=model_name
        )
        if len(model_paths) > 0:
            return model_paths

        # Else download the models from Fuel and then try again
        if collection:
            download_uri = "%s/%s/%s/collections/%s" % (
                server,
                server_version,
                owner,
                collection,
            )
        elif model_name:
            download_uri = "%s/%s/%s/models/%s" % (
                server,
                server_version,
                owner,
                model_name,
            )
        download_command = 'ign fuel download -v 3 -t model -j %s -u "%s"' % (
            os.cpu_count(),
            download_uri,
        )
        os.system(download_command)

        model_paths = scenario_gazebo.get_local_cache_model_paths(
            owner=owner, name=model_name
        )
        if 0 == len(model_paths):
            logger.error(
                'URI "%s" is not valid and does not contain any models that are \
                          owned by the owner of the collection'
                % download_uri
            )
            pass

        return model_paths

    def random_model(
        self,
        min_scale=0.125,
        max_scale=0.175,
        min_mass=0.05,
        max_mass=0.25,
        min_friction=0.75,
        max_friction=1.5,
        decimation_fraction_of_visual=0.25,
        decimation_min_faces=40,
        decimation_max_faces=200,
        max_faces=40000,
        max_vertices=None,
        component_min_faces_fraction=0.1,
        component_max_volume_fraction=0.35,
        fix_mtl_texture_paths=True,
        skip_blacklisted=True,
        return_sdf_path=True,
    ) -> str:

        # Loop until a model is found, checked for validity, configured and returned
        # If any of these steps fail, sample another model and try again
        # Note: Due to this behaviour, the function could stall if all models are invalid
        # TODO: Add a simple timeout to random sampling of valid model (# of attempts or time-based)
        while True:

            # Get path to a random model from the collection
            model_path = self.get_random_model_path()

            # Check if the model is already blacklisted and skip if desired
            if skip_blacklisted and self.is_blacklisted(model_path):
                continue

            # Check is the model is already configured
            if self.is_configured(model_path):
                # If so, break the loop
                break

            # Process the model and break loop only if it is valid
            if self.process_model(
                model_path,
                decimation_fraction_of_visual=decimation_fraction_of_visual,
                decimation_min_faces=decimation_min_faces,
                decimation_max_faces=decimation_max_faces,
                max_faces=max_faces,
                max_vertices=max_vertices,
                component_min_faces_fraction=component_min_faces_fraction,
                component_max_volume_fraction=component_max_volume_fraction,
                fix_mtl_texture_paths=fix_mtl_texture_paths,
            ):
                break

        # Apply randomization
        self.randomize_configured_model(
            model_path,
            min_scale=min_scale,
            max_scale=max_scale,
            min_friction=min_friction,
            max_friction=max_friction,
            min_mass=min_mass,
            max_mass=max_mass,
        )

        if return_sdf_path:
            # Return path to the configured SDF file
            return self.get_configured_sdf_path(model_path)
        else:
            # Return path to the model directory
            return model_path

    def process_all_models(
        self,
        decimation_fraction_of_visual=0.025,
        decimation_min_faces=8,
        decimation_max_faces=400,
        max_faces=40000,
        max_vertices=None,
        component_min_faces_fraction=0.1,
        component_max_volume_fraction=0.35,
        fix_mtl_texture_paths=True,
    ):
        if self._unique_cache:
            model_paths = self._model_paths
        else:
            model_paths = self._class_model_paths

        blacklist_model_counter = 0
        for i in range(len(model_paths)):
            if not self.process_model(
                model_paths[i],
                decimation_fraction_of_visual=decimation_fraction_of_visual,
                decimation_min_faces=decimation_min_faces,
                decimation_max_faces=decimation_max_faces,
                max_faces=max_faces,
                max_vertices=max_vertices,
                component_min_faces_fraction=component_min_faces_fraction,
                component_max_volume_fraction=component_max_volume_fraction,
                fix_mtl_texture_paths=fix_mtl_texture_paths,
            ):
                blacklist_model_counter += 1

            print('Processed model %i/%i "%s"' % (i, len(model_paths), model_paths[i]))

        print("Number of blacklisted models: %i" % blacklist_model_counter)

    def process_model(
        self,
        model_path,
        decimation_fraction_of_visual=0.25,
        decimation_min_faces=40,
        decimation_max_faces=200,
        max_faces=40000,
        max_vertices=None,
        component_min_faces_fraction=0.1,
        component_max_volume_fraction=0.35,
        fix_mtl_texture_paths=True,
    ) -> bool:

        # Parse the SDF of the model
        sdf = parse_sdf(self.get_sdf_path(model_path))

        # Process the model(s) contained in the SDF
        for model in sdf.models:

            # Process the link(s) of each model
            for link in model.links:

                # Get rid of the existing collisions prior to simplifying it
                link.collisions.clear()

                # Values for the total inertial properties of current link
                # These values will be updated for each body that the link contains
                total_mass = 0.0
                total_inertia = [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]
                common_centre_of_mass = [0.0, 0.0, 0.0]

                # Go through the visuals and process them
                for visual in link.visuals:

                    # Get path to the mesh of the link's visual
                    mesh_path = self.get_mesh_path(model_path, visual)

                    # If desired, fix texture path in 'mtl' files for '.obj' mesh format
                    if fix_mtl_texture_paths:
                        self.fix_mtl_texture_paths(
                            model_path, mesh_path, model.attributes["name"]
                        )

                    # Load the mesh (without materials)
                    mesh = trimesh.load(mesh_path, force="mesh", skip_materials=True)

                    # Check if model has too much geometry (blacklist if needed)
                    if not self.check_excessive_geometry(
                        mesh, model_path, max_faces=max_faces, max_vertices=max_vertices
                    ):
                        return False

                    # Check if model has disconnected geometry/components (blacklist if needed)
                    if not self.check_disconnected_components(
                        mesh,
                        model_path,
                        component_min_faces_fraction=component_min_faces_fraction,
                        component_max_volume_fraction=component_max_volume_fraction,
                    ):
                        return False

                    # Compute inertial properties for this mesh
                    (
                        total_mass,
                        total_inertia,
                        common_centre_of_mass,
                    ) = self.sum_inertial_properties(
                        mesh, total_mass, total_inertia, common_centre_of_mass
                    )

                    # Add decimated collision geometry to the SDF
                    self.add_collision(
                        mesh,
                        link,
                        model_path,
                        fraction_of_visual=decimation_fraction_of_visual,
                        min_faces=decimation_min_faces,
                        max_faces=decimation_max_faces,
                    )

                    # Write original scale (size) into the SDF
                    # This is used for later reference during randomization (for scale limits)
                    self.write_original_scale(mesh, model_path)

                # Make sure the link has valid inertial properties (blacklist if needed)
                if not self.check_inertial_properties(
                    model_path, total_mass, total_inertia
                ):
                    return False

                # Write inertial properties to the SDF of the link
                self.write_inertial_properties(
                    link, total_mass, total_inertia, common_centre_of_mass
                )

        # Write the configured SDF into a file
        sdf.export_xml(self.get_configured_sdf_path(model_path))
        return True

    def add_collision(
        self,
        mesh,
        link,
        model_path,
        fraction_of_visual=0.05,
        min_faces=8,
        max_faces=750,
        friction=1.0,
    ):

        # Determine name of path to the collistion geometry
        collision_name = (
            link.attributes["name"] + "_collision_" + str(len(link.collisions))
        )
        collision_mesh_path = self.get_collision_mesh_path(model_path, collision_name)

        # Determine number of faces to keep after the decimation
        face_count = min(
            max(fraction_of_visual * len(mesh.faces), min_faces), max_faces
        )

        # Simplify mesh via decimation
        collision_mesh = mesh.simplify_quadratic_decimation(face_count)

        # Export the collision mesh to the appropriate location
        os.makedirs(os.path.dirname(collision_mesh_path), exist_ok=True)
        collision_mesh.export(
            collision_mesh_path, file_type=self.__collision_mesh_file_type
        )

        # Create collision SDF element
        collision = create_sdf_element("collision")

        # Add collision geometry to the SDF
        collision.geometry.mesh = create_sdf_element("mesh")
        collision.geometry.mesh.uri = os.path.relpath(
            collision_mesh_path, start=model_path
        )

        # Add surface friction to the SDF of collision (default to 1 and randomize later)
        collision.surface = create_sdf_element("surface")
        collision.surface.friction = create_sdf_element("friction", "surface")
        collision.surface.friction.ode = create_sdf_element("ode", "collision")
        collision.surface.friction.ode.mu = friction
        collision.surface.friction.ode.mu2 = friction

        # Add it to the SDF of the link
        collision_name = os.path.basename(collision_mesh_path).split(".")[0]
        link.add_collision(collision_name, collision)

    def sum_inertial_properties(
        self, mesh, total_mass, total_inertia, common_centre_of_mass, density=1.0
    ) -> Tuple[float, float, float]:

        # Arbitrary density is used here
        # The mass will be randomized once it is fully computed for a link
        mesh.density = density

        # Tmp variable to store the mass of all previous geometry, used to determine centre of mass
        mass_of_others = total_mass

        # For each additional mesh, simply add the mass and inertia
        total_mass += mesh.mass
        total_inertia += mesh.moment_inertia

        # Compute a common centre of mass between all previous geometry and the new mesh
        common_centre_of_mass = [
            mass_of_others * common_centre_of_mass[0] + mesh.mass * mesh.center_mass[0],
            mass_of_others * common_centre_of_mass[1] + mesh.mass * mesh.center_mass[1],
            mass_of_others * common_centre_of_mass[2] + mesh.mass * mesh.center_mass[2],
        ] / total_mass

        return total_mass, total_inertia, common_centre_of_mass

    def randomize_configured_model(
        self,
        model_path,
        min_scale=0.05,
        max_scale=0.25,
        min_mass=0.1,
        max_mass=3.0,
        min_friction=0.75,
        max_friction=1.5,
    ):

        # Get path to the configured SDF file
        configured_sdf_path = self.get_configured_sdf_path(model_path)

        # Parse the configured SDF that needs to be randomized
        sdf = parse_sdf(configured_sdf_path)

        # Process the model(s) contained in the SDF
        for model in sdf.models:

            # Process the link(s) of each model
            for link in model.links:

                # Randomize scale of the link
                self.randomize_scale(
                    model_path, link, min_scale=min_scale, max_scale=max_scale
                )

                # Randomize inertial properties of the link
                self.randomize_inertial(link, min_mass=min_mass, max_mass=max_mass)

                # Randomize friction of the link
                self.randomize_friction(
                    link, min_friction=min_friction, max_friction=max_friction
                )

        # Overwrite the configured SDF file with randomized values
        sdf.export_xml(configured_sdf_path)

    def randomize_scale(self, model_path, link, min_scale=0.05, max_scale=0.25):

        # Note: This function currently supports only scaling of links with single mesh geometry
        if len(link.visuals) > 1:
            return False

        # Get a random scale for the size of mesh
        random_scale = self.np_random.uniform(min_scale, max_scale)

        # Determine a scale factor that will result in such scale for the size of mesh
        original_mesh_scale = self.read_original_scale(model_path)
        scale_factor = random_scale / original_mesh_scale

        # Determine scale factor for inertial properties based on random scale and current scale
        current_scale = link.visuals[0].geometry.mesh.scale.value[0]
        inertial_scale_factor = scale_factor / current_scale

        # Write scale factor into SDF for visual and collision geometry
        link.visuals[0].geometry.mesh.scale = [scale_factor] * 3
        link.collisions[0].geometry.mesh.scale = [scale_factor] * 3

        # Recompute inertial properties according to the scale
        link.inertial.pose.x *= inertial_scale_factor
        link.inertial.pose.y *= inertial_scale_factor
        link.inertial.pose.z *= inertial_scale_factor

        # Mass is scaled n^3
        link.mass = link.mass.value * inertial_scale_factor**3

        # Inertia is scaled n^5
        inertial_scale_factor_n5 = inertial_scale_factor**5
        link.inertia.ixx = link.inertia.ixx.value * inertial_scale_factor_n5
        link.inertia.iyy = link.inertia.iyy.value * inertial_scale_factor_n5
        link.inertia.izz = link.inertia.izz.value * inertial_scale_factor_n5
        link.inertia.ixy = link.inertia.ixy.value * inertial_scale_factor_n5
        link.inertia.ixz = link.inertia.ixz.value * inertial_scale_factor_n5
        link.inertia.iyz = link.inertia.iyz.value * inertial_scale_factor_n5

    def randomize_inertial(
        self, link, min_mass=0.1, max_mass=3.0
    ) -> Tuple[float, float]:

        random_mass = self.np_random.uniform(min_mass, max_mass)
        mass_scale_factor = random_mass / link.mass.value

        link.mass = random_mass
        link.inertia.ixx = link.inertia.ixx.value * mass_scale_factor
        link.inertia.iyy = link.inertia.iyy.value * mass_scale_factor
        link.inertia.izz = link.inertia.izz.value * mass_scale_factor
        link.inertia.ixy = link.inertia.ixy.value * mass_scale_factor
        link.inertia.ixz = link.inertia.ixz.value * mass_scale_factor
        link.inertia.iyz = link.inertia.iyz.value * mass_scale_factor

    def randomize_friction(self, link, min_friction=0.75, max_friction=1.5):

        for collision in link.collisions:
            random_friction = self.np_random.uniform(min_friction, max_friction)

            collision.surface.friction.ode.mu = random_friction
            collision.surface.friction.ode.mu2 = random_friction

    def write_inertial_properties(self, link, mass, inertia, centre_of_mass):

        link.mass = mass

        link.inertia.ixx = inertia[0][0]
        link.inertia.iyy = inertia[1][1]
        link.inertia.izz = inertia[2][2]
        link.inertia.ixy = inertia[0][1]
        link.inertia.ixz = inertia[0][2]
        link.inertia.iyz = inertia[1][2]

        link.inertial.pose = [
            centre_of_mass[0],
            centre_of_mass[1],
            centre_of_mass[2],
            0.0,
            0.0,
            0.0,
        ]

    def write_original_scale(self, mesh, model_path):

        file = open(self.get_original_scale_path(model_path), "w")
        file.write(str(mesh.scale))
        file.close()

    def read_original_scale(self, model_path) -> float:

        file = open(self.get_original_scale_path(model_path), "r")
        original_scale = file.read()
        file.close()

        return float(original_scale)

    def check_excessive_geometry(
        self, mesh, model_path, max_faces=40000, max_vertices=None
    ) -> bool:

        if max_faces is not None:
            num_faces = len(mesh.faces)
            if num_faces > max_faces:
                self.blacklist_model(
                    model_path, reason="Excessive geometry (%d faces)" % num_faces
                )
                return False

        if max_vertices is not None:
            num_vertices = len(mesh.vertices)
            if num_vertices > max_vertices:
                self.blacklist_model(
                    model_path, reason="Excessive geometry (%d vertices)" % num_vertices
                )
                return False

        return True

    def check_disconnected_components(
        self,
        mesh,
        model_path,
        component_min_faces_fraction=0.05,
        component_max_volume_fraction=0.1,
    ) -> bool:

        # Get a list of all connected componends inside the mesh
        # Consider components only with `component_min_faces_fraction` percent faces
        min_faces = round(component_min_faces_fraction * len(mesh.faces))
        connected_components = trimesh.graph.connected_components(
            mesh.face_adjacency, min_len=min_faces
        )

        # If more than 1 objects were detected, consider also relative volume of the meshes
        if len(connected_components) > 1:
            total_volume = mesh.volume

            large_component_counter = 0
            for component in connected_components:
                submesh = mesh.copy()
                mask = np.zeros(len(mesh.faces), dtype=np.bool)
                mask[component] = True
                submesh.update_faces(mask)

                volume_fraction = submesh.volume / total_volume
                if volume_fraction > component_max_volume_fraction:
                    large_component_counter += 1

                if large_component_counter > 1:
                    self.blacklist_model(
                        model_path,
                        reason="Disconnected components (%d instances)"
                        % len(connected_components),
                    )
                    return False

        return True

    def check_inertial_properties(self, model_path, mass, inertia) -> bool:

        if (
            mass < 1e-10
            or inertia[0][0] < 1e-10
            or inertia[1][1] < 1e-10
            or inertia[2][2] < 1e-10
        ):
            self.blacklist_model(model_path, reason="Invalid inertial properties")
            return False

        return True

    def get_random_model_path(self) -> str:

        if self._unique_cache:
            return self.np_random.choice(self._model_paths)
        else:
            return self.np_random.choice(self._class_model_paths)

    def get_collision_mesh_path(self, model_path, collision_name) -> str:

        return os.path.join(
            model_path,
            self.__collision_mesh_dir,
            collision_name + "." + self.__collision_mesh_file_type,
        )

    def get_sdf_path(self, model_path) -> str:

        return os.path.join(model_path, self.__sdf_base_name)

    def get_configured_sdf_path(self, model_path) -> str:

        return os.path.join(model_path, self.__configured_sdf_base_name)

    def get_blacklisted_path(self, model_path) -> str:

        return os.path.join(model_path, self.__blacklisted_base_name)

    def get_mesh_path(self, model_path, visual_or_collision) -> str:

        # TODO: This might need fixing for certain collections/models
        mesh_uri = visual_or_collision.geometry.mesh.uri.value
        return os.path.join(model_path, mesh_uri)

    def get_original_scale_path(self, model_path) -> str:

        return os.path.join(model_path, self.__original_scale_base_name)

    def blacklist_model(self, model_path, reason="Unknown"):

        if self._enable_blacklisting:
            bl_file = open(self.get_blacklisted_path(model_path), "w")
            bl_file.write(reason)
            bl_file.close()
        logger.warn(
            '%s model "%s". Reason: %s.'
            % (
                "Blacklisting" if self._enable_blacklisting else "Skipping",
                model_path,
                reason,
            )
        )

    def is_blacklisted(self, model_path) -> bool:

        return os.path.isfile(self.get_blacklisted_path(model_path))

    def is_configured(self, model_path) -> bool:

        return os.path.isfile(self.get_configured_sdf_path(model_path))

    def fix_mtl_texture_paths(self, model_path, mesh_path, model_name):

        # The `.obj` files use mtl
        if mesh_path.endswith(".obj"):

            # Find all textures located in the model path, used later to relative linking
            texture_files = glob.glob(os.path.join(model_path, "**", "textures", "*.*"))

            # Find location of mtl file, if any
            mtllib_file = None
            with open(mesh_path, "r") as file:
                for line in file:
                    if "mtllib" in line:
                        mtllib_file = line.split(" ")[-1].strip()
                        break

            if mtllib_file is not None:
                mtllib_file = os.path.join(os.path.dirname(mesh_path), mtllib_file)

                fin = open(mtllib_file, "r")
                data = fin.read()
                for line in data.splitlines():
                    if "map_" in line:
                        # Find the name of the texture/map in the mtl
                        map_file = line.split(" ")[-1].strip()

                        # Find the first match of the texture/map file
                        for texture_file in texture_files:
                            if os.path.basename(
                                texture_file
                            ) == map_file or os.path.basename(
                                texture_file
                            ) == os.path.basename(
                                map_file
                            ):

                                # Make the file unique to the model (unless it already is)
                                if model_name in texture_file:
                                    new_texture_file_name = texture_file
                                else:
                                    new_texture_file_name = texture_file.replace(
                                        map_file, model_name + "_" + map_file
                                    )

                                os.rename(texture_file, new_texture_file_name)

                                # Apply the correct relative path
                                data = data.replace(
                                    map_file,
                                    os.path.relpath(
                                        new_texture_file_name,
                                        start=os.path.dirname(mesh_path),
                                    ),
                                )
                                break
                fin.close()

                # Write in the correct data
                fout = open(mtllib_file, "w")
                fout.write(data)
                fout.close()
