import matplotlib.pyplot as plt
import torch
import numpy as np
import smplx
from human_body_prior.body_model.body_model_vposer import BodyModelWithPoser
from human_body_prior.tools.model_loader import load_vposer


class VPoserModel():
    def __init__(
        self,
        model_type='smpl',
        vposer_model_path="./vposer_v1_0",
        ext='npz',
        gender='neutral',
        create_body_pose=True,
        plot_joints=True,
        num_betas=10,
        sample_shape=False,
        sample_expression=False,
        num_expression_coeffs=10,
        use_face_contour=False,
        use_vposer=True
    ):
        self.vposer_model_path = vposer_model_path
        self.model_type = model_type
        self.ext = ext
        self.gender = gender
        self.plot_joints = plot_joints
        self.num_betas = num_betas
        self.sample_shape = sample_shape
        self.sample_expression = sample_expression
        self.num_expression_coeffs = num_expression_coeffs
        self.create_body_pose = create_body_pose
        self.use_vposer = use_vposer

        self.create_model()

    def create_model(self):

        self.model = BodyModelWithPoser(
            bm_path="./models/smplx/SMPLX_MALE.npz",
            batch_size=1,
            poser_type="vposer",
            smpl_exp_dir=self.vposer_model_path
        )
        return self.model

    def get_vposer_latens(self):
        return self.model.poZ_body

    def get_pose(self):
        return self.model.pose_body


class SMPLyModel():
    def __init__(
        self,
        model_folder,
        model_type='smplx',
        ext='npz',
        gender='male',
        create_body_pose=True,
        plot_joints=True,
        num_betas=10,
        sample_shape=False,
        sample_expression=True,
        num_expression_coeffs=10,
        plotting_module='pyrender',
        use_face_contour=False,
        use_vposer_init=False,
        device=torch.device('cpu')
    ):
        self.model_folder = model_folder
        self.model_type = model_type
        self.ext = ext
        self.gender = gender
        self.plot_joints = plot_joints
        self.num_betas = num_betas
        self.sample_shape = sample_shape
        self.sample_expression = sample_expression
        self.num_expression_coeffs = num_expression_coeffs
        self.plotting_module = plotting_module
        self.use_face_contour = use_face_contour
        self.create_body_pose = create_body_pose
        self.use_vposer_init = use_vposer_init
        self.device = device

    def create_model(self):
        initial_pose = None
        if self.use_vposer_init:
            # sample a valid human shape via vposer
            vp, ps = load_vposer("./vposer_v1_0")
            vp = vp.to(device=self.device)
            self.vp = vp
            self.vposer_sample = torch.from_numpy(
                np.random.randn(1, 32).astype(np.float32)).to(device=self.device)

            # sample SMPL body pose from vposer
            initial_pose = self.vp.decode(
                self.vposer_sample, output_type='aa').view(-1, 63)

        self.model = smplx.create(
            self.model_folder,
            model_type=self.model_type,
            gender=self.gender,
            body_pose=initial_pose,
            use_face_contour=self.use_face_contour,
            create_body_pose=self.create_body_pose,
            num_betas=self.num_betas,
            return_verts=True,
            num_expression_coeffs=self.num_expression_coeffs,
            ext=self.ext)
        return self.model