import sys
sys.path.append('../')
import constants as cnst

import torch
import torchvision
import dataset_loaders
from my_utils.visualize_flame_overlay import OverLayViz
import numpy as np
try:
    from my_utils.flm_dynamic_fit_overlay import camera_dynamic, camera_ringnetpp
except ModuleNotFoundError:
    camera_ringnetpp = None

import imageio
import os
import tqdm


def get_unique_shuffle_indices(n):
    """
    :param n: an integer n
    :return: a shuffled index numpy array that displaces all the indices from their original position
    """
    indices = np.arange(0, n)
    np.random.shuffle(indices)
    while (indices == np.arange(0, n)).any():
        np.random.shuffle(indices)

    return indices


def get_images_from_flame_params(flame_params, pose, model, step, alpha, input_indices):
    batch_size = 16
    model.eval()
    img_tensor = []
    flame_params = torch.from_numpy(flame_params)
    if pose is not None:
        pose = torch.from_numpy(pose)
    input_indices = torch.from_numpy(input_indices)
    pose_this_batch = None
    with torch.no_grad():
        for b_id in range(0, flame_params.shape[0], batch_size):
            # import ipdb; ipdb.set_trace()
            flame_param_this_batch = flame_params[b_id:b_id + batch_size].cuda()
            input_indices_this_batch = input_indices[b_id:b_id + batch_size].cuda()
            if pose is not None:
                pose_this_batch = pose[b_id:b_id + batch_size].cuda()
            # import ipdb; ipdb.set_trace()
            img_batch = model(flame_param_this_batch, pose_this_batch, step=step, alpha=alpha,
                              input_indices=input_indices_this_batch)[-1]
            img_tensor.append(torch.clamp(img_batch.cpu(), -1, 1))

    img_tensor = torch.cat(img_tensor, dim=0)
    return img_tensor


def requires_grad(model, flag=True):
    for p in model.parameters():
        p.requires_grad = flag


def accumulate(model1, model2, decay=0.999):
    '''
    Implementation of running Averages
        input:
            model1: nn.Module
            model2: nn.Module

        model1_params = model1_params*decay + (1-decay)*model2_params
    '''
    par1 = dict(model1.named_parameters())
    par2 = dict(model2.named_parameters())

    for k in par1.keys():
        par1[k].data.mul_(decay).add_(1 - decay, par2[k].data)


def adjust_lr(optimizer, lr, style_gn2):
    if not style_gn2:
        for group in optimizer.param_groups:
            mult = group.get('mult', 1)
            group['lr'] = lr * mult


class VisualizationSaver():
    def __init__(self, gen_i, gen_j, sampling_flame_labels, dataset, input_indices, overlay_mesh=False):
        self.gen_i = gen_i
        self.gen_j = gen_j
        self.sampling_flame_labels = sampling_flame_labels
        self.overlay_mesh = overlay_mesh
        self.overlay_visualizer = OverLayViz()
        self.dataset = dataset
        self.input_indices = input_indices
        self.cam_t = np.array([0., 0., 2.5])

    def set_flame_params(self, pose, sampling_flame_labels, input_indices):
        self.pose = pose
        self.sampling_flame_labels = sampling_flame_labels
        self.input_indices = input_indices

    def save_samples(self, i, model, step, alpha, resolution, fid, run_id):
        images = []
        # camera_params = camera_dynamic((resolution, resolution), self.cam_t)
        flength = 5000
        cam_t = np.array([0., 0., 0])
        camera_params = camera_ringnetpp((512, 512), trans=cam_t, focal=flength)

        with torch.no_grad():
            for img_idx in range(self.gen_i):
                flame_param_this_batch = self.sampling_flame_labels[img_idx * self.gen_j:(img_idx + 1) * self.gen_j]
                if self.pose is not None:
                    pose_this_batch = self.pose[img_idx * self.gen_j:(img_idx + 1) * self.gen_j]
                else:
                    pose_this_batch = None
                idx_this_batch = self.input_indices[img_idx * self.gen_j:(img_idx + 1) * self.gen_j]
                img_tensor = model(flame_param_this_batch.clone(), pose_this_batch, step=step, alpha=alpha,
                                   input_indices=idx_this_batch)[-1]

                img_tensor = self.overlay_visualizer.range_normalize_images(
                    dataset_loaders.fast_image_reshape(img_tensor, height_out=256, width_out=256,
                                                       non_diff_allowed=True))

                images.append(img_tensor.data.cpu())

        torchvision.utils.save_image(
            torch.cat(images, 0),
            f'{cnst.output_root}sample/{str(run_id)}/{str(i + 1).zfill(6)}_res{resolution}x{resolution}_fid_{fid:.2f}.png',
            nrow=self.gen_i,
            normalize=True,
            range=(0, 1))


def save_set_of_images(path, prefix, images, show_prog_bar=False, name_list=None):
    os.makedirs(path, exist_ok=True)
    if len(images.shape) == 4:  # Color channels present
        if images.shape[-1] == 1:
            images = images[:, :, :, 0]
        elif images.shape[-1] != 1 and images.shape[1] == 3:
            images = images.transpose((0, 2, 3, 1))

    if images.dtype =='float32':
        if np.min(images) < 0 or np.max(images) > 1:
            print('Warning! Not in the range of 0 - 1')

        images = (np.clip(images, 0, 1)*255).astype('uint8')


    if show_prog_bar:
        progbar = tqdm.tqdm(range(images.shape[0]))
        progbar.set_description('Saving images')
        for i in progbar:
            if name_list is None:
                cur_name = i
            else:
                cur_name = name_list[i]
            imageio.imwrite(os.path.join(path, prefix + f'{cur_name}.png'), images[i])
    else:
        for i, img in enumerate(images):
            if name_list is None:
                cur_name = i
            else:
                cur_name = name_list[i]
            imageio.imwrite(os.path.join(path, prefix + f'{cur_name}.png'), img)
