from pathlib import Path

import numpy as np
import torch
from einops import einsum, rearrange, repeat
from jaxtyping import Float
from plyfile import PlyData, PlyElement
from scipy.spatial.transform import Rotation as R
from torch import Tensor

from ..misc.sh_rotation import rotate_sh


def construct_list_of_attributes(num_rest: int) -> list[str]:
    attributes = ["x", "y", "z", "nx", "ny", "nz"]
    for i in range(3):
        attributes.append(f"f_dc_{i}")
    for i in range(num_rest):
        attributes.append(f"f_rest_{i}")
    attributes.append("opacity")
    for i in range(3):
        attributes.append(f"scale_{i}")
    for i in range(4):
        attributes.append(f"rot_{i}")
    return attributes


def export_ply(
    extrinsics: Float[Tensor, "4 4"],
    means: Float[Tensor, "gaussian 3"],
    scales: Float[Tensor, "gaussian 3"],
    rotations: Float[Tensor, "gaussian 4"],
    harmonics: Float[Tensor, "gaussian 3 d_sh"],
    opacities: Float[Tensor, " gaussian"],
    path: Path,
    align_to_view: bool = False,
):
    path = Path(path)
    with torch.autocast(device_type=means.device.type, enabled=False):
        means = means.float()
        scales = scales.float()
        rotations = rotations.float()
        harmonics = harmonics.float()
        opacities = opacities.float()
        extrinsics = extrinsics.to(means)
        rotations = rotations / rotations.norm(dim=-1, keepdim=True).clamp_min(1e-8)

        if align_to_view:
            view_rotation = extrinsics[:3, :3].transpose(-1, -2)
            view_origin = extrinsics[:3, 3]
            means = einsum(
                view_rotation,
                means - view_origin,
                "i j, ... j -> ... i",
            )
            harmonics = rotate_sh(harmonics, view_rotation)
            rotations = R.from_quat(rotations.detach().cpu().numpy()).as_matrix()
            rotations = view_rotation.detach().cpu().numpy() @ rotations
            rotations = R.from_matrix(rotations).as_quat()
        else:
            rotations = rotations.detach().cpu().numpy()

        x, y, z, w = rearrange(rotations, "g xyzw -> xyzw g")
        rotations = np.stack((w, x, y, z), axis=-1)
        means = means.detach().cpu().numpy()
        normals = np.zeros_like(means)
        harmonics_dc = harmonics[..., 0].detach().cpu().contiguous().numpy()
        harmonics_rest = (
            harmonics[..., 1:].flatten(start_dim=1).detach().cpu().contiguous().numpy()
        )
        logit_opacities = (
            torch.logit(opacities[..., None].clamp(1e-6, 1 - 1e-6))
            .detach()
            .cpu()
            .numpy()
        )
        log_scales = scales.clamp_min(1e-12).log().detach().cpu().numpy()

    num_rest = 3 * (harmonics.shape[-1] - 1)
    dtype_full = [
        (attribute, "f4") for attribute in construct_list_of_attributes(num_rest)
    ]
    elements = np.empty(means.shape[0], dtype=dtype_full)
    attributes = (
        means,
        normals,
        harmonics_dc,
        harmonics_rest,
        logit_opacities,
        log_scales,
        rotations,
    )
    attributes = np.concatenate(attributes, axis=1)
    elements[:] = list(map(tuple, attributes))
    path.parent.mkdir(exist_ok=True, parents=True)
    PlyData([PlyElement.describe(elements, "vertex")]).write(path)


def infer_gaussian_grid(
    num_gaussians: int,
    num_views: int,
    image_height: int,
    image_width: int,
) -> tuple[int, int, int]:
    if num_gaussians % num_views != 0:
        raise ValueError(
            f"Cannot split {num_gaussians} Gaussians over {num_views} views."
        )

    gaussians_per_view = num_gaussians // num_views
    for downsample in (1, 2, 4, 8, 16, 32):
        if image_height % downsample != 0 or image_width % downsample != 0:
            continue
        height = image_height // downsample
        width = image_width // downsample
        pixels = height * width
        if pixels > 0 and gaussians_per_view % pixels == 0:
            return height, width, gaussians_per_view // pixels

    raise ValueError(
        "Could not infer Gaussian grid layout from "
        f"num_gaussians={num_gaussians}, num_views={num_views}, "
        f"image_shape=({image_height}, {image_width})."
    )


def select_gaussians(
    tensor: Tensor,
    mask: Tensor,
    num_views: int,
    height: int,
    width: int,
    points_per_pixel: int,
) -> Tensor:
    tensor = rearrange(
        tensor,
        "() (v h w spp) ... -> h w spp v ...",
        v=num_views,
        h=height,
        w=width,
        spp=points_per_pixel,
    )
    return tensor[mask]


def local_to_world_rotations(
    rotations: Tensor,
    context_extrinsics: Tensor,
    num_views: int,
    height: int,
    width: int,
    points_per_pixel: int,
    mask: Tensor | None,
) -> Tensor:
    if mask is None:
        local_rotations = rotations[0]
        c2w_rotations = repeat(
            context_extrinsics[0, :, :3, :3],
            "v i j -> (v h w spp) i j",
            h=height,
            w=width,
            spp=points_per_pixel,
        )
    else:
        local_rotations = select_gaussians(
            rotations,
            mask,
            num_views,
            height,
            width,
            points_per_pixel,
        )
        c2w_rotations = repeat(
            context_extrinsics[0, :, :3, :3],
            "v i j -> h w spp v i j",
            h=height,
            w=width,
            spp=points_per_pixel,
        )[mask]

    local_rotation_mats = R.from_quat(
        local_rotations.detach().cpu().numpy()
    ).as_matrix()
    world_rotation_mats = c2w_rotations.detach().cpu().numpy() @ local_rotation_mats
    world_rotations = R.from_matrix(world_rotation_mats).as_quat()
    return torch.from_numpy(world_rotations).to(
        device=rotations.device,
        dtype=rotations.dtype,
    )


def save_gaussian_ply(
    gaussians,
    visualization_dump,
    example,
    save_path,
    save_all_gaussians: bool = True,
    align_to_view: bool = False,
):
    del visualization_dump

    v, _, image_height, image_width = example["context"]["image"].shape[1:]
    h, w, points_per_pixel = infer_gaussian_grid(
        gaussians.means.shape[1],
        v,
        image_height,
        image_width,
    )

    if gaussians.scales is None or gaussians.rotations is None:
        raise ValueError("PLY export requires Gaussians.scales and Gaussians.rotations.")

    if save_all_gaussians:
        mask = None
        selected_means = gaussians.means[0]
        selected_scales = gaussians.scales[0]
        selected_harmonics = gaussians.harmonics[0]
        selected_opacities = gaussians.opacities[0]
    else:
        mask = torch.zeros(
            (h, w, points_per_pixel, v),
            dtype=torch.bool,
            device=gaussians.means.device,
        )
        gaussian_trim = 8
        mask[gaussian_trim:-gaussian_trim, gaussian_trim:-gaussian_trim, :, :] = True
        selected_means = select_gaussians(
            gaussians.means, mask, v, h, w, points_per_pixel
        )
        selected_scales = select_gaussians(
            gaussians.scales, mask, v, h, w, points_per_pixel
        )
        selected_harmonics = select_gaussians(
            gaussians.harmonics, mask, v, h, w, points_per_pixel
        )
        selected_opacities = select_gaussians(
            gaussians.opacities, mask, v, h, w, points_per_pixel
        )

    world_rotations = local_to_world_rotations(
        gaussians.rotations,
        example["context"]["extrinsics"],
        v,
        h,
        w,
        points_per_pixel,
        mask,
    )

    export_ply(
        example["context"]["extrinsics"][0, v // 2],
        selected_means,
        selected_scales,
        world_rotations,
        selected_harmonics,
        selected_opacities,
        save_path,
        align_to_view=align_to_view,
    )