import numpy as np
import torch
from .humor_render_tools.tools import viz_smpl_seq
from smplx.utils import Struct
from .video import Video
import os
from multiprocessing import Pool
from tqdm import tqdm
from multiprocessing import Process

THIS_FOLDER = os.path.dirname(os.path.abspath(__file__))
FACE_PATH = os.path.join(THIS_FOLDER, "humor_render_tools/smplh.faces")
FACES = torch.from_numpy(np.int32(np.load(FACE_PATH)))


class HumorRenderer:
    def __init__(self, fps=30.0, **kwargs):
        self.kwargs = kwargs
        self.fps = fps

    def __call__(self, vertices, output, render_pair=False,text=None,colors=[],**kwargs):
        params = self.kwargs | kwargs
        fps = self.fps
        if "fps" in params:
            fps = params.pop("fps")

        if render_pair:
            render_overlaid(vertices, output,fps,colors,**params)
        else:
            render(vertices, output, fps,colors,**params)

        fname = f'{output}.mp4'
        return fname


def render_overlaid(vertices, out_path, fps,colors=[], progress_bar=tqdm,**kwargs):
    assert isinstance(vertices, list) and len(vertices) == 2
    # Put the vertices at the floor level
    # F X N x 3 ===> [F X N x 3, F X N x 3] 
    ground = vertices[0][..., 2].min()
    vertices[0][..., 2] -= ground 
    vertices[1][..., 2] -= ground
    verts0 = vertices[0]
    verts1 = vertices[1]
    
    import pyrender

    # remove title if it exists
    kwargs.pop("title", None)

    # vertices: SMPL-H vertices
    # verts = np.load("interval_2_verts.npy")
    out_folder = os.path.splitext(out_path)[0]

    verts0 = torch.from_numpy(verts0)
    body_pred0 = Struct(v=verts0, f=FACES)
    verts1 = torch.from_numpy(verts1)
    body_pred1 = Struct(v=verts1, f=FACES)


    # out_folder, body_pred, start, end, fps, kwargs = args
    viz_smpl_seq(
        pyrender, out_folder, [body_pred0, body_pred1], fps=fps,progress_bar=progress_bar,vertex_color_list=colors, **kwargs
    )
    video = Video(out_folder, fps=fps)
    video.save(out_path)

def render(vertices, out_path, fps, colors=[], progress_bar=tqdm, **kwargs):
    # Put the vertices at the floor level
    ground = vertices[..., 2].min()
    vertices[..., 2] -= ground

    import pyrender

    # remove title if it exists
    kwargs.pop("title", None)

    # vertices: SMPL-H vertices
    # verts = np.load("interval_2_verts.npy")
    out_folder = os.path.splitext(out_path)[0]

    verts = torch.from_numpy(vertices)
    body_pred = Struct(v=verts, f=FACES)

    # out_folder, body_pred, start, end, fps, kwargs = args
    viz_smpl_seq(
        pyrender, out_folder, body_pred, fps=fps,progress_bar=progress_bar,vertex_color=colors, **kwargs
    )

    video = Video(out_folder, fps=fps)
    video.save(out_path)
    import shutil
    shutil.rmtree(out_folder)

def render_offset(args):
    import pyrender

    out_folder, body_pred, start, end, fps, kwargs = args
    viz_smpl_seq(
        pyrender, out_folder, body_pred, start=start, end=end, fps=fps, **kwargs
    )
    return 0


def render_multiprocess(vertices, out_path, fps, **kwargs):
    # WIP: does not work yet
    # import ipdb

    # ipdb.set_trace()
    # remove title if it exists
    kwargs.pop("title", None)

    # vertices: SMPL-H vertices
    # verts = np.load("interval_2_verts.npy")
    out_folder = os.path.splitext(out_path)[0]

    verts = torch.from_numpy(vertices)
    body_pred = Struct(v=verts, f=FACES)

    # faster rendering
    # by rendering part of the sequence in parallel
    # still work in progress, use one process for now
    n_processes = 1

    verts_lst = np.array_split(verts, n_processes)
    len_split = [len(x) for x in verts_lst]
    starts = [0] + np.cumsum([x for x in len_split[:-1]]).tolist()
    ends = np.cumsum([x for x in len_split]).tolist()
    out_folders = [out_folder for _ in range(n_processes)]
    fps_s = [fps for _ in range(n_processes)]
    kwargs_s = [kwargs for _ in range(n_processes)]
    body_pred_s = [body_pred for _ in range(n_processes)]

    arguments = [out_folders, body_pred_s, starts, ends, fps_s, kwargs_s]
    # sanity
    # lst = [verts[start:end] for start, end in zip(starts, ends)]
    # assert (torch.cat(lst) == verts).all()

    processes = []
    for _, args in zip(range(n_processes), zip(*arguments)):
        process = Process(target=render_offset, args=(args,))
        process.start()
        processes.append(process)

    for process in processes:
        process.join()

    if False:
        # start 4 worker processes
        with Pool(processes=n_processes) as pool:
            # print "[0, 1, 4,..., 81]"
            # print same numbers in arbitrary order
            print(f"0/{n_processes} rendered")
            i = 0
            for _ in pool.imap_unordered(render_offset, zip(*arguments)):
                i += 1
                print(f"i/{n_processes} rendered")

    video = Video(out_folder, fps=fps)
    video.save(out_path)