# copyright (c) 2018- polygoniq xyz s.r.o.

import typing
import collections
import bpy
import logging
from . import object_render_estimator

logger = logging.getLogger(f"polygoniq.{__name__}")


DECIMATE_MODIFIER_NAME = "memsaver_decimate"


def get_meshes_used_in_object(obj: bpy.types.Object) -> typing.Iterable[bpy.types.Mesh]:
    if obj.type == 'EMPTY':
        if obj.instance_type == 'COLLECTION':
            if obj.instance_collection is None:
                return  # warn?

            for instanced_object in obj.instance_collection.objects:
                yield from get_meshes_used_in_object(instanced_object)

    elif obj.type == 'MESH':
        if obj.data is None:
            return  # warn?

        yield obj.data


def generate_mesh_objects_map(
    objects: typing.Iterable[bpy.types.Object],
) -> typing.DefaultDict[bpy.types.Mesh, typing.List[bpy.types.Object]]:
    mesh_objects_map: typing.DefaultDict[bpy.types.Mesh, typing.List[bpy.types.Object]] = (
        collections.defaultdict(list)
    )

    for obj in objects:
        meshes = get_meshes_used_in_object(obj)
        for mesh in meshes:
            mesh_objects_map[mesh].append(obj)

    return mesh_objects_map


def update_object_decimation_ratio_map(
    objects_decimation_ratio_map: typing.DefaultDict[bpy.types.Object, float],
    mesh_objects_map: typing.DefaultDict[bpy.types.Mesh, typing.List[bpy.types.Object]],
    scene: bpy.types.Scene,
    camera: bpy.types.Object,
    objects: typing.Iterable[bpy.types.Object],
    full_quality_distance: float,
    lowest_quality_distance: float,
    lowest_quality_decimation_ratio: float,
    lowest_face_count: float,
) -> None:
    decimation_interpolation_denominator = lowest_quality_distance - full_quality_distance
    if decimation_interpolation_denominator <= 0:
        decimation_interpolation_denominator = 1

    for mesh, objects in mesh_objects_map.items():
        if len(mesh.polygons) < lowest_face_count:
            continue  # ignore meshes with lower face count than the one given

        mesh_min_distance = float("inf")

        for obj in objects:
            if obj.library is not None:  # linked, non-editable object
                # set min distance to 0 to avoid any decimation, we won't be able to add the same
                # decimation to all objects using this mesh, so let's not decimate at all
                mesh_min_distance = 0.0
                break

            _, _, obj_min_distance = object_render_estimator.get_object_2d_bounds(
                scene, camera, obj
            )
            if obj_min_distance is not None:
                mesh_min_distance = min(mesh_min_distance, obj_min_distance)

        # object_render_estimator will return negative values of the object intersects the camera
        # we assume non-negative values so let's clamp to 0
        mesh_min_distance = max(0.0, mesh_min_distance)

        # Everything closer than full_quality_distance has decimation ratio 1.0
        # Everything further than lowest_quality_distance has decimation ratio set to
        # lowest_quality_decimation_ratio
        # Between minimum_distance and lowest_quality_distance interpolate linearly from 1.0 to
        # lowest_quality_decimation_ratio

        decimation_ratio = 1.0
        if mesh_min_distance <= full_quality_distance:
            decimation_ratio = 1.0
        elif mesh_min_distance <= lowest_quality_distance:
            nominator = mesh_min_distance - full_quality_distance
            interpolation_factor = nominator / decimation_interpolation_denominator
            assert interpolation_factor >= 0.0
            assert interpolation_factor <= 1.0
            decimation_ratio = (
                1.0 - interpolation_factor
            ) * 1.0 + interpolation_factor * lowest_quality_decimation_ratio
        else:
            decimation_ratio = lowest_quality_decimation_ratio

        assert decimation_ratio >= 0.0
        assert decimation_ratio <= 1.0
        for obj in objects:
            objects_decimation_ratio_map[obj] = max(
                objects_decimation_ratio_map[obj], decimation_ratio
            )


def get_objects_decimation_ratio_map_current_frame(
    scene: bpy.types.Scene,
    camera: bpy.types.Object,
    objects: typing.Iterable[bpy.types.Object],
    minimum_distance: float,
    lowest_quality_distance: float,
    lowest_quality_decimation_ratio: float,
    lowest_face_count: float,
) -> typing.DefaultDict[bpy.types.Object, float]:
    mesh_objects_map = generate_mesh_objects_map(objects)
    objects_decimation_ratio_map: typing.DefaultDict[bpy.types.Object, float] = (
        collections.defaultdict(float)
    )
    update_object_decimation_ratio_map(
        objects_decimation_ratio_map,
        mesh_objects_map,
        scene,
        camera,
        objects,
        minimum_distance,
        lowest_quality_distance,
        lowest_quality_decimation_ratio,
        lowest_face_count,
    )
    return objects_decimation_ratio_map


def get_objects_decimation_ratio_map_animation_mode(
    scene: bpy.types.Scene,
    camera: bpy.types.Object,
    objects: typing.Iterable[bpy.types.Object],
    minimum_distance: float,
    lowest_quality_distance: float,
    lowest_quality_decimation_ratio: float,
    lowest_face_count: float,
) -> typing.DefaultDict[bpy.types.Object, float]:
    previous_frame_current = scene.frame_current
    try:
        mesh_objects_map = generate_mesh_objects_map(objects)
        objects_decimation_ratio_map: typing.DefaultDict[bpy.types.Object, float] = (
            collections.defaultdict(float)
        )
        current_frame = scene.frame_start
        while current_frame <= scene.frame_end:
            scene.frame_current = current_frame
            bpy.context.view_layer.update()

            update_object_decimation_ratio_map(
                objects_decimation_ratio_map,
                mesh_objects_map,
                scene,
                camera,
                objects,
                minimum_distance,
                lowest_quality_distance,
                lowest_quality_decimation_ratio,
                lowest_face_count,
            )
            current_frame += 1

        return objects_decimation_ratio_map
    finally:
        scene.frame_current = previous_frame_current
        bpy.context.view_layer.update()


def revert_to_original(obj: bpy.types.Object) -> None:
    if DECIMATE_MODIFIER_NAME in obj.modifiers:
        obj.modifiers.remove(obj.modifiers[DECIMATE_MODIFIER_NAME])


def set_decimation_ratio(obj: bpy.types.Object, decimation_ratio: float) -> None:
    if decimation_ratio == 1.0:
        # no decimation
        revert_to_original(obj)
    else:
        memsaver_decimation = None
        if DECIMATE_MODIFIER_NAME in obj.modifiers:
            memsaver_decimation = obj.modifiers[DECIMATE_MODIFIER_NAME]
        else:
            memsaver_decimation = obj.modifiers.new(DECIMATE_MODIFIER_NAME, 'DECIMATE')

        if memsaver_decimation is None:
            logger.error(
                f"Failed to create memsaver decimation modifier for object {obj.name} "
                f"of type {obj.type}, skipping!"
            )
            return
        memsaver_decimation.ratio = decimation_ratio