import torch
import numpy as np

from .vision3d import rigidTransformRT
from .utils_smplx import getSMPLXOut


def renderFromSMPLXParams(smplx_kwargs, smplx_params, K, RT, save_folder, img_hw, save_ext='.png', BS_render=64):
    """Render SMPLX params to images. Requires pytorch3d rendering (optional)."""
    try:
        from drawer import render_3d_overlay_pytorch3d_split
    except ImportError:
        raise ImportError("Rendering requires 'drawer' module. This is optional for fitting.")
    assert K.shape == (3, 3)
    with torch.no_grad():
        smplx_out, smplx_model = getSMPLXOut(smplx_kwargs, 'cuda', smplx_params, BS=512, run_with_verts=True)
    if RT is not None:
        assert RT.shape == (3, 4)
        smplx_out['vertices'] = rigidTransformRT(smplx_out['vertices'], RT[None])
    render_3d_overlay_pytorch3d_split(smplx_out['vertices'], smplx_model.faces, K, None, BS_render, False, save_folder, None, save_ext, np.arange(len(smplx_out['vertices'])), img_hw=img_hw)


def replaceKeypointsWeightsWithHandModelSuccess(kpts_with_weights, rhand_idxs, lhand_idxs, success_r, success_l):
    # Model-based hand pose estimation seems more robust than 2D detectors for hands, so zero-out hand confidences if model-based estimation failed
    assert len(success_l) == len(success_r) == kpts_with_weights.shape[0]
    assert kpts_with_weights.ndim == 3 and kpts_with_weights.shape[2] == 3

    rhands_failed = kpts_with_weights[~success_r]
    lhands_failed = kpts_with_weights[~success_l]
    rhands_failed[...,rhand_idxs,2] = 0.0
    lhands_failed[...,lhand_idxs,2] = 0.0
    kpts_with_weights[~success_r] = rhands_failed
    kpts_with_weights[~success_l] = lhands_failed
    return kpts_with_weights