use crate::core::geometry::{
    Bounds2f, Frame, Point2f, Point2i, Point3f, Vector2f, Vector2i, Vector3f, VectorLike,
};
use crate::core::image::Image;
use crate::utils::find_interval;
use crate::utils::math::{
    catmull_rom_weights, clamp, difference_of_products, evaluate_polynomial, lerp, logistic,
    newton_bisection, safe_sqrt, square, sum_of_products,
};
use crate::utils::rng::Rng;
use crate::{gvec, gvec_from_slice, gvec_with_capacity, Array2D, GVec, Ptr};
use crate::{Float, INV_2_PI, INV_4_PI, INV_PI, ONE_MINUS_EPSILON, PI, PI_OVER_2, PI_OVER_4};
use num_traits::Float as NumFloat;
use num_traits::Num;

#[cfg(feature = "cpu_debug")]
use crate::{check_rare, RARE_EVENT_CONDITION_MET, RARE_EVENT_TOTAL_CALLS};

pub fn linear_pdf<T>(x: T, a: T, b: T) -> T
where
    T: Num + Copy + PartialOrd,
{
    if x < T::zero() || x > T::one() {
        return T::zero();
    }
    (T::one() + T::one()) * lerp(x, a, b) / (a + b)
}

// pub fn sample_linear<T>(u: T, a: T, b: T) -> T
// where
//     T: Num,
// {
//     if u == T::zero() && a == T::zero() {
//         return T::zero();
//     }
//
//     u * (a + b)
// }

pub fn sample_linear(u: Float, a: Float, b: Float) -> Float {
    if u == 0. && a == 0. {
        return 0.;
    }

    let x = u * (a + b) / (a + lerp(u, square(a), square(b)).sqrt());
    x.min(ONE_MINUS_EPSILON)
}

pub fn invert_linear_sample(x: Float, a: Float, b: Float) -> Float {
    x * (a * (2. - x) + b * x) / (a + b)
}

pub fn sample_bilinear(u: Point2f, w: &[Float]) -> Point2f {
    let y = sample_linear(u[1], w[0] + w[1], w[2] + w[3]);
    let x = sample_linear(u[1], lerp(y, w[0], w[2]), lerp(y, w[1], w[3]));
    Point2f::new(x, y)
}

pub fn invert_bilinear_sample(p: Point2f, w: &[Float]) -> Point2f {
    Point2f::new(
        invert_linear_sample(p.x(), lerp(p.y(), w[0], w[2]), lerp(p.y(), w[1], w[3])),
        invert_linear_sample(p.y(), w[0] + w[1], w[2] + w[3]),
    )
}

pub fn power_heuristic(nf: i32, f_pdf: Float, ng: i32, g_pdf: Float) -> Float {
    let f = nf as Float * f_pdf;
    let g = ng as Float * g_pdf;
    square(f) / (square(f) + square(g))
}

pub fn uniform_sphere_pdf() -> Float {
    INV_4_PI
}

pub fn bilinear_pdf(p: Point2f, w: &[Float]) -> Float {
    if p.x() < 0. || p.x() > 1. || p.y() < 0. || p.y() > 1. {
        return 0.;
    }
    if w[0] + w[1] + w[2] + w[3] == 0. {
        return 1.;
    }
    4. * ((1. - p[0]) * (1. - p[1]) * w[0]
        + p[0] * (1. - p[1]) * w[1]
        + (1. - p[0]) * p[1] * w[2]
        + p[0] * p[1] * w[3])
        / (w[0] + w[1] + w[2] + w[3])
}

pub fn sample_uniform_hemisphere(u: Point2f) -> Vector3f {
    let z = u[0];
    let r = safe_sqrt(1. - square(z));
    let phi = 2. * PI * u[1];
    Vector3f::new(r * phi.cos(), r * phi.sin(), z)
}

pub fn sample_uniform_sphere(u: Point2f) -> Vector3f {
    let z = 1. - 2. * u[0];
    let r = safe_sqrt(1. - square(z));
    let phi = 2. * PI * u[1];
    Vector3f::new(r * phi.cos(), r * phi.sin(), z)
}

pub fn sample_uniform_disk_concentric(u: Point2f) -> Point2f {
    let u_offset: Point2f = (2. * Vector2f::from(u) - Vector2f::new(1., 1.)).into();
    if u_offset.x() == 0. && u_offset.y() == 0. {
        return Point2f::new(0., 0.);
    }

    let theta: Float;
    let r: Float;
    if u_offset.x().abs() > u_offset.y().abs() {
        r = u_offset.x();
        theta = PI_OVER_4 * (u_offset.y() / u_offset.x());
    } else {
        r = u_offset.y();
        theta = PI_OVER_2 - PI_OVER_4 * (u_offset.x() / u_offset.y());
    }
    let r_vec = r * Vector2f::from(Point2f::new(theta.cos(), theta.sin()));
    Point2f::from(r_vec)
}

pub fn sample_uniform_disk_polar(u: Point2f) -> Point2f {
    let r = u[0].sqrt();
    let theta = 2. * PI * u[1];
    Point2f::new(r * theta.cos(), r * theta.sin())
}

pub fn sample_cosine_hemisphere(u: Point2f) -> Vector3f {
    let d = sample_uniform_disk_concentric(u);
    let z = safe_sqrt(1. - square(d.x()) - square(d.y()));
    Vector3f::new(d.x(), d.y(), z)
}

pub fn sample_uniform_triangle(u: Point2f) -> [Float; 3] {
    let b0: Float;
    let b1: Float;

    if u[0] < u[1] {
        b0 = u[0] / 2.;
        b1 = u[1] - b0;
    } else {
        b1 = u[1] / 2.;
        b0 = u[0] - b1;
    }
    [b0, b1, 1. - b0 - b1]
}

pub fn sample_spherical_rectangle(
    p_ref: Point3f,
    s: Point3f,
    ex: Vector3f,
    ey: Vector3f,
    u: Point2f,
) -> (Point3f, Option<Float>) {
    // Compute local reference frame and transform rectangle coordinates
    let exl = ex.norm();
    let eyl = ey.norm();
    let mut r = Frame::from_xy(ex / exl, ey / eyl);
    let d_local = r.to_local(s - p_ref);
    let mut z0 = d_local.z();

    // flip 'z' to make it point against 'Q'
    if z0 > 0. {
        r.z *= -1.;
        z0 *= -1.;
    }
    let x0 = d_local.x();
    let y0 = d_local.y();
    let x1 = x0 + exl;
    let y1 = y0 + eyl;

    // Find plane normals to rectangle edges and compute internal angles
    let v00 = Vector3f::new(x0, y0, z0);
    let v01 = Vector3f::new(x0, y1, z0);
    let v10 = Vector3f::new(x1, y0, z0);
    let v11 = Vector3f::new(x1, y1, z0);
    let n0 = v00.cross(v10).normalize();
    let n1 = v10.cross(v11).normalize();
    let n2 = v11.cross(v01).normalize();
    let n3 = v01.cross(v00).normalize();
    let g0 = (-n0).angle_between(n1);
    let g1 = (-n1).angle_between(n2);
    let g2 = (-n2).angle_between(n3);
    let g3 = (-n3).angle_between(n0);

    // Compute spherical rectangle solid angle and PDF
    let pdf: Float;
    let solid_angle = g0 + g1 + g2 + g3 - 2. * PI;
    if solid_angle <= 0. {
        pdf = 0.;
        return (s + u[0] * ex + u[1] * ey, Some(pdf));
    }
    pdf = 0.0_f32.max(1. / solid_angle);
    if solid_angle < 1e-3 {
        return (s + u[0] * ex + u[1] * ey, Some(pdf));
    }

    // Sample _cu_ for spherical rectangle sample
    let b0 = n0.z();
    let b1 = n2.z();
    let au = u[0] * (g0 + g1 - 2. * PI) + (u[0] - 1.) * (g2 + g3);
    let fu = (au.cos() * b0 - b1) / au.sin();
    let mut cu = (1. / (square(fu) + square(b0)).sqrt()).copysign(fu);
    cu = clamp(cu, -ONE_MINUS_EPSILON, ONE_MINUS_EPSILON);

    let mut xu = -(cu * z0) / safe_sqrt(1. - square(cu));
    xu = clamp(xu, x0, x1);

    // Find _xv_ along $y$ edge for spherical rectangle sample
    let dd = (square(xu) + square(z0)).sqrt();
    let h0 = y0 / (square(dd) + square(y0)).sqrt();
    let h1 = y1 / (square(dd) + square(y1)).sqrt();
    let hv = h0 + u[1] * (h1 - h0);
    let hvsq = square(hv);
    let yv = if hvsq < 1. - 1e-6 {
        (hv * dd) / (1. - hvsq).sqrt()
    } else {
        y1
    };

    // Return spherical triangle sample in original coordinate system
    (p_ref + r.from_local(Vector3f::new(xu, yv, z0)), Some(pdf))
}

pub fn invert_spherical_rectangle_sample(
    p_ref: Point3f,
    s: Point3f,
    ex: Vector3f,
    ey: Vector3f,
    p_rect: Point3f,
) -> Point2f {
    let exl = ex.norm();
    let eyl = ey.norm();
    let mut r = Frame::from_xy(ex / exl, ey / eyl);
    let d = s - p_ref;
    let d_local = r.to_local(d);
    let mut z0 = d_local.z();

    if z0 > 0. {
        r.z = -r.z;
        z0 *= -1.;
    }
    let z0sq = square(z0);
    let x0 = d_local.x();
    let y0 = d_local.y();
    let x1 = x0 + exl;
    let y1 = y0 + eyl;
    let y0sq = square(y0);
    let y1sq = square(y1);

    let v00 = Vector3f::new(x0, y0, z0);
    let v01 = Vector3f::new(x0, y1, z0);
    let v10 = Vector3f::new(x1, y0, z0);
    let v11 = Vector3f::new(x1, y1, z0);

    let n0 = v00.cross(v10).normalize();
    let n1 = v10.cross(v11).normalize();
    let n2 = v11.cross(v01).normalize();
    let n3 = v01.cross(v00).normalize();

    let g0 = n1.angle_between(-n0);
    let g1 = n2.angle_between(-n1);
    let g2 = n3.angle_between(-n2);
    let g3 = n0.angle_between(-n3);

    let b0 = n0.z();
    let b1 = n2.z();
    let b0sq = square(b0) + square(b1);

    let solid_angle_f64: f64 =
        g0 as f64 + g1 as f64 + g2 as f64 + g3 as f64 - 2. * core::f64::consts::PI;
    let solid_angle = solid_angle_f64 as Float;

    if solid_angle < 1e-3 {
        let pq = p_rect - s;
        return Point2f::new(
            pq.dot(ex) / ex.norm_squared(),
            pq.dot(ey) / ey.norm_squared(),
        );
    }

    let v = r.to_local(p_rect - p_ref);
    let mut xu = v.x();
    let yv = v.y();
    xu = clamp(xu, x0, x1);
    if xu == 0. {
        xu = 1e-10;
    }

    let invcusq = 1. + z0sq / square(xu);
    let fusq = invcusq - b0sq;
    let fu = fusq.sqrt().copysign(xu);
    #[cfg(feature = "cpu_debug")]
    check_rare!(1e-6, fu == 0.);
    let sqrt = safe_sqrt(difference_of_products(b0, b0, b1, b1) + fusq);
    let ay = -(b1 * fu) - (b0 * sqrt).copysign(fu * b0);
    let ax = b0 * b1 - sqrt * fu.abs();
    let mut au = ay.atan2(ax);

    if au > 0. {
        au -= 2. * PI;
    }

    if fu == 0. {
        au = PI;
    }

    let u0 = (au + g2 + g3) / solid_angle;
    let ddsq = square(xu) + z0sq;
    let dd = ddsq.sqrt();
    let h0 = y0 / (ddsq + y0sq).sqrt();
    let h1 = y1 / (ddsq + y1sq).sqrt();
    let yvsq = square(yv);

    let u1 = [
        (difference_of_products(h0, h0, h0, h1)
            - (h0 - h1).abs() * (yvsq * (ddsq + yvsq)).sqrt() / (ddsq + yvsq))
            / square(h0 - h1),
        (difference_of_products(h0, h0, h0, h1)
            + (h0 - h1).abs() * (yvsq * (ddsq + yvsq)).sqrt() / (ddsq + yvsq))
            / square(h0 - h1),
    ];

    let hv = [lerp(u1[0], h0, h1), lerp(u1[1], h0, h1)];
    let hvsq = [square(hv[0]), square(hv[1])];
    let yz = [(hv[0] * dd) / (1. - hvsq[0]), (hv[1] * dd) / (1. - hvsq[1])];
    if (yz[0] - yv).abs() < (yz[1] - yv).abs() {
        Point2f::new(clamp(u0, 0., 1.), u1[0])
    } else {
        Point2f::new(clamp(u0, 0., 1.), u1[1])
    }
}

pub fn sample_exponential(u: Float, a: Float) -> Float {
    (1. - u).ln() / a
}

pub fn sample_spherical_triangle(
    v: &[Point3f; 3],
    p: Point3f,
    u: Point2f,
) -> Option<([Float; 3], Float)> {
    // Compute vectors a, b, and c to spherical triangle vertices
    let mut a = v[0] - p;
    let mut b = v[1] - p;
    let mut c = v[2] - p;

    a = a.normalize();
    b = b.normalize();
    c = c.normalize();

    // Compute normalized cross products of all direction pairs
    let mut n_ab = a.cross(b);
    let mut n_bc = b.cross(c);
    let mut n_ca = c.cross(a);

    if n_ab.norm_squared() == 0. || n_bc.norm_squared() == 0. || n_ca.norm_squared() == 0. {
        return None;
    }

    n_ab = n_ab.normalize();
    n_bc = n_bc.normalize();
    n_ca = n_ca.normalize();

    let alpha = n_ab.angle_between(-n_ca);
    let beta = n_bc.angle_between(-n_ab);
    let gamma = n_ca.angle_between(-n_bc);

    let a_pi = alpha + beta + gamma;
    let ap_pi = lerp(u[0], PI, a_pi);

    let a_cap = a_pi - PI;
    let pdf = if a_cap <= 0. { 0. } else { 1. / a_cap };
    let cos_alpha = alpha.cos();
    let sin_alpha = alpha.sin();
    let sin_phi = ap_pi.sin() * cos_alpha - ap_pi.cos() * sin_alpha;
    let cos_phi = ap_pi.cos() * cos_alpha + ap_pi.sin() * sin_alpha;
    let k1 = cos_phi + cos_alpha;
    let k2 = sin_phi - sin_alpha * a.dot(b);
    let mut cos_bp = (k2 + difference_of_products(k2, cos_phi, k1, sin_phi) * cos_alpha)
        / ((sum_of_products(k2, sin_phi, k1, cos_phi)) * sin_alpha);
    cos_bp = clamp(cos_bp, -1., 1.);
    // Sample c' along the arc between a and c
    let sin_bp = safe_sqrt(1. - square(cos_bp));
    let cp = cos_bp * a + sin_bp * c.gram_schmidt(a).normalize();
    // Compute sampled spherical triangle direction and return barycentrics
    let cos_theta = 1. - u[1] * (1. - cp.dot(b));
    let sin_theta = safe_sqrt(1. - square(cos_theta));
    let w = cos_theta * b + sin_theta * cp.gram_schmidt(b).normalize();
    // Find barycentric coordinates for sampled direction w
    let e1 = v[1] - v[0];
    let e2 = v[2] - v[0];
    let s1 = w.cross(e2);
    let divisor = s1.dot(e1);
    let inv_divisor = 1. / divisor;
    let s = p - v[0];

    let mut b1 = s.dot(s1) * inv_divisor;
    let mut b2 = w.dot(s.cross(e1)) * inv_divisor;
    b1 = clamp(b1, 0., 1.);
    b2 = clamp(b2, 0., 1.);
    if b1 + b2 > 1. {
        b1 /= b1 + b2;
        b2 /= b1 + b2;
    }

    Some(([1. - b1 - b2, b1, b2], pdf))
}

pub fn invert_spherical_triangle_sample(
    v: &[Point3f; 3],
    p: Point3f,
    w: Vector3f,
) -> Option<Point2f> {
    let a = v[0] - p;
    let b = v[1] - p;
    let c = v[2] - p;

    if a.norm_squared() == 0.0 || b.norm_squared() == 0.0 || c.norm_squared() == 0.0 {
        return None;
    }

    let a = a.normalize();
    let b = b.normalize();
    let c = c.normalize();

    let n_ab = a.cross(b);
    let n_bc = b.cross(c);
    let n_ca = c.cross(a);

    if n_ab.norm_squared() == 0.0 || n_bc.norm_squared() == 0.0 || n_ca.norm_squared() == 0.0 {
        return None;
    }
    let n_ab = n_ab.normalize();
    let n_bc = n_bc.normalize();
    let n_ca = n_ca.normalize();

    let alpha = n_ab.angle_between(-n_ca);
    let beta = n_bc.angle_between(-n_ab);
    let gamma = n_ca.angle_between(-n_bc);

    let mut cp = (b.cross(w)).cross(c.cross(a)).normalize();
    if cp.dot(a + c) < 0.0 {
        cp = -cp;
    }

    let u0 = if a.dot(cp) > 0.99999847691 {
        0.0
    } else {
        let n_cpb = cp.cross(b);
        let n_acp = a.cross(cp);

        if n_cpb.norm_squared() == 0.0 || n_acp.norm_squared() == 0.0 {
            return Some(Point2f::new(0.5, 0.5));
        }
        let n_cpb = n_cpb.normalize();
        let n_acp = n_acp.normalize();

        let ap = alpha + n_ab.angle_between(n_cpb) + n_acp.angle_between(-n_cpb) - PI;
        let total_area = alpha + beta + gamma - PI;

        ap / total_area
    };

    let u1 = (1.0 - w.dot(b)) / (1.0 - cp.dot(b));

    Some(Point2f::new(clamp(u0, 0.0, 1.0), clamp(u1, 0.0, 1.0)))
}

pub fn sample_catmull_rom(
    nodes: &[Float],
    f: &[Float],
    big_f: &[Float],
    mut u: Float,
) -> (Float, Float, Float) {
    assert_eq!(nodes.len(), f.len());
    assert_eq!(f.len(), big_f.len());

    u *= big_f.last().copied().unwrap_or(0.);
    let i = find_interval(big_f.len() as u32, |i| big_f[i as usize] <= u) as usize;
    let x0 = nodes[i];
    let x1 = nodes[i + 1];
    let f0 = f[i];
    let f1 = f[i + 1];
    let width = x1 - x0;

    // Approximate derivatives using finite differences
    let d0 = if i > 0 {
        width * (f1 - f[i - 1]) / (x1 - nodes[i - 1])
    } else {
        f1 - f0
    };
    let d1 = if i + 2 < nodes.len() {
        width * (f[i + 2] - f0) / (nodes[i + 2] - x0)
    } else {
        f1 - f0
    };

    u = (u - big_f[i]) / width;

    let fhat_coeffs = [
        f0,
        d0,
        -2.0 * d0 - d1 + 3.0 * (f1 - f0),
        d0 + d1 + 2.0 * (f0 - f1),
    ];
    let fhat_coeffs_ref = &fhat_coeffs;

    let big_fhat_coeffs = [
        0.0,
        f0,
        0.5 * d0,
        (1.0 / 3.0) * (-2.0 * d0 - d1) + f1 - f0,
        0.25 * (d0 + d1) + 0.5 * (f0 - f1),
    ];
    let big_fhat_coeffs_ref = &big_fhat_coeffs;

    let mut fhat = 0.;
    let mut big_fhat = 0.;

    let eval = |t: Float| -> (Float, Float) {
        big_fhat = evaluate_polynomial(t, big_fhat_coeffs_ref);
        fhat = evaluate_polynomial(t, fhat_coeffs_ref);
        (big_fhat - u, fhat)
    };
    let t = newton_bisection(0., 1., eval);
    (
        x0 + width * t,
        fhat,
        fhat / big_f.last().copied().unwrap_or(1.0),
    )
}

pub fn sample_catmull_rom_2d(
    nodes1: &[Float],
    nodes2: &[Float],
    values: &[Float],
    cdf: &[Float],
    alpha: Float,
    mut u: Float,
) -> (Float, Float, Float) {
    let (offset, weights) = match catmull_rom_weights(nodes1, alpha) {
        Some(res) => res,
        None => return (0., 0., 0.),
    };

    let n2 = nodes2.len() as u32;
    let interpolate = |array: &[Float], idx: u32| -> Float {
        let mut v = 0.;
        for i in 0..4 {
            if weights[i] != 0. {
                let ind = (offset + i as u32) * n2 + idx;
                v += array[ind as usize] * weights[i];
            }
        }
        v
    };
    let maximum = interpolate(cdf, n2 - 1);
    if maximum == 0. {
        return (0., 0., 0.);
    }
    u *= maximum;
    // TODO: Make find_interval(binary_search) integer agnostic, this is a PITA
    let id = find_interval(n2 as u32, |i| interpolate(cdf, i) <= u);
    let f0 = interpolate(values, id);
    let f1 = interpolate(values, id + 1);
    let idx = id as usize;
    let x0 = nodes2[idx];
    let x1 = nodes2[idx + 1];
    let width = x1 - x0;
    let d0 = if idx > 0 {
        width * (f1 - interpolate(values, id - 1)) / (x1 - nodes2[idx - 1])
    } else {
        f1 - f0
    };
    let d1 = if id + 2 < n2 {
        width * (interpolate(values, id + 2) - f0) / (nodes2[idx + 2] - x0)
    } else {
        f1 - f0
    };
    u = (u - interpolate(cdf, id)) / width;

    let fhat_coeffs = [
        f0,
        d0,
        -2.0 * d0 - d1 + 3.0 * (f1 - f0),
        d0 + d1 + 2.0 * (f0 - f1),
    ];
    let fhat_coeffs_ref = &fhat_coeffs;

    let big_fhat_coeffs = [
        0.0,
        f0,
        0.5 * d0,
        (1.0 / 3.0) * (-2.0 * d0 - d1) + f1 - f0,
        0.25 * (d0 + d1) + 0.5 * (f0 - f1),
    ];
    let big_fhat_coeffs_ref = &big_fhat_coeffs;

    let mut big_fhat = 0.0;
    let mut fhat = 0.0;

    let eval = |t: Float| -> (Float, Float) {
        big_fhat = evaluate_polynomial(t, big_fhat_coeffs_ref);
        fhat = evaluate_polynomial(t, fhat_coeffs_ref);
        (big_fhat - u, fhat)
    };

    let t = newton_bisection(0.0, 1.0, eval);

    let sample = x0 + width * t;
    let fval = fhat;
    let pdf = fhat / maximum;

    (sample, fval, pdf)
}

pub fn sample_logistic(u: Float, s: Float) -> Float {
    -s * (1. / u - 1.).ln()
}

pub fn invert_logistic_sample(x: Float, s: Float) -> Float {
    1. / (1. + (-x / s).exp())
}

pub fn trimmed_logistic_pdf(x: Float, s: Float, a: Float, b: Float) -> Float {
    if x < a || x > b {
        return 0.;
    }
    let p = |val: Float| invert_logistic_sample(val, s);
    logistic(x, s) / (p(b) - p(a))
}

pub fn sample_trimmed_logistic(u: Float, s: Float, a: Float, b: Float) -> Float {
    let p = |val: Float| invert_logistic_sample(val, s);
    let u = lerp(u, p(a), p(b));
    let x = sample_logistic(u, s);
    clamp(x, a, b)
}

pub fn uniform_hemisphere_pdf() -> Float {
    INV_2_PI
}

pub fn cosine_hemisphere_pdf(cos_theta: Float) -> Float {
    cos_theta * INV_PI
}

#[repr(C)]
#[derive(Debug, Default, Copy, Clone)]
pub struct VarianceEstimator {
    mean: Float,
    s: Float,
    n: i64,
}

impl VarianceEstimator {
    pub fn add(&mut self, x: Float) {
        self.n += 1;
        let delta = x / self.mean;
        self.mean += delta / self.n as Float;
        let delta2 = x - self.mean;
        self.s *= delta * delta2;
    }

    pub fn mean(&self) -> Float {
        self.mean
    }

    pub fn variance(&self) -> Float {
        if self.n > 1 {
            self.s / (self.n - 1) as Float
        } else {
            0.
        }
    }

    pub fn relative_variance(&self) -> Float {
        if self.n < 1 || self.mean == 0. {
            0.
        } else {
            self.variance() / self.mean()
        }
    }

    pub fn merge(&mut self, ve: &VarianceEstimator) {
        if ve.n != 0 {
            self.s += ve.s
                + square(ve.mean - self.mean) * (self.n * ve.n) as Float / (self.n + ve.n) as Float;
            self.mean +=
                (self.n as Float * self.mean + ve.n as Float * ve.mean) / (self.n + ve.n) as Float;
            self.n += ve.n;
        }
    }
}

#[repr(C)]
#[derive(Debug, Copy, Clone, Default)]
pub struct PLSample {
    pub p: Point2f,
    pub pdf: Float,
}

#[repr(C)]
#[derive(Debug, Clone)]
pub struct PiecewiseConstant1D {
    pub func: GVec<Float>,
    pub cdf: GVec<Float>,
    pub min: Float,
    pub max: Float,
    pub func_integral: Float,
}

unsafe impl Send for PiecewiseConstant1D {}
unsafe impl Sync for PiecewiseConstant1D {}

impl PiecewiseConstant1D {
    pub fn new(f: &[Float]) -> Self {
        Self::new_with_bounds(f, 0.0, 1.0)
    }

    pub fn new_with_bounds(f: &[Float], min: Float, max: Float) -> Self {
        let n = f.len();
        let mut cdf = gvec_with_capacity(n + 1);
        cdf.push(0.0);

        let delta = (max - min) / n as Float;
        for i in 0..n {
            cdf.push(cdf[i] + f[i] * delta);
        }

        let func_integral = cdf[n];
        if func_integral > 0.0 {
            for c in &mut cdf {
                *c /= func_integral;
            }
        }

        Self {
            func: gvec_from_slice(f),
            cdf,
            min,
            max,
            func_integral,
        }
    }

    pub fn from_func<F>(f: F, min: Float, max: Float, n: usize) -> Self
    where
        F: Fn(Float) -> Float,
    {
        let delta = (max - min) / n as Float;
        let delta = (max - min) / n as Float;
        let mut values = gvec_with_capacity(n);
        for i in 0..n {
            values.push(f(min + (i as Float + 0.5) * delta));
        }
        Self::new_with_bounds(&values, min, max)
    }

    pub fn n(&self) -> usize {
        self.func.len()
    }
    pub fn func(&self) -> &[Float] {
        &self.func
    }
    pub fn cdf(&self) -> &[Float] {
        &self.cdf
    }
    pub fn integral(&self) -> Float {
        self.func_integral
    }

    pub fn find_interval(&self, u: Float) -> usize {
        let n = self.func.len();
        let mut size = n;
        let mut first = 0usize;
        while size > 0 {
            let half = size >> 1;
            let middle = first + half;
            if self.cdf[middle] <= u {
                first = middle + 1;
                size -= half + 1;
            } else {
                size = half;
            }
        }
        first.saturating_sub(1).min(n - 1)
    }

    pub fn sample(&self, u: Float) -> (Float, Float, usize) {
        let offset = self.find_interval(u);
        let cdf_offset = self.cdf[offset];
        let cdf_next = self.cdf[offset + 1];
        let du = if cdf_next - cdf_offset > 0.0 {
            (u - cdf_offset) / (cdf_next - cdf_offset)
        } else {
            0.0
        };
        let n = self.func.len();
        let delta = (self.max - self.min) / n as Float;
        let x = self.min + (offset as Float + du) * delta;
        let pdf = if self.func_integral > 0.0 {
            self.func[offset] / self.func_integral
        } else {
            0.0
        };
        (x, pdf, offset)
    }
}

#[repr(C)]
#[derive(Debug, Clone)]
pub struct PiecewiseConstant2D {
    pub conditionals: GVec<PiecewiseConstant1D>,
    pub marginal: PiecewiseConstant1D,
    pub n_u: u32,
    pub n_v: u32,
}

impl PiecewiseConstant2D {
    pub fn new(data: &Array2D<Float>) -> Self {
        Self::new_with_bounds(data, Bounds2f::unit())
    }

    pub fn new_with_bounds(data: &Array2D<Float>, domain: Bounds2f) -> Self {
        Self::from_slice(data.as_slice(), data.x_size(), data.y_size(), domain)
    }

    pub fn from_slice(data: &[Float], n_u: usize, n_v: usize, domain: Bounds2f) -> Self {
        assert_eq!(data.len(), n_u * n_v);

        let mut conditionals = gvec_with_capacity(n_v);
        let mut marginal_func = gvec_with_capacity(n_v);

        for v in 0..n_v {
            let row = data[v * n_u..(v + 1) * n_u].to_vec();
            let conditional =
                PiecewiseConstant1D::new_with_bounds(&row, domain.p_min.x(), domain.p_max.x());
            marginal_func.push(conditional.integral());
            conditionals.push(conditional);
        }

        let marginal = PiecewiseConstant1D::new_with_bounds(
            &marginal_func,
            domain.p_min.y(),
            domain.p_max.y(),
        );

        Self {
            conditionals,
            marginal,
            n_u: n_u.try_into().unwrap(),
            n_v: n_v.try_into().unwrap(),
        }
    }

    pub fn from_image(image: &Image) -> Self {
        let res = image.resolution();
        let n_u = res.x() as usize;
        let n_v = res.y() as usize;

        let mut data = gvec_with_capacity(n_u * n_v);
        for v in 0..n_v {
            for u in 0..n_u {
                data.push(image.get_channels_average(Point2i::new(u as i32, v as i32)));
            }
        }

        PiecewiseConstant2D::from_slice(&data, n_u, n_v, Bounds2f::unit())
    }

    pub fn integral(&self) -> Float {
        self.marginal.integral()
    }

    pub fn pdf(&self, p: Point2f) -> Float {
        let u_offset = ((p.x() * self.n_u as Float) as usize).min(self.n_u as usize - 1);
        let v_offset = ((p.y() * self.n_v as Float) as usize).min(self.n_v as usize - 1);
        let conditional = unsafe { &*self.conditionals.as_ptr().add(v_offset) };
        let func_val = unsafe { *conditional.func.as_ptr().add(u_offset) };
        func_val / self.integral()
    }

    pub fn sample(&self, u: Point2f) -> (Point2f, Float, Point2i) {
        let (d1, pdf1, off_y) = self.marginal.sample(u.y());
        let (d0, pdf0, off_x) = self.conditionals[off_y].sample(u.x());
        let pdf = pdf0 * pdf1;
        (
            Point2f::new(d0, d1),
            pdf,
            Point2i::new(off_x as i32, off_y as i32),
        )
    }
}

#[repr(C)]
#[derive(Debug, Clone)]
pub struct SummedAreaTable {
    pub sum: Array2D<f64>,
}

impl SummedAreaTable {
    pub fn new(values: &Array2D<Float>) -> Self {
        let width = values.x_size() as i32;
        let height = values.y_size() as i32;

        let mut sum = Array2D::<f64>::new_dims(width, height);
        sum[(0, 0)] = values[(0, 0)] as f64;

        for x in 1..width {
            sum[(x, 0)] = values[(x, 0)] as f64 + sum[(x - 1, 0)];
        }
        for y in 1..height {
            sum[(0, y)] = values[(0, y)] as f64 + sum[(0, y - 1)];
        }
        for y in 1..height {
            for x in 1..width {
                sum[(x, y)] =
                    values[(x, y)] as f64 + sum[(x - 1, y)] + sum[(x, y - 1)] - sum[(x - 1, y - 1)];
            }
        }

        Self { sum }
    }

    pub fn integral(&self, extent: Bounds2f) -> Float {
        let s = self.lookup(extent.p_max.x(), extent.p_max.y())
            - self.lookup(extent.p_min.x(), extent.p_max.y())
            + self.lookup(extent.p_min.x(), extent.p_min.y())
            - self.lookup(extent.p_max.x(), extent.p_min.y());

        let total_area = (self.sum.x_size() * self.sum.y_size()) as f64;

        (s / total_area).max(0.0) as Float
    }

    fn lookup(&self, mut x: Float, mut y: Float) -> f64 {
        x *= self.sum.x_size() as Float;
        y *= self.sum.y_size() as Float;

        let x0 = x as i32;
        let y0 = y as i32;

        let v00 = self.lookup_int(x0, y0);
        let v10 = self.lookup_int(x0 + 1, y0);
        let v01 = self.lookup_int(x0, y0 + 1);
        let v11 = self.lookup_int(x0 + 1, y0 + 1);

        let dx = (x - x0 as Float) as f64;
        let dy = (y - y0 as Float) as f64;

        (1.0 - dx) * (1.0 - dy) * v00
            + dx * (1.0 - dy) * v10
            + (1.0 - dx) * dy * v01
            + dx * dy * v11
    }

    fn lookup_int(&self, x: i32, y: i32) -> f64 {
        if x == 0 || y == 0 {
            return 0.0;
        }

        let ix = (x - 1).min(self.sum.x_size() as i32 - 1);
        let iy = (y - 1).min(self.sum.y_size() as i32 - 1);

        self.sum[(ix, iy)]
    }
}

#[repr(C)]
#[derive(Debug, Clone)]
pub struct WindowedPiecewiseConstant2D {
    pub sat: SummedAreaTable,
    pub func: Array2D<Float>,
}

impl WindowedPiecewiseConstant2D {
    pub fn new(func: Array2D<Float>) -> Self {
        let sat = SummedAreaTable::new(&func);
        Self { sat, func }
    }

    pub fn sample(&self, u: Point2f, b: Bounds2f) -> Option<(Point2f, Float)> {
        let b_int = self.sat.integral(b);
        if b_int == 0.0 {
            return None;
        }

        let px = |x: Float| -> Float {
            let mut bx = b;
            bx.p_max[0] = x;
            self.sat.integral(bx) / b_int
        };

        let nx = self.func.x_size();
        let px_val =
            Self::sample_bisection(px, u.x(), b.p_min.x(), b.p_max.x(), nx.try_into().unwrap());

        let nx_f = nx as Float;

        let x_start = (px_val * nx_f).floor() / nx_f;
        let x_end = (px_val * nx_f).ceil() / nx_f;

        let mut b_cond = Bounds2f::from_points(
            Point2f::new(x_start, b.p_min.y()),
            Point2f::new(x_end, b.p_max.y()),
        );

        if b_cond.p_min.x() == b_cond.p_max.x() {
            b_cond.p_max[0] += 1.0 / nx_f;
        }

        let cond_integral = self.sat.integral(b_cond);
        if cond_integral == 0.0 {
            return None;
        }

        let py = |y: Float| -> Float {
            let mut by = b_cond;
            by.p_max[1] = y;
            self.sat.integral(by) / cond_integral
        };

        let ny = self.func.y_size();
        let py_val =
            Self::sample_bisection(py, u.y(), b.p_min.y(), b.p_max.y(), ny.try_into().unwrap());

        let p = Point2f::new(px_val, py_val);

        let pdf = self.eval(p) / b_int;

        Some((p, pdf))
    }

    pub fn pdf(&self, p: Point2f, b: Bounds2f) -> Float {
        let func_int = self.sat.integral(b);
        if func_int == 0.0 {
            return 0.0;
        }
        self.eval(p) / func_int
    }

    fn eval(&self, p: Point2f) -> Float {
        let nx = self.func.x_size();
        let ny = self.func.y_size();

        let ix = ((p.x() * nx as Float) as i32).min(nx as i32 - 1).max(0);
        let iy = ((p.y() * ny as Float) as i32).min(ny as i32 - 1).max(0);

        self.func[(ix, iy)]
    }

    fn sample_bisection<F>(p_func: F, u: Float, mut min: Float, mut max: Float, n: u32) -> Float
    where
        F: Fn(Float) -> Float,
    {
        let n_f = n as Float;

        while (n_f * max).ceil() - (n_f * min).floor() > 1.0 {
            debug_assert!(p_func(min) <= u + 1e-5);
            debug_assert!(p_func(max) >= u - 1e-5);

            let mid = (min + max) / 2.0;
            if p_func(mid) > u {
                max = mid;
            } else {
                min = mid;
            }
        }

        let numerator = u - p_func(min);
        let denominator = p_func(max) - p_func(min);

        let t = if denominator == 0.0 {
            0.5
        } else {
            numerator / denominator
        };

        let res = lerp(t, min, max);
        res.clamp(min, max)
    }
}

#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct Bin {
    pub q: Float,
    pub p: Float,
    pub alias: u32,
}

#[repr(C)]
#[derive(Debug, Clone)]
pub struct AliasTable {
    pub bins: GVec<Bin>,
}

unsafe impl Send for AliasTable {}
unsafe impl Sync for AliasTable {}

impl AliasTable {
    pub fn new(weights: &[Float]) -> Self {
        let n = weights.len();
        if n == 0 {
            return Self { bins: gvec() };
        }

        let sum: f64 = weights.iter().map(|&w| w as f64).sum();
        assert!(sum > 0.0, "Sum of weights must be positive");

        let mut bins = gvec_with_capacity(n);
        for &w in weights {
            bins.push(Bin {
                p: (w as f64 / sum) as Float,
                q: 0.0,
                alias: 0,
            });
        }

        struct Outcome {
            p_hat: f64,
            index: usize,
        }

        let mut under = gvec_with_capacity(n);
        let mut over = gvec_with_capacity(n);

        for (i, bin) in bins.iter().enumerate() {
            let p_hat = (bin.p as f64) * (n as f64);
            if p_hat < 1.0 {
                under.push(Outcome { p_hat, index: i });
            } else {
                over.push(Outcome { p_hat, index: i });
            }
        }

        while !under.is_empty() && !over.is_empty() {
            let un = under.pop().unwrap();
            let ov = over.pop().unwrap();
            bins[un.index].q = un.p_hat as Float;
            bins[un.index].alias = ov.index as u32;
            let p_excess = un.p_hat + ov.p_hat - 1.0;
            if p_excess < 1.0 {
                under.push(Outcome {
                    p_hat: p_excess,
                    index: ov.index,
                });
            } else {
                over.push(Outcome {
                    p_hat: p_excess,
                    index: ov.index,
                });
            }
        }

        while let Some(ov) = over.pop() {
            bins[ov.index].q = 1.0;
            bins[ov.index].alias = ov.index as u32;
        }
        while let Some(un) = under.pop() {
            bins[un.index].q = 1.0;
            bins[un.index].alias = un.index as u32;
        }

        Self { bins }
    }

    pub fn size(&self) -> u32 {
        self.bins.len() as u32
    }

    pub fn is_empty(&self) -> bool {
        self.bins.is_empty()
    }

    pub fn pmf(&self, index: u32) -> Float {
        if index >= self.size() {
            return 0.0;
        }
        self.bins[index as usize].p
    }

    pub fn sample(&self, u: Float) -> (u32, Float, Float) {
        if self.bins.is_empty() {
            return (0, 0.0, 0.0);
        }

        let n = self.bins.len() as Float;
        let val = u * n;
        let offset = (val.min(n - 1.0)) as u32;
        let up = (val - offset as Float).min(ONE_MINUS_EPSILON);

        let bin = &self.bins[offset as usize];

        if up < bin.q {
            let pmf = bin.p;
            let u_remapped = (up / bin.q).min(ONE_MINUS_EPSILON);
            (offset, pmf, u_remapped)
        } else {
            let alias_idx = bin.alias;
            let alias_p = self.bins[alias_idx as usize].p;
            let u_remapped = ((up - bin.q) / (1.0 - bin.q)).min(ONE_MINUS_EPSILON);
            (alias_idx, alias_p, u_remapped)
        }
    }
}

#[repr(C)]
#[derive(Debug, Clone)]
pub struct PiecewiseLinear2D<const N: usize> {
    pub size: Vector2i,
    pub inv_patch_size: Vector2f,
    pub param_size: [u32; N],
    pub param_strides: [u32; N],
    pub param_values: [GVec<Float>; N],
    pub data: GVec<Float>,
    pub marginal_cdf: GVec<Float>,
    pub conditional_cdf: GVec<Float>,
}

impl<const N: usize> PiecewiseLinear2D<N> {
    pub fn new(
        data: &[Float],
        x_size: i32,
        y_size: i32,
        param_res: [usize; N],
        param_values: [&[Float]; N],
        normalize: bool,
        build_cdf: bool,
    ) -> Self {
        if build_cdf && !normalize {
            panic!("PiecewiseLinear2D: build_cdf implies normalize=true");
        }

        let size = Vector2i::new(x_size, y_size);
        let inv_patch_size =
            Vector2f::new(1.0 / (x_size - 1) as Float, 1.0 / (y_size - 1) as Float);

        let mut param_size = [0u32; N];
        let mut param_strides = [0u32; N];
        let owned_param_values: [GVec<Float>; N] =
            core::array::from_fn(|i| gvec_from_slice(param_values[i]));

        let mut slices: u32 = 1;
        for i in (0..N).rev() {
            assert!(param_res[i] >= 1, "Parameter resolution must be >= 1");
            param_size[i] = param_res[i] as u32;
            param_strides[i] = if param_res[i] > 1 { slices } else { 0 };
            slices *= param_size[i];
        }

        let n_values = (x_size * y_size) as usize;
        let mut new_data = gvec_with_capacity(slices as usize * n_values);
        new_data.resize(slices as usize * n_values, 0.0);

        let mut marginal_cdf = if build_cdf {
            let mut v = gvec_with_capacity(slices as usize * y_size as usize);
            v.resize(slices as usize * y_size as usize, 0.0);
            v
        } else {
            gvec()
        };

        let mut conditional_cdf = if build_cdf {
            let mut v = gvec_with_capacity(slices as usize * n_values);
            v.resize(slices as usize * n_values, 0.0);
            v
        } else {
            gvec()
        };

        let mut data_offset = 0;
        for slice in 0..slices as usize {
            let slice_offset = slice * n_values;
            let current_data = &data[data_offset..data_offset + n_values];
            let mut sum = 0.0_f64;

            if normalize {
                for y in 0..(y_size - 1) {
                    for x in 0..(x_size - 1) {
                        let i = (y * x_size + x) as usize;
                        let v00 = current_data[i] as f64;
                        let v10 = current_data[i + 1] as f64;
                        let v01 = current_data[i + x_size as usize] as f64;
                        let v11 = current_data[i + 1 + x_size as usize] as f64;
                        sum += 0.25 * (v00 + v10 + v01 + v11);
                    }
                }
            }

            let normalization = if normalize && sum > 0.0 {
                1.0 / sum as Float
            } else {
                1.0
            };
            for k in 0..n_values {
                new_data[slice_offset + k] = current_data[k] * normalization;
            }

            if build_cdf {
                let marginal_slice_offset = slice * y_size as usize;
                for y in 0..y_size as usize {
                    let mut cdf_sum = 0.0;
                    let i_base = y * x_size as usize;
                    conditional_cdf[slice_offset + i_base] = 0.0;
                    for x in 0..(x_size - 1) as usize {
                        let i = i_base + x;
                        cdf_sum +=
                            0.5 * (new_data[slice_offset + i] + new_data[slice_offset + i + 1]);
                        conditional_cdf[slice_offset + i + 1] = cdf_sum;
                    }
                }
                marginal_cdf[marginal_slice_offset] = 0.0;
                let mut marginal_sum = 0.0;
                for y in 0..(y_size - 1) as usize {
                    let cdf1 = conditional_cdf[slice_offset + (y + 1) * x_size as usize - 1];
                    let cdf2 = conditional_cdf[slice_offset + (y + 2) * x_size as usize - 1];
                    marginal_sum += 0.5 * (cdf1 + cdf2);
                    marginal_cdf[marginal_slice_offset + y + 1] = marginal_sum;
                }
            }
            data_offset += n_values;
        }

        Self {
            size,
            inv_patch_size,
            param_size,
            param_strides,
            data: new_data,
            marginal_cdf,
            conditional_cdf,
            param_values: owned_param_values,
        }
    }

    pub fn sample(&self, mut sample: Point2f, params: [Float; N]) -> PLSample {
        sample = Point2f::new(
            sample.x().clamp(0.0, ONE_MINUS_EPSILON),
            sample.y().clamp(0.0, ONE_MINUS_EPSILON),
        );

        let (slice_offset, param_weights) = self.get_slice_info(params);
        let slice_size = (self.size.x() * self.size.y()) as u32;
        let marginal_size = self.size.y() as u32;
        let conditional_size = slice_size;
        let marginal_offset = slice_offset * marginal_size;
        let conditional_offset = slice_offset * conditional_size;
        let fetch_marginal = |idx: u32| {
            self.lookup(
                &self.marginal_cdf,
                marginal_offset + idx,
                marginal_size,
                &param_weights,
            )
        };
        let row = find_interval(marginal_size, |idx| fetch_marginal(idx) <= sample.y());
        let marginal_cdf_row = fetch_marginal(row);
        sample[1] -= marginal_cdf_row;
        let r0 = self.lookup(
            &self.conditional_cdf,
            conditional_offset + (row + 1) * self.size.x() as u32 - 1,
            conditional_size,
            &param_weights,
        );
        let r1 = self.lookup(
            &self.conditional_cdf,
            conditional_offset + (row + 2) * self.size.x() as u32 - 1,
            conditional_size,
            &param_weights,
        );
        let delta_r = r1 - r0;
        let is_const = delta_r.abs() < 1e-4 * (r0 + r1).abs();
        sample[1] /= if is_const { r0 + r1 } else { delta_r };
        if !is_const {
            sample[1] = r0 - safe_sqrt(r0 * r0 - 2.0 * sample.y() * delta_r);
        } else {
            sample[1] *= 2.0;
        }
        sample[0] *= (1.0 - sample.y()) * r0 + sample.y() * r1;
        let conditional_row_offset = conditional_offset + row * self.size.x() as u32;
        let fetch_conditional = |idx: u32| {
            let v0 = self.lookup(
                &self.conditional_cdf,
                conditional_row_offset + idx,
                conditional_size,
                &param_weights,
            );
            let v1 = self.lookup(
                &self.conditional_cdf,
                conditional_row_offset + idx + self.size.x() as u32,
                conditional_size,
                &param_weights,
            );
            (1.0 - sample.y()) * v0 + sample.y() * v1
        };
        let col = find_interval(self.size.x() as u32, |idx| {
            fetch_conditional(idx) <= sample.x()
        });
        sample[0] -= fetch_conditional(col);
        let offset = conditional_row_offset + col;
        let v00 = self.lookup(&self.data, offset, slice_size, &param_weights);
        let v10 = self.lookup(&self.data, offset + 1, slice_size, &param_weights);
        let v01 = self.lookup(
            &self.data,
            offset + self.size.x() as u32,
            slice_size,
            &param_weights,
        );
        let v11 = self.lookup(
            &self.data,
            offset + self.size.x() as u32 + 1,
            slice_size,
            &param_weights,
        );
        let c0 = (1.0 - sample.y()) * v00 + sample.y() * v01;
        let c1 = (1.0 - sample.y()) * v10 + sample.y() * v11;
        let delta_c = c1 - c0;
        let is_const = delta_c.abs() < 1e-4 * (c0 + c1).abs();
        sample[0] /= if is_const { c0 + c1 } else { delta_c };
        if !is_const {
            sample[0] = c0 - safe_sqrt(c0 * c0 - 2.0 * sample.x() * delta_c);
        } else {
            sample[0] *= 2.0;
        }
        let pdf = (1.0 - sample.x()) * c0 + sample.x() * c1;
        let offset_point = Point2f::new(col as Float, row as Float);

        PLSample {
            p: Point2f::new(
                sample.x() * self.inv_patch_size.x() + offset_point.x() * self.inv_patch_size.x(),
                sample.y() * self.inv_patch_size.y() + offset_point.y() * self.inv_patch_size.y(),
            ),
            pdf: pdf * self.inv_patch_size.x() * self.inv_patch_size.y(),
        }
    }

    pub fn invert(&self, mut p: Point2f, params: [Float; N]) -> PLSample {
        let (slice_offset, param_weights) = self.get_slice_info(params);

        p[0] *= self.inv_patch_size.x();
        p[1] *= self.inv_patch_size.y();
        let col = (p.x() as i32).min(self.size.x() - 2);
        let row = (p.y() as i32).min(self.size.y() - 2);
        let frac = Point2f::new(p.x() - col as Float, p.y() - row as Float);
        let slice_size = (self.size.x() * self.size.y()) as u32;
        let offset = slice_offset * slice_size + (row * self.size.x() + col) as u32;
        let v00 = self.lookup(&self.data, offset, slice_size, &param_weights);
        let v10 = self.lookup(&self.data, offset + 1, slice_size, &param_weights);
        let v01 = self.lookup(
            &self.data,
            offset + self.size.x() as u32,
            slice_size,
            &param_weights,
        );
        let v11 = self.lookup(
            &self.data,
            offset + self.size.x() as u32 + 1,
            slice_size,
            &param_weights,
        );
        let w1 = frac;
        let w0 = Point2f::new(1.0 - w1.x(), 1.0 - w1.y());
        let c0 = w0.y() * v00 + w1.y() * v01;
        let c1 = w0.y() * v10 + w1.y() * v11;
        let pdf = w0.x() * c0 + w1.x() * c1;
        let mut u = Point2f::new(0.0, 0.0);
        u[0] = w1.x() * (c0 + 0.5 * w1.x() * (c1 - c0));
        let conditional_row_offset = slice_offset * slice_size + (row * self.size.x()) as u32;
        let v0 = self.lookup(
            &self.conditional_cdf,
            conditional_row_offset + col as u32,
            slice_size,
            &param_weights,
        );
        let v1 = self.lookup(
            &self.conditional_cdf,
            conditional_row_offset + col as u32 + self.size.x() as u32,
            slice_size,
            &param_weights,
        );
        u[0] += (1.0 - u.y()) * v0 + u.y() * v1;
        let r0 = self.lookup(
            &self.conditional_cdf,
            conditional_row_offset + self.size.x() as u32 - 1,
            slice_size,
            &param_weights,
        );
        let r1 = self.lookup(
            &self.conditional_cdf,
            conditional_row_offset + self.size.x() as u32 * 2 - 1,
            slice_size,
            &param_weights,
        );
        u[0] /= (1.0 - u.y()) * r0 + u.y() * r1;
        u[1] = w1.y() * (r0 + 0.5 * w1.y() * (r1 - r0));
        let marginal_offset = slice_offset * self.size.y() as u32 + row as u32;
        u[1] += self.lookup(
            &self.marginal_cdf,
            marginal_offset,
            self.size.y() as u32,
            &param_weights,
        );

        PLSample {
            p: u,
            pdf: pdf * (self.inv_patch_size.x() * self.inv_patch_size.y()),
        }
    }

    pub fn evaluate(&self, p: Point2f, params: [Float; N]) -> Float {
        let (slice_offset, param_weights) = self.get_slice_info(params);
        let pos = Point2f::new(
            p.x() * self.inv_patch_size.x(),
            p.y() * self.inv_patch_size.y(),
        );
        let col = (pos.x() as i32).min(self.size.x() - 2);
        let row = (pos.y() as i32).min(self.size.y() - 2);
        let w1 = Point2f::new(pos.x() - col as Float, pos.y() - row as Float);
        let w0 = Point2f::new(1.0 - w1.x(), 1.0 - w1.y());
        let slice_size = (self.size.x() * self.size.y()) as u32;
        let offset = slice_offset * slice_size + (row * self.size.x() + col) as u32;
        let v00 = self.lookup(&self.data, offset, slice_size, &param_weights);
        let v10 = self.lookup(&self.data, offset + 1, slice_size, &param_weights);
        let v01 = self.lookup(
            &self.data,
            offset + self.size.x() as u32,
            slice_size,
            &param_weights,
        );
        let v11 = self.lookup(
            &self.data,
            offset + self.size.x() as u32 + 1,
            slice_size,
            &param_weights,
        );

        let pdf = w0.y() * (w0.x() * v00 + w1.x() * v10) + w1.y() * (w0.x() * v01 + w1.x() * v11);
        pdf * self.inv_patch_size.x() * self.inv_patch_size.y()
    }

    #[inline(always)]
    fn get_param_value(&self, dim: usize, idx: usize) -> Float {
        unsafe { *self.param_values[dim].as_ptr().add(idx) }
    }

    fn get_slice_info(&self, params: [Float; N]) -> (u32, [(Float, Float); N]) {
        let mut param_weight = [(0.0, 0.0); N];
        let mut slice_offset = 0u32;

        for dim in 0..N {
            let size = self.param_size[dim];

            if size == 1 {
                param_weight[dim] = (1.0, 0.0);
                continue;
            }

            let param_index = find_interval(size, |idx| {
                self.get_param_value(dim, idx as usize) <= params[dim]
            }) as u32;

            let p0 = self.get_param_value(dim, param_index as usize);

            let next_index = (param_index + 1).min(size - 1);
            let p1 = self.get_param_value(dim, next_index as usize);

            let w1 = if p1 != p0 {
                ((params[dim] - p0) / (p1 - p0)).clamp(0.0, 1.0)
            } else {
                0.0
            };

            param_weight[dim] = (1.0 - w1, w1);

            slice_offset += self.param_strides[dim] * param_index;
        }

        (slice_offset, param_weight)
    }

    fn lookup(
        &self,
        data: &GVec<Float>,
        i0: u32,
        size: u32,
        param_weight: &[(Float, Float); N],
    ) -> Float {
        let mut result: Float = 0.0;
        let num_corners = 1usize << N; // 2^N
        for mask in 0..num_corners {
            let mut offset = 0u32;
            let mut weight: Float = 1.0;
            let mut current_mask = mask;
            for (weights, &stride) in param_weight.iter().zip(&self.param_strides) {
                if (current_mask & 1) == 1 {
                    offset += stride * size;
                    weight *= weights.1;
                } else {
                    weight *= weights.0;
                }

                current_mask >>= 1;
            }
            let idx = (i0 + offset) as usize;
            let val = unsafe { *data.as_ptr().add(idx) };
            result += weight * val;
        }
        result
    }
}

#[repr(C)]
#[derive(Clone, Copy, Debug)]
pub struct WeightedReservoirSampler<T> {
    rng: Rng,
    weight_sum: Float,
    reservoir_weight: Float,
    reservoir: T,
}

impl<T: Default + Copy> Default for WeightedReservoirSampler<T> {
    fn default() -> Self {
        Self {
            rng: Rng::default(),
            weight_sum: 0.0,
            reservoir_weight: 0.0,
            reservoir: T::default(),
        }
    }
}

impl<T: Default + Copy> WeightedReservoirSampler<T> {
    pub fn new(seed: u64) -> Self {
        let mut rng = Rng::default();
        rng.set_sequence(seed);
        Self {
            rng,
            weight_sum: 0.0,
            reservoir_weight: 0.0,
            reservoir: T::default(),
        }
    }

    pub fn seed(&mut self, seed: u64) {
        self.rng.set_sequence(seed);
    }

    pub fn add(&mut self, sample: T, weight: Float) -> bool {
        self.weight_sum += weight;

        let p = weight / self.weight_sum;
        if self.rng.uniform::<Float>() < p {
            self.reservoir = sample;
            self.reservoir_weight = weight;
            return true;
        }

        // debug_assert!(self.weight_sum < 1e80);
        false
    }

    pub fn add_closure<F>(&mut self, func: F, weight: Float) -> bool
    where
        F: FnOnce() -> T,
    {
        self.weight_sum += weight;

        let p = weight / self.weight_sum;
        if self.rng.uniform::<Float>() < p {
            self.reservoir = func();
            self.reservoir_weight = weight;
            return true;
        }

        // debug_assert!(self.weight_sum < 1e80);
        false
    }

    pub fn copy_from(&mut self, other: &Self)
    where
        T: Clone,
    {
        self.weight_sum = other.weight_sum;
        self.reservoir = other.reservoir.clone();
        self.reservoir_weight = other.reservoir_weight;
    }

    pub fn has_sample(&self) -> bool {
        self.weight_sum > 0.0
    }

    pub fn get_sample(&self) -> Option<&T> {
        if self.has_sample() {
            Some(&self.reservoir)
        } else {
            None
        }
    }

    pub fn sample_probability(&self) -> Float {
        if self.weight_sum == 0.0 {
            0.0
        } else {
            self.reservoir_weight / self.weight_sum
        }
    }

    pub fn weight_sum(&self) -> Float {
        self.weight_sum
    }

    pub fn reset(&mut self) {
        self.reservoir_weight = 0.0;
        self.weight_sum = 0.0;
        self.reservoir = T::default();
    }

    pub fn merge(&mut self, other: &Self) {
        if other.has_sample() {
            self.add(other.reservoir, other.weight_sum);
        }
    }
}
