from typing import Optional

import cv2
import numpy as np
import torch

from configs.train_config import TrainConfig
from models.model import HifiFace


def inference(source_face: str, target_face: str, model_path: str, model_idx: Optional[int], swapped_face: str):
    opt = TrainConfig()
    opt.use_ddp = False

    device = "cpu"
    checkpoint = (model_path, model_idx)
    model = HifiFace(opt.identity_extractor_config, is_training=False, device=device, load_checkpoint=checkpoint)
    model.eval()

    src = cv2.cvtColor(cv2.imread(source_face), cv2.COLOR_BGR2RGB)
    src = cv2.resize(src, (256, 256))
    src = src.transpose(2, 0, 1)
    src = torch.from_numpy(src).unsqueeze(0).to(device).float()
    src = src / 255.0

    tgt = cv2.cvtColor(cv2.imread(target_face), cv2.COLOR_BGR2RGB)
    tgt = cv2.resize(tgt, (256, 256))
    tgt = tgt.transpose(2, 0, 1)
    tgt = torch.from_numpy(tgt).unsqueeze(0).to(device).float()
    tgt = tgt / 255.0

    with torch.no_grad():
        result_face = model.forward(src, tgt).cpu()
        result_face = torch.clamp(result_face, 0, 1) * 255
    result_face = result_face.numpy()[0].astype(np.uint8)
    result_face = result_face.transpose(1, 2, 0)

    result_face = cv2.cvtColor(result_face, cv2.COLOR_BGR2RGB)
    cv2.imwrite(swapped_face, result_face)


if __name__ == "__main__":
    source_face = "/home/xuehongyang/data/female_1.jpg"
    target_face = "/home/xuehongyang/data/female_2.jpg"
    model_path = "/data/checkpoints/hififace/baseline_1k_ddp_with_cyc_1681278017147"
    model_idx = 80000
    swapped_face = "/home/xuehongyang/data/male_1_to_male_2.jpg"
    inference(source_face, target_face, model_path, model_idx, swapped_face)