import os, shutil
from zipfile import ZipFile
import numpy as np
import gradio as gr
from PIL import Image
from rembg import new_session, remove
from logging import getLogger, StreamHandler, DEBUG
from utils.functions import complete, clean, clean_by_name, get_random_name

logger = getLogger(__name__)
handler = StreamHandler()
handler.setLevel(DEBUG)
logger.setLevel(DEBUG)
logger.addHandler(handler)
logger.propagate = False
        
def run_rembg(img):
    """
    Remove background from an image using U2-Net algorithm.

    Args:
        img: numpy array, input image data.

    Returns:
        A cropped PIL image object.
    """
    output = remove(img)
    output_pil = Image.fromarray(output)

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

    return output_pil

def run_rembg_withmodel(img, model):    
    logger.debug(f"model name : {model}")
    logger.debug(f"model name : {type(model)}")
    session = new_session(model)
    logger.debug(f"session name : {session}")
    output = remove(img, session=session)
    output_pil = Image.fromarray(output)

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

    return output_pil
    
def from_zip(inputs):
    """
    Read images from a zip file and output a processed zip file.

    Args:
        inputs: list of file objects, input zip file.

    Returns:
        A tuple of output zip file name and a completion message.
    """
    work_dir = get_random_name()
    os.makedirs(work_dir, 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)
                        # logger.debug(f"image name : {image_name}")
                        category_dir = image_name.split("/")[0]
                        # image_name = image_name.split("/")[1]
                        os.makedirs(f"{work_dir}/{category_dir}", exist_ok=True)
                        image_data_dict[image_name] = image_array

                except Exception as e:
                    logger.info(f"Exception : {e}")
    
    
    for image_name, image_data in image_data_dict.items():

        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"{work_dir}/{image_name}")
    
    shutil.make_archive(work_dir, "zip", work_dir)
    shutil.rmtree(work_dir)
    
    return f"{work_dir}.zip", complete(work_dir)


def from_image_files(images, text_class_name):

    if not text_class_name=="":
        dir_name = text_class_name
    else:
        dir_name = get_random_name()

    os.makedirs(dir_name, exist_ok=True)

    for image in images:
        image_name = image.name
        # logger.debug(f"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[:image_name.find("_", image_name.find("_") + 1)] + ".png"
        # logger.debug(f"save image name : {image_name}")
        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("complete")+"+"+dir_name




if __name__=="__main__":

       
    with gr.Blocks() as demo:

        
        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_images.png' width="90%" height="90%">
                    """
                )        
                gr.Markdown(
                    """
                    ### Output Zip File
                    <img src='file/assets/output_zip.png' width="85%" height="85%">
                    """
                )        
            
            with gr.Row():
                with gr.Column():
                    text_class_name = gr.Textbox(label="Class Name", value="", placeholder="cat")
                    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, text_class_name],
                outputs=[image_output, text_output]
                )
            text_output.change(
                fn=clean_by_name,
                inputs=text_output,
                outputs=text_output

            )


        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%">
                    Zip file can include multiple directories.
                    """
                )        
                gr.Markdown(
                    """
                    ### Output Zip File
                    <img src='file/assets/output_zip.png' width="85%" height="85%">
                    If input has multiple directories, output has the same multiple diretocories.
                    """
                )        
            
            with gr.Row():
                image_input = gr.File(file_count="multiple")
                image_output = gr.File()
                text_output = gr.Textbox(visible=False, value="idle_state")
            
            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("ImageWithModel"):
            gr.Markdown(
                """
                <center><h1>Image Matting using different models</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():
                with gr.Column():
                    model_name = gr.Textbox(label="Model Name", value="", placeholder="u2net")
                    image_input = gr.Image(type="numpy")                
                image_output = gr.Image(type="pil")
            
            btn = gr.Button("Run!")


            btn.click(
                fn=run_rembg_withmodel,
                inputs=[image_input, model_name],
                outputs=image_output,
                api_name="imageMattingWithModel"
                )      



        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"
    )