# Copyright 2024 Toyota Research Institute.  All rights reserved.

import numpy as np
from numba import njit, prange
import torch

from anydata.geometry.cameras.base import CameraBase
from anydata.utils.decorators import iterate1
from anydata.utils.types import is_tensor, is_numpy


def calculate_normals(depth, camera: CameraBase, to_world=False, euclidean=False):
    """Calculate normals from a pointcloud map or from a depth map + intrinsics"""
    # Create pointcloud map if intrinsics are provided
    points = camera.reconstruct_depth_map(depth, euclidean=euclidean, to_world=to_world)
    # Prepare points for cross-product
    p0 = points[:, :, :-1, :-1]
    p1 = points[:, :,  1:, :-1]
    p2 = points[:, :, :-1,  1:]
    # Calculate and normalize normals
    normals = torch.cross(p1 - p0, p2 - p0, 1)
    normals = normals / normals.norm(dim=1, keepdim=True)
    # Pad normals
    normals = torch.nn.functional.pad(normals, [0, 1, 0, 1], mode='replicate')
    # Substitute nan values with zero
    normals[torch.isnan(normals)] = 0.0
    # Return normals
    return normals


def calc_dot_prod(pts: torch.Tensor, nrm: torch.Tensor):
    """Calculate dot product between points (pts) and normals (nrm)"""
    pts = pts / pts.norm(dim=1, keepdim=True)
    nrm = nrm / nrm.norm(dim=1, keepdim=True)
    dots = torch.sum(pts * nrm, dim=1, keepdim=True)
    return dots


@iterate1
@iterate1
def inv2depth(inv_depth: torch.Tensor):
    """Invert an inverse depth map to produce a depth map"""
    if is_tensor(inv_depth):
        depth = 1. / inv_depth.clamp(min=1e-8, max=None)
    elif is_numpy(inv_depth):
        depth = 1. / inv_depth.clip(min=1e-8, max=None)
    else:
        raise NotImplementedError('Wrong inverse depth format.')
    depth[inv_depth <= 0.] = 0.
    return depth


@iterate1
@iterate1
def depth2inv(depth: torch.Tensor):
    """Invert a depth map to produce an inverse depth map"""
    if is_tensor(depth):
        inv_depth = 1. / depth.clamp(min=1e-8, max=None)
    elif is_numpy(depth):
        inv_depth = 1. / depth.clip(min=1e-8, max=None)
    else:
        raise NotImplementedError('Wrong depth format')
    inv_depth[depth <= 0.] = 0.
    return inv_depth


def norm_from_info(data, encode_info):
    min_val, max_val = float(encode_info['min'][0]), float(encode_info['max'][0])
    data = (data - min_val) / max(max_val - min_val, 1e-8)  # Clip denominator
    return data * (2 ** int(encode_info['bits'][0]) - 1)


def unnorm_from_info(data, encode_info):
    min_val, max_val = float(encode_info['min'][0]), float(encode_info['max'][0])
    data = data / (2 ** int(encode_info['bits'][0]) - 1)
    return data * (max_val - min_val) + min_val


def encode_1bit(data: np.ndarray):
    return np.round(data).astype(bool)


def decode_1bit(data: np.ndarray, encode_info: dict = None):
    return data


def encode_8bit(data: np.ndarray):
    return np.round(data).astype(np.uint8)


def decode_8bit(data: np.ndarray, encode_info: dict = None):
    return data.astype(np.float32)


def encode_24bit(data_array: np.ndarray):
    data_uint32 = np.round(data_array).astype(np.uint32)
    data_uint8 = data_uint32.view(np.uint8).reshape(*data_uint32.shape, 4)
    data_uint24 = data_uint8[..., :3]  # Use first 3 channels for RGB encoding
    return data_uint24


def decode_24bit(depth_rgb: np.ndarray, encode_info: dict = None):
    pad_vals = [(0, 0)] * len(depth_rgb.shape[:-1]) + [(0, 1)]  # Pad only the last dimension to 4 channels
    depth_rgb_full = np.pad(depth_rgb, pad_vals, mode='constant')  # Pad to 4 channels
    depth_uint32_reconstructed = depth_rgb_full.view(np.uint32).reshape(depth_rgb.shape[:-1])
    decoded = depth_uint32_reconstructed.astype(np.float32)
    return unnorm_from_info(decoded, encode_info)


@njit(parallel=True)
def encode_18bit_numba(depth_flat: np.ndarray):
    """Pack 4x18-bit values into 9 bytes using Numba"""
    assert depth_flat.shape[0] % 4 == 0, "Input length must be a multiple of 4 for packing"
    num_groups = depth_flat.shape[0] // 4
    packed: np.ndarray = np.zeros((num_groups, 9), dtype=np.uint8)
    
    for i in prange(num_groups):
        # Get 4 values (or fewer for last group)
        v0, v1, v2, v3 = depth_flat[i*4:(i+1)*4]
        # Pack into 9 bytes
        packed[i, 0] = v0 & 0xFF
        packed[i, 1] = (v0 >> 8) & 0xFF
        packed[i, 2] = ((v0 >> 16) & 0x03) | ((v1 & 0x3F) << 2)
        packed[i, 3] = (v1 >> 6) & 0xFF
        packed[i, 4] = ((v1 >> 14) & 0x0F) | ((v2 & 0x0F) << 4)
        packed[i, 5] = (v2 >> 4) & 0xFF
        packed[i, 6] = ((v2 >> 12) & 0x3F) | ((v3 & 0x03) << 6)
        packed[i, 7] = (v3 >> 2) & 0xFF
        packed[i, 8] = (v3 >> 10) & 0xFF
    return packed


def encode_18bit(depth: np.ndarray):
    """Fast wrapper for 18-bit encoding"""
    assert depth.shape[-2] % 2 == 0 and depth.shape[-1] % 2 == 0, "Height and width must be even"
    base_shape = depth.shape[:-2] + (depth.shape[-2] // 2, depth.shape[-1] // 2)
    depth_u32 = np.round(depth).astype(np.uint32).flatten()
    assert np.max(depth_u32) < 2**18, "Depth values are too large to encode in 18 bits"
    encoded_flat = encode_18bit_numba(depth_u32)
    encoded = encoded_flat.reshape(base_shape + (encoded_flat.shape[-1],))
    return encoded


@njit(parallel=True)
def decode_18bit_numba(packed: np.ndarray) -> np.ndarray:
    """Decode 18-bit packed values using Numba"""
    assert packed.shape[0] % 9 == 0, "Input length must be a multiple of 9 for unpacking"
    num_groups = packed.shape[0] // 9
    decoded = np.zeros((num_groups, 4), dtype=np.uint32)
    
    for i in prange(num_groups):
        # Extract 4 values from 9 bytes
        p0, p1, p2, p3, p4, p5, p6, p7, p8 = packed[i*9:(i+1)*9]
        # Decode 4 values
        decoded[i, 0] = p0 | (p1 << 8) | ((p2 & 0x03) << 16)
        decoded[i, 1] = ((p2 >> 2) & 0x3F) | (p3 << 6) | ((p4 & 0x0F) << 14)
        decoded[i, 2] = ((p4 >> 4) & 0x0F) | (p5 << 4) | ((p6 & 0x3F) << 12)
        decoded[i, 3] = ((p6 >> 6) & 0x03) | (p7 << 2) | (p8 << 10)
    return decoded


def decode_18bit(packed: np.ndarray, encode_info: dict = None):
    """Fast wrapper for 18-bit decoding"""
    base_shape = packed.shape[:-1]
    packed_flat = packed.flatten()
    decoded_flat = decode_18bit_numba(packed_flat)
    decoded = decoded_flat.reshape(base_shape[:-2] + (base_shape[-2] * 2, base_shape[-1] * 2))
    decoded = decoded.astype(np.float32)
    return unnorm_from_info(decoded, encode_info)


def get_encode_info(data, bits=None):
    """Automatically determine bits and return normalized depth"""

    # Store interval range
    data_min, data_max = data.min(), data.max()
    encode_info = {
        'min': np.array([data_min]), 
        'max': np.array([data_max]),
    }

    if bits is None:
        # Automatically determine number of bits to use
        interval = data_max - data_min
        if interval <= 100:
            encode_info['bits'] = np.array([18])
        elif interval <= 300:
            encode_info['bits'] = np.array([24])
        else:
            encode_info['bits'] = np.array([32])
    else:
        # Use value given
        encode_info['bits'] = np.array([bits])

    # Normalize depth map to [0,1]
    if int(encode_info['bits']) in [18, 24]:
        data = norm_from_info(data, encode_info)

    return data, encode_info


def encode_dense(dense: np.ndarray, bits: int = None):
    """Encode dense map into a compact format based on the specified bit"""

    # Normalize and get bits used
    dense, encode_info = get_encode_info(dense, bits=bits)
    bits = int(encode_info['bits'][0])

    leftover_h = dense.shape[-2] % 4
    if leftover_h > 0:
        dense = dense[..., :-leftover_h, :]
    leftover_w = dense.shape[-1] % 4
    if leftover_w > 0:
        dense = dense[..., :, :-leftover_w]

    # Encode with provided bits
    if bits == 1:
        return encode_1bit(dense), encode_info
    elif bits == 8:
        return encode_8bit(dense), encode_info
    elif bits == 18:
        return encode_18bit(dense), encode_info
    elif bits == 24:
        return encode_24bit(dense), encode_info
    elif bits == 32:
        return dense.astype(np.float32), encode_info
    else:
        raise ValueError(f"Unsupported bit: {bits}. Supported values are 8, 18, 24, and 32.")


def decode_dense(encoded: np.ndarray, encode_info: dict = None):
    """Decode depth map from a compact format based on the specified bit depth"""
    if encode_info is None:
        return encoded
    bits = int(encode_info['bits'][0])
    if bits == 1:
        decoded = decode_1bit(encoded, encode_info)
    elif bits == 8:
        decoded = decode_8bit(encoded, encode_info)
    elif bits == 18:
        decoded = decode_18bit(encoded, encode_info)
    elif bits == 24:
        decoded = decode_24bit(encoded, encode_info)
    elif bits == 32:
        decoded = encoded.astype(np.float32)  # No decoding needed for 32-bit depth
    else:
        raise ValueError(f"Unsupported bit depth: {bits}. Supported values are 8, 18, 24, and 32.")

    return decoded