import torch
from argparse import Namespace
from torchvision.transforms import transforms

from configs import paths_config
from models.e4e.psp import pSp
from scripts.latent_creators.base_latent_creator import BaseLatentCreator
from utils.log_utils import log_image_from_w


class E4ELatentCreator(BaseLatentCreator):

    def __init__(self, use_wandb=False):
        self.e4e_inversion_pre_process = transforms.Compose([
            transforms.Resize((256, 256)),
            transforms.ToTensor(),
            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])

        super().__init__('e4e', self.e4e_inversion_pre_process, use_wandb=use_wandb)

        e4e_model_path = paths_config.e4e
        ckpt = torch.load(e4e_model_path, map_location='cpu')
        opts = ckpt['opts']
        opts['batch_size'] = 1
        opts['checkpoint_path'] = e4e_model_path
        opts = Namespace(**opts)
        self.e4e_inversion_net = pSp(opts)
        self.e4e_inversion_net.eval()
        self.e4e_inversion_net = self.e4e_inversion_net.cuda()

    def run_projection(self, fname, image):
        _, e4e_image_latent = self.e4e_inversion_net(image, randomize_noise=False, return_latents=True,
                                                     resize=False,
                                                     input_code=False)

        if self.use_wandb:
            log_image_from_w(e4e_image_latent, self.old_G, 'First e4e inversion')

        return e4e_image_latent


if __name__ == '__main__':
    e4e_latent_creator = E4ELatentCreator()
    e4e_latent_creator.create_latents()