from utils import quat2rot
import numpy as np

class Gaussian:
    def __init__(self, positions, scales, rotations, features_dc, features_rest, opacity):
        self.num_gaussians = positions.shape[0]
        self.positions = positions

        self.scales = scales
        self.rotations = rotations

        # self.scales = np.exp(scales)
        
        # rot_length = np.repeat(np.linalg.norm(rotations, axis=1), 4).reshape(-1, 4)
        # self.rotations = rotations / rot_length

        self.features_dc = features_dc
        self.features_rest = features_rest

        self.opacity = opacity

    def empty_with_cap(cap: int) -> "Gaussian":
        positions = np.zeros((cap, 3))
        scales = np.zeros((cap, 3))
        rotations = np.zeros((cap, 4))
        features_dc = np.zeros((cap, 1, 3))
        features_rest = np.zeros((cap, 15, 3))
        opacity = np.zeros((cap, 1))

        gaussian = Gaussian(positions, scales, rotations, features_dc, features_rest, opacity)
        gaussian.num_gaussians = 0
        return gaussian

    @property
    def sh(self):
        return np.concatenate((self.features_dc, self.features_rest), axis=1)
    
    def calc(self, x, id: int, scale: float):
        x = x * scale
        S = np.diag(self.scales[id])
        R = quat2rot(self.rotations[id])
        M = S @ R

        Sigma = M.T @ M
        Sigma = np.linalg.inv(Sigma)

        return np.dot(x, Sigma @ x)
        # M = R @ S
        # return np.dot(x, M @ M.T @ x)

    def clip_to_box(self, min: np.ndarray, max: np.ndarray):
        indices = ((self.positions > min) & (self.positions < max)).all(axis=1)
        self.apply_filter(indices)
    
    def apply_filter(self, filter):
        self.positions = self.positions[filter]
        self.scales = self.scales[filter]
        self.rotations = self.rotations[filter]
        self.features_dc = self.features_dc[filter]
        self.features_rest = self.features_rest[filter]
        self.opacity = self.opacity[filter]
        self.num_gaussians = self.positions.shape[0]

    def add(self, position, scale, rotation, features_dc, features_rest, opacity):
        n = self.num_gaussians

        self.positions[n] = position
        self.scales[n] = scale
        self.rotations[n] = rotation
        self.features_dc[n] = features_dc
        self.features_rest[n] = features_rest
        self.opacity[n] = opacity

        self.num_gaussians += 1
        return n
    
    def concat(self, other: "Gaussian"):
        self.positions = np.concatenate((self.positions, other.positions), axis=0)
        self.scales = np.concatenate((self.scales, other.scales), axis=0)
        self.rotations = np.concatenate((self.rotations, other.rotations), axis=0)
        self.features_dc = np.concatenate((self.features_dc, other.features_dc), axis=0)
        self.features_rest = np.concatenate((self.features_rest, other.features_rest), axis=0)
        self.opacity = np.concatenate((self.opacity, other.opacity), axis=0)
        self.num_gaussians = self.num_gaussians + other.num_gaussians

    def replace(self, id1: int, id2: int,
                position, scale, rotation, sh, opacity):
        # 用新的 gaussian 替换 id1 位置的 gaussian
        self.positions[id1] = position
        self.scales[id1] = scale
        self.rotations[id1] = rotation
        self.features_dc[id1] = sh[0, :]
        self.features_rest[id1] = sh[1:, :]
        self.opacity[id1] = opacity

        # 把最后一个 gaussian 放到 id2 的位置上
        n = self.num_gaussians
        self.positions[id2] = self.positions[n - 1]
        self.scales[id2] = self.scales[n - 1]
        self.rotations[id2] = self.rotations[n - 1]
        self.features_dc[id2] = self.features_dc[n - 1]
        self.features_rest[id2] = self.features_rest[n - 1]
        self.opacity[id2] = self.opacity[n - 1]
        
        self.num_gaussians -= 1
        pass

    def copy(self) -> "Gaussian":
        positions = np.copy(self.positions)
        scales = np.copy(self.scales)
        rotations = np.copy(self.rotations)
        features_dc = np.copy(self.features_dc)
        features_rest = np.copy(self.features_rest)
        opacity = np.copy(self.opacity)

        return Gaussian(positions, scales, rotations, features_dc, features_rest, opacity)