import numpy as np
import torch
from plyfile import PlyData, PlyElement
from typing import NamedTuple
import smplx
import tqdm
import cv2 as cv
import os

from scipy.spatial.transform import Rotation as R


class GaussianAttributes(NamedTuple):
    xyz: np.ndarray
    opacities: np.ndarray
    features_dc: np.ndarray
    features_extra: np.ndarray
    scales: np.ndarray
    rot: np.ndarray

def load_gaussians_from_ply(path):
    max_sh_degree = 3
    plydata = PlyData.read(path)

    xyz = np.stack((np.asarray(plydata.elements[0]["x"]),
                    np.asarray(plydata.elements[0]["y"]),
                    np.asarray(plydata.elements[0]["z"])), axis=1)
    opacities = np.asarray(plydata.elements[0]["opacity"])[..., np.newaxis]

    features_dc = np.zeros((xyz.shape[0], 3, 1))
    features_dc[:, 0, 0] = np.asarray(plydata.elements[0]["f_dc_0"])
    features_dc[:, 1, 0] = np.asarray(plydata.elements[0]["f_dc_1"])
    features_dc[:, 2, 0] = np.asarray(plydata.elements[0]["f_dc_2"])

    extra_f_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("f_rest_")]
    extra_f_names = sorted(extra_f_names, key=lambda x: int(x.split('_')[-1]))
    assert len(extra_f_names) == 3 * (max_sh_degree + 1) ** 2 - 3
    features_extra = np.zeros((xyz.shape[0], len(extra_f_names)))
    for idx, attr_name in enumerate(extra_f_names):
        features_extra[:, idx] = np.asarray(plydata.elements[0][attr_name])
    # Reshape (P,F*SH_coeffs) to (P, F, SH_coeffs except DC)
    features_extra = features_extra.reshape((features_extra.shape[0], 3, (max_sh_degree + 1) ** 2 - 1))

    scale_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("scale_")]
    scale_names = sorted(scale_names, key=lambda x: int(x.split('_')[-1]))
    scales = np.zeros((xyz.shape[0], len(scale_names)))
    for idx, attr_name in enumerate(scale_names):
        scales[:, idx] = np.asarray(plydata.elements[0][attr_name])

    rot_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("rot")]
    rot_names = sorted(rot_names, key=lambda x: int(x.split('_')[-1]))
    rots = np.zeros((xyz.shape[0], len(rot_names)))
    for idx, attr_name in enumerate(rot_names):
        rots[:, idx] = np.asarray(plydata.elements[0][attr_name])

    return GaussianAttributes(xyz, opacities, features_dc, features_extra, scales, rots)


def construct_list_of_attributes(_features_dc, _features_rest, _scaling, _rotation):
    l = ['x', 'y', 'z', 'nx', 'ny', 'nz']
    # All channels except the 3 DC
    for i in range(_features_dc.shape[1] * _features_dc.shape[2]):
        l.append('f_dc_{}'.format(i))
    for i in range(_features_rest.shape[1] * _features_rest.shape[2]):
        l.append('f_rest_{}'.format(i))
    l.append('opacity')
    for i in range(_scaling.shape[1]):
        l.append('scale_{}'.format(i))
    for i in range(_rotation.shape[1]):
        l.append('rot_{}'.format(i))
    return l


def select_gaussians(gaussian_attrs, select_mask_or_idx):
    return GaussianAttributes(
        xyz=gaussian_attrs.xyz[select_mask_or_idx],
        opacities=gaussian_attrs.opacities[select_mask_or_idx],
        features_dc=gaussian_attrs.features_dc[select_mask_or_idx],
        features_extra=gaussian_attrs.features_extra[select_mask_or_idx],
        scales=gaussian_attrs.scales[select_mask_or_idx],
        rot=gaussian_attrs.rot[select_mask_or_idx]
    )


def combine_gaussians(gaussian_attrs_list):
    return GaussianAttributes(
        xyz=np.concatenate([gau.xyz for gau in gaussian_attrs_list], axis=0),
        opacities=np.concatenate([gau.opacities for gau in gaussian_attrs_list], axis=0),
        features_dc=np.concatenate([gau.features_dc for gau in gaussian_attrs_list], axis=0),
        features_extra=np.concatenate([gau.features_extra for gau in gaussian_attrs_list], axis=0),
        scales=np.concatenate([gau.scales for gau in gaussian_attrs_list], axis=0),
        rot=np.concatenate([gau.rot for gau in gaussian_attrs_list], axis=0),
    )


def apply_transformation_to_gaussians(gaussian_attrs, spatial_transformation, color_transformation=None):
    xyzs = np.copy(gaussian_attrs.xyz)
    xyzs = np.matmul(xyzs, spatial_transformation[:3, :3].transpose()) + spatial_transformation[:3, 3].reshape([1, 3])

    gaussian_rotmats = R.from_quat(gaussian_attrs.rot[:, (1, 2, 3, 0)]).as_matrix()
    new_rots = []
    for rotmat in gaussian_rotmats:
        rotmat = np.matmul(spatial_transformation[:3, :3], rotmat)
        rotq = R.from_matrix(rotmat).as_quat()
        rotq = np.array([rotq[3], rotq[0], rotq[1], rotq[2]])
        new_rots.append(rotq)
    new_rots = np.stack(new_rots, axis=0)
    if color_transformation is not None:
        if color_transformation.shape[0] == 3 and color_transformation.shape[1] == 3:
            new_clrs = np.matmul(gaussian_attrs.features_dc[:, :, 0], color_transformation)[:, :, np.newaxis]
        elif color_transformation.shape[0] == 4 and color_transformation.shape[1] == 4:
            clrs = gaussian_attrs.features_dc[:, :, 0]
            clrs = np.concatenate([clrs, np.ones_like(clrs[:, :1])], axis=1)
            new_clrs = np.matmul(clrs, color_transformation)
            new_clrs = new_clrs[:, :3, np.newaxis]
    else:
        new_clrs = gaussian_attrs.features_dc

    return GaussianAttributes(
        xyz=xyzs,
        opacities=gaussian_attrs.opacities,
        features_dc=new_clrs,
        features_extra=gaussian_attrs.features_extra,
        scales=gaussian_attrs.scales,
        rot=new_rots,
    )


def update_gaussian_attributes(
        orig_gaussian,
        new_xyz=None, new_rgb=None, new_rot=None, new_opacity=None, new_scale=None):
    return GaussianAttributes(
        xyz=orig_gaussian.xyz if new_xyz is None else new_xyz,
        opacities=orig_gaussian.opacities if new_opacity is None else new_opacity,
        features_dc=orig_gaussian.features_dc if new_rgb is None else new_rgb,
        features_extra=orig_gaussian.features_extra,
        scales=orig_gaussian.scales if new_scale is None else new_scale,
        rot=orig_gaussian.rot if new_rot is None else new_rot,
    )


def save_gaussians_as_ply(path, gaussian_attrs: GaussianAttributes):
    os.makedirs(os.path.dirname(path), exist_ok=True)

    xyz = gaussian_attrs.xyz
    normals = np.zeros_like(xyz)
    features_dc = gaussian_attrs.features_dc
    features_rest = gaussian_attrs.features_extra
    opacities = gaussian_attrs.opacities
    scale = gaussian_attrs.scales
    rotation = gaussian_attrs.rot

    dtype_full = [(attribute, 'f4') for attribute in construct_list_of_attributes(features_dc, features_rest, scale, rotation)]

    elements = np.empty(xyz.shape[0], dtype = dtype_full)
    attributes = np.concatenate((xyz, normals, features_dc.reshape(features_dc.shape[0], -1),
                                 features_rest.reshape(features_rest.shape[0], -1),
                                 opacities, scale, rotation), axis=1)
    elements[:] = list(map(tuple, attributes))
    el = PlyElement.describe(elements, 'vertex')
    PlyData([el]).write(path)
    return