
import bpy
import math
import numpy as np

from pathlib import Path


def alpha_to_color(pixels_data, color):
    """Convert Alpha to WhiteBG"""
    new_pixels_data = []
    for i in pixels_data:
        height, width, array_d = i.shape 
        mask = i[:,:,3:]
        background = np.array([color[0], color[1], color[2] ,1], dtype=np.float32)
        background = np.tile(background, (height*width))
        background = np.reshape(background, (height,width,4))
        new_pixels_data.append(i * mask + background * (1 - mask))
        # print(new_pixels_data)#Dbg
    return new_pixels_data

def create_array(height, width):
    return np.zeros((height*width*4), dtype=np.float32)

def read_pixels_data(img, source_height, source_width):
    img_w, img_h = img.size

    if img_w != source_width :
        scale = abs(img_w/source_width)
        img.scale(int(img_w/scale), int(img_h/scale))
        img_w, img_h = img.size

    array = create_array(img_h, img_w)
    img.pixels.foreach_get(array)
    array = array.reshape(img_h, img_w, 4)
    
    if array.shape[0] != source_height:
        #print('ARRAY SHAPE', array.shape[:], source_height)
        missing_height = int(abs(source_height-img_h)/2)
        empty_array = create_array(missing_height, source_width)
        empty_array = empty_array.reshape(missing_height, source_width, 4)
        array = np.vstack((empty_array, array, empty_array))

    return array.reshape(source_height, source_width, 4)

def create_final(output_name, pixels_data, final_height, final_width):
    #print('output_name: ', output_name)

    new_img = bpy.data.images.get(output_name)
    if new_img:
        bpy.data.images.remove(new_img)

    new_img = bpy.data.images.new(output_name, final_width, final_height)
    new_img.generated_color=(0,0,0,0)

    #print('pixels_data: ', pixels_data)
    new_img.pixels.foreach_set(pixels_data)

    return new_img

def guess_input_format(img_list):
    for i in img_list:
        if i.size[0] == i.size[1]:
            return i.size

def format_files(files, catalog_data):
    img_dict = {}
    for k, v in catalog_data.items():
        if '/' not in k:
            continue
        img_dict[v['name']] = [f for f in files if v['name'] in f]

    return img_dict

def mosaic_export(
    files, catalog_data, row=2, columns=2, auto_calculate=True,
    bg_color=(0.18, 0.18, 0.18,), resize_output=100,
):

    img_dict = format_files(files, catalog_data)
    
    for cat, files_list in img_dict.items():

        if not files_list:
            continue

        for i in bpy.data.images:
            bpy.data.images.remove(i)
            
        img_list = []
        
        chars = Path(files_list[0]).parts[-4]
        output_dir = str(Path(files_list[0]).parent.parent)

        ext = 'jpg'
        output_name = f'{chars}_{cat}.{ext}'

        for img in files_list:
            img_list.append(bpy.data.images.load(img, check_existing=True))

        for i in img_list:
            i.colorspace_settings.name = 'Raw'

        if auto_calculate:
            rows = int(math.sqrt(len(img_list)))
            columns = math.ceil(len(img_list)/rows)

        if rows*columns < len(img_list):
            raise AttributeError('Grid too small for number of images')
        
        src_w, src_h = img_list[0].size
        final_w = src_w * columns
        final_h = src_h * rows

        img_pixels = [read_pixels_data(img, src_h, src_w) for img in img_list]
        
        #Check if there is enough "data" to create an horizontal stack
        ##It not, create empty array
        h_stack = []
        total_len = rows*columns
        if len(img_pixels) < total_len:
            for i in range(total_len-len(img_pixels)):
                img_pixels.append(create_array(src_h, src_w).reshape(src_h, src_w, 4))
        
        img_pixels = alpha_to_color(img_pixels, bg_color)
        for i in range(0,len(img_pixels),columns):
            h_stack.append(np.hstack(img_pixels[i:i+columns]))
        if rows > 1:
            combined_stack = np.vstack(h_stack[::-1])
        else:
            combined_stack = np.hstack((h_stack[:]))

        combined_img = create_final(output_name, combined_stack.flatten(), final_h, final_w)

        if resize_output != 100:
            w, h = combined_img.size
            combined_img.scale(w*(resize_output*.01), h*(resize_output*.01))
        

        combined_img.filepath_raw = '/'.join([output_dir, output_name])
        combined_img.file_format = 'JPEG'
        combined_img.save()

        print(f"""
        Image saved: {combined_img.filepath_raw}
        """)
