from zipfile import ZipFile
import numpy as np

import os, shutil

import gradio as gr
from PIL import Image
import sys, os
from rembg import remove
from utils import functions


def complete():
    return "complete"


def clean(text_output):
    if text_output=="complete":
        print("clean up")
        os.remove("objects.zip")
        return ""
    else:
        pass

def clean_by_name(text_output):
    text_output, dir_name = text_output.split("+")
    if text_output=="complete":
        print("clean up")
        os.remove(f"{dir_name}.zip")
        return ""
    else:
        pass
    
    

def run_rembg(img):
    output = remove(img)
    output_pil = Image.fromarray(output)

    # Remove margins 
    cropped_image = output_pil.crop(output_pil.getbbox())

    return cropped_image


def from_zip(inputs):
    os.makedirs("objects", exist_ok=True)

    image_data_dict = {}
    with ZipFile(inputs[0].name, "r") as zip_file:
        image_names = zip_file.namelist()

        prefix = ""
        for name in image_names:
            if prefix=="":
                prefix = name.split("/")[0]
            else:
                break

        image_files = []      
          
        for image_name in image_names:
            if image_name[-3:] in "pngjpg":
                try:
                    with zip_file.open(image_name) as f:
                        image = Image.open(f)
                        image_files.append(image_name)
                        
                        image_array = np.array(image)
                        image_name = image_name.split("/")[1]
                        image_data_dict[image_name] = image_array

                except Exception as e:
                    print("Exception : ", e)
    
    
    for image_name, image_data in image_data_dict.items():
        print(type(image_data))
        output = remove(image_data)        
        output_pil = Image.fromarray(output)
        # Remove margins 
        cropped_image = output_pil.crop(output_pil.getbbox())

        image_name = image_name.replace("jpg", "png")
        cropped_image.save(f"objects/{image_name}")
    
    shutil.make_archive("objects", "zip", "objects")
    shutil.rmtree("objects")
    
    return "objects.zip", complete()


def from_image_files(images):
    dir_name = functions.get_random_name()
    os.makedirs(dir_name, exist_ok=True)

    for image in images:
        image_name = image.name

    
        # 読み込み
        image_data = np.array(Image.open(image_name))

        output = remove(image_data)        
        output_pil = Image.fromarray(output)
        # Remove margins 
        cropped_image = output_pil.crop(output_pil.getbbox())

        image_name = image_name.split("/")[-1]
        image_name = image_name.replace("jpg", "png")
        cropped_image.save(f"{dir_name}/{image_name}")
    
    shutil.make_archive(f"{dir_name}", "zip", f"{dir_name}")
    shutil.rmtree(f"{dir_name}")
    
    return f"{dir_name}.zip", complete()+"+"+dir_name




if __name__=="__main__":

       
    with gr.Blocks() as demo:

        with gr.Tab("Zip"):
            gr.Markdown(
            """
            <center><h1>Image Matting using U<sup>2</sup>-Net</h1></center>
            """
            )        
            with gr.Row():
                gr.Markdown(
                    """
                    ### Input Zip File
                    <img src='file/assets/input_zip.png' width="85%" height="85%">
                    """
                )        
                gr.Markdown(
                    """
                    ### Output Zip File
                    <img src='file/assets/output_zip.png' width="85%" height="85%">
                    """
                )        
            
            with gr.Row():
                image_input = gr.File(file_count="multiple")
                image_output = gr.File()
                text_output = gr.Textbox(visible=False)
            
            btn = gr.Button("Run!")
            
            btn.click(
                fn=from_zip,
                inputs=image_input,
                outputs=[image_output, text_output]
                )
            text_output.change(
                fn=clean,
                inputs=text_output,
                outputs=text_output

            )


        with gr.Tab("Image"):
            gr.Markdown(
                """
                <center><h1>Image Matting using U<sup>2</sup>-Net</h1></center>
                """
            )        
            with gr.Row():
                gr.Markdown(
                    """
                    ### Input Image
                    <img src='file/assets/cat-3038243_1280.jpg' width="50%" height="50%">
                    """
                )        
                
                gr.Markdown(
                    """
                    ### Output Image
                    <img src='file/assets/objects/cat-3038243_1280.png' width="20%" height="20%">
                    """
                )        
            with gr.Row():
                image_input = gr.Image(type="numpy")
                image_output = gr.Image(type="pil")
            
            btn = gr.Button("Run!")


            btn.click(
                fn=run_rembg,
                inputs=image_input,
                outputs=image_output,
                api_name="imageMatting"
                )
        

        
        with gr.Tab("Images"):
            gr.Markdown(
            """
            <center><h1>Image Matting using U<sup>2</sup>-Net</h1></center>
            """
            )        
            with gr.Row():
                gr.Markdown(
                    """
                    ### Input Image Files
                    <img src='file/assets/input_zip.png' width="85%" height="85%">
                    """
                )        
                gr.Markdown(
                    """
                    ### Output Zip File
                    <img src='file/assets/output_zip.png' width="85%" height="85%">
                    """
                )        
            
            with gr.Row():
                image_input = gr.File(file_count="multiple")
                image_output = gr.File()
                text_output = gr.Textbox(visible=False)
            
            btn = gr.Button("Run!")
            
            btn.click(
                fn=from_image_files,
                inputs=image_input,
                outputs=[image_output, text_output]
                )
            text_output.change(
                fn=clean_by_name,
                inputs=text_output,
                outputs=text_output

            )




        gr.Markdown(
            """
            ---
            <b>Acknowledgments</b>
            - Library
              - Library Git hub : [danielgatis/rembg](https://github.com/danielgatis/rembg)
              - Cloned on 2023/3/12  
            - Algorithm        
              - Library Git hub : [U<sup>2</sup>-Net](https://github.com/xuebinqin/U-2-Net)  
            - Image
              - Cat Image from [Pixabay](https://pixabay.com/images/id-3038243/)
            """
        )

    demo.launch(
        favicon_path="./assets/ハサミのフリーアイコン.png"
    )