use crate::Float;
use crate::core::texture::{
    GPUFloatTexture, GPUSpectrumTexture, TextureEvalContext, TextureMapping2D, TextureMapping3D,
    TextureMapping3DTrait,
};
use crate::spectra::{SampledSpectrum, SampledWavelengths};
use crate::utils::{Ptr, math::square};
use num_traits::Float as NumFloat;

fn checkerboard(
    ctx: &TextureEvalContext,
    map2d: Ptr<TextureMapping2D>,
    map3d: Ptr<TextureMapping3D>,
) -> Float {
    let d = |x: Float| -> Float {
        let y = x / 2. - (x / 2.).floor() - 0.5;
        return x / 2. + y * (1. - 2. * y.abs());
    };

    let bf = |x: Float, r: Float| -> Float {
        if (x.floor() - r) == (x + r).floor() {
            return 1. - 2. * (x.floor() as i32 & 1) as Float;
        }
        (d(x + r) - 2. * d(x) + d(x - r)) / square(r)
    };

    if !map2d.is_null() {
        assert!(map3d.is_null());
        let c = map2d.map(&ctx);
        let ds = 1.5 * c.dsdx.abs().max(c.dsdy.abs());
        let dt = 1.5 * c.dtdx.abs().max(c.dtdy.abs());
        // Integrate product of 2D checkerboard function and triangle filter
        0.5 - bf(c.st[0], ds) * bf(c.st[1], dt) / 2.
    } else {
        assert!(!map3d.is_null());
        let c = map3d.map(&ctx);
        let dx = 1.5 * c.dpdx.x().abs().max(c.dpdy.x().abs());
        let dy = 1.5 * c.dpdx.y().abs().max(c.dpdy.y().abs());
        let dz = 1.5 * c.dpdx.z().abs().max(c.dpdy.z().abs());
        0.5 - bf(c.p.x(), dx) * bf(c.p.y(), dy) * bf(c.p.z(), dz)
    }
}

#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct FloatCheckerboardTexture {
    pub map2d: Ptr<TextureMapping2D>,
    pub map3d: Ptr<TextureMapping3D>,
    pub tex: [Ptr<GPUFloatTexture>; 2],
}

impl FloatCheckerboardTexture {
    pub fn evaluate(&self, ctx: &TextureEvalContext) -> Float {
        let w = checkerboard(&ctx, self.map2d, self.map3d);

        let mut t0 = 0.0;
        let mut t1 = 0.0;

        if w != 1.0 {
            if let Some(tex) = self.tex[0].get() {
                t0 = tex.evaluate(ctx);
            }
        }

        if w != 0.0 {
            if let Some(tex) = self.tex[1].get() {
                t1 = tex.evaluate(ctx);
            }
        }

        (1.0 - w) * t0 + w * t1
    }
}

#[repr(C)]
#[derive(Clone, Copy, Debug)]
pub struct SpectrumCheckerboardTexture {
    pub map2d: Ptr<TextureMapping2D>,
    pub map3d: Ptr<TextureMapping3D>,
    pub tex: [Ptr<GPUSpectrumTexture>; 2],
}

impl SpectrumCheckerboardTexture {
    pub fn evaluate(
        &self,
        ctx: &TextureEvalContext,
        lambda: &SampledWavelengths,
    ) -> SampledSpectrum {
        let w = checkerboard(ctx, self.map2d, self.map3d);
        let mut t0 = SampledSpectrum::new(0.);
        let mut t1 = SampledSpectrum::new(0.);
        if w != 1.0 {
            if let Some(tex) = self.tex[0].get() {
                t0 = tex.evaluate(ctx, lambda);
            }
        }

        if w != 0.0 {
            if let Some(tex) = self.tex[1].get() {
                t1 = tex.evaluate(ctx, lambda);
            }
        }

        t0 * (1.0 - w) + t1 * w
    }
}