"""ctypes binding of the C++ mesh renderer (mesh_render.cpp; build with `make` in this folder).

A textured triangle mesh is ray-traced through a thin lens with radial distortion, a disk light,
soft shadows, a background plane and an optional rectangular occluder. Outputs: rgb, a sheet mask
and uv (texture pixel seen by every output pixel).
"""

import ctypes as C
import sys
from dataclasses import dataclass
from pathlib import Path

import cv2
import numpy as np

LIB = Path(__file__).resolve().parent / ("mesh_render.dll" if sys.platform.startswith("win") else "mesh_render.so")


@dataclass
class Optics:
    F: float = 50.0               # focal distance
    L: float = 66.7               # lens-sensor distance
    R: float = 10.0               # aperture radius


@dataclass
class Distortion:
    use: bool = False
    k1: float = 0.0
    cx: float = 0.0
    cy: float = 0.0


@dataclass
class Lighting:
    use: bool = False
    x: float = 0.0
    y: float = -100.0
    z: float = -100.0
    intensity: float = 1.0
    diameter: float = 20.0
    light_mix_koef: float = 0.8


@dataclass
class BackgroundShadow:
    use: bool = False
    bg_z: float = 0.0
    bottom_shadow_koef: float = 0.0


@dataclass
class RectOccluder:
    use: bool = False             # not read by the renderer: disable with a far, empty occluder
    cx: float = 0.0
    cy: float = 0.0
    cz: float = 30.0
    w: float = 20.0
    h: float = 20.0
    yaw: float = 0.0
    pitch: float = 0.0
    roll: float = 0.0
    circle_segments: int = 32


class Vector3d(C.Structure):
    _fields_ = [("x", C.c_float), ("y", C.c_float), ("z", C.c_float)]


class Vector2d(C.Structure):
    _fields_ = [("x", C.c_float), ("y", C.c_float)]


class Triangle(C.Structure):
    _fields_ = [("v0", C.POINTER(Vector3d)), ("v1", C.POINTER(Vector3d)), ("v2", C.POINTER(Vector3d)),
                ("imagePoint0", C.POINTER(Vector2d)), ("imagePoint1", C.POINTER(Vector2d)),
                ("imagePoint2", C.POINTER(Vector2d))]


class Mesh(C.Structure):
    _fields_ = [
        ("triangles", C.POINTER(Triangle)), ("points3d", C.POINTER(Vector3d)), ("points2d", C.POINTER(Vector2d)),
        ("triangleCount", C.c_int), ("pointCount", C.c_int),
        ("use_light_info", C.c_bool), ("light_x", C.c_float), ("light_y", C.c_float), ("light_z", C.c_float),
        ("light_intensivity", C.c_float), ("use_shadow_info", C.c_bool), ("light_diameter", C.c_float),
        ("use_bg_shadow", C.c_bool), ("bg_z", C.c_float),
        ("occ_cx", C.c_float), ("occ_cy", C.c_float), ("occ_cz", C.c_float), ("occ_w", C.c_float),
        ("occ_h", C.c_float), ("occ_yaw", C.c_float), ("occ_pitch", C.c_float), ("occ_roll", C.c_float),
        ("occ_circle_segments", C.c_int),
        ("camera_tilt_x_rad", C.c_float), ("use_distortion_k1", C.c_bool), ("k1", C.c_float),
        ("dist_norm", C.c_float), ("cx", C.c_float), ("cy", C.c_float),
        ("bottom_shadow_koef", C.c_float), ("light_mix_koef", C.c_float),
    ]


def _load():
    dll = C.CDLL(str(LIB))
    dll.create_mesh.argtypes = [C.c_int, C.c_int, C.c_int, C.c_int]
    dll.create_mesh.restype = C.POINTER(Mesh)
    dll.render.argtypes = [C.c_void_p, C.c_int, C.c_int, C.c_int, C.POINTER(Mesh), C.c_void_p, C.c_int, C.c_int,
                           C.c_int, C.c_bool, C.c_void_p, C.c_int, C.c_bool, C.c_void_p,
                           C.c_float, C.c_float, C.c_float]
    dll.render.restype = C.c_int
    dll.delete_mesh.argtypes = [C.POINTER(Mesh)]
    dll.delete_mesh.restype = None
    dll.mesh_set_rect_occluder.argtypes = [C.POINTER(Mesh)] + [C.c_float] * 8 + [C.c_int]
    dll.mesh_set_rect_occluder.restype = None
    return dll


class MeshRenderer:
    """A (grid_w + 1) x (grid_h + 1) node mesh over an input_width x input_height texture.
    points2d / points3d: numpy views of the node texture coordinates and 3-D positions (camera frame)."""

    def __init__(self, input_width, input_height, grid_w, grid_h):
        self.dll = _load()
        self.mesh = self.dll.create_mesh(int(input_width), int(input_height), int(grid_w), int(grid_h))
        if not self.mesh:
            raise RuntimeError("create_mesh returned NULL")
        m = self.mesh.contents
        self.points2d = np.ctypeslib.as_array(C.cast(m.points2d, C.POINTER(C.c_float)), shape=(m.pointCount, 2))
        self.points3d = np.ctypeslib.as_array(C.cast(m.points3d, C.POINTER(C.c_float)), shape=(m.pointCount, 3))

    def configure(self, optics, distortion, lighting, bg_shadow, occluder):
        self.optics = optics
        m = self.mesh.contents
        m.use_distortion_k1 = bool(distortion.use)
        m.k1 = float(distortion.k1)
        m.dist_norm = float(optics.L)
        m.cx, m.cy = float(distortion.cx), float(distortion.cy)
        m.use_light_info = bool(lighting.use)
        m.light_x, m.light_y, m.light_z = float(lighting.x), float(lighting.y), float(lighting.z)
        m.light_intensivity = float(lighting.intensity)
        m.light_diameter = float(lighting.diameter)
        m.light_mix_koef = float(lighting.light_mix_koef)
        m.use_bg_shadow = bool(bg_shadow.use)
        m.bg_z = float(bg_shadow.bg_z)
        m.bottom_shadow_koef = float(bg_shadow.bottom_shadow_koef)
        m.use_shadow_info = bool(bg_shadow.use)
        m.camera_tilt_x_rad = 0.0
        o = occluder
        self.dll.mesh_set_rect_occluder(self.mesh, float(o.cx), float(o.cy), float(o.cz), float(o.w), float(o.h),
                                        float(o.yaw), float(o.pitch), float(o.roll), int(o.circle_segments))

    def render(self, texture, out_size, background):
        """texture uint8 (H, W, 3), background BGR uint8 of out_size (W, H).
        -> rgb uint8 (H, W, 3), mask uint8 (H, W) 0/255, uv float32 (H, W, 2) in texture px."""
        w, h = int(out_size[0]), int(out_size[1])
        rgba = np.zeros((h, w * 4), np.uint8)
        rgba.reshape(h, w, 4)[...] = cv2.cvtColor(background, cv2.COLOR_BGR2RGBA)
        mask = np.zeros((h, w), np.uint8)
        uv = np.zeros((h, w, 2), np.float32)
        ret = self.dll.render(texture.ctypes.data_as(C.c_void_p), int(texture.shape[1]), int(texture.strides[0]),
                              int(texture.shape[0]), self.mesh, rgba.ctypes.data_as(C.c_void_p), w, w * 4, h,
                              True, mask.ctypes.data_as(C.c_void_p), int(mask.strides[0] // mask.itemsize),
                              True, uv.ctypes.data_as(C.c_void_p),
                              C.c_float(self.optics.R), C.c_float(self.optics.L), C.c_float(self.optics.F))
        if ret != 0:
            raise RuntimeError("render() failed with code {}".format(ret))
        rgb = cv2.flip(np.ascontiguousarray(rgba.reshape(h, w, 4)[..., :3]), -1)
        return rgb, mask[::-1, ::-1], cv2.flip(uv, -1)

    def close(self):
        if self.mesh:
            self.dll.delete_mesh(self.mesh)
            self.mesh = None