from __future__ import annotations

import sys
from pathlib import Path

import torch
import torch.nn.functional as F

from diffusers import (
    StableDiffusionPipeline,
    UNet2DConditionModel,
    DPMSolverMultistepScheduler,
)

from setup_models import setup_arc2face

BASE_MODEL = "stable-diffusion-v1-5/stable-diffusion-v1-5"


def load_arc2face(device="cuda"):

    # Ensure code and pretrained weights exist.
    repo_path, model_dir = setup_arc2face()

    repo_path = Path(repo_path).resolve()
    model_dir = Path(model_dir).resolve()

    if str(repo_path) not in sys.path:
        sys.path.insert(0, str(repo_path))

    from arc2face import (
        CLIPTextModelWrapper,
        project_face_embs,
    )

    device = torch.device(device)
    dtype = (
        torch.float16
        if device.type == "cuda"
        else torch.float32
    )

    # Load Arc2Face's custom CLIP encoder.
    encoder = CLIPTextModelWrapper.from_pretrained(
        str(model_dir),
        subfolder="encoder",
        torch_dtype=dtype,
    )

    # Load the fine-tuned Arc2Face UNet.
    unet = UNet2DConditionModel.from_pretrained(
        str(model_dir),
        subfolder="arc2face",
        torch_dtype=dtype,
    )

    # Construct Stable Diffusion pipeline.
    pipe = StableDiffusionPipeline.from_pretrained(
        BASE_MODEL,
        text_encoder=encoder,
        unet=unet,
        torch_dtype=dtype,
        safety_checker=None,
    )

    pipe.scheduler = DPMSolverMultistepScheduler.from_config(
        pipe.scheduler.config
    )

    pipe = pipe.to(device)
    pipe.set_progress_bar_config(disable=True)

    return pipe, dtype


@torch.inference_mode()
def generate_image(
    pipe,
    identity_embedding,
    seed: int,
    dtype,
):

    identity_embedding = F.normalize(
        identity_embedding.to(
            device=pipe.device,
            dtype=dtype,
        ),
        p=2,
        dim=-1,
    )

    from arc2face import project_face_embs

    # Project the 512-D ArcFace embedding into
    # Arc2Face's identity-conditioned text embedding.
    prompt_embeds = project_face_embs(
        pipe,
        identity_embedding,
    )

    generator = torch.Generator(
        device=pipe.device
    ).manual_seed(int(seed))

    result = pipe(
        prompt_embeds=prompt_embeds,
        num_inference_steps=25,
        guidance_scale=3.0,
        generator=generator,
    )

    return result.images[0].convert("RGB")