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)