#!/usr/bin/env python3
import os
import math
import numpy as np
import cv2
import tifffile

import mitsuba as mi
import drjit as dr
mi.set_variant("llvm_ad_spectral")

def wavelen_to_wavenum(wavelength):
    return dr.two_pi / (wavelength * 1e-3)

def poisson_pmf(g, m):
    return dr.exp(-g + m * dr.log(dr.maximum(g, 1e-20)) - dr.lgamma(m + 1.0))

def sample_poisson(g, u, max_m=64):
    max_m = g + 3+dr.sqrt(g)
    g = dr.maximum(g, 0.0)
    u = dr.clip(u, 0.0, 1.0 - 1e-7)

    m   = mi.UInt32(0)
    p   = dr.exp(-g)
    cdf = p

    def cond(m, p, cdf):
        return (m < mi.UInt32(max_m)) & (cdf < u)

    def body(m, p, cdf):
        m_next = m + mi.UInt32(1)
        p      = p * g / mi.Float(m_next)
        cdf    = cdf + p
        return m_next, p, cdf

    m, p, cdf = dr.while_loop((m, p, cdf), cond, body)
    return dr.minimum(m, mi.UInt32(max_m))

def fresnel_conductor_spectral(cos_i, eta, k):
    cos2 = dr.square(cos_i); sin2 = 1.0 - cos2
    eta2, k2 = dr.square(eta), dr.square(k)
    t0 = eta2 - k2 - sin2
    a2pb2 = dr.sqrt(dr.square(t0) + 4.0 * eta2 * k2)
    a = dr.sqrt(0.5 * (a2pb2 + t0))
    t1, t2 = a2pb2 + cos2, 2.0 * cos_i * a
    Rs = (t1 - t2) / (t1 + t2)
    t3, t4 = cos2 * a2pb2 + dr.square(sin2), t2 * sin2
    Rp = Rs * (t3 - t4) / (t3 + t4)
    return 0.5 * (Rs + Rp)

class Fractal:
    def __init__(self, sigmah, T, gamma, T_exp, gamma_exp):
        self.sigmah2 = dr.square(sigmah)
        self.T = T
        self.gamma = gamma
        self.T_exp = T_exp
        self.gamma_exp = gamma_exp

    def _T_m(self, m):
        m = dr.maximum(mi.Float(m), 1.0)
        return dr.select(m == 1.0, self.T, self.T * dr.power(m, self.T_exp))

    def _gamma_m(self, m):
        m = dr.maximum(mi.Float(m), 1.0)
        return dr.select(m == 1.0, self.gamma, self.gamma * dr.power(m, self.gamma_exp))

    def _sigma2_normalized(self, k, m):
        T, gamma = self._T_m(m), self._gamma_m(m)
        x = 1 + k*k * T
        p = dr.power(x, (gamma-1)/2)
        return 1/(1 - (1/p))

    def psd(self, wi: mi.Vector3f, wo: mi.Vector3f, k: float, m):
        T, gamma = self._T_m(m), self._gamma_m(m)
        z = (mi.Vector2f(wi.x, wi.y) + mi.Vector2f(wo.x, wo.y)) * k
        x = 1 + T * dr.dot(z, z)
        p = dr.power(x, (gamma+1)/2)
        f = 1/p
        sigma2_norm = self._sigma2_normalized(k, m)
        return sigma2_norm * dr.inv_two_pi * k*k * (gamma - 1) * T * f

    def alpha(self, wi: mi.Vector3f, wo: mi.Vector3f, k: float):
        a = dr.square((dr.abs(wi.z) + dr.abs(wo.z)) * k) * self.sigmah2
        return dr.exp(-a)

    def sample(self, wi: mi.Vector3f, k: mi.Float, u2: mi.Point2f, m):
        gamma, T = self._gamma_m(m), self._T_m(m)
        s = dr.sqrt(dr.maximum(0.0, 1.0 - dr.square(wi.z)))
        phi_i = dr.select(s > 0.0, dr.atan2(wi.y, wi.x), mi.Float(0.0))
        sqrtT = dr.sqrt(T)
        k2T = (k * k) * T
        M = 1.0 - dr.power(1.0 + k2T * dr.square(1.0 + s), -(gamma - 1.0) * 0.5)
        inner = dr.maximum(dr.power(1.0 - M * u2.x, -2.0 / (gamma - 1.0)) - 1.0, 0.0)
        f = dr.select(sqrtT > 0.0, dr.sqrt(inner) / sqrtT, mi.Float(0.0))
        f_k = dr.select(k != 0.0, f / k, mi.Float(0.0))
        num = dr.square(f_k) + dr.square(s) - 1.0
        EPS = 1e-8
        denom = 2.0 * f_k * s
        mask  = dr.abs(denom) > EPS
        denom_safe = dr.select(mask, denom, mi.Float(1.0))
        cos_arg = num / denom_safe
        cos_arg = dr.select(mask, cos_arg, mi.Float(1.0))
        cos_arg = dr.clip(cos_arg, -1.0, 1.0)
        phi_max = dr.select((f == 0.0) | (s == 0.0), mi.Float(dr.pi), dr.acos(cos_arg))
        phi_f = phi_i + (2.0 * u2.y - 1.0) * phi_max
        zeta_k_x = dr.select(k != 0.0, (f * dr.cos(phi_f)) / k, mi.Float(0.0))
        zeta_k_y = dr.select(k != 0.0, (f * dr.sin(phi_f)) / k, mi.Float(0.0))
        wo_x, wo_y = zeta_k_x - wi.x, zeta_k_y - wi.y
        r2 = wo_x * wo_x + wo_y * wo_y
        z = dr.sqrt(dr.maximum(0.0, 1.0 - r2))
        wo_z = dr.select(wi.z >= 0.0, z, -z)
        wo = mi.Vector3f(wo_x, wo_y, wo_z)
        psd_val = self.psd(wi, wo, k, m)
        w   = mi.Float(dr.inv_pi) * phi_max
        w_s = dr.maximum(w, EPS)
        pdf_raw = (z * psd_val) / w_s
        pdf     = dr.select(w > 1e-2, pdf_raw, mi.Float(0.0))
        return wo, pdf, psd_val, w

class FullGHS(mi.BSDF):
    def __init__(self, props):
        super().__init__(props)
        self.eta   = props.get('eta',   mi.load_dict({"type": "uniform", "value": 0.0}))
        self.kappa = props.get('kappa', mi.load_dict({"type": "uniform", "value": 1.0}))
        self.t      = props.get('t', mi.Float(1e-12))
        self.sigmah = props.get('sigmah', mi.Float(1e-9))
        self.gamma  = props.get('gamma', 3.0)
        T_exp      = props.get('T_exp', mi.Float(-0.8))
        gamma_exp = props.get('gamma_exp', mi.Float(0.05))
        self.profile = Fractal(self.sigmah, self.t, self.gamma, T_exp, gamma_exp)
        self.m_flags = (mi.BSDFFlags.DeltaReflection |
                        mi.BSDFFlags.GlossyReflection |
                        mi.BSDFFlags.FrontSide)

    def flags(self, _=mi.BSDFFlags.All, **__):
        return self.m_flags

    def sample(self, ctx, si, sample1, sample2, active):
        wi = si.wi
        eta  = self.eta.eval(si, active)
        kappa = self.kappa.eval(si, active)
        cos_theta_i = mi.Frame3f.cos_theta(wi)
        k = wavelen_to_wavenum(si.wavelengths)
        hero_k = k[0]  # Wavelengths are randomly shuffled

        g = dr.square(hero_k * self.sigmah * (cos_theta_i + cos_theta_i))  # could undersample for diffuse

        m = sample_poisson(g, sample1)
        alpha = dr.exp(-dr.square(k * self.sigmah * (cos_theta_i + cos_theta_i)))
        pdf_specular = dr.exp(-g)
        is_specular = (m < 1) & active
        bs = mi.BSDFSample3f(); bs.eta = 1.0
        wo_spec   = mi.reflect(wi)
        weight_spec = alpha * fresnel_conductor_spectral(cos_theta_i, eta, kappa)

        wo_scatter, pdf_scatter, psd_scatter, _ = self.profile.sample(wi, hero_k, sample2, m)
        pdf_scatter  *= (1.0 - pdf_specular)
        h = wi + wo_scatter
        h = dr.normalize(h * dr.sign(wi.z))
        F_scatter = fresnel_conductor_spectral(dr.dot(wi, h), eta, kappa)
        weight_scatter = (1-alpha) * dr.abs(wo_scatter.z) * psd_scatter * F_scatter

        bs.wo  = dr.select(is_specular, wo_spec, wo_scatter)
        bs.pdf = dr.select(is_specular, pdf_specular, pdf_scatter)
        bs.sampled_type = dr.select(
            is_specular,
            mi.UInt32(+mi.BSDFFlags.DeltaReflection),
            mi.UInt32(+mi.BSDFFlags.GlossyReflection),
        )
        bs.sampled_component = mi.UInt32(0)

        EPS = 1e-8
        value = dr.select(is_specular, weight_spec, weight_scatter)
        value  = value / dr.maximum(bs.pdf, EPS)
        return bs, value

    def eval(self, ctx, si, wo, active):
        cos_o = mi.Frame3f.cos_theta(wo)
        cos_i = mi.Frame3f.cos_theta(si.wi)
        active = active & (cos_o > 0)
        k = wavelen_to_wavenum(si.wavelengths)
        hero_k = k[0]
        eta   = self.eta.eval(si, active)
        kappa = self.kappa.eval(si, active)
        alpha = self.profile.alpha(si.wi, wo, k)
        g = dr.square(hero_k * self.sigmah * (cos_i + cos_o))

        # Loop through 99% of weights
        sd3    = 3.0 * dr.sqrt(dr.maximum(g, 0.0))
        min_mf = dr.maximum(1.0, g - sd3)
        max_mf = dr.minimum(g + sd3, 128.0)
        state = (
            mi.UInt32(0),
            dr.zeros(mi.Float, dr.width(g))
        )

        def cond(m, psd_scatter):
            return m < mi.UInt32(128)

        def body(m, psd_scatter):
            m_f      = mi.Float(m)
            active_m = (m_f >= min_mf) & (m_f <= max_mf)
            w_m      = dr.select(active_m, poisson_pmf(g, m), 0.0)
            psd_m = self.profile.psd(si.wi, wo, hero_k, m=m)
            psd_scatter += w_m * psd_m
            return m + mi.UInt32(1), psd_scatter

        _, psd_scatter = dr.while_loop(state, cond, body)

        F = fresnel_conductor_spectral(cos_i, eta, kappa)
        f_s = (1.0 - alpha) * dr.abs(cos_o) * psd_scatter * F
        out = dr.select(active, f_s, mi.Spectrum(0))
        out = dr.select(dr.isfinite(out), out, mi.Spectrum(0))
        return out

    def pdf(self, ctx, si, wo, active):
        cos_o = mi.Frame3f.cos_theta(wo)
        cos_i = mi.Frame3f.cos_theta(si.wi)
        active = active & (cos_o > 0)
        k = wavelen_to_wavenum(si.wavelengths)
        hero_k = k[0]
        alpha = self.profile.alpha(si.wi, wo, hero_k)
        g = dr.square(hero_k * self.sigmah * (cos_i + cos_o))

        # Loop through 99% of weights
        sd3    = 3.0 * dr.sqrt(dr.maximum(g, 0.0))
        min_mf = dr.maximum(1.0, g - sd3)
        max_mf = dr.minimum(g + sd3, 128.0)
        state = (
            mi.UInt32(0),
            dr.zeros(mi.Float, dr.width(g))
        )

        def cond(m, psd_scatter):
            return m < mi.UInt32(128)

        def body(m, psd_scatter):
            m_f      = mi.Float(m)
            active_m = (m_f >= min_mf) & (m_f <= max_mf)
            w_m      = dr.select(active_m, poisson_pmf(g, m), 0.0)
            psd_m = self.profile.psd(si.wi, wo, hero_k, m=m)
            psd_scatter += w_m * psd_m
            return m + mi.UInt32(1), psd_scatter

        _, psd_scatter = dr.while_loop(state, cond, body)

        f_s = (1.0 - alpha) * dr.abs(cos_o) * psd_scatter
        out = dr.select(active, f_s, 0.0)
        out = dr.select(dr.isfinite(out), out, 0.0)
        return out


mi.register_bsdf("fullghs", FullGHS)

ARUCO_DICTS = {
    "DICT_4X4_50": cv2.aruco.DICT_4X4_50,
    "DICT_5X5_50": cv2.aruco.DICT_5X5_50,
    "DICT_6X6_50": cv2.aruco.DICT_6X6_50,
    "DICT_ARUCO_ORIGINAL": cv2.aruco.DICT_ARUCO_ORIGINAL,
}

def camera_from_aruco(image_path: str,
                      focal_mm: float,
                      sensor_w_mm: float,
                      sensor_h_mm: float,
                      marker_inches: float,
                      aruco_dict_name: str,
                      expected_ids: set[int] | None,
                      assume_zero_distortion: bool = True):
    img = cv2.imread(image_path, cv2.IMREAD_COLOR)
    if img is None:
        raise FileNotFoundError(image_path)
    h, w = img.shape[:2]

    fx = (focal_mm / sensor_w_mm) * w
    fy = (focal_mm / sensor_h_mm) * h
    cx, cy = w * 0.5, h * 0.5
    K  = np.array([[fx, 0, cx],
                   [ 0, fy, cy],
                   [ 0,  0,  1]], dtype=np.float64)
    dist = np.zeros(5) if assume_zero_distortion else np.array([0,0,0,0,0], np.float64)

    dictionary = cv2.aruco.getPredefinedDictionary(ARUCO_DICTS[aruco_dict_name])
    params = cv2.aruco.DetectorParameters()
    params.adaptiveThreshWinSizeMin = 3
    params.adaptiveThreshWinSizeMax = 50
    params.adaptiveThreshWinSizeStep = 5
    params.adaptiveThreshConstant   = 7
    params.minMarkerPerimeterRate   = 0.02
    params.maxMarkerPerimeterRate   = 4.0
    params.cornerRefinementMethod   = cv2.aruco.CORNER_REFINE_SUBPIX

    detector = cv2.aruco.ArucoDetector(dictionary, params)
    corners, ids, _ = detector.detectMarkers(img)
    if ids is None or len(ids) == 0:
        raise RuntimeError("No ArUco markers detected.")

    ids_list = ids.flatten().tolist()
    if expected_ids and (expected_ids & set(ids_list)):
        use_id = list(expected_ids & set(ids_list))[0]
        i = ids_list.index(use_id)
    else:
        # fallback: first detected
        i = 0

    marker_m = marker_inches * 0.0254
    rvecs, tvecs, _ = cv2.aruco.estimatePoseSingleMarkers(corners, marker_m, K, dist)
    rvec, tvec = rvecs[i].reshape(3,1), tvecs[i].reshape(3,1)

    half = marker_m * 0.5
    obj_corners = np.array([[-half,  half, 0.0],
                            [ half,  half, 0.0],
                            [ half, -half, 0.0],
                            [-half, -half, 0.0]], dtype=np.float64)
    img_corners = corners[i][0].astype(np.float64)
    ok, rvec_ref, tvec_ref = cv2.solvePnP(obj_corners, img_corners, K, dist,
                                          rvec, tvec, useExtrinsicGuess=True,
                                          flags=cv2.SOLVEPNP_ITERATIVE)
    if ok:
        rvec, tvec = rvec_ref, tvec_ref

    R_cam_mark, _ = cv2.Rodrigues(rvec)
    t_cam_mark    = tvec
    return K, R_cam_mark, t_cam_mark, h, w

def np_to_mi(T44: np.ndarray):
    return mi.ScalarTransform4f(T44.astype(np.float32))

def build_scene_dict(K, R_cam_mark, t_cam_mark, img_w, img_h,
                     res_scale: float,
                     spp: int,
                     ior_eta_path: str,
                     ior_k_path: str,
                     t_val: float,
                     sigmah: float,
                     gamma: float,
                     T_exp: float,
                     gamma_exp: float,
                     Angle: str):
    fx = float(K[0, 0])
    fov_x_deg = 2.0 * math.degrees(math.atan((img_w * 0.5) / max(1e-9, fx)))
    fov_x_deg = 2.0 * math.degrees(math.atan((img_w * 0.5) / max(1e-9, fx)))

    Rcw = np.asarray(R_cam_mark, dtype=np.float64)
    tcw = np.asarray(t_cam_mark).reshape(3,1)
    Rwc, twc = Rcw.T, -Rcw.T @ tcw

    A = np.diag([1.0, -1.0, 1.0])
    R_world_from_mi = Rwc @ A
    T_cam = np.eye(4)
    T_cam[:3,:3] = R_world_from_mi
    T_cam[:3, 3] = twc[:,0]

    mm_to_m = 1e-3
    inches_to_mm = 25.4
    ground_half = 5.0

    L1, W1, H1 = 6*inches_to_mm*mm_to_m, 6*inches_to_mm*mm_to_m, 1*mm_to_m

    if Angle == "grazing":
        sample1_cy_mm = 7 * inches_to_mm
    elif Angle == "medium":
        sample1_cy_mm = -1 * inches_to_mm
    else:
        sample1_cy_mm = 0 * inches_to_mm

    T_sample1 = np.diag([L1*0.5, W1*0.5, H1*0.5, 1.0])
    T_sample1[:3, 3] = [
        0.0,
        sample1_cy_mm * mm_to_m,
        H1 * 0.5
    ]

    sample2_len_mm = 6 * inches_to_mm
    sample2_wid_mm = 5
    sample2_hgt_mm = 8 * inches_to_mm
    sample2_cx_mm  = 0.0
    if Angle == "grazing":
        sample2_cy_mm  = 13*inches_to_mm + 6 * inches_to_mm / 2 
    else:
        sample2_cy_mm  = 6 * inches_to_mm / 2
    sample2_cz_mm  = sample2_hgt_mm / 2

    L2 = sample2_len_mm * mm_to_m
    W2 = sample2_wid_mm * mm_to_m
    H2 = sample2_hgt_mm * mm_to_m
    CX = sample2_cx_mm * mm_to_m
    CY = sample2_cy_mm * mm_to_m
    CZ = sample2_cz_mm * mm_to_m

    T_sample2 = np.eye(4, dtype=np.float64)
    T_sample2[0, 0] = L2 * 0.5
    T_sample2[1, 1] = W2 * 0.5
    T_sample2[2, 2] = H2 * 0.5
    T_sample2[:3, 3] = [CX, CY, CZ]

    def Rx(deg):
        th = math.radians(deg)
        c, s = math.cos(th), math.sin(th)
        R = np.array([[1,0,0],[0,c,-s],[0,s,c]], dtype=np.float64)
        return R
    z_eps = 1e-6
    S_face = np.diag([L2*0.5, H2*0.5, 1.0, 1.0])

    R_face_negY = np.eye(4, dtype=np.float64); R_face_negY[:3,:3] = Rx(+90)
    T_face_negY = np.eye(4, dtype=np.float64); T_face_negY[:3, 3] = [CX, CY - (W2 * 0.5) - z_eps, CZ]
    T_negY = T_face_negY @ R_face_negY @ S_face

    # Light position (world coords)
    light_pos = [12*inches_to_mm*mm_to_m, -9*inches_to_mm*mm_to_m, 9.5*inches_to_mm*mm_to_m]

    # Resolution
    low_w = 1024
    low_h = int(1024 * 1.5)

    scene = {
        "type": "scene",
        "integrator": {"type": "path", "max_depth": 3},
        "sensor": {
            "type": "perspective",
            "to_world": np_to_mi(T_cam),
            "fov_axis": "x", "fov": fov_x_deg,
            "film": {"type":"hdrfilm","width":low_w,"height":low_h,
                     "rfilter":{"type":"gaussian"},
                     "pixel_format":"rgb"},
            "sampler": {"type":"multijitter","sample_count": 64}
        },
        "light": {
            "type": "point", "position": light_pos,
            "intensity": {"type":"blackbody","temperature":2700.0}
        },
        "ground": {
            "type":"rectangle",
            "to_world": np_to_mi(np.array([
                [ground_half,0,0,0],[0,ground_half,0,0],[0,0,1,0],[0,0,0,1]
            ],dtype=np.float64)),
            "bsdf":{"type":"diffuse","reflectance":{"type":"rgb","value":[0.7,0.7,0.8]}}
        },
        "sample_box": {
            "type":"cube",
            "to_world": np_to_mi(T_sample1),
            "bsdf":{
                "type":"fullghs",
                "t": t_val,
                "sigmah": sigmah,
                "gamma": gamma,
                "T_exp": T_exp,
                "gamma_exp": gamma_exp,
                "eta": {"type":"spectrum","filename": ior_eta_path},
                "kappa": {"type":"spectrum","filename": ior_k_path}
            }
        },
        "checker_board": {
            "type": "cube",
            "to_world": np_to_mi(T_sample2),
            "bsdf": { "type": "diffuse",
                      "reflectance": { "type": "rgb", "value": [0.0, 0.0, 0.0] } }
        },
        "checker_side_negY": {
            "type": "rectangle",
            "to_world": np_to_mi(T_negY),
            "bsdf": {
                "type": "diffuse",
                "reflectance": {
                    "type": "bitmap",
                    "filename": "checker_6x8.png",
                    "wrap_mode": "clamp",
                    "filter_type": "nearest"
                }
            }
        }
    }
    return scene

def spectrum_file(path):
    return {"type": "spectrum", "filename": path}

def build_scene_dict_ggx(K, R_cam_mark, t_cam_mark, img_w, img_h,
                     res_scale: float,
                     spp: int,
                     ior_eta_path: str,
                     ior_k_path: str,
                     alpha: float,
                     Angle: str):

    fx = float(K[0, 0])
    fov_x_deg = 2.0 * math.degrees(math.atan((img_w * 0.5) / max(1e-9, fx)))
    fov_x_deg = 2.0 * math.degrees(math.atan((img_w * 0.5) / max(1e-9, fx)))

    Rcw = np.asarray(R_cam_mark, dtype=np.float64)
    tcw = np.asarray(t_cam_mark).reshape(3,1)
    Rwc, twc = Rcw.T, -Rcw.T @ tcw

    A = np.diag([1.0, -1.0, 1.0])
    R_world_from_mi = Rwc @ A
    T_cam = np.eye(4)
    T_cam[:3,:3] = R_world_from_mi
    T_cam[:3, 3] = twc[:,0]

    mm_to_m = 1e-3
    inches_to_mm = 25.4
    ground_half = 5.0

    L1, W1, H1 = 6*inches_to_mm*mm_to_m, 6*inches_to_mm*mm_to_m, 1*mm_to_m

    # move it forward by 7 inches
    if Angle == "grazing":
        sample1_cy_mm = 7 * inches_to_mm
    elif Angle == "medium":
        sample1_cy_mm = -1 * inches_to_mm
    else:
        sample1_cy_mm = 0 * inches_to_mm

    T_sample1 = np.diag([L1*0.5, W1*0.5, H1*0.5, 1.0])
    T_sample1[:3, 3] = [
        0.0,
        sample1_cy_mm * mm_to_m,
        H1 * 0.5
    ]

    sample2_len_mm = 6 * inches_to_mm
    sample2_wid_mm = 5
    sample2_hgt_mm = 8 * inches_to_mm
    sample2_cx_mm  = 0.0
    if Angle == "grazing":
        sample2_cy_mm  = 13*inches_to_mm + 6 * inches_to_mm / 2 
    else:
        sample2_cy_mm  = 6 * inches_to_mm / 2
    sample2_cz_mm  = sample2_hgt_mm / 2

    L2 = sample2_len_mm * mm_to_m
    W2 = sample2_wid_mm * mm_to_m
    H2 = sample2_hgt_mm * mm_to_m
    CX = sample2_cx_mm * mm_to_m
    CY = sample2_cy_mm * mm_to_m
    CZ = sample2_cz_mm * mm_to_m

    T_sample2 = np.eye(4, dtype=np.float64)
    T_sample2[0, 0] = L2 * 0.5
    T_sample2[1, 1] = W2 * 0.5
    T_sample2[2, 2] = H2 * 0.5
    T_sample2[:3, 3] = [CX, CY, CZ]

    def Rx(deg):
        th = math.radians(deg)
        c, s = math.cos(th), math.sin(th)
        R = np.array([[1,0,0],[0,c,-s],[0,s,c]], dtype=np.float64)
        return R
    z_eps = 1e-6
    S_face = np.diag([L2*0.5, H2*0.5, 1.0, 1.0])

    R_face_negY = np.eye(4, dtype=np.float64); R_face_negY[:3,:3] = Rx(+90)
    T_face_negY = np.eye(4, dtype=np.float64); T_face_negY[:3, 3] = [CX, CY - (W2 * 0.5) - z_eps, CZ]
    T_negY = T_face_negY @ R_face_negY @ S_face

    light_pos = [12*inches_to_mm*mm_to_m, -9*inches_to_mm*mm_to_m, 9.5*inches_to_mm*mm_to_m]

    # Resolution
    low_w = 1024
    low_h = int(1024 * 1.5)

    scene = {
        "type": "scene",
        "integrator": {"type": "path", "max_depth": 3},
        "sensor": {
            "type": "perspective",
            "to_world": np_to_mi(T_cam),
            "fov_axis": "x", "fov": fov_x_deg,
            "film": {"type":"hdrfilm","width":low_w,"height":low_h,
                     "rfilter":{"type":"gaussian"},
                     "pixel_format":"rgb"},
            "sampler": {"type":"multijitter","sample_count": 64}
        },
        "light": {
            "type": "point", "position": light_pos,
            "intensity": {"type":"blackbody","temperature":2700.0}
        },
        "ground": {
            "type":"rectangle",
            "to_world": np_to_mi(np.array([
                [ground_half,0,0,0],[0,ground_half,0,0],[0,0,1,0],[0,0,0,1]
            ],dtype=np.float64)),
            "bsdf":{"type":"diffuse","reflectance":{"type":"rgb","value":[0.7,0.7,0.8]}}
        },
        "sample_box": {
            "type":"cube",
            "to_world": np_to_mi(T_sample1),
            "bsdf":{
                "type":"roughconductor", "distribution":"ggx",
                "alpha": float(alpha), "sample_visible": True,
                "eta": spectrum_file(ior_eta_path),   # spectrum from ior/<mat>_eta.txt
                "k":   spectrum_file(ior_k_path)      # spectrum from ior/<mat>_k.txt
            }
        },
        "checker_board": {
            "type": "cube",
            "to_world": np_to_mi(T_sample2),
            "bsdf": { "type": "diffuse",
                      "reflectance": { "type": "rgb", "value": [0.0, 0.0, 0.0] } }
        },
        "checker_side_negY": {
            "type": "rectangle",
            "to_world": np_to_mi(T_negY),
            "bsdf": {
                "type": "diffuse",
                "reflectance": {
                    "type": "bitmap",
                    "filename": "checker_6x8.png",
                    "wrap_mode": "clamp",
                    "filter_type": "nearest"
                }
            }
        }
    }
    return scene

def build_cc_scene_dict(K, R_cam_mark, t_cam_mark, img_w, img_h,
                        res_scale: float,
                        spp: int,
                        cc_png_path: str,
                        cc_len_in: float = 6.0,
                        cc_wid_in: float = 4.0,
                        cc_thk_mm: float = 2.0):
    """Colour-checker-only scene (plus ground + checkerboard), matching your reference cell."""

    fx = float(K[0, 0])
    fov_x_deg = 2.0 * math.degrees(math.atan((img_w * 0.5) / max(1e-9, fx)))

    Rcw = np.asarray(R_cam_mark, dtype=np.float64)
    tcw = np.asarray(t_cam_mark).reshape(3, 1)
    Rwc, twc = Rcw.T, -Rcw.T @ tcw

    A = np.diag([1.0, -1.0, 1.0])
    R_world_from_mi = Rwc @ A
    T_cam = np.eye(4, dtype=np.float64)
    T_cam[:3, :3] = R_world_from_mi
    T_cam[:3, 3]  = twc[:, 0]

    mm_to_m      = 1e-3
    inches_to_mm = 25.4
    ground_half  = 5.0

    L = (cc_len_in * inches_to_mm) * mm_to_m
    W = (cc_wid_in * inches_to_mm) * mm_to_m
    H = (cc_thk_mm) * mm_to_m

    T_panel = np.eye(4, dtype=np.float64)
    T_panel[0, 0] = L * 0.5
    T_panel[1, 1] = W * 0.5
    T_panel[2, 2] = 1.0
    T_panel[2, 3] = H * 0.5

    sample2_len_mm = 6 * inches_to_mm
    sample2_wid_mm = 5
    sample2_hgt_mm = 8 * inches_to_mm
    sample2_cx_mm  = 0.0
    sample2_cy_mm  = 6 * inches_to_mm / 2
    sample2_cz_mm  = sample2_hgt_mm / 2

    L2 = sample2_len_mm * mm_to_m
    W2 = sample2_wid_mm * mm_to_m
    H2 = sample2_hgt_mm * mm_to_m
    CX = sample2_cx_mm * mm_to_m
    CY = sample2_cy_mm * mm_to_m
    CZ = sample2_cz_mm * mm_to_m

    T_sample2 = np.eye(4, dtype=np.float64)
    T_sample2[0, 0] = L2 * 0.5
    T_sample2[1, 1] = W2 * 0.5
    T_sample2[2, 2] = H2 * 0.5
    T_sample2[:3, 3] = [CX, CY, CZ]

    def Rx(deg):
        th = math.radians(deg)
        c, s = math.cos(th), math.sin(th)
        return np.array([[1,0,0],[0,c,-s],[0,s,c]], dtype=np.float64)

    z_eps = 1e-6
    S_face = np.diag([L2*0.5, H2*0.5, 1.0, 1.0])

    R_face_negY = np.eye(4, dtype=np.float64); R_face_negY[:3, :3] = Rx(+90)
    T_face_negY = np.eye(4, dtype=np.float64); T_face_negY[:3, 3] = [CX, CY - (W2 * 0.5) - z_eps, CZ]
    T_negY = T_face_negY @ R_face_negY @ S_face

    light_pos = [
        12 * inches_to_mm * mm_to_m,
        -9 * inches_to_mm * mm_to_m,
        9.5 * inches_to_mm * mm_to_m
    ]

    # Resolution
    low_w = 1024
    low_h = int(1024 * 1.5)

    scene = {
        "type": "scene",
        "integrator": {"type": "path", "max_depth": 3},
        "sensor": {
            "type": "perspective",
            "to_world": np_to_mi(T_cam),
            "fov_axis": "x", "fov": fov_x_deg,
            "film": {"type": "hdrfilm", "width": low_w, "height": low_h,
                     "rfilter": {"type": "gaussian"},
                     "pixel_format": "rgb"},
            "sampler": {"type":"multijitter","sample_count": 64}
        },

        "light": {
            "type": "point",
            "position": light_pos,
            "intensity": {"type": "blackbody", "temperature": 2700.0}
        },

        "ground": {
            "type": "rectangle",
            "to_world": np_to_mi(np.array([
                [ground_half, 0.0,       0.0, 0.0],
                [0.0,        ground_half,0.0, 0.0],
                [0.0,        0.0,       1.0, 0.0],
                [0.0,        0.0,       0.0, 1.0]
            ], dtype=np.float64)),
            "bsdf": {"type": "diffuse",
                     "reflectance": {"type": "rgb", "value": [0.7, 0.7, 0.8]}}
        },

        "spyder_panel": {
            "type": "rectangle",
            "to_world": np_to_mi(T_panel),
            "bsdf": {
                "type": "diffuse",
                "reflectance": {
                    "type": "bitmap",
                    "filename": cc_png_path,
                    "wrap_mode": "clamp",
                    "filter_type": "nearest"
                }
            }
        },

        "checker_board": {
            "type": "cube",
            "to_world": np_to_mi(T_sample2),
            "bsdf": {"type": "diffuse",
                     "reflectance": {"type": "rgb", "value": [0.0, 0.0, 0.0]}}
        },
        "checker_side_negY": {
            "type": "rectangle",
            "to_world": np_to_mi(T_negY),
            "bsdf": {
                "type": "diffuse",
                "reflectance": {
                    "type": "bitmap",
                    "filename": "checker_6x8.png",
                    "wrap_mode": "clamp",
                    "filter_type": "nearest"
                }
            }
        }
    }
    return scene

def main():
    # Camera inputs
    ANGLES = ["normal", "medium", "grazing"]
    RENDER_TYPE = 'full'
    for ANGLE in ANGLES:
        IMAGE_PATH      = f"final_reference/{ANGLE}/aruco.jpg"
        DICT_ID_NAME    = "DICT_4X4_50"
        EXPECTED_IDS    = {0}
        MARKER_INCHES   = 2.0
        FOCAL_MM        = 51.3 if ANGLE == "grazing" else 49.1
        SENSOR_W_MM     = 24.0 / 1.5
        SENSOR_H_MM     = 35.9 / 1.5
        ASSUME_ZERO_DISTORTION = True


        # Material Parameters
        sigma_hs = np.array([0.808, 0.280, 0.138, 0.091, 0.027])
        lcs = np.array([16.6646, 25.6949, 19.0496, 35.1056, 198.73])
        gammas = np.array([3.66, 2.34, 2.11, 2.04, 1.949])
        lc_exps = np.array([-0.7710, -0.8307, -0.9338, -0.9745, -1.0273])
        gamma_exps = np.array([0.3009, 0.0234, 0.0097, 0.0031, -0.0052])
        alphas = [0.121, 0.038, 0.021, 0.01, 0.001]  # smoothed GGX roughness values (profilometer)
        # alphas = [0.141, 0.0683, 0.059, 0.049, 0.020]  # non-smooth GGX roughness values (profilometer)
        print(len(alphas))
        RENDER_MATS = np.array(['SS', 'Brass', 'Brass', 'Brass', 'SS'])
        print(len(RENDER_MATS))
        IOR_DIR   = "ior"

        SCALE     = 1e-2
        SPP       = 2 ** 7
        RES_SCALE = 0.25

        BASE_DIR = "renders"
        OUTDIR = os.path.join(BASE_DIR, ANGLE)
        if RENDER_TYPE == 'ggx':
            OUTDIR = os.path.join(OUTDIR, "ggx")
        elif RENDER_TYPE == 'raw_ggx':
            alphas = [0.141, 0.0683, 0.059, 0.049, 0.020]
            OUTDIR = os.path.join(OUTDIR, "raw_ggx")
        elif RENDER_TYPE == 'first':
            OUTDIR = os.path.join(OUTDIR, "first")
        else:
            OUTDIR = os.path.join(OUTDIR, "ghs")

        PNG_DIR   = os.path.join(OUTDIR, "png")
        TIFF_DIR  = os.path.join(OUTDIR, "tiff")
        os.makedirs(PNG_DIR, exist_ok=True)
        os.makedirs(TIFF_DIR, exist_ok=True)

        MATERIAL_FILES = {
            "Al":    (os.path.join(IOR_DIR, "al_eta.spd"),    os.path.join(IOR_DIR, "al_k.spd")),
            "Brass": (os.path.join(IOR_DIR, "brass_eta.spd"), os.path.join(IOR_DIR, "brass_k.spd")),
            "SS":    (os.path.join(IOR_DIR, "ss_eta.spd"),    os.path.join(IOR_DIR, "ss_k.spd")),
        }

        K, R_cam_mark, t_cam_mark, h, w = camera_from_aruco(
            image_path=IMAGE_PATH,
            focal_mm=FOCAL_MM,
            sensor_w_mm=SENSOR_W_MM,
            sensor_h_mm=SENSOR_H_MM,
            marker_inches=MARKER_INCHES,
            aruco_dict_name=DICT_ID_NAME,
            expected_ids=EXPECTED_IDS,
            assume_zero_distortion=ASSUME_ZERO_DISTORTION
        )

        combo_iter = zip(RENDER_MATS, gammas, sigma_hs, lcs, lc_exps, gamma_exps, alphas)
        for i, (mat_name, gamma, sigma_h, lc, lc_exp, gamma_exp, alpha) in enumerate(combo_iter):
            if mat_name not in MATERIAL_FILES:
                raise ValueError(f"Unknown material '{mat_name}'")
            eta_path, k_path = MATERIAL_FILES[mat_name]
            if not (os.path.isfile(eta_path) and os.path.isfile(k_path)):
                raise FileNotFoundError(f"Missing IOR files for {mat_name}: {eta_path}, {k_path}")

            if RENDER_TYPE == 'ggx' or RENDER_TYPE == 'raw_ggx':
                scene_dict = build_scene_dict_ggx(
                    K=K, R_cam_mark=R_cam_mark, t_cam_mark=t_cam_mark,
                    img_w=w, img_h=h,
                    res_scale=RES_SCALE, spp=SPP,
                    ior_eta_path=eta_path, ior_k_path=k_path,
                    alpha=float(alpha),
                    Angle=ANGLE
                )
                png_path  = os.path.join(PNG_DIR,  f"{i}_alpha_{alpha}.png")
                tiff_path = os.path.join(TIFF_DIR, f"{i}_alpha_{alpha}.tiff")
            else:
                scene_dict = build_scene_dict(
                    K=K, R_cam_mark=R_cam_mark, t_cam_mark=t_cam_mark,
                    img_w=w, img_h=h,
                    res_scale=RES_SCALE, spp=SPP,
                    ior_eta_path=eta_path, ior_k_path=k_path,
                    t_val=float(dr.square(lc)), sigmah=float(sigma_h), gamma=float(gamma),
                    T_exp=float(2*lc_exp), gamma_exp=float(gamma_exp),
                    Angle=ANGLE
                )

                # export scene_dict
                with open(f"scene_dump_{ANGLE}.txt", "w") as f:
                    f.write(str(scene_dict))
                png_path  = os.path.join(PNG_DIR,  f"{i}.png")
                tiff_path = os.path.join(TIFF_DIR, f"{i}.tiff")

            scene = mi.load_dict(scene_dict)
            img   = mi.render(scene, spp=int(SPP), seed=1234)


            out = np.array(SCALE * img)
            out = cv2.flip(out, 1)
            out = np.nan_to_num(out, nan=0.0, posinf=0.0, neginf=0.0)
            if np.any(out < 0):
                print("less than zero!")
                out[out < 0] = 0

            bmp = mi.Bitmap(out)
            bmp_tiff = bmp.convert(
                pixel_format=mi.Bitmap.PixelFormat.RGB,
                component_format=mi.Struct.Type.Float16,
                srgb_gamma=False
            )

            mi.util.write_bitmap(png_path, bmp)
            tifffile.imwrite(tiff_path, np.array(bmp_tiff, np.float16), photometric='rgb')
            print(f"wrote {tiff_path}")

if __name__ == "__main__":
    main()
