"""Partial point clouds of MOVi objects observed by a simulated depth camera (data.partial_view).

A pinhole depth camera ray-casts the posed visual meshes of a scene (data_generation/objects/<shape>/
visual_geometry.obj, the meshes rendered by Kubric). The points where rays hit object i, i.e. its segment of the
depth image back-projected to 3D, form its observation. Observed points are stored in the object frame (asset units)
and moved with the ground-truth pose to every frame of a sample, so each point keeps its identity over time: the
input is partial, the supervision is the ground-truth motion of the same surface points, and the camera is used only
at the observation frame.

Cameras follow the Kubric convention: they look along their -z axis, K maps camera coordinates to pixels and R is the
camera-to-world matrix. With the original Kubric camera, the ray cast reproduces the `image_positions` and
`visibility` recorded by Kubric (scripts/check_depth_camera.py).

Config (data.partial_view):
    enable: true
    camera_file: data_generation/scene_cameras_MoviB.json   # Kubric camera of every scene (export_scene_cameras.py)
    eval_camera: pulled_back    # scene camera direction, moved back to `distance`; or scene (the Kubric camera)
    train_camera: random        # azimuth, elevation and distance drawn per sample; or pulled_back / scene
    train_azimuth_jitter_deg: 0.0  # train split, pulled_back camera: rotate the direction about the vertical
                                # through `target` by ~ U(-j, j) degrees, keeping its elevation
    distance: 12.0              # metres from `target`
    target: [0.0, 0.0, 1.5]
    fov_y_deg: 58.0
    resolution: [480, 640]      # (H, W) of the evaluation camera
    train_resolutions: [[480, 640]]
    random_azimuth_deg: [0.0, 360.0]
    random_elevation_deg: [20.0, 70.0]
    random_distance: [10.0, 14.0]
    observe_index: 3            # sample frame observed by the camera (the last warm-up frame)
    max_points: 512             # at most this many points per object (farthest-point subsample)
    min_points: 16              # objects with fewer observed points count as unobserved
    depth_noise_std: 0.0        # train split: depth noise std (m) at 1 m, scaled by depth^2
    boundary_erosion_px: 0      # train split: drop pixels within this distance of another segment
    occluder_prob: 0.0          # train split: probability that 1-3 random discs hide the pixels behind them
    occluder_radius: [0.03, 0.12]  # disc radius as a fraction of the image height
    random_max_points: null     # train split: [lo, hi]; keep the first n ~ U{lo..hi} FPS points of each object
    full_object_prob: 0.0       # train split: probability of using the whole collision meshes (vertices in FPS
                                # order, <= max_points) instead of the observed points
    completion: null            # CompletionNet checkpoint (src/data/shape_completion.py): add predicted points on
                                # the hidden side; the union is FPS-subsampled to max_points
    completion_input_points: 256   # observed points read by the completion network (an FPS prefix)
"""
import json
from functools import lru_cache
from pathlib import Path
from typing import Dict, List, Optional, Sequence, Tuple

import numpy as np

DEFAULTS = {
    'eval_camera': 'pulled_back',
    'train_camera': 'random',
    'train_azimuth_jitter_deg': 0.0,
    'distance': 12.0,
    'target': [0.0, 0.0, 1.5],
    'fov_y_deg': 58.0,
    'resolution': [480, 640],
    'train_resolutions': [[480, 640]],
    'random_azimuth_deg': [0.0, 360.0],
    'random_elevation_deg': [20.0, 70.0],
    'random_distance': [10.0, 14.0],
    'observe_index': 3,
    'max_points': 512,
    'min_points': 16,
    'depth_noise_std': 0.0,
    'boundary_erosion_px': 0,
    'occluder_prob': 0.0,
    'occluder_radius': [0.03, 0.12],
    'random_max_points': None,
    'full_object_prob': 0.0,
    'completion': None,
    'completion_input_points': 256,
}


def quat_to_rotmat(q: np.ndarray) -> np.ndarray:
    """Rotation matrices (..., 3, 3) of unit quaternions (..., 4) in (w, x, y, z) order."""
    q = np.asarray(q, dtype=np.float64)
    w, x, y, z = q[..., 0], q[..., 1], q[..., 2], q[..., 3]
    R = np.empty(q.shape[:-1] + (3, 3), dtype=np.float64)
    R[..., 0, 0] = 1 - 2 * (y * y + z * z)
    R[..., 0, 1] = 2 * (x * y - w * z)
    R[..., 0, 2] = 2 * (x * z + w * y)
    R[..., 1, 0] = 2 * (x * y + w * z)
    R[..., 1, 1] = 1 - 2 * (x * x + z * z)
    R[..., 1, 2] = 2 * (y * z - w * x)
    R[..., 2, 0] = 2 * (x * z - w * y)
    R[..., 2, 1] = 2 * (y * z + w * x)
    R[..., 2, 2] = 1 - 2 * (x * x + y * y)
    return R


def farthest_point_orders(clouds: Sequence[np.ndarray], k: int) -> List[np.ndarray]:
    """Batched farthest-point sampling of clouds (N_i, 3): min(k, N_i) indices per cloud in sampling order, starting
    at the point farthest from the centroid (ties: lowest index), so the first m indices are an m-point FPS subset.
    Distances are computed in float32."""
    sizes = np.array([c.shape[0] for c in clouds])
    n, rows = int(sizes.max()), np.arange(len(clouds))
    X, Y, Z = (np.zeros((len(clouds), n), dtype=np.float32) for _ in range(3))
    valid = np.zeros((len(clouds), n), dtype=bool)
    for i, c in enumerate(clouds):
        c = np.asarray(c, dtype=np.float32)
        X[i, :sizes[i]], Y[i, :sizes[i]], Z[i, :sizes[i]] = c[:, 0], c[:, 1], c[:, 2]
        valid[i, :sizes[i]] = True
    d = np.zeros((len(clouds), n), dtype=np.float32)
    buf = np.empty_like(d)
    for A in (X, Y, Z):
        np.subtract(A, (A.sum(axis=1) / sizes)[:, None].astype(np.float32), out=buf)
        d += buf * buf
    d[~valid] = -np.inf
    sel = np.zeros((len(clouds), min(int(k), n)), dtype=np.int64)
    sel[:, 0] = np.argmax(d, axis=1)
    d.fill(np.inf)
    d[~valid] = -np.inf
    dd = np.empty_like(d)
    for i in range(1, sel.shape[1]):
        last = sel[:, i - 1]
        np.subtract(X, X[rows, last][:, None], out=buf)
        np.multiply(buf, buf, out=dd)
        for A in (Y, Z):
            np.subtract(A, A[rows, last][:, None], out=buf)
            buf *= buf
            dd += buf
        np.minimum(d, dd, out=d)
        d[rows, last] = -np.inf
        sel[:, i] = np.argmax(d, axis=1)
    return [sel[i, :min(int(k), int(sizes[i]))] for i in range(len(clouds))]


def farthest_point_order(P: np.ndarray, k: int) -> np.ndarray:
    """farthest_point_orders of a single cloud."""
    return farthest_point_orders([P], k)[0]


def kabsch(A: np.ndarray, B: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
    """Rigid transforms (R (..., 3, 3), t (..., 3)) minimising |R a + t - b| from points A (N, 3) to B (..., N, 3)."""
    A = np.asarray(A, dtype=np.float64)
    B = np.asarray(B, dtype=np.float64)
    ca = A.mean(axis=0)
    cb = B.mean(axis=-2)
    H = np.einsum('ni,...nj->...ij', A - ca, B - cb[..., None, :])
    U, _, Vt = np.linalg.svd(H)
    d = np.sign(np.linalg.det(np.einsum('...ij,...jk->...ik', U, Vt)))
    D = np.zeros(H.shape)
    D[..., 0, 0] = 1.0
    D[..., 1, 1] = 1.0
    D[..., 2, 2] = d
    R = np.einsum('...ji,...jk,...lk->...il', Vt, D, U)       # V D U^T
    t = cb - np.einsum('...ij,j->...i', R, ca)
    return R, t


def look_at(eye: Sequence[float], target: Sequence[float], fov_y_deg: float, height: int, width: int):
    """Kubric-convention camera at `eye` aimed at `target` (z up): K (pixels) and camera-to-world R (4, 4)."""
    eye = np.asarray(eye, dtype=np.float64)
    target = np.asarray(target, dtype=np.float64)
    f = 0.5 * height / np.tan(np.radians(fov_y_deg) / 2)
    K = np.array([[f, 0.0, -width / 2], [0.0, -f, -height / 2], [0.0, 0.0, -1.0]])
    back = eye - target
    back /= np.linalg.norm(back)
    right = np.cross([0.0, 0.0, 1.0], back)
    right /= np.linalg.norm(right)
    up = np.cross(back, right)
    R = np.eye(4)
    R[:3, 0], R[:3, 1], R[:3, 2], R[:3, 3] = right, up, back, eye
    return K, R


def view_frame(points: np.ndarray, camera: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """View frame of an observed cloud X (n, 3; world, z up, floor at z = 0) seen from camera centre (3,).

    Returns the origin (3,), i.e. the centroid; the rotation R (3, 3) whose columns are the frame axes in world
    coordinates (z up, x horizontal towards the camera); and the extra inputs (centroid height, sine of the camera
    elevation, camera distance / 10). View coordinates are (X - origin) @ R."""
    c = np.asarray(points, np.float64).mean(axis=0)
    d = np.asarray(camera, np.float64) - c
    h = float(np.hypot(d[0], d[1]))
    x = np.array([d[0] / h, d[1] / h, 0.0]) if h > 1e-9 else np.array([1.0, 0.0, 0.0])
    R = np.stack([x, np.cross([0.0, 0.0, 1.0], x), [0.0, 0.0, 1.0]], axis=1)
    dist = float(np.linalg.norm(d))
    return c, R, np.array([c[2], d[2] / max(dist, 1e-9), dist / 10.0])


@lru_cache(maxsize=None)
def visual_mesh(objects_dir: str, shape: str) -> Tuple[np.ndarray, np.ndarray]:
    """Vertices (asset units) and faces of a KuBasic visual mesh."""
    import trimesh
    m = trimesh.load(Path(objects_dir) / shape / 'visual_geometry.obj', force='mesh', process=False)
    return np.asarray(m.vertices, dtype=np.float64), np.asarray(m.faces, dtype=np.int64)


@lru_cache(maxsize=None)
def collision_points(objects_dir: str, shape: str, k: int) -> np.ndarray:
    """At most k vertices (asset units) of a KuBasic collision mesh in FPS order (the full-object input)."""
    import trimesh
    V = np.asarray(trimesh.load(Path(objects_dir) / shape / 'collision_geometry.obj', force='mesh',
                                process=False).vertices, dtype=np.float64)
    return V[farthest_point_order(V, min(k, len(V)))]


class PartialView:
    """Observations of a scene by the configured depth camera (see the module docstring)."""

    def __init__(self, config: dict, repo_root: Path):
        unknown = sorted(set(config) - set(DEFAULTS) - {'enable', 'camera_file'})
        if unknown:
            raise ValueError(f'data.partial_view: unknown keys {unknown}')
        self.cfg = {**DEFAULTS, **{k: v for k, v in config.items() if k != 'enable'}}
        for key in ('eval_camera', 'train_camera'):
            if self.cfg[key] not in ('pulled_back', 'scene', 'random'):
                raise ValueError(f"data.partial_view.{key} must be pulled_back, scene or random; got {self.cfg[key]}")
        if self.cfg['eval_camera'] == 'random':
            raise ValueError('data.partial_view.eval_camera must be deterministic (pulled_back or scene)')
        self.objects_dir = str(repo_root / 'data_generation' / 'objects')
        cam_file = self.cfg.get('camera_file')
        if not cam_file:
            raise ValueError('data.partial_view.camera_file is required (see data_generation/export_scene_cameras.py)')
        cam_path = Path(cam_file) if Path(cam_file).is_absolute() else repo_root / cam_file
        self.cameras = json.loads(cam_path.read_text())['scenes']
        self.observe_index = int(self.cfg['observe_index'])
        self.max_points = int(self.cfg['max_points'])
        self.min_points = int(self.cfg['min_points'])
        rmp = self.cfg['random_max_points']
        if rmp is not None and not (len(rmp) == 2 and self.min_points <= int(rmp[0]) <= int(rmp[1]) <= self.max_points):
            raise ValueError(f'data.partial_view.random_max_points must be [lo, hi] with min_points <= lo <= hi <= '
                             f'max_points; got {rmp}')
        self.completer = None
        if self.cfg['completion']:
            from .shape_completion import Completer
            path = Path(self.cfg['completion'])
            self.completer = Completer(path if path.is_absolute() else repo_root / path,
                                       int(self.cfg['completion_input_points']))
        # Import Open3D in the process that builds the dataset: DataLoader workers forked after PyTorch has started its
        # OpenMP threads hang when they import it themselves.
        import open3d  # noqa: F401

    def describe(self, train: bool) -> str:
        c = self.cfg
        cam = c['train_camera'] if train else c['eval_camera']
        res = c['train_resolutions'] if train else [c['resolution']]
        aug = ''
        if train:
            aug = (f"; augmentation: depth noise {c['depth_noise_std']}, erosion {c['boundary_erosion_px']} px, "
                   f"occluders p={c['occluder_prob']} r={c['occluder_radius']}, random points {c['random_max_points']}, "
                   f"whole objects p={c['full_object_prob']}")
        comp = (f"; completion {c['completion']} (reads {c['completion_input_points']} points)"
                if c['completion'] else '')
        if train and cam == 'pulled_back' and float(c['train_azimuth_jitter_deg']) > 0:
            cam += f" (azimuth +-{c['train_azimuth_jitter_deg']} deg)"
        return (f"partial view: {cam} camera, resolutions {res}, fov_y {c['fov_y_deg']} deg, distance "
                f"{c['random_distance'] if cam == 'random' else c['distance']} m, observed at sample frame "
                f"{self.observe_index}, <= {self.max_points} points per object, objects with < {self.min_points} "
                f"points not observed{aug}{comp}")

    def camera(self, scene: str, train: bool, rng: np.random.Generator):
        """(K, R, height, width) of the camera that observes `scene`."""
        c = self.cfg
        kind = c['train_camera'] if train else c['eval_camera']
        if kind == 'scene':
            rec = self.cameras[scene]
            H, W = rec.get('resolution', [256, 256])
            return np.asarray(rec['K'], np.float64), np.asarray(rec['R'], np.float64), int(H), int(W)
        target = np.asarray(c['target'], dtype=np.float64)
        if train:
            H, W = c['train_resolutions'][rng.integers(len(c['train_resolutions']))]
        else:
            H, W = c['resolution']
        if kind == 'pulled_back':
            direction = np.asarray(self.cameras[scene]['R'], np.float64)[:3, 3] - target
            direction /= np.linalg.norm(direction)
            jitter = float(c['train_azimuth_jitter_deg'])
            if train and jitter > 0:
                a = np.radians(rng.uniform(-jitter, jitter))
                direction = np.array([np.cos(a) * direction[0] - np.sin(a) * direction[1],
                                      np.sin(a) * direction[0] + np.cos(a) * direction[1], direction[2]])
            dist = float(c['distance'])
        else:
            az = np.radians(rng.uniform(*c['random_azimuth_deg']))
            el = np.radians(rng.uniform(*c['random_elevation_deg']))
            direction = np.array([np.cos(el) * np.cos(az), np.cos(el) * np.sin(az), np.sin(el)])
            dist = float(rng.uniform(*c['random_distance']))
        K, R = look_at(target + dist * direction, target, float(c['fov_y_deg']), int(H), int(W))
        return K, R, int(H), int(W)

    def whole_objects(self, shapes: Sequence[str]) -> List[np.ndarray]:
        """Every object as its whole collision mesh (vertices, FPS order, <= max_points), the full-object input."""
        return [collision_points(self.objects_dir, s, self.max_points) for s in shapes]

    def observe(self, scene: str, shapes: Sequence[str], sizes: np.ndarray, positions: np.ndarray,
                quaternions: np.ndarray, train: bool, rng: np.random.Generator,
                camera: Optional[tuple] = None) -> List[Optional[np.ndarray]]:
        """Observed points of every object in its own frame (asset units, (n_i, 3), FPS order), None if unobserved.

        positions (N, 3) and quaternions (N, 4; w, x, y, z) are the object poses at the observation frame and sizes
        (N,) the metadata sizes: point u of object i is at positions[i] + sizes[i] R(quaternions[i]) u in the world.
        camera is (K, R, height, width) as returned by camera(), drawn here if None. With completion enabled, the
        predicted points are added.
        """
        import open3d as o3d
        if train and float(self.cfg['full_object_prob']) > 0 and rng.random() < float(self.cfg['full_object_prob']):
            return self.whole_objects(shapes)
        K, Rc, H, W = self.camera(scene, train, rng) if camera is None else camera
        Rq = quat_to_rotmat(quaternions)
        scene_rc = o3d.t.geometry.RaycastingScene()
        Kinv = np.linalg.inv(K)
        cover = np.zeros((H, W), dtype=bool)          # rays are cast only inside the objects' projected boxes
        for i, shape in enumerate(shapes):
            V, F = visual_mesh(self.objects_dir, shape)
            Vw = positions[i] + sizes[i] * V @ Rq[i].T
            scene_rc.add_triangles(o3d.core.Tensor(Vw.astype(np.float32)), o3d.core.Tensor(F.astype(np.uint32)))
            Xc = (Vw - Rc[:3, 3]) @ Rc[:3, :3]
            front = Xc[:, 2] < 0
            if front.any():
                h = Xc[front] @ K.T
                uv = h[:, :2] / h[:, 2:3]
                x0, y0 = np.clip(np.floor(uv.min(axis=0)), 0, [W, H]).astype(int)
                x1, y1 = np.clip(np.ceil(uv.max(axis=0)), 0, [W, H]).astype(int)
                cover[y0:y1, x0:x1] = True
        out: List[Optional[np.ndarray]] = [None] * len(shapes)
        if not cover.any():
            return out
        y0, y1 = np.flatnonzero(cover.any(axis=1))[[0, -1]] + [0, 1]
        x0, x1 = np.flatnonzero(cover.any(axis=0))[[0, -1]] + [0, 1]
        inside = cover[y0:y1, x0:x1].reshape(-1)
        u, v = np.meshgrid(np.arange(x0, x1) + 0.5, np.arange(y0, y1) + 0.5)
        d = np.stack([u.reshape(-1)[inside], v.reshape(-1)[inside], np.ones(int(inside.sum()))], -1) @ Kinv.T @ Rc[:3, :3].T
        d /= np.linalg.norm(d, axis=1, keepdims=True)
        rays = np.concatenate([np.broadcast_to(Rc[:3, 3], d.shape), d], axis=1).astype(np.float32)
        d = rays[:, 3:].astype(np.float64)                             # directions as cast (float32)
        ans = scene_rc.cast_rays(o3d.core.Tensor(rays), nthreads=1)
        t_hit = ans['t_hit'].numpy().astype(np.float64)
        ids = ans['geometry_ids'].numpy().astype(np.int64)
        ids[~np.isfinite(t_hit)] = -1
        if train and int(self.cfg['boundary_erosion_px']) > 0:
            img = np.full(inside.shape, -1, dtype=np.int64)            # pixels outside all boxes are background
            img[inside] = ids
            img = img.reshape(y1 - y0, x1 - x0)
            keep = np.ones_like(img, dtype=bool)
            for _ in range(int(self.cfg['boundary_erosion_px'])):
                edge = np.zeros_like(keep)
                edge[1:] |= img[1:] != img[:-1]
                edge[:-1] |= img[:-1] != img[1:]
                edge[:, 1:] |= img[:, 1:] != img[:, :-1]
                edge[:, :-1] |= img[:, :-1] != img[:, 1:]
                keep &= ~edge
                img = np.where(keep, img, -2)
            ids = np.where(keep.reshape(-1)[inside], ids, -1)
        if train and float(self.cfg['occluder_prob']) > 0 and rng.random() < float(self.cfg['occluder_prob']):
            pu, pv = u.reshape(-1)[inside], v.reshape(-1)[inside]          # pixel centres of the cast rays
            for _ in range(int(rng.integers(1, 4))):                        # 1-3 discs in front of all objects
                cx, cy, r = rng.uniform(x0, x1), rng.uniform(y0, y1), rng.uniform(*self.cfg['occluder_radius']) * H
                ids[(pu - cx) ** 2 + (pv - cy) ** 2 < r * r] = -1
        if train and float(self.cfg['depth_noise_std']) > 0:
            hits = np.isfinite(t_hit)
            t_hit[hits] += rng.normal(size=int(hits.sum())) * float(self.cfg['depth_noise_std']) * t_hit[hits] ** 2
        seen, clouds = [], []
        for i in range(len(shapes)):
            hit = ids == i
            if int(hit.sum()) < self.min_points:
                continue
            Xw = rays[0, :3].astype(np.float64) + t_hit[hit, None] * d[hit]
            U = (Xw - positions[i]) @ Rq[i] / sizes[i]                 # object frame, asset units
            if U.shape[0] > 2 * self.max_points:                         # random subsample of dense segments before FPS
                U = U[np.sort(rng.choice(U.shape[0], 2 * self.max_points, replace=False))]
            seen.append(i)
            clouds.append(U)
        if clouds:
            rmp = self.cfg['random_max_points'] if train else None
            for i, U, order in zip(seen, clouds, farthest_point_orders(clouds, self.max_points)):
                if rmp is not None:                                         # an FPS prefix still covers the segment
                    order = order[:max(self.min_points, int(rng.integers(int(rmp[0]), int(rmp[1]) + 1)))]
                out[i] = U[order]
        if self.completer is not None and seen:
            pred = self.completer.complete([positions[i] + sizes[i] * out[i] @ Rq[i].T for i in seen], Rc[:3, 3])
            union = [np.concatenate([out[i], (Y - positions[i]) @ Rq[i] / sizes[i]]) for i, Y in zip(seen, pred)]
            for i, U, order in zip(seen, union, farthest_point_orders(union, self.max_points)):
                out[i] = U[order]
        return out

    def images(self, scene: str, shapes: Sequence[str], sizes: np.ndarray, positions: np.ndarray,
               quaternions: np.ndarray):
        """Depth image and segmentation of the evaluation camera, for display. depth (H, W): distance along the
        optical axis (m), including the floor (z = 0), inf where nothing is hit; segment (H, W): index of the object
        hit, -1 for floor and background. Same ray cast as observe(), over the whole image and without training
        augmentations."""
        import open3d as o3d
        K, Rc, H, W = self.camera(scene, False, np.random.default_rng(0))
        Rq = quat_to_rotmat(quaternions)
        scene_rc = o3d.t.geometry.RaycastingScene()
        for i, shape in enumerate(shapes):
            V, F = visual_mesh(self.objects_dir, shape)
            Vw = positions[i] + sizes[i] * V @ Rq[i].T
            scene_rc.add_triangles(o3d.core.Tensor(Vw.astype(np.float32)), o3d.core.Tensor(F.astype(np.uint32)))
        u, v = np.meshgrid(np.arange(W) + 0.5, np.arange(H) + 0.5)
        d = np.stack([u.reshape(-1), v.reshape(-1), np.ones(H * W)], -1) @ np.linalg.inv(K).T @ Rc[:3, :3].T
        d /= np.linalg.norm(d, axis=1, keepdims=True)
        rays = np.concatenate([np.broadcast_to(Rc[:3, 3], d.shape), d], axis=1).astype(np.float32)
        ans = scene_rc.cast_rays(o3d.core.Tensor(rays))
        t_hit = ans['t_hit'].numpy().astype(np.float64)
        segment = ans['geometry_ids'].numpy().astype(np.int64)
        segment[~np.isfinite(t_hit)] = -1
        down = d[:, 2] < -1e-9                                           # floor (z = 0) where no object is hit
        t_floor = np.where(down, -Rc[2, 3] / np.where(down, d[:, 2], -1.0), np.inf)
        t = np.where(np.isfinite(t_hit), t_hit, t_floor)
        forward = -Rc[:3, 2]                                             # Kubric cameras look along -z
        depth = t * (d @ forward)
        return depth.reshape(H, W), segment.reshape(H, W)