import numpy as np
import torch
import PIL
from tqdm import tqdm
import random
from PIL import Image



def generate_composite_images(model, z, decision_boundaries, lambdas, latent_space='W', negative_colors=None):
    """
    The regenerate_images function takes a model, z, and decision_boundary as input.  It then
    constructs an inverse rotation/translation matrix and passes it to the generator.  The generator
    expects this matrix as an inverse to avoid potentially failing numerical operations in the network.
    The function then generates images using G(z_0, label) where z_0 is a linear combination of z and the decision boundary.
    
    :param model: Pass in the model to be used for image generation
    :param z: Generate the starting point of the line
    :param decision_boundary: Generate images along the direction of the decision boundary
    :param min_epsilon: Set the minimum value of lambda
    :param max_epsilon: Set the maximum distance from the original image to generate
    :param count: Determine the number of images that are generated
    :return: A list of images and a list of lambdas
    :doc-author: Trelent
    """
    device = torch.device('cpu')
    G = model.to(device) # type: ignore
    
    # Labels.
    label = torch.zeros([1, G.c_dim], device=device)

    z = torch.from_numpy(z.copy()).to(device)
    repetitions = 16 
    z_0 = z
    
    if negative_colors:
        for decision_boundary, lmbd, neg_boundary in zip(decision_boundaries, lambdas, negative_colors):
            decision_boundary = torch.from_numpy(decision_boundary.copy()).to(device)
            if neg_boundary != 'None':
                neg_boundary = torch.from_numpy(neg_boundary.copy()).to(device)
            
                z_0 = z_0 + int(lmbd) * (decision_boundary - (neg_boundary.T * decision_boundary) * neg_boundary)
            else:
                z_0 = z_0 + int(lmbd) * decision_boundary
    else:            
        for decision_boundary, lmbd in zip(decision_boundaries, lambdas):
            decision_boundary = torch.from_numpy(decision_boundary.copy()).to(device)
            z_0 = z_0 + int(lmbd) * decision_boundary
        
        
    if latent_space == 'Z':
        W_0 = G.mapping(z_0, label, truncation_psi=1).to(torch.float32)
        # W = G.mapping(z, label, truncation_psi=1).to(torch.float32)
    else:
        W_0 = z_0.expand((repetitions, -1)).unsqueeze(0).to(torch.float32)
        # W = z.expand((repetitions, -1)).unsqueeze(0).to(torch.float32)
        
    # if layers:
    #     W_f = torch.empty_like(W).copy_(W).to(torch.float32)
    #     W_f[:, layers, :] = W_0[:, layers, :]
    #     img = G.synthesis(W_f, noise_mode='const')
    # else:
    img = G.synthesis(W_0, noise_mode='const')
                
    img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
    img = PIL.Image.fromarray(img[0].cpu().numpy(), 'RGB')
            
    return img
    
  
def generate_original_image(z, model, latent_space='W'):
    """
    The generate_original_image function takes in a latent vector and the model,
    and returns an image generated from that latent vector.
    
    
    :param z: Generate the image
    :param model: Generate the image
    :return: A pil image
    :doc-author: Trelent
    """
    repetitions = 16
    
    device = torch.device('cpu')
    G = model.to(device) # type: ignore
    # Labels.
    label = torch.zeros([1, G.c_dim], device=device)
    if latent_space == 'Z':
        z = torch.from_numpy(z.copy()).to(device)
        img = G(z, label, truncation_psi=1, noise_mode='const')
    else:
        W = torch.from_numpy(np.repeat(z, repetitions, axis=0).reshape(1, repetitions, z.shape[1]).copy()).to(device)
        img = G.synthesis(W, noise_mode='const')
    
    img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
    return PIL.Image.fromarray(img[0].cpu().numpy(), 'RGB')