import os 
import cv2
import gradio as gr
import requests
from PIL import Image
from PIL import UnidentifiedImageError
import io
from io import BytesIO
import AnimeGANv3_src # Import the actual AnimeGANv3_src module



class AnimeGANv3:
    def __init__(self):
        # Ensure the output directory exists
        os.makedirs('output', exist_ok=True)
        
    def inference(self, img_path, style, face_detection):
        print(img_path, style, face_detection)
        try:
            # Read an image from the specified path
            img = cv2.imread(img_path)
            if img is None:
                raise ValueError(f"Unable to read image at {img_path}")
                
            # Convert the image to RGB 
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            
            # Determine the style code based on the style parameter
            style_code = {
                "AnimeGANv3_Disney v1.0": "D",
                
            }
            f = style_code.get(style, "U") # Default to "U" if style is not found
            
            try:
                det_face = face_detection == "Yes"
                # Ust the convert method from the AnimeGANv3_src module
                output = AnimeGANv3_src.Convert(img, f, det_face)
                
                # Save the output image to a file
                save_path = f"output/out.{img_path.rsplit('.')[-1]}"
                cv2.imwrite(save_path, cv2.cvtColor(output, cv2.COLOR_RGB2BGR))
                
                return output, save_path
            except RuntimeError as error:
                print('global exception', error)
                return None, None
        except Exception as error:
            print('global exception', error)
            return None, None





# Define the DESCRIPTION markdown with an embedded image
DESCRIPTION = f"""
# [AnimeGANv3_Disney v1.0 Implementation](https://github.com/TachibanaYoshino/AnimeGANv3)

"""

# Create an instance of the AnimeGANv3 class
anime_gan_disney = AnimeGANv3()

# Define the Gradio interface
iface = gr.Interface(
    fn=anime_gan_disney.inference,
    inputs=[
        gr.components.Image(type="filepath", label="Input"),
        gr.components.Dropdown(choices=[
            'AnimeGANv3_Disney v1.0',
            # ... (other choices)
        ], label='AnimeGANv3 Style'),
        gr.components.Radio(choices=["Yes", "No"], label='Extract face'),
    ],
    outputs=[
        gr.components.Image(type="numpy", label="Output (The whole image)"),
        gr.components.File(label="Download the output image")
    ],
    description=DESCRIPTION,
    allow_flagging="never",
    examples=[
        ['samples/Angela_Aguilar_0.jpg', 'AnimeGANv3_Disney v1.0', "Yes"],
       # ['samples/Angela_Aguilar_.jpg', 'AnimeGANv3_Disney v1.0', "No"],
        ['samples/Avantika_Vandanapu.jpg', 'AnimeGANv3_Disney v1.0', "Yes"],
        #['samples/Avantika_Vandanapu_.jpg', 'AnimeGANv3_Disney v1.0', "No"],
        #['samples/Ester_Exposito.jpg', 'AnimeGANv3_Disney v1.0', "No"],
        ['samples/Ester_Exposito_.jpg', 'AnimeGANv3_Disney v1.0', "Yes"],
        ['samples/Giorgia_Whigham.jpg', 'AnimeGANv3_Disney v1.0', "Yes"],
        #['samples/Giorgia_Whigham_.jpg', 'AnimeGANv3_Disney v1.0', "Yes"],
        #['samples/Laura_Osma.jpg', 'AnimeGANv3_Disney v1.0', "Yes"],
        ['samples/Laura_Osma_.jpg', 'AnimeGANv3_Disney v1.0', "No"],
       # ['samples/Michelle_Pellicer.jpg', 'AnimeGANv3_Disney v1.0', "No"],
        ['samples/Michelle_Pellicer_.jpeg', 'AnimeGANv3_Disney v1.0', "No"],
        #['samples/Olivia_Rodrigo.jpg', 'AnimeGANv3_Disney v1.0', "No"],
        ['samples/Olivia_Rodrigo_.jpg', 'AnimeGANv3_Disney v1.0', "No"],
        #['output/output_Gemelas.png', 'AnimeGANv3_Disney v1.0', "No"],
    ]
)
    
# Launch the Gradio interface
iface.launch()