use crate::core::image::Image;
use crate::utils::arena::Arena;
use crate::utils::backend::GpuAllocator;
use crate::utils::containers::Array2D;
use shared::core::geometry::{Bounds2f, Point2i, Vector2f, Vector2i};
use shared::utils::sampling::{
    AliasTable, Bin, DevicePiecewiseConstant1D, DevicePiecewiseConstant2D, DeviceSummedAreaTable,
    DeviceWindowedPiecewiseConstant2D, PiecewiseLinear2D,
};
use shared::utils::{gpu_array_from_fn, Ptr};
use shared::Float;
use std::sync::Arc;

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

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

    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 values: Vec<Float> = (0..n)
            .map(|i| f(min + (i as Float + 0.5) * delta))
            .collect();
        Self::new_with_bounds(values, min, max)
    }

    pub fn new_with_bounds(f: Vec<Float>, min: Float, max: Float) -> Self {
        let n = f.len();
        let mut cdf = Vec::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: f, cdf, min, max }
    }

    // Accessors
    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 {
        // func_integral is the un-normalized sum. After normalization cdf[n] == 1.0,
        // so we reconstruct from the last CDF entry before normalization.
        // But since we normalized in-place, we need to store it. Let's compute it.
        let n = self.func.len();
        let delta = (self.max - self.min) / n as Float;
        self.func.iter().sum::<Float>() * delta
    }

    /// Host-side sampling (for scene construction, not rendering).
    /// During rendering, use the device struct via arena-uploaded Ptrs.
    pub fn sample_host(&self, u: Float) -> (Float, Float, usize) {
        let n = self.func.len();
        let offset = self.find_interval_host(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 delta = (self.max - self.min) / n as Float;
        let x = self.min + (offset as Float + du) * delta;
        let func_integral = self.integral();
        let pdf = if func_integral > 0.0 {
            self.func[offset] / func_integral
        } else {
            0.0
        };
        (x, pdf, offset)
    }

    fn find_interval_host(&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)
    }
}

#[derive(DeviceRepr)]
#[device(name = "DevicePiecewiseConstant1D")]
pub struct PiecewiseConstant1D {
    pub func: Vec<Float>,
    pub cdf: Vec<Float>,
    pub min: Float,
    pub max: Float,
    pub n: u32,
    #[device(expr = "self.integral()")]
    pub func_integral: Float,
}

#[derive(Clone, Debug, DeviceRepr)]
#[device(name = "DevicePiecewiseConstant2D")]
pub struct PiecewiseConstant2D {
    pub conditionals: Vec<PiecewiseConstant1D>,
    #[device(flatten)]
    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 = Vec::with_capacity(n_v);
        let mut marginal_func = Vec::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_v }
    }

    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 = Vec::with_capacity(n_u * n_v);
        for v in 0..n_v {
            for u in 0..n_u {
                data.push(image.get_channels(Point2i::new(u as i32, v as i32)).average());
            }
        }

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

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

struct PiecewiseLinear2DStorage<const N: usize> {
    data: Vec<Float>,
    marginal_cdf: Vec<Float>,
    conditional_cdf: Vec<Float>,
    param_values: [Vec<Float>; N],
}

pub struct PiecewiseLinear2DHost<const N: usize> {
    size: Vector2i,
    inv_patch_size: Vector2f,
    param_size: [u32; N],
    param_strides: [u32; N],
    storage: Arc<PiecewiseLinear2DStorage<N>>,
}


impl<const N: usize> PiecewiseLinear2DHost<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::new: build_cdf implies normalize=true");
        }

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

        let mut param_size = [0u32; N];
        let mut param_strides = [0u32; N];
        let param_values = gpu_array_from_fn(|i| param_values[i].to_vec());

        let mut slices: u32 = 1;
        for i in (0..N).rev() {
            if param_res[i] < 1 {
                panic!("PiecewiseLinear2D::new: 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 = vec![0.0; slices as usize * n_values];
        let mut marginal_cdf = if build_cdf {
            vec![0.0; slices as usize * y_size as usize]
        } else {
            Vec::new()
        };
        let mut conditional_cdf = if build_cdf {
            vec![0.0; slices as usize * n_values]
        } else {
            Vec::new()
        };

        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.;

            // Construct conditional CDF
            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;
                // Construct marginal CDF
                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;
                    }
                }
                // Construct marginal CDF
                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;
        }

        let view = PiecewiseLinear2D {
            size,
            inv_patch_size,
            param_size,
            param_strides,
            param_values: param_values.each_ref().map(|x| x.as_ptr().into()),
            data: Ptr::null(),
            marginal_cdf: marginal_cdf.as_ptr().into(),
            conditional_cdf: conditional_cdf.as_ptr().into(),
        };

        let storage = Arc::new(PiecewiseLinear2DStorage {
            data: new_data,
            marginal_cdf,
            conditional_cdf,
            param_values,
        });

        let mut final_view = view;
        final_view.data = storage.data.as_ptr().into();
        final_view.marginal_cdf = storage.marginal_cdf.as_ptr().into();
        final_view.conditional_cdf = storage.conditional_cdf.as_ptr().into();
        for i in 0..N {
            final_view.param_values[i] = storage.param_values[i].as_ptr().into();
        }

        Self {
            view: final_view,
            _storage: storage,
        }
    }
}

impl<const N: usize> DeviceRepr for PiecewiseLinear2DHost<N> {
    type Target = PiecewiseLinear2D<N>;

    fn upload_value<A: GpuAllocator>(&self, arena: &Arena<A>) -> PiecewiseLinear2D<N> {
        let s = &self.storage;

        let (data_ptr, _) = arena.alloc_slice(&s.data);
        let (marginal_ptr, _) = arena.alloc_slice(&s.marginal_cdf);
        let (conditional_ptr, _) = arena.alloc_slice(&s.conditional_cdf);

        let param_ptrs: [Ptr<Float>; N] = std::array::from_fn(|i| {
            let (ptr, _) = arena.alloc_slice(&s.param_values[i]);
            ptr
        });

        PiecewiseLinear2D {
            size: self.size,
            inv_patch_size: self.inv_patch_size,
            param_size: self.param_size,
            param_strides: self.param_strides,
            param_values: param_ptrs,
            data: data_ptr,
            marginal_cdf: marginal_ptr,
            conditional_cdf: conditional_ptr,
        }
    }
}

#[derive(Debug, Clone)]
pub struct AliasTableHost {
    bins: Vec<Bin>,
}

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

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

        let mut bins = Vec::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 = Vec::with_capacity(n);
        let mut over = Vec::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) -> usize { self.bins.len() }
    pub fn is_empty(&self) -> bool { self.bins.is_empty() }
}

impl DeviceRepr for AliasTableHost {
    type Target = AliasTable;

    fn upload_value<A: GpuAllocator>(&self, arena: &Arena<A>) -> AliasTable {
        if self.bins.is_empty() {
            return AliasTable { bins: Ptr::null(), size: 0 };
        }
        let (bins_ptr, _) = arena.alloc_slice(&self.bins);
        AliasTable {
            bins: bins_ptr,
            size: self.bins.len() as u32,
        }
    }
}

#[derive(Clone, Debug, DeviceRepr)]
#[device(name = "DeviceSummedAreaTable")]
pub struct SummedAreaTable {
    #[device(flatten)]
    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 }
    }
}

#[derive(Clone, Debug, DeviceRepr)]
#[device(name = "DeviceWindowedPiecewiseConstant2D")]
pub struct WindowedPiecewiseConstant2D {
    #[device(flatten)]
    sat: SummedAreaTable,
    #[device(flatten)]
    func: Array2D<Float>,
}

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

struct PiecewiseLinear2DStorage<const N: usize> {
    data: Vec<Float>,
    marginal_cdf: Vec<Float>,
    conditional_cdf: Vec<Float>,
    param_values: [Vec<Float>; N],
}

pub struct PiecewiseLinear2DHost<const N: usize> {
    size: Vector2i,
    inv_patch_size: Vector2f,
    param_size: [u32; N],
    param_strides: [u32; N],
    storage: Arc<PiecewiseLinear2DStorage<N>>,
}

impl<const N: usize> PiecewiseLinear2DHost<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: [Vec<Float>; N] = gpu_array_from_fn(|i| param_values[i].to_vec());

        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 = vec![0.0; slices as usize * n_values];
        let mut marginal_cdf = if build_cdf {
            vec![0.0; slices as usize * y_size as usize]
        } else {
            Vec::new()
        };
        let mut conditional_cdf = if build_cdf {
            vec![0.0; slices as usize * n_values]
        } else {
            Vec::new()
        };

        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;
        }

        let storage = Arc::new(PiecewiseLinear2DStorage {
            data: new_data,
            marginal_cdf,
            conditional_cdf,
            param_values: owned_param_values,
        });

        Self {
            size,
            inv_patch_size,
            param_size,
            param_strides,
            storage,
        }
    }
}

impl<const N: usize> DeviceRepr for PiecewiseLinear2DHost<N> {
    type Target = PiecewiseLinear2D<N>;

    fn upload_value<A: GpuAllocator>(&self, arena: &Arena<A>) -> PiecewiseLinear2D<N> {
        let s = &self.storage;

        let (data_ptr, _) = arena.alloc_slice(&s.data);
        let (marginal_ptr, _) = arena.alloc_slice(&s.marginal_cdf);
        let (conditional_ptr, _) = arena.alloc_slice(&s.conditional_cdf);

        let param_ptrs: [Ptr<Float>; N] = std::array::from_fn(|i| {
            let (ptr, _) = arena.alloc_slice(&s.param_values[i]);
            ptr
        });

        PiecewiseLinear2D {
            size: self.size,
            inv_patch_size: self.inv_patch_size,
            param_size: self.param_size,
            param_strides: self.param_strides,
            param_values: param_ptrs,
            data: data_ptr,
            marginal_cdf: marginal_ptr,
            conditional_cdf: conditional_ptr,
        }
    }
}