# coding: utf-8

"""
Make video template
"""

import os
import cv2
import numpy as np
import pickle
from rich.progress import track
from .utils.cropper import Cropper

from .utils.io import load_driving_info
from .utils.camera import get_rotation_matrix
from .utils.helper import mkdir, basename
from .utils.rprint import rlog as log
from .config.crop_config import CropConfig
from .config.inference_config import InferenceConfig
from .live_portrait_wrapper import LivePortraitWrapper

class TemplateMaker:

    def __init__(self, inference_cfg: InferenceConfig, crop_cfg: CropConfig):
        self.live_portrait_wrapper: LivePortraitWrapper = LivePortraitWrapper(cfg=inference_cfg)
        self.cropper = Cropper(crop_cfg=crop_cfg)

    def make_motion_template(self, video_fp: str, output_path: str, **kwargs):
        """ make video template (.pkl format)
        video_fp: driving video file path
        output_path: where to save the pickle file
        """

        driving_rgb_lst = load_driving_info(video_fp)
        driving_rgb_lst = [cv2.resize(_, (256, 256)) for _ in driving_rgb_lst]
        driving_lmk_lst = self.cropper.get_retargeting_lmk_info(driving_rgb_lst)
        I_d_lst = self.live_portrait_wrapper.prepare_driving_videos(driving_rgb_lst)

        n_frames = I_d_lst.shape[0]

        templates = []


        for i in track(range(n_frames), description='Making templates...', total=n_frames):
            I_d_i = I_d_lst[i]
            x_d_i_info = self.live_portrait_wrapper.get_kp_info(I_d_i)
            R_d_i = get_rotation_matrix(x_d_i_info['pitch'], x_d_i_info['yaw'], x_d_i_info['roll'])
            # collect s_d, R_d, δ_d and t_d for inference
            template_dct = {
                'n_frames': n_frames,
                'frames_index': i,
            }
            template_dct['scale'] = x_d_i_info['scale'].cpu().numpy().astype(np.float32)
            template_dct['R_d'] = R_d_i.cpu().numpy().astype(np.float32)
            template_dct['exp'] = x_d_i_info['exp'].cpu().numpy().astype(np.float32)
            template_dct['t'] = x_d_i_info['t'].cpu().numpy().astype(np.float32)

            templates.append(template_dct)

        mkdir(output_path)
        # Save the dictionary as a pickle file
        pickle_fp = os.path.join(output_path, f'{basename(video_fp)}.pkl')
        with open(pickle_fp, 'wb') as f:
            pickle.dump([templates, driving_lmk_lst], f)
        log(f"Template saved at {pickle_fp}")