import matplotlib.pyplot as plt
import time
import numpy as np

from IPython import display as ipythondisplay
from string import Formatter


def display_model(model):
    import tensorflow as tf
    tf.keras.utils.plot_model(model, to_file="tmp.png", show_shapes=True)
    return ipythondisplay.Image("tmp.png")


def plot_sample(x, y, vae, backend='tf'):
    """Plot original and reconstructed images side by side.
    
    Args:
        x: Input images array of shape [B, H, W, C] (TF) or [B, C, H, W] (PT)
        y: Labels array of shape [B] where 1 indicates a face
        vae: VAE model (TensorFlow or PyTorch)
        framework: 'tf' or 'pt' indicating which framework to use
    """
    plt.figure(figsize=(2, 1))

    if backend == 'tf':
        idx = np.where(y == 1)[0][0]
        _, _, _, recon = vae(x)
        recon = np.clip(recon, 0, 1)

    elif backend == 'pt':
        import torch 
        y = y.detach().cpu().numpy()
        face_indices = np.where(y == 1)[0]
        idx = face_indices[0] if len(face_indices) > 0 else 0

        with torch.inference_mode():
            _, _, _, recon = vae(x)
        recon = torch.clamp(recon, 0, 1)
        recon = recon.permute(0, 2, 3, 1).detach().cpu().numpy()
        x = x.permute(0, 2, 3, 1).detach().cpu().numpy()

    else:
        raise ValueError("framework must be 'tf' or 'pt'")

    plt.subplot(1, 2, 1)
    plt.imshow(x[idx])
    plt.grid(False)

    plt.subplot(1, 2, 2) 
    plt.imshow(recon[idx])
    plt.grid(False)

    if backend == 'pt':
        plt.show()


class LossHistory:
    def __init__(self, smoothing_factor=0.0):
        self.alpha = smoothing_factor
        self.loss = []

    def append(self, value):
        self.loss.append(
            self.alpha * self.loss[-1] + (1 - self.alpha) * value
            if len(self.loss) > 0
            else value
        )

    def get(self):
        return self.loss


class PeriodicPlotter:
    def __init__(self, sec, xlabel="", ylabel="", scale=None):
        self.xlabel = xlabel
        self.ylabel = ylabel
        self.sec = sec
        self.scale = scale

        self.tic = time.time()

    def plot(self, data):
        if time.time() - self.tic > self.sec:
            plt.cla()

            if self.scale is None:
                plt.plot(data)
            elif self.scale == "semilogx":
                plt.semilogx(data)
            elif self.scale == "semilogy":
                plt.semilogy(data)
            elif self.scale == "loglog":
                plt.loglog(data)
            else:
                raise ValueError("unrecognized parameter scale {}".format(self.scale))

            plt.xlabel(self.xlabel)
            plt.ylabel(self.ylabel)
            ipythondisplay.clear_output(wait=True)
            ipythondisplay.display(plt.gcf())

            self.tic = time.time()


def create_grid_of_images(xs, size=(5, 5)):
    """Combine a list of images into a single image grid by stacking them into an array of shape `size`"""

    grid = []
    counter = 0
    for i in range(size[0]):
        row = []
        for j in range(size[1]):
            row.append(xs[counter])
            counter += 1
        row = np.hstack(row)
        grid.append(row)
    grid = np.vstack(grid)
    return grid