import bpy, os, sys, math, mathutils, importlib
import numpy as np
from . rectpack import newPacker, PackingMode, PackingBin

def postpack():

    cv_installed = False

    cv2 = importlib.util.find_spec("cv2")

    if cv2 is None:
        print("CV2 not found - Ignoring postpacking")
        return 0
    else:
        cv2 = importlib.__import__("cv2")
        cv_installed = True

    if cv_installed:

        lightmap_directory = os.path.join(os.path.dirname(bpy.data.filepath), bpy.context.scene.TLM_EngineProperties.tlm_lightmap_savedir)

        packedAtlas = {}

        #TODO - TEST WITH ONLY 1 ATLAS AT FIRST (1 Atlas for each, but only 1 bin (no overflow))
        #PackedAtlas = Packer
        #Each atlas has bins
        #Each bins has rects
        #Each rect corresponds to a pack_object

        scene = bpy.context.scene
        
        sceneProperties = scene.TLM_SceneProperties

        end = "_baked"

        if sceneProperties.tlm_denoise_use:

            end = "_denoised"

        if sceneProperties.tlm_filtering_use:

            end = "_filtered"

        formatEnc = ".hdr"

        image_channel_depth = cv2.IMREAD_ANYDEPTH
        linear_straight = False
        
        if sceneProperties.tlm_encoding_use and scene.TLM_EngineProperties.tlm_bake_mode != "Background":

            if sceneProperties.tlm_encoding_device == "CPU":

                if sceneProperties.tlm_encoding_mode_a == "HDR":

                    if sceneProperties.tlm_format == "EXR":

                        formatEnc = ".exr"

                if sceneProperties.tlm_encoding_mode_a == "RGBM":

                    formatEnc = "_encoded.png"
                    image_channel_depth = cv2.IMREAD_UNCHANGED

            else:

                if sceneProperties.tlm_encoding_mode_b == "HDR":

                    if sceneProperties.tlm_format == "EXR":

                        formatEnc = ".exr"

                if sceneProperties.tlm_encoding_mode_b == "LogLuv":

                    formatEnc = "_encoded.png"
                    image_channel_depth = cv2.IMREAD_UNCHANGED
                    linear_straight = True

                if sceneProperties.tlm_encoding_mode_b == "RGBM":

                    formatEnc = "_encoded.png"
                    image_channel_depth = cv2.IMREAD_UNCHANGED

                if sceneProperties.tlm_encoding_mode_b == "RGBD":

                    formatEnc = "_encoded.png"
                    image_channel_depth = cv2.IMREAD_UNCHANGED

        packer = {}

        for atlas in bpy.context.scene.TLM_PostAtlasList: #For each atlas

            packer[atlas.name] = newPacker(PackingMode.Offline, PackingBin.BFF, rotation=False)

            if scene.TLM_EngineProperties.tlm_setting_supersample == "2x":
                supersampling_scale = 2
            elif scene.TLM_EngineProperties.tlm_setting_supersample == "4x":
                supersampling_scale = 4
            else:
                supersampling_scale = 1

            atlas_resolution = int(int(atlas.tlm_atlas_lightmap_resolution) / int(scene.TLM_EngineProperties.tlm_resolution_scale) * int(supersampling_scale))

            packer[atlas.name].add_bin(atlas_resolution, atlas_resolution, 1)

            #AtlasList same name prevention?
            rect = []

            #For each object that targets the atlas
            for obj in bpy.context.scene.objects:
                if obj.TLM_ObjectProperties.tlm_mesh_lightmap_use:
                    if obj.TLM_ObjectProperties.tlm_postpack_object:
                        if obj.TLM_ObjectProperties.tlm_postatlas_pointer == atlas.name:

                            res = int(int(obj.TLM_ObjectProperties.tlm_mesh_lightmap_resolution) / int(scene.TLM_EngineProperties.tlm_resolution_scale) * int(supersampling_scale))

                            rect.append((res, res, obj.name))

            for r in rect:
                packer[atlas.name].add_rect(*r)

            packedAtlas[atlas.name] = np.zeros((atlas_resolution,atlas_resolution, 3), dtype="float32")

            #Continue here...overwrite value if using 8-bit encoding
            if sceneProperties.tlm_encoding_use:
                if sceneProperties.tlm_encoding_device == "CPU":
                    if sceneProperties.tlm_encoding_mode_a == "RGBM":
                        packedAtlas[atlas.name] = np.zeros((atlas_resolution,atlas_resolution, 4), dtype=np.uint8)
                    if sceneProperties.tlm_encoding_mode_a == "RGBD":
                        packedAtlas[atlas.name] = np.zeros((atlas_resolution,atlas_resolution, 4), dtype=np.uint8)
                
                if sceneProperties.tlm_encoding_device == "GPU":
                    if sceneProperties.tlm_encoding_mode_b == "RGBM":
                        packedAtlas[atlas.name] = np.zeros((atlas_resolution,atlas_resolution, 4), dtype=np.uint8)
                    if sceneProperties.tlm_encoding_mode_b == "RGBD":
                        packedAtlas[atlas.name] = np.zeros((atlas_resolution,atlas_resolution, 4), dtype=np.uint8)
                    if sceneProperties.tlm_encoding_mode_b == "LogLuv":
                        packedAtlas[atlas.name] = np.zeros((atlas_resolution,atlas_resolution, 4), dtype=np.uint8)

            packer[atlas.name].pack()

            for idy, rect in enumerate(packer[atlas.name].rect_list()):

                aob = rect[5]

                src = cv2.imread(os.path.join(lightmap_directory, aob + end + formatEnc), image_channel_depth) #"_baked.hdr"

                print("Obj name is: " + aob)

                x,y,w,h = rect[1],rect[2],rect[3],rect[4]

                print(src.shape)
                print(packedAtlas[atlas.name].shape)

                packedAtlas[atlas.name][y:h+y, x:w+x] = src
                
                obj = bpy.data.objects[aob]

                for idx, layer in enumerate(obj.data.uv_layers):

                    if not obj.TLM_ObjectProperties.tlm_use_default_channel:
                        uv_channel = obj.TLM_ObjectProperties.tlm_uv_channel
                    else:
                        uv_channel = "UVMap_Lightmap"

                    if layer.name == uv_channel:
                        obj.data.uv_layers.active_index = idx

                        print("UVLayer set to: " + str(obj.data.uv_layers.active_index))

                for uv_verts in obj.data.uv_layers.active.data:

                    #SET UV LAYER TO 

                    atlasRes = atlas_resolution
                    texRes = rect[3]
                    x,y,w,z = x,y,texRes,texRes
                    
                    ratio = atlasRes/texRes
                    
                    if x == 0:
                        x_offset = 0
                    else:
                        x_offset = 1/(atlasRes/x)

                    if y == 0:
                        y_offset = 0
                    else:
                        y_offset = 1/(atlasRes/y)
                    
                    vertex_x = (uv_verts.uv[0] * 1/(ratio)) + x_offset
                    vertex_y = (1 - ((uv_verts.uv[1] * 1/(ratio)) + y_offset))
                    
                    uv_verts.uv[0] = vertex_x
                    uv_verts.uv[1] = vertex_y

                scaleUV(obj.data.uv_layers.active, (1, -1), getBoundsCenter(obj.data.uv_layers.active))
                print(getCenter(obj.data.uv_layers.active))

            cv2.imwrite(os.path.join(lightmap_directory, atlas.name + end + formatEnc), packedAtlas[atlas.name])
            print("Written: " + str(os.path.join(lightmap_directory, atlas.name + end + formatEnc)))

            #Change the material for each material, slot
            for obj in bpy.context.scene.objects:
                if obj.TLM_ObjectProperties.tlm_mesh_lightmap_use:
                    if obj.TLM_ObjectProperties.tlm_postpack_object:
                        if obj.TLM_ObjectProperties.tlm_postatlas_pointer == atlas.name:
                            for slot in obj.material_slots:
                                nodetree = slot.material.node_tree

                                for node in nodetree.nodes:

                                    if node.name == "TLM_Lightmap":

                                        existing_image = node.image

                                        atlasImage = bpy.data.images.load(os.path.join(lightmap_directory, atlas.name + end + formatEnc), check_existing=True)

                                        if linear_straight:
                                            if atlasImage.colorspace_settings.name != 'Linear':
                                                atlasImage.colorspace_settings.name = 'Linear'

                                        node.image = atlasImage

                                        os.remove(os.path.join(lightmap_directory, obj.name + end + formatEnc))
                                        existing_image.user_clear()

            #Add dilation map here...
            for obj in bpy.context.scene.objects:
                if obj.TLM_ObjectProperties.tlm_mesh_lightmap_use:
                    if obj.TLM_ObjectProperties.tlm_postpack_object:
                        if obj.TLM_ObjectProperties.tlm_postatlas_pointer == atlas.name:
                            if atlas.tlm_atlas_dilation:
                                for slot in obj.material_slots:
                                    nodetree = slot.material.node_tree

                                    for node in nodetree.nodes:

                                        if node.name == "TLM_Lightmap":

                                            existing_image = node.image

                                            atlasImage = bpy.data.images.load(os.path.join(lightmap_directory, atlas.name + end + formatEnc), check_existing=True)

                                            img = cv2.imread(atlasImage.filepath_raw, image_channel_depth)

                                            kernel = np.ones((5,5), dtype="float32")

                                            img_dilation = cv2.morphologyEx(img, cv2.MORPH_OPEN, kernel)

                                            cv2.imshow('Dilation', img_dilation)
                                            cv2.waitKey(0)
                                            
                                            print("TODO: Adding dilation for: " + obj.name)
                                            #TODO MASKING OPTION!



    else:

        print("OpenCV not installed. Skipping postpacking process.")

def getCenter(uv_layer):

    total_x, total_y = 0,0
    len = 0

    for uv_verts in uv_layer.data:
        total_x += uv_verts.uv[0]
        total_y += uv_verts.uv[1]

        len += 1 

    center_x = total_x / len
    center_y = total_y / len

    return (center_x, center_y)

def getBoundsCenter(uv_layer):

    min_x = getCenter(uv_layer)[0]
    max_x = getCenter(uv_layer)[0]
    min_y = getCenter(uv_layer)[1]
    max_y = getCenter(uv_layer)[1]

    len = 0

    for uv_verts in uv_layer.data:

        if uv_verts.uv[0] < min_x:
            min_x = uv_verts.uv[0]
        if uv_verts.uv[0] > max_x:
            max_x = uv_verts.uv[0]
        if uv_verts.uv[1] < min_y:
            min_y = uv_verts.uv[1]
        if uv_verts.uv[1] > max_y:
            max_y = uv_verts.uv[1]

    center_x = (max_x - min_x) / 2 + min_x
    center_y = (max_y - min_y) / 2 + min_y

    return (center_x, center_y)


def scale2D(v, s, p):
    return (p[0] + s[0]*(v[0] - p[0]), p[1] + s[1]*(v[1] - p[1]))

def scaleUV( uvMap, scale, pivot ):
    for uvIndex in range( len(uvMap.data) ):
        uvMap.data[uvIndex].uv = scale2D(uvMap.data[uvIndex].uv, scale, pivot)