import subprocess
import shutil
import os
import gradio as gr
import torchvision.transforms as T
import sys
import spaces
from PIL import Image

subprocess.run(["git", "clone", "https://github.com/AIRI-Institute/HairFastGAN"], check=True)
os.chdir("HairFastGAN")

subprocess.run(["git", "clone", "https://huggingface.co/AIRI-Institute/HairFastGAN"], check=True)

os.chdir("HairFastGAN")
subprocess.run(["git", "lfs", "pull"], check=True)
os.chdir("..")

shutil.move("HairFastGAN/pretrained_models", "pretrained_models")
shutil.move("HairFastGAN/input", "input")

shutil.rmtree("HairFastGAN")

items = os.listdir()

for item in items:
    print(item)
    shutil.move(item, os.path.join('..', item))

os.chdir("..")

shutil.rmtree("HairFastGAN")

from hair_swap import HairFast, get_parser

hair_fast = HairFast(get_parser().parse_args([]))

def resize(image_path):
    img = Image.open("image_path")
    square_size = 1024
    
    left = (img.width - square_size) / 2
    top = (img.height - square_size) / 2
    right = (img.width + square_size) / 2
    bottom = (img.height + square_size) / 2
    
    img_cropped = img.crop((left, top, right, bottom))
    return img_cropped

@spaces.GPU
def swap_hair(source, target_1, target_2, progress=gr.Progress(track_tqdm=True)):
    target_2 = target_2 if target_2 else target_1
    final_image = hair_fast.swap(source, target_1, target_2)
    return T.functional.to_pil_image(final_image)
    
with gr.Blocks() as demo:
    gr.Markdown("## HairFastGan")
    with gr.Row():
        source = gr.Image(label="Photo that you want to replace the hair", type="filepath")
        target_1 = gr.Image(label="Reference hair you want to get", type="filepath")
        target_2 = gr.Image(label="Reference color hair you want to get (optional)", type="filepath")
    btn = gr.Button("Get the haircut")
    output = gr.Image(label="Your result")
    gr.Examples(examples=[("michael_cera-min.png", "leo_square-min.png", "pink_hair_celeb-min.png")], inputs=[source, target_1, target_2], outputs=output)
    source.upload(fn=resize, input=source, output=source)
    target_1.upload(fn=resize, input=target_1, output=target_1)
    target_2.upload(fn=resize, input=target_2, output=target_2)
    btn.click(fn=swap_hair, inputs=[source, target_1, target_2], outputs=[output])

demo.launch()