import gradio as gr
import datetime
import shutil
import numpy as np
import os, subprocess, re, glob, cv2
import os.path as osp
import insightface
from insightface.app import FaceAnalysis
from insightface.data import get_image as ins_get_image
import matplotlib.pyplot as plt
import gdown
from PIL import Image
from rich.progress import Progress, track

#/.insightface/models/buffalo_l.zip
app = FaceAnalysis(name='buffalo_l',providers=['CUDAExecutionProvider']) # , root=str('models/insightface')
app.prepare(ctx_id=0, det_size=(640, 640))
# Import the library
import gdown

# Google Drive file ID
file_id = '1MLAMn5pdHsx7lafeA3KjPechXiqYiYEM'

# Construct the download URL
url = f'https://drive.google.com/uc?id={file_id}'
#url = 'https://drive.google.com/file/d/1MLAMn5pdHsx7lafeA3KjPechXiqYiYEM/view?usp=drive_link'

# Download the file
gdown.download(url, 'inswapper_128.onnx', quiet=False)

swapper = insightface.model_zoo.get_model('inswapper_128.onnx',download=False,download_zip=False)
os.system("mkdir output")
os.system("mkdir frame")
os.system("mkdir allpic")


def facecompare(face1,face2):
  feat1 = face1.normed_embedding
  feat2 = face2.normed_embedding
  similarity = np.dot(feat1, feat2)
  threshold = 0.3
  if similarity > threshold:
      return True
  else:
      return False

def faceswapdir(face_image,path=".",face_id=0,fcfrmid=0):
  img1 = cv2.imread(face_image)
  lst = os.listdir(path)
  lst.sort()
  img_lst = [cv2.imread(f"{path}/{img}") for img in track(lst,description='Read frames')]
  base_face = app.get(img1)[0]
  face_lst =[app.get(face)[face_id] if app.get(face) else None  for face in track(img_lst,description='Extract faces ')]
  out_img = [swapper.get(swap.copy(), face_lst[i], base_face, paste_back=True)  if face_lst[i] is not None  and  facecompare(face_lst[i],face_lst[fcfrmid]) else img_lst[i] for i,swap in enumerate(track(img_lst, description='Swap frames '))]
  swapped_img1_rgb = [cv2.cvtColor(swapped, cv2.COLOR_BGR2RGB) for swapped in track(out_img,description='Refine swapped faces ')]

  for i,swapped2 in enumerate(swapped_img1_rgb):
    Image.fromarray(np.uint8(swapped2)).save(f"output/{i+1:06}.jpg")

def faceswapdirold(face_image,path=".",face_id=0):
  img1 = cv2.imread(face_image)
  lst = os.listdir(path)
  lst.sort()
  img_lst = [cv2.imread(f"{path}/{img}") for img in lst]
  base_face = app.get(img1)[0]
  face_lst =[app.get(face)[0] if app.get(face) else None  for face in img_lst]
  out_img = [swapper.get(swap.copy(), face_lst[i], base_face, paste_back=True)  if face_lst[i] is not None  else img_lst[i] for i,swap in enumerate(track(img_lst, description="Swap frames "))]
  swapped_img1_rgb = [cv2.cvtColor(swapped, cv2.COLOR_BGR2RGB) for swapped in out_img]

  for i,swapped2 in enumerate(swapped_img1_rgb):
    Image.fromarray(np.uint8(swapped2)).save(f"output/{i+1:06}.jpg")

def faceswapfnvidold(path,path2,num1=0):         
    shutil.copyfile(path.name, 'video.mp4')
    shutil.copyfile(path2.name, 'baseimg.jpg')  
    
    os.system('ffmpeg -i video.mp4 -q:a 0 -map a ia.mp3')  
    os.system('ffmpeg -hwaccel auto -i video.mp4 -q:v 18 -pix_fmt rgb24 -vf fps=25  frame/%06d.jpg')
    faceswapdirold('baseimg.jpg','frame',num1)
    os.system('ffmpeg -hwaccel auto -r 25 -i output/%06d.jpg -vcodec libx264 -crf 18  -pix_fmt yuv420p   ov.mp4')
    os.system('ffmpeg -i ov.mp4 -i ia.mp3 -c:v copy -c:a libmp3lame  out_video.mp4')    
    
    return ['out_video.mp4']
    
def faceswapfnvid(path,num1,path2):         
    shutil.copyfile(path.name, 'video.mp4')
    shutil.copyfile(path2.name, 'baseimg.jpg')  
    
    os.system('ffmpeg -i video.mp4 -q:a 0 -map a ia.mp3')  
    os.system('ffmpeg -hwaccel auto -i video.mp4 -q:v 18 -pix_fmt rgb24 -vf fps=25  frame/%06d.jpg')
    faceswapdir('baseimg.jpg','frame',0,num1)
    os.system('ffmpeg -hwaccel auto -r 25 -i output/%06d.jpg -vcodec libx264 -crf 18  -pix_fmt yuv420p   ov.mp4')
    os.system('ffmpeg -i ov.mp4 -i ia.mp3 -c:v copy -c:a libmp3lame  out_video.mp4')    
    
    return ['out_video.mp4']

def faceswapfn(path,path2):         
    shutil.copyfile(path.name, 'allpic/targ.jpg')
    shutil.copyfile(path2.name, 'baseimg.jpg')  
    
    faceswapdir('baseimg.jpg','allpic',0,0)
    
    return ['output/000001.jpg']

#demo = gr.Interface(fn=faceswapfn, title="Subtitle the Video" ,inputs="file", outputs="file")
gr.set_static_paths(paths=["/home/user/app"])

with gr.Blocks() as demo:
    with gr.Tab('Video Faceswap old'):
        gr.Markdown("Upload Video old")
        with gr.Row():
            inp = [gr.File(file_types=['.mp4'],label="Target Video"),gr.File(file_types=['.jpg','.jpeg','.png'],label="Base Face Image"),gr.Number(value=0)]
            out = gr.File(label="Swapped Video")
        btn = gr.Button("Face Swap")
        btn.click(fn=faceswapfnvidold, inputs=inp, outputs=out)
    with gr.Tab('Video Faceswap'):
        gr.Markdown("Upload Video")
        with gr.Row():
            inp = [gr.File(file_types=['.mp4'],label="Target Video"),gr.Number(value=0),gr.File(file_types=['.jpg','.jpeg','.png'],label="Base Face Image")]
            out = gr.File(label="Swapped Video")
        btn = gr.Button("Face Swap")
        btn.click(fn=faceswapfnvid, inputs=inp, outputs=out)
    with gr.Tab('Image Faceswap'):
        gr.Markdown("Upload Image")
        with gr.Row():
            inp = [gr.File(file_types=['.jpg','.jpeg','.png'], label="Target Image"),gr.File(file_types=['.jpg','.jpeg','.png'], label="Base Face Image")]
            out = gr.File(label="Swapped Image")
        btn = gr.Button("Face Swap")
        btn.click(fn=faceswapfn, inputs=inp, outputs=out)
    
demo.launch(share=True,allowed_paths=['out_video.mp4','output/000001.jpg'])