import os
import cv2
import numpy as np
import tempfile
from tqdm import tqdm
from subprocess import call

import trimesh
from scipy.io import wavfile
import librosa

class Struct(object):
    def __init__(self, **kwargs):
        for key, val in kwargs.items():
            setattr(self, key, val)

def add_image_text(img, text, color=(0,0,255), w=800, h=800):
    font = cv2.FONT_HERSHEY_SIMPLEX
    img = np.require(img, dtype='f4', requirements=['O', 'W'])
    img.flags.writeable = True
    img1 = img.copy()
    img1 = cv2.putText(img1, '%s' % (text), (50, 100), font, 2, color, 2, 1)
    img1 = cv2.rectangle(img1, (0, 0), (w, h), color, thickness=3)
    return img1

class RenderTool:
    def __init__(self, out_path):
        path = os.path.join(os.getcwd(), 'visualise/smplx/SMPLX_NEUTRAL.npz')
        model_data = np.load(path, allow_pickle=True)
        data_struct = Struct(**model_data)
        self.f = data_struct.f  # faces
        self.out_path = out_path
        if not os.path.exists(self.out_path):
            os.makedirs(self.out_path)

    def _render_sequences(self, cur_wav_file, v_list, j=-1, run_in_parallel=False):
        symbol = '/'
        print("Render {} {} sequence.".format(cur_wav_file.split(symbol)[-2], cur_wav_file.split(symbol)[-1]))
        
        directory = os.path.join(self.out_path, cur_wav_file.split(symbol)[2].split(symbol)[0])
        if not os.path.exists(directory):
            os.makedirs(directory)
        
        if j == -1:
            video_fname = os.path.join(directory, '%s.mp4' % cur_wav_file.split(symbol)[-1].split('.')[-2].split(symbol)[-1])
        elif j == -2:
            video_fname = os.path.join(directory, cur_wav_file.split(symbol)[-3]+'--%s.mp4' % cur_wav_file.split(symbol)[-1].split('.')[-2].split(symbol)[-1])
        else:
            video_fname = os.path.join(directory, str(j)+'_%s.mp4' % cur_wav_file.split(symbol)[-1].split('.')[-2].split(symbol)[-1])

        self._render_sequences_helper(video_fname, cur_wav_file, v_list)

    def _render_sequences_helper(self, video_fname, cur_wav_file, v_list):
        num_frames = v_list[0].shape[0]

        # Prepare output frames folder
        frames_dir = os.path.join(os.path.dirname(video_fname), "frames_tmp")
        os.makedirs(frames_dir, exist_ok=True)

        center = np.mean(v_list[0][0], axis=0)

        for i_frame in tqdm(range(num_frames)):
            cur_img_list = []
            for i in range(len(v_list)):
                mesh = trimesh.Trimesh(vertices=v_list[i][i_frame], faces=self.f)
                scene = trimesh.Scene(mesh)
                png = scene.save_image(resolution=[800, 800], visible=True)
                
                img_array = np.asarray(bytearray(png), dtype=np.uint8)
                img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)

                cur_img_list.append(img)

            if len(cur_img_list) == 1:
                final_img = cur_img_list[0]
            else:
                final_img = np.hstack(cur_img_list)

            frame_path = os.path.join(frames_dir, f"frame_{i_frame:04d}.png")
            cv2.imwrite(frame_path, final_img)

        # Create video from frames
        tmp_audio_file = tempfile.NamedTemporaryFile('w', suffix='.wav', dir=os.path.dirname(video_fname))
        tmp_audio_file.close()

        audio, sr = librosa.load(cur_wav_file, sr=16000)
        wavfile.write(tmp_audio_file.name, sr, audio)

        tmp_video_file = tempfile.NamedTemporaryFile('w', suffix='.mp4', dir=os.path.dirname(video_fname))
        tmp_video_file.close()

        cmd_frames_to_video = (
            f"ffmpeg -framerate 30 -i {frames_dir}/frame_%04d.png -c:v libx264 -pix_fmt yuv420p {tmp_video_file.name}"
        ).split()
        call(cmd_frames_to_video)

        cmd_merge_audio = (
            f"ffmpeg -i {tmp_video_file.name} -i {tmp_audio_file.name} -vcodec copy -acodec aac {video_fname}"
        ).split()
        call(cmd_merge_audio)

        # Cleanup
        os.remove(tmp_audio_file.name)
        os.remove(tmp_video_file.name)
        for f in os.listdir(frames_dir):
            os.remove(os.path.join(frames_dir, f))
        os.rmdir(frames_dir)