# Copyright (C) 2022 Thomas Hoppe (h0bB1T). All rights reserved.
#
# Unauthorized copying of this file via any medium is strictly prohibited.
# Proprietary and confidential.

import bpy

from typing import Union

from bpy.types import Operator
from bpy.props import StringProperty, FloatProperty, IntProperty

import numpy as np

from ..utils.tools import Measure
from ..utils.dev import dbg
from ..utils.blender import drop_node
from .helpers.drop_node import DropNodeOperator
from .helpers.pbr_map_calculator import PbrQuantizer
from .helpers.setup_pbr_map import create_pbr_map, setup_pbr_map


class UI_OT_create_pbr_map(DropNodeOperator, Operator):
    bl_idname = 'snw.create_pbr_map'
    bl_label = 'Create Map'
    bl_description = 'Create PBR map node from PBR set'
    bl_options = {'REGISTER'}


    name: StringProperty() # type: ignore
    diffuse: StringProperty() # type: ignore
    metallic: StringProperty() # type: ignore
    specular: StringProperty() # type: ignore
    roughness: StringProperty() # type: ignore
    glossiness: StringProperty() # type: ignore
    metallic_const: FloatProperty() # type: ignore
    specular_const: FloatProperty() # type: ignore
    roughness_const: FloatProperty() # type: ignore
    colors: IntProperty() # type: ignore
    max_cluster_size: FloatProperty() # type: ignore
    order: StringProperty() # type: ignore


    @staticmethod
    def __load_blender_image(
        name_or_color: Union[str, float], 
        rgb: bool, 
        width: int = 0, 
        height: int = 0
        ):
        """
        Returns (pixels, (w, h)), either (R,G,B) or gray in [0..255] format. If image must be loaded,
        it is added to imgs_to_remove for later cleanup.
        """
        if isinstance(name_or_color, str):
            dbg(f'Load {name_or_color}')
            # Load and downscale for faster processing.
            img = bpy.data.images.load(name_or_color)

            pixels = img.pixels[:]
            img.scale(512, 512)

            # Extract pixel information.
            i_pixels = img.pixels[:]
            chnls = img.channels
            pixels = []
            scl = 511 if img.depth > 32 else 255 # 8/16 bit images, not always, have no idea yet ..
            if rgb:
                assert(chnls >= 3)
                r = np.rint(np.array(i_pixels, dtype=float)[0::chnls] * scl).astype(int)
                g = np.rint(np.array(i_pixels, dtype=float)[1::chnls] * scl).astype(int)
                b = np.rint(np.array(i_pixels, dtype=float)[2::chnls] * scl).astype(int)
                pixels = np.column_stack((r, g, b)).tolist()
            else:
                pixels = np.rint(np.array(i_pixels, dtype=float)[0::chnls] * 255).astype(int).tolist()

            # Remove from blenders image list.
            bpy.data.images.remove(img)

            return (pixels, (512, 512))
        else:
            # Shouldn't be RGB.
            assert(rgb == False)
            c = round(name_or_color * 255)
            return ((width * height) * [c,], (width, height))


    def execute(self, context: bpy.types.Context):
        with Measure('Load images ..'):
            diffuse_map, size = UI_OT_create_pbr_map.__load_blender_image(self.diffuse, True)

            metallic = self.metallic if self.metallic else self.metallic_const
            metallic_map, _ = UI_OT_create_pbr_map.__load_blender_image(metallic, False, *size)

            specular = self.specular if self.specular else self.specular_const
            specular_map, _ = UI_OT_create_pbr_map.__load_blender_image(specular, False, *size)

            if self.roughness:
                roughness_map, _ = UI_OT_create_pbr_map.__load_blender_image(self.roughness, False, *size)
            elif self.glossiness:
                roughness_map, _ = UI_OT_create_pbr_map.__load_blender_image(self.glossiness, False, *size)
                roughness_map = [ 255 - i for i in roughness_map ]
            else:
                roughness_map, _ = UI_OT_create_pbr_map.__load_blender_image(self.roughness_const, False, *size)

        with Measure('Create PBR map ..'):
            quantizer = PbrQuantizer()
            pbr_map = quantizer.process(diffuse_map, metallic_map, specular_map, roughness_map, self.colors, self.max_cluster_size)
            pbr_map.reorder(self.order)

        # Build the node that is dropped.
        nb = create_pbr_map(f'{self.name}.{self.order}.PBRMap')
        setup_pbr_map(nb, pbr_map)
        return drop_node(nb.tree)
