import numpy as np
import cv2
import os
import time
import imageio
from tqdm import tqdm
from PIL import Image
import os
import sys
sys.path.insert(1, os.path.join(sys.path[0], '..'))
import torch
import torchvision
from torchvision.utils import make_grid
from torch import Tensor
from torchvision.transforms.functional import to_tensor

# ----------------------------------------------------------------------------------------------
def savenp2sheet(imgs, savepath, nrow=None):
    """ save multiple imgs (in numpy array type) to a img sheet.
        img sheet is one row.

    imgs: 
        np array of size [N, H, W, 3] or List[array] with array size = [H,W,3] 
    """
    if imgs.ndim == 4:
        img_list = [imgs[i] for i in range(imgs.shape[0])]
        imgs = img_list
    
    imgs_new = []
    for i, img in enumerate(imgs):
        if img.ndim == 3 and img.shape[0] == 3:
            img = np.transpose(img,(1,2,0))
        
        assert(img.ndim == 3 and img.shape[-1] == 3), img.shape # h,w,3
        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        imgs_new.append(img)
    n = len(imgs)
    if nrow is not None:
        n_cols = nrow
    else:
        n_cols=int(n**0.5)
    n_rows=int(np.ceil(n/n_cols))
    print(n_cols)
    print(n_rows)

    imgsheet = cv2.vconcat([cv2.hconcat(imgs_new[i*n_cols:(i+1)*n_cols]) for i in range(n_rows)])
    cv2.imwrite(savepath, imgsheet)
    print(f'saved in {savepath}')

# ----------------------------------------------------------------------------------------------
def save_np_to_img(img, path, norm=True):
    if norm:
        img = (img + 1) / 2 * 255
    img = img.astype(np.uint8)
    image = Image.fromarray(img)
    image.save(path, q=95)

# ----------------------------------------------------------------------------------------------
def npz_to_imgsheet_5d(data_path, res_dir, nrow=None,):
    if isinstance(data_path, str):
        imgs = np.load(data_path)['arr_0'] # NTHWC
    elif isinstance(data_path, np.ndarray):
        imgs = data_path
    else:
        raise Exception
    
    if os.path.isdir(res_dir):
        res_path = os.path.join(res_dir, f'samples.jpg')
    else:
        assert(res_dir.endswith('.jpg'))
        res_path = res_dir
    imgs = np.concatenate([imgs[i] for i in range(imgs.shape[0])], axis=0)
    savenp2sheet(imgs, res_path, nrow=nrow)

# ----------------------------------------------------------------------------------------------
def npz_to_imgsheet_4d(data_path, res_path, nrow=None,):
    if isinstance(data_path, str):
        imgs = np.load(data_path)['arr_0'] # NHWC
    elif isinstance(data_path, np.ndarray):
        imgs = data_path
    else:
        raise Exception
    print(imgs.shape)
    savenp2sheet(imgs, res_path, nrow=nrow)


# ----------------------------------------------------------------------------------------------
def tensor_to_imgsheet(tensor, save_path):
    """ 
        save a batch of videos in one image sheet with shape of [batch_size * num_frames].
        data: [b,c,t,h,w]
    """
    assert(tensor.dim() == 5)
    b,c,t,h,w = tensor.shape
    imgs = [tensor[bi,:,ti, :, :] for bi in range(b) for ti in range(t)]
    torchvision.utils.save_image(imgs, save_path, normalize=True, nrow=t)


# ----------------------------------------------------------------------------------------------
def npz_to_frames(data_path, res_dir, norm, num_frames=None, num_samples=None):
    start = time.time()
    arr = np.load(data_path)
    imgs = arr['arr_0'] # [N, T, H, W, 3]
    print('original data shape: ', imgs.shape)

    if num_samples is not None:
        imgs = imgs[:num_samples, :, :, :, :]
        print('after sample selection: ', imgs.shape)
    
    if num_frames is not None:
        imgs = imgs[:, :num_frames, :, :, :]
        print('after frame selection: ', imgs.shape)

    for vid in tqdm(range(imgs.shape[0]), desc='Video'):
        video_dir = os.path.join(res_dir, f'video{vid:04d}')
        os.makedirs(video_dir, exist_ok=True)
        for fid in range(imgs.shape[1]):
            frame = imgs[vid, fid, :, :, :] #HW3
            save_np_to_img(frame, os.path.join(video_dir, f'frame{fid:04d}.jpg'), norm=norm)
    print('Finish')
    print(f'Total time = {time.time()- start}')

# ----------------------------------------------------------------------------------------------
def npz_to_gifs(data_path, res_dir, duration=0.2, start_idx=0, num_videos=None, mode='gif'):
    os.makedirs(res_dir, exist_ok=True)
    if isinstance(data_path, str):
        imgs = np.load(data_path)['arr_0'] # NTHWC
    elif isinstance(data_path, np.ndarray):
        imgs = data_path
    else:
        raise Exception

    for i in range(imgs.shape[0]):
        frames = [imgs[i,j,:,:,:] for j in range(imgs[i].shape[0])] # [(h,w,3)]
        if mode == 'gif':
            imageio.mimwrite(os.path.join(res_dir, f'samples_{start_idx+i}.gif'), frames, format='GIF', duration=duration)
        elif mode == 'mp4':
            frames = [torch.from_numpy(frame) for frame in frames]
            frames = torch.stack(frames, dim=0).to(torch.uint8) # [T, H, W, C]
            torchvision.io.write_video(os.path.join(res_dir, f'samples_{start_idx+i}.mp4'),
                frames, fps=0.5, video_codec='h264', options={'crf': '10'})
        if i+ 1 == num_videos:
            break

# ----------------------------------------------------------------------------------------------
def fill_with_black_squares(video, desired_len: int) -> Tensor:
    if len(video) >= desired_len:
        return video

    return torch.cat([
        video,
        torch.zeros_like(video[0]).unsqueeze(0).repeat(desired_len - len(video), 1, 1, 1),
    ], dim=0)

# ----------------------------------------------------------------------------------------------
def load_num_videos(data_path, num_videos):
    # data_path can be either data_path of np array 
    if isinstance(data_path, str):
        videos = np.load(data_path)['arr_0'] # NTHWC
    elif isinstance(data_path, np.ndarray):
        videos = data_path
    else:
        raise Exception

    if num_videos is not None:
        videos = videos[:num_videos, :, :, :, :]
    return videos

# ----------------------------------------------------------------------------------------------
def npz_to_video_grid(data_path, out_path, num_frames=None, fps=8, num_videos=None, nrow=None, verbose=True):
    if isinstance(data_path, str):
        videos = load_num_videos(data_path, num_videos)
    elif isinstance(data_path, np.ndarray):
        videos = data_path
    else:
        raise Exception
    n,t,h,w,c = videos.shape

    videos_th = []
    for i in range(n):
        video = videos[i, :,:,:,:]
        images = [video[j, :,:,:] for j in range(t)]
        images = [to_tensor(img) for img in images]
        video = torch.stack(images)
        videos_th.append(video)
    
    if num_frames is None:
        num_frames = videos.shape[1]
    if verbose:
        videos = [fill_with_black_squares(v, num_frames) for v in tqdm(videos_th, desc='Adding empty frames')] # NTCHW
    else:
        videos = [fill_with_black_squares(v, num_frames) for v in videos_th] # NTCHW

    frame_grids = torch.stack(videos).permute(1, 0, 2, 3, 4) # [T, N, C, H, W]
    if nrow is None:
        nrow = int(np.ceil(np.sqrt(n)))
    if verbose:
        frame_grids = [make_grid(fs, nrow=nrow) for fs in tqdm(frame_grids, desc='Making grids')]
    else:
        frame_grids = [make_grid(fs, nrow=nrow) for fs in frame_grids]

    if os.path.dirname(out_path) != "":
        os.makedirs(os.path.dirname(out_path), exist_ok=True)
    frame_grids = (torch.stack(frame_grids) * 255).to(torch.uint8).permute(0, 2, 3, 1) # [T, H, W, C]
    torchvision.io.write_video(out_path, frame_grids, fps=fps, video_codec='h264', options={'crf': '10'})

# ----------------------------------------------------------------------------------------------
def npz_to_gif_grid(data_path, out_path, n_cols=None, num_videos=20):
    arr = np.load(data_path)
    imgs = arr['arr_0'] # [N, T, H, W, 3]
    imgs = imgs[:num_videos]
    n, t, h, w, c = imgs.shape
    assert(n == num_videos)
    n_cols = n_cols if n_cols else imgs.shape[0]
    n_rows = np.ceil(imgs.shape[0] / n_cols).astype(np.int8)
    H, W = h * n_rows, w * n_cols
    grid = np.zeros((t, H, W, c), dtype=np.uint8)

    for i in range(n_rows):
        for j in range(n_cols):
            if i*n_cols+j < imgs.shape[0]:
                grid[:, i*h:(i+1)*h, j*w:(j+1)*w, :] = imgs[i*n_cols+j, :, :, :, :]
    
    videos = [grid[i] for i in range(grid.shape[0])] # grid: TH'W'C
    imageio.mimwrite(out_path, videos, format='GIF', duration=0.5,palettesize=256)


# ----------------------------------------------------------------------------------------------
def torch_to_video_grid(videos, out_path, num_frames, fps, num_videos=None, nrow=None, verbose=True):
    """
    videos: -1 ~ 1, torch.Tensor, BCTHW
    """
    n,t,h,w,c = videos.shape
    videos_th = [videos[i, ...] for i in range(n)]
    if verbose:
        videos = [fill_with_black_squares(v, num_frames) for v in tqdm(videos_th, desc='Adding empty frames')] # NTCHW
    else:
        videos = [fill_with_black_squares(v, num_frames) for v in videos_th] # NTCHW

    frame_grids = torch.stack(videos).permute(1, 0, 2, 3, 4) # [T, N, C, H, W]
    if nrow is None:
        nrow = int(np.ceil(np.sqrt(n)))
    if verbose:
        frame_grids = [make_grid(fs, nrow=nrow) for fs in tqdm(frame_grids, desc='Making grids')]
    else:
        frame_grids = [make_grid(fs, nrow=nrow) for fs in frame_grids]

    if os.path.dirname(out_path) != "":
        os.makedirs(os.path.dirname(out_path), exist_ok=True)
    frame_grids = ((torch.stack(frame_grids) + 1) / 2 * 255).to(torch.uint8).permute(0, 2, 3, 1) # [T, H, W, C]
    torchvision.io.write_video(out_path, frame_grids, fps=fps, video_codec='h264', options={'crf': '10'})