import matplotlib.pyplot as plt
import math
import utils
from . import parse

save_ind = 0

def visualize(image, title, colorbar=False, show_plot=True, **kwargs):
    plt.title(title)
    plt.imshow(image, **kwargs)
    if colorbar:
        plt.colorbar()
    if show_plot:
        plt.show()

def visualize_arrays(image_title_pairs, colorbar_index=-1, show_plot=True, figsize=None, **kwargs):
    if figsize is not None:
        plt.figure(figsize=figsize)
    num_subplots = len(image_title_pairs)
    for idx, image_title_pair in enumerate(image_title_pairs):
        plt.subplot(1, num_subplots, idx+1)
        if isinstance(image_title_pair, (list, tuple)):
            image, title = image_title_pair
        else:
            image, title = image_title_pair, None
        
        if title is not None:
            plt.title(title)
        
        plt.imshow(image, **kwargs)
        if idx == colorbar_index:
            plt.colorbar()
            
    if show_plot:
        plt.show()

def visualize_masked_latents(latents_all, masked_latents, timestep_T=False, timestep_0=True):
    if timestep_T:
        # from T to 0
        latent_idx = 0

        plt.subplot(1, 2, 1)
        plt.title("latents_all (t=T)")
        plt.imshow((latents_all[latent_idx, 0, :3].cpu().permute(1,2,0).numpy().astype(float) / 1.5).clip(0., 1.), cmap="gray")

        plt.subplot(1, 2, 2)
        plt.title("mask latents (t=T)")
        plt.imshow((masked_latents[latent_idx, 0, :3].cpu().permute(1,2,0).numpy().astype(float) / 1.5).clip(0., 1.), cmap="gray")

        plt.show()

    if timestep_0:
        latent_idx = -1
        plt.subplot(1, 2, 1)
        plt.title("latents_all (t=0)")
        plt.imshow((latents_all[latent_idx, 0, :3].cpu().permute(1,2,0).numpy().astype(float) / 1.5).clip(0., 1.), cmap="gray")

        plt.subplot(1, 2, 2)
        plt.title("mask latents (t=0)")
        plt.imshow((masked_latents[latent_idx, 0, :3].cpu().permute(1,2,0).numpy().astype(float) / 1.5).clip(0., 1.), cmap="gray")

        plt.show()

# This function has not been adapted to new `saved_attn`.
def visualize_attn(token_map, cross_attention_probs_tensors, stage_id, block_id, visualize_step_start=10, input_ca_has_condition_only=False):
    """
    Visualize cross attention: `stage_id`th downsampling block, mean over all timesteps starting from step start, `block_id`th Transformer block, second item (conditioned), mean over heads, show each token
    cross_attention_probs_tensors:
    One of `cross_attention_probs_down_tensors`, `cross_attention_probs_mid_tensors`, and `cross_attention_probs_up_tensors`
    stage_id: index of downsampling/mid/upsaming block
    block_id: index of the transformer block
    """
    
    plt.figure(figsize=(20, 8))

    for token_id in range(len(token_map)):
        token = token_map[token_id]
        plt.subplot(1, len(token_map), token_id + 1)
        plt.title(token)
        attn = cross_attention_probs_tensors[stage_id][visualize_step_start:].mean(dim=0)[block_id]
        
        if not input_ca_has_condition_only:
            assert attn.shape[0] == 2, f"Expect to have 2 items (uncond and cond), but found {attn.shape[0]} items"
            attn = attn[1]
        else:
            assert attn.shape[0] == 1, f"Expect to have 1 item (cond only), but found {attn.shape[0]} items"
            attn = attn[0]
        
        attn = attn.mean(dim=0)[:, token_id]
        H = W = int(math.sqrt(attn.shape[0]))
        attn = attn.reshape((H, W))
        plt.imshow(attn.cpu().numpy())
        
    plt.show()

# This function has not been adapted to new `saved_attn`.
def visualize_across_timesteps(token_id, cross_attention_probs_tensors, stage_id, block_id, visualize_step_start=10, input_ca_has_condition_only=False):
    """
    Visualize cross attention for one token, across timesteps: `stage_id`th downsampling block, mean over all timesteps starting from step start, `block_id`th Transformer block, second item (conditioned), mean over heads, show each token
    cross_attention_probs_tensors:
    One of `cross_attention_probs_down_tensors`, `cross_attention_probs_mid_tensors`, and `cross_attention_probs_up_tensors`
    stage_id: index of downsampling/mid/upsaming block
    block_id: index of the transformer block
    
    `visualize_step_start` is not used. We visualize all timesteps.
    """
    plt.figure(figsize=(50, 8))
    
    attn_stage = cross_attention_probs_tensors[stage_id]
    num_inference_steps = attn_stage.shape[0]

    for t in range(num_inference_steps):
        plt.subplot(1, num_inference_steps, t + 1)
        plt.title(f"t: {t}")

        attn = attn_stage[t][block_id]
        
        if not input_ca_has_condition_only:
            assert attn.shape[0] == 2, f"Expect to have 2 items (uncond and cond), but found {attn.shape[0]} items"
            attn = attn[1]
        else:
            assert attn.shape[0] == 1, f"Expect to have 1 item (cond only), but found {attn.shape[0]} items"
            attn = attn[0]
        
        attn = attn.mean(dim=0)[:, token_id]
        H = W = int(math.sqrt(attn.shape[0]))
        attn = attn.reshape((H, W))
        plt.imshow(attn.cpu().numpy())
        plt.axis("off")
        plt.tight_layout()

    plt.show()

def visualize_bboxes(bboxes, H, W):
    num_boxes = len(bboxes)
    for ind, bbox in enumerate(bboxes):
        plt.subplot(1, num_boxes, ind + 1)
        fg_mask = utils.proportion_to_mask(bbox, H, W)
        plt.title(f"transformed bbox ({ind})")
        plt.imshow(fg_mask.cpu().numpy())
    plt.show()

def display(image, save_prefix="", ind=None):
    global save_ind
    if save_prefix != "":
        save_prefix = save_prefix + "_"
    ind = f"{ind}_" if ind is not None else ""
    path = f"{parse.img_dir}/{save_prefix}{ind}{save_ind}.png"
    
    print(f"Saved to {path}")
    
    image.save(path)
    save_ind = save_ind + 1