# -*- coding: utf-8 -*-
from os.path import dirname, abspath, join

import torch
from huggingface_hub import hf_hub_url, cached_download
from omegaconf import OmegaConf

from .model import VQGanGumbelVAE


def get_vae(pretrained=True, dwt=False, cache_dir="models"):
    config = OmegaConf.load(
        join(dirname(abspath(__file__)), "vqgan.gumbelf8-sber.config.yml")
    )
    vae = VQGanGumbelVAE(config, dwt=dwt)
    if pretrained:
        repo_id = "shonenkov/rudalle-utils"
        if dwt:
            filename = "vqgan.gumbelf8-sber-dwt.model.ckpt"
        else:
            filename = "vqgan.gumbelf8-sber.model.ckpt"
        cache_dir = join(cache_dir, "vae")
        config_file_url = hf_hub_url(repo_id=repo_id, filename=filename)
        cached_download(config_file_url, cache_dir=cache_dir, force_filename=filename)
        checkpoint = torch.load(join(cache_dir, filename), map_location="cpu")
        if dwt:
            vae.load_state_dict(checkpoint["state_dict"])
        else:
            vae.model.load_state_dict(checkpoint["state_dict"], strict=False)
    return vae
