use crate::materials::*;
use enum_dispatch::enum_dispatch;
use std::ops::Deref;

use crate::Float;
use crate::bxdfs::{
    CoatedConductorBxDF, CoatedDiffuseBxDF, ConductorBxDF, DielectricBxDF, DiffuseBxDF,
};
use crate::core::bsdf::BSDF;
use crate::core::bssrdf::BSSRDF;
use crate::core::bxdf::BxDF;
use crate::core::geometry::{Frame, Normal3f, Point2f, Point3f, Vector2f, Vector3f, VectorLike};
use crate::core::image::{DeviceImage, WrapMode, WrapMode2D};
use crate::core::interaction::{Interaction, InteractionTrait, ShadingGeom, SurfaceInteraction};
use crate::core::scattering::TrowbridgeReitzDistribution;
use crate::core::spectrum::{Spectrum, SpectrumTrait};
use crate::core::texture::{
    GPUFloatTexture, GPUSpectrumTexture, TextureEvalContext, TextureEvaluator,
};
use crate::materials::*;
use crate::spectra::{SampledSpectrum, SampledWavelengths};
use crate::utils::Ptr;
use crate::utils::hash::hash_float;
use crate::utils::math::clamp;

#[repr(C)]
#[derive(Clone, Debug, Copy)]
pub struct MaterialEvalContext {
    pub texture: TextureEvalContext,
    pub wo: Vector3f,
    pub ns: Normal3f,
    pub dpdus: Vector3f,
}

impl Deref for MaterialEvalContext {
    type Target = TextureEvalContext;

    fn deref(&self) -> &Self::Target {
        &self.texture
    }
}

impl From<&SurfaceInteraction> for MaterialEvalContext {
    fn from(si: &SurfaceInteraction) -> Self {
        Self {
            texture: TextureEvalContext::from(si),
            wo: si.common.wo,
            ns: si.shading.n,
            dpdus: si.shading.dpdu,
        }
    }
}

#[repr(C)]
#[derive(Clone, Debug, Default)]
pub struct NormalBumpEvalContext {
    pub p: Point3f,
    pub uv: Point2f,
    pub n: Normal3f,
    pub shading: ShadingGeom,
    pub dpdx: Vector3f,
    pub dpdy: Vector3f,
    // All 0
    pub dudx: Float,
    pub dudy: Float,
    pub dvdx: Float,
    pub dvdy: Float,
    pub face_index: i32,
}

impl From<&SurfaceInteraction> for NormalBumpEvalContext {
    fn from(si: &SurfaceInteraction) -> Self {
        Self {
            p: si.p(),
            uv: si.common.uv,
            n: si.n(),
            shading: si.shading.clone(),
            dudx: si.dudx,
            dudy: si.dudy,
            dvdx: si.dvdx,
            dvdy: si.dvdy,
            dpdx: si.dpdx,
            dpdy: si.dpdy,
            face_index: si.face_index,
        }
    }
}

impl From<&NormalBumpEvalContext> for TextureEvalContext {
    fn from(ctx: &NormalBumpEvalContext) -> Self {
        Self {
            p: ctx.p,
            uv: ctx.uv,
            n: ctx.n,
            dpdx: ctx.dpdx,
            dpdy: ctx.dpdy,
            dudx: ctx.dudx,
            dudy: ctx.dudy,
            dvdx: ctx.dvdx,
            dvdy: ctx.dvdy,
            face_index: ctx.face_index,
        }
    }
}

pub fn normal_map(normal_map: &DeviceImage, ctx: &NormalBumpEvalContext) -> (Vector3f, Vector3f) {
    let wrap = WrapMode2D::from(WrapMode::Repeat);
    let uv = Point2f::new(ctx.uv[0], 1. - ctx.uv[1]);
    let r = normal_map.bilerp_channel_with_wrap(uv, 0, wrap);
    let g = normal_map.bilerp_channel_with_wrap(uv, 1, wrap);
    let b = normal_map.bilerp_channel_with_wrap(uv, 2, wrap);
    let mut ns = Vector3f::new(2.0 * r - 1.0, 2.0 * g - 1.0, 2.0 * b - 1.0);
    ns = ns.normalize();
    let frame = Frame::from_xz(ctx.shading.dpdu.normalize(), Vector3f::from(ctx.shading.n));
    ns = frame.from_local(ns);
    let ulen = ctx.shading.dpdu.norm();
    let vlen = ctx.shading.dpdv.norm();
    let dpdu = ctx.shading.dpdu.gram_schmidt(ns).normalize() * ulen;
    let dpdv = ctx.shading.dpdu.cross(dpdu).normalize() * vlen;
    (dpdu, dpdv)
}

pub fn bump_map<T: TextureEvaluator>(
    tex_eval: &T,
    displacement: &GPUFloatTexture,
    ctx: &NormalBumpEvalContext,
) -> (Vector3f, Vector3f) {
    debug_assert!(tex_eval.can_evaluate(&[Ptr::from(displacement)], &[]));
    let mut du = 0.5 * (ctx.dudx.abs() + ctx.dudy.abs());
    if du == 0.0 {
        du = 0.0005;
    }

    let mut dv = 0.5 * (ctx.dvdx.abs() + ctx.dvdy.abs());
    if dv == 0.0 {
        dv = 0.0005;
    }
    let mut shifted_ctx = TextureEvalContext::from(ctx);
    shifted_ctx.p = ctx.p + ctx.shading.dpdu * du;
    shifted_ctx.uv = ctx.uv + Vector2f::new(du, 0.0);
    let u_displace = tex_eval.evaluate_float(displacement, &shifted_ctx);
    shifted_ctx.p = ctx.p + ctx.shading.dpdv * dv;
    shifted_ctx.uv = ctx.uv + Vector2f::new(0.0, dv);
    let v_displace = tex_eval.evaluate_float(displacement, &shifted_ctx);
    let center_ctx = TextureEvalContext::from(ctx);
    let displace = tex_eval.evaluate_float(displacement, &center_ctx);

    let d_displace_du = (u_displace - displace) / du;
    let d_displace_dv = (v_displace - displace) / dv;
    let n_vec = Vector3f::from(ctx.shading.n);
    let dndu_vec = Vector3f::from(ctx.shading.dndu);
    let dndv_vec = Vector3f::from(ctx.shading.dndv);

    let dpdu = ctx.shading.dpdu + n_vec * d_displace_du + dndu_vec * displace;
    let dpdv = ctx.shading.dpdv + n_vec * d_displace_dv + dndv_vec * displace;
    (dpdu, dpdv)
}

#[enum_dispatch]
pub trait MaterialTrait {
    fn get_bsdf<T: TextureEvaluator>(
        &self,
        tex_eval: &T,
        ctx: &MaterialEvalContext,
        lambda: &SampledWavelengths,
    ) -> BSDF;

    fn get_bssrdf<T: TextureEvaluator>(
        &self,
        tex_eval: &T,
        ctx: &MaterialEvalContext,
        lambda: &SampledWavelengths,
    ) -> Option<BSSRDF>;

    fn can_evaluate_textures(&self, tex_eval: &dyn TextureEvaluator) -> bool;
    fn get_normal_map(&self) -> Option<&DeviceImage>;
    fn get_displacement(&self) -> Ptr<GPUFloatTexture>;
    fn has_subsurface_scattering(&self) -> bool;
}

#[derive(Clone, Copy, Debug)]
#[enum_dispatch(MaterialTrait)]
pub enum Material {
    CoatedDiffuse(CoatedDiffuseMaterial),
    CoatedConductor(CoatedConductorMaterial),
    Conductor(ConductorMaterial),
    Dielectric(DielectricMaterial),
    Diffuse(DiffuseMaterial),
    DiffuseTransmission(DiffuseTransmissionMaterial),
    Hair(HairMaterial),
    Measured(MeasuredMaterial),
    Subsurface(SubsurfaceMaterial),
    ThinDielectric(ThinDielectricMaterial),
    Mix(MixMaterial),
}