import numpy as np
from utils import quat2rot

def aabb_size(scale, rotation, scale_factor=1):
    scale = scale_factor * scale
    if scale.shape != (3, 3):
        scale_matrix = np.diag(scale)
    else:
        scale_matrix = scale
    
    if rotation.shape != (3, 3):
        rotation = rotation / np.linalg.norm(rotation)
        rotation_matrix = quat2rot(rotation)
        # r, x, y, z = rotation[0], rotation[1], rotation[2], rotation[3]
        # rotation_matrix = np.array([
        #     [1. - 2. * (y * y + z * z), 2. * (x * y - r * z), 2. * (x * z + r * y)],
        #     [2. * (x * y + r * z), 1. - 2. * (x * x + z * z), 2. * (y * z - r * x)],
        #     [2. * (x * z - r * y), 2. * (y * z + r * x), 1. - 2. * (x * x + y * y)]
        # ], dtype=np.float32)
    else:
        rotation_matrix = rotation

    M = scale_matrix @ rotation_matrix
    R = M.T @ M

    x_span = np.sqrt(R[0][0])
    y_span = np.sqrt(R[1][1])
    z_span = np.sqrt(R[2][2])

    return x_span, y_span, z_span

def process_aabb(position, scale, rotation, scale_factor=1):
    x_span, y_span, z_span = aabb_size(scale, rotation, scale_factor)

    xmin = position[0] - x_span
    xmax = position[0] + x_span
    ymin = position[1] - y_span
    ymax = position[1] + y_span
    zmin = position[2] - z_span
    zmax = position[2] + z_span
    
    aabb = np.array([xmin, xmax, ymin, ymax, zmin, zmax])
    if np.isnan(aabb).any():
        print("nan detected!")
        print(scale, rotation, aabb)

    return aabb

def solve_scale(rotation, aabb):
    dest = np.array([
        (aabb[1] - aabb[0]) * (aabb[1] - aabb[0]) / 4,
        (aabb[3] - aabb[2]) * (aabb[3] - aabb[2]) / 4,
        (aabb[5] - aabb[4]) * (aabb[5] - aabb[4]) / 4,
    ])

    rotation = rotation / np.linalg.norm(rotation)
    r, x, y, z = rotation[0], rotation[1], rotation[2], rotation[3]
    rotation_matrix = np.array([
        [1. - 2. * (y * y + z * z), 2. * (x * y - r * z), 2. * (x * z + r * y)],
        [2. * (x * y + r * z), 1. - 2. * (x * x + z * z), 2. * (y * z - r * x)],
        [2. * (x * z - r * y), 2. * (y * z + r * x), 1. - 2. * (x * x + y * y)]
    ], dtype=np.float32)

    coeff = rotation_matrix.T * rotation_matrix.T
    solution = np.linalg.solve(coeff, dest)
    print(solution)
    return np.sqrt(solution)


def intersect(aabb1, aabb2):
    new_aabb = np.zeros_like(aabb1)
    new_aabb[0] = max(aabb1[0], aabb2[0])
    new_aabb[1] = min(aabb1[1], aabb2[1])
    new_aabb[2] = max(aabb1[2], aabb2[2])
    new_aabb[3] = min(aabb1[3], aabb2[3])
    new_aabb[4] = max(aabb1[4], aabb2[4])
    new_aabb[5] = min(aabb1[5], aabb2[5])

    if new_aabb[0] < new_aabb[1] and new_aabb[2] < new_aabb[3] and new_aabb[4] < new_aabb[5]:
        return new_aabb
    else:
        return None

def aabb_merge(aabb1, aabb2):
    new_aabb = np.zeros_like(aabb1)
    new_aabb[0] = min(aabb1[0], aabb2[0])
    new_aabb[1] = max(aabb1[1], aabb2[1])
    new_aabb[2] = min(aabb1[2], aabb2[2])
    new_aabb[3] = max(aabb1[3], aabb2[3])
    new_aabb[4] = min(aabb1[4], aabb2[4])
    new_aabb[5] = max(aabb1[5], aabb2[5])
    return new_aabb

def center(aabb):
    return np.array([
        aabb[0] + aabb[1],
        aabb[2] + aabb[3],
        aabb[4] + aabb[5]
    ]) / 2

def min_interval(aabb):
    return min(aabb[1] - aabb[0], min(aabb[3] - aabb[2], aabb[5] - aabb[4]))