from __future__ import annotations

import os
import pathlib
import sys

import huggingface_hub
import numpy as np
import torch
import torch.nn as nn

if os.getenv("SYSTEM") == "spaces":
    os.system("sed -i '14,21d' StyleSwin/op/fused_act.py")
    os.system("sed -i '12,19d' StyleSwin/op/upfirdn2d.py")

current_dir = pathlib.Path(__file__).parent
submodule_dir = current_dir / "StyleSwin"
sys.path.insert(0, submodule_dir.as_posix())

from models.generator import Generator


class Model:
    MODEL_NAMES = [
        "CelebAHQ_256",
        "FFHQ_256",
        "LSUNChurch_256",
        "CelebAHQ_1024",
        "FFHQ_1024",
    ]

    def __init__(self):
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        self._download_all_models()
        self.model_name = self.MODEL_NAMES[3]
        self.model = self._load_model(self.model_name)

        self.std = torch.FloatTensor([0.229, 0.224, 0.225])[None, :, None, None].to(self.device)
        self.mean = torch.FloatTensor([0.485, 0.456, 0.406])[None, :, None, None].to(self.device)

    def _load_model(self, model_name: str) -> nn.Module:
        size = int(model_name.split("_")[1])
        channel_multiplier = 1 if size == 1024 else 2
        model = Generator(size, style_dim=512, n_mlp=8, channel_multiplier=channel_multiplier)
        ckpt_path = huggingface_hub.hf_hub_download("public-data/StyleSwin", f"models/{model_name}.pt")
        ckpt = torch.load(ckpt_path)
        model.load_state_dict(ckpt["g_ema"])
        model.to(self.device)
        model.eval()
        return model

    def set_model(self, model_name: str) -> None:
        if model_name == self.model_name:
            return
        self.model_name = model_name
        self.model = self._load_model(model_name)

    def _download_all_models(self):
        for name in self.MODEL_NAMES:
            self._load_model(name)

    def generate_z(self, seed: int) -> torch.Tensor:
        seed = int(np.clip(seed, 0, np.iinfo(np.uint32).max))
        z = np.random.RandomState(seed).randn(1, 512)
        return torch.from_numpy(z).float().to(self.device)

    def postprocess(self, tensors: torch.Tensor) -> np.ndarray:
        assert tensors.dim() == 4
        tensors = tensors * self.std + self.mean
        tensors = (tensors * 255).clamp(0, 255).to(torch.uint8)
        return tensors.permute(0, 2, 3, 1).cpu().numpy()

    @torch.inference_mode()
    def generate_image(self, seed: int) -> np.ndarray:
        z = self.generate_z(seed)
        out, _ = self.model(z)
        out = self.postprocess(out)
        return out[0]

    def set_model_and_generate_image(self, model_name: str, seed: int) -> np.ndarray:
        self.set_model(model_name)
        return self.generate_image(seed)