import os
import cv2
import gradio as gr
import AnimeGANv3_src
import numpy as np
from moviepy.editor import VideoFileClip, ImageSequenceClip

class AnimeGANv3:
    def __init__(self):
        # Ensure the output directory exists
        os.makedirs('output', exist_ok=True)

    def process_frame(self, frame, style_code, det_face):
        """Process a single frame with AnimeGANv3."""
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        output = AnimeGANv3_src.Convert(frame_rgb, style_code, det_face)
        return output

    def inference(self, video_path, style, if_face=None):
        print(video_path, style, if_face)
        try:
            # Map style names to codes
            style_codes = {
                "AnimeGANv3_Arcane": "A",
                "AnimeGANv3_Trump v1.0": "T",
                "AnimeGANv3_Shinkai": "S",
                "AnimeGANv3_PortraitSketch": "P",
                "AnimeGANv3_Hayao": "H",
                "AnimeGANv3_Disney v1.0": "D",
                "AnimeGANv3_JP_face v1.0": "J",
                "AnimeGANv3_Kpop v2.0": "K",
            }
            style_code = style_codes.get(style, "U")
            det_face = if_face == "Yes"

            # Open the video
            cap = cv2.VideoCapture(video_path)
            if not cap.isOpened():
                raise Exception("Could not open video file")

            # Get video properties
            fps = cap.get(cv2.CAP_PROP_FPS)
            frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            print(f"Processing {frame_count} frames at {fps} FPS")

            # Extract and process frames
            frames = []
            while True:
                ret, frame = cap.read()
                if not ret:
                    break
                processed_frame = self.process_frame(frame, style_code, det_face)
                frames.append(processed_frame)
            cap.release()

            # Convert frames to video
            save_path = "output/out.mp4"
            clip = ImageSequenceClip(frames, fps=fps)
            clip.write_videofile(save_path, codec="libx264", audio=False)

            return save_path
        except Exception as error:
            print('Error:', error)
            return None

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

# Define the Gradio interface
title = "AnimeGANv3: Video to Anime Converter"
description = r"""Upload a video to convert it into anime style using AnimeGANv3.<br>
Select a style and choose whether to optimize for faces.<br>
<a href='https://github.com/TachibanaYoshino/AnimeGANv3' target='_blank'><b>AnimeGANv3 GitHub</b></a> | 
<a href='https://www.patreon.com/Asher_Chan' target='_blank'><b>Patreon</b></a>"""

iface = gr.Interface(
    fn=anime_gan.inference,
    inputs=[
        gr.Video(label="Input Video"),
        gr.Dropdown(choices=[
            'AnimeGANv3_Hayao',
            'AnimeGANv3_Shinkai',
            'AnimeGANv3_Arcane',
            'AnimeGANv3_Trump v1.0',
            'AnimeGANv3_Disney v1.0',
            'AnimeGANv3_PortraitSketch',
            'AnimeGANv3_JP_face v1.0',
            'AnimeGANv3_Kpop v2.0',
        ], label='AnimeGANv3 Style', default='AnimeGANv3_Arcane'),
        gr.Radio(choices=["Yes", "No"], label='Extract face', default="No"),
    ],
    outputs=[
        gr.Video(label="Output Video")
    ],
    title=title,
    description=description,
    allow_flagging="never"
)

# Launch the interface
iface.launch()