from gaussian import Gaussian
from utils import quat2rot
from aabb import process_aabb, intersect, min_interval
import numpy as np

def similarity(gaussian: Gaussian, id1: int, id2: int):
    geometry_sim = geometry_similarity(gaussian, id1, id2)
    color_sim = color_similarity(gaussian.sh[id1], gaussian.sh[id2])
    color_max = np.max(np.abs(gaussian.sh))
    # return geometry_sim * geometry_coeff + color_sim * color_coeff
    return geometry_sim, color_sim / (color_max * color_max * 48)
    
def color_similarity(c1: np.ndarray, c2: np.ndarray):
    return np.sum(c1 * c2)

def geometry_similarity(gaussian: Gaussian, id1: int, id2: int):
    def geometry_similarity_using_overlap(gaussian: Gaussian, id1: int, id2: int):
        aabb1 = process_aabb(gaussian.positions[id1], gaussian.scales[id1], gaussian.rotations[id1], scale_factor=2)
        aabb2 = process_aabb(gaussian.positions[id2], gaussian.scales[id2], gaussian.rotations[id2], scale_factor=2)
        # 在两个 aabb 里 stratified sample 点，对两个 Gaussian 的乘积做积分
        # 然后除以总sample数作为重叠程度的 metric

        intersected_aabb = intersect(aabb1, aabb2)

        if intersected_aabb is None:
            return 0

        x_stride = (intersected_aabb[1] - intersected_aabb[0]) / (4. + 1e-3)
        y_stride = (intersected_aabb[3] - intersected_aabb[2]) / (4. + 1e-3)
        z_stride = (intersected_aabb[5] - intersected_aabb[4]) / (4. + 1e-3)
        
        score = 0
        for i in range(int(np.floor((intersected_aabb[1] - intersected_aabb[0]) / x_stride))):
            x = intersected_aabb[0] + i * x_stride
            for j in range(int(np.floor((intersected_aabb[3] - intersected_aabb[2]) / y_stride))):
                y = intersected_aabb[2] + j * y_stride
                for k in range(int(np.floor((intersected_aabb[5] - intersected_aabb[4]) / z_stride))):
                    z = intersected_aabb[4] + k * z_stride
                    pos = np.array([x, y, z])
                    score += gaussian.calc(pos, id1, 10) * gaussian.calc(pos, id2, 10) * gaussian.opacity[id1] * gaussian.opacity[id2]

        scale1 = gaussian.scales[id1][0] * gaussian.scales[id1][1] * gaussian.scales[id1][2]
        scale2 = gaussian.scales[id2][0] * gaussian.scales[id2][1] * gaussian.scales[id2][2]

        # print(f"calculating pair ({id1},{id2}), used {64} samples.")
        return score / (scale1 * scale2)

    def geometry_similarity_using_position(gaussian: Gaussian, id1: int, id2: int):
        def position_similarity(p1: np.ndarray, p2: np.ndarray):
            return np.linalg.norm(p1 - p2)
        
        def scale_similarity(s1: np.ndarray, s2: np.ndarray):
            return np.linalg.norm(s1 - s2)

        def rotation_similarity(q1: np.ndarray, q2: np.ndarray):
            return 1 - np.dot(q1, q2)
            # r1 = quat2rot(q1)
            # r2 = quat2rot(q2)
            # eigenvalues, eigenvectors = np.linalg.eig(r1 @ r2.T)
            # return np.linalg.norm(eigenvectors @ np.diag(np.log(eigenvalues)) @ np.linalg.inv(eigenvectors)) 
        position_coeff = 1.
        scale_coeff = 1.
        rotation_coeff = 1.
        return position_similarity(gaussian.positions[id1], gaussian.positions[id2]) * position_coeff + \
            scale_similarity(gaussian.scales[id1], gaussian.scales[id2]) * scale_coeff + \
            rotation_similarity(gaussian.rotations[id1], gaussian.rotations[id2]) * rotation_coeff

    return geometry_similarity_using_overlap(gaussian, id1, id2)