use crate::Float;
use crate::bxdfs::DiffuseBxDF;
use crate::core::bsdf::BSDF;
use crate::core::bssrdf::BSSRDF;
use crate::core::bxdf::{BxDF, BxDFFlags};
use crate::core::camera::{Camera, CameraTrait};
use crate::core::geometry::{
    Normal3f, Point2f, Point3f, Point3fi, Ray, RayDifferential, Vector3f, VectorLike,
};
use crate::core::image::Image;
use crate::core::light::{Light, LightTrait};
use crate::core::material::{
    Material, MaterialEvalContext, MaterialTrait, NormalBumpEvalContext, bump_map, normal_map,
};
use crate::core::medium::{Medium, MediumInterface, PhaseFunction};
use crate::core::sampler::{Sampler, SamplerTrait};
use crate::core::shape::Shape;
use crate::core::texture::{GPUFloatTexture, UniversalTextureEvaluator};
use crate::spectra::{SampledSpectrum, SampledWavelengths};
use crate::utils::Ptr;
use crate::utils::math::{clamp, difference_of_products, square};
use enum_dispatch::enum_dispatch;

#[repr(C)]
#[derive(Default, Copy, Clone, Debug)]
pub struct InteractionBase {
    pub pi: Point3fi,
    pub n: Normal3f,
    pub time: Float,
    pub wo: Vector3f,
    pub uv: Point2f,
    pub medium_interface: MediumInterface,
    pub medium: Ptr<Medium>,
}

impl InteractionBase {
    pub fn new_surface_geom(
        pi: Point3fi,
        n: Normal3f,
        uv: Point2f,
        wo: Vector3f,
        time: Float,
    ) -> Self {
        Self {
            pi,
            n,
            uv,
            wo: wo.normalize(),
            time,
            medium_interface: MediumInterface::default(),
            medium: Ptr::null(),
        }
    }

    pub fn new_medium(p: Point3f, wo: Vector3f, time: Float, medium: Ptr<Medium>) -> Self {
        Self {
            pi: Point3fi::new_from_point(p),
            n: Normal3f::zero(),
            uv: Point2f::default(),
            wo: wo.normalize(),
            time,
            medium_interface: MediumInterface::default(),
            medium,
        }
    }

    pub fn new_minimal(pi: Point3fi, n: Normal3f) -> Self {
        Self {
            pi,
            n,
            ..Default::default()
        }
    }

    pub fn new_boundary(p: Point3f, time: Float, medium_interface: MediumInterface) -> Self {
        Self {
            pi: Point3fi::new_from_point(p),
            time,
            medium_interface,
            ..Default::default()
        }
    }
}

#[enum_dispatch]
pub trait InteractionTrait {
    fn get_common(&self) -> &InteractionBase;
    fn get_common_mut(&mut self) -> &mut InteractionBase;

    fn p(&self) -> Point3f {
        self.get_common().pi.into()
    }
    fn pi(&self) -> Point3fi {
        self.get_common().pi
    }

    fn time(&self) -> Float {
        self.get_common().time
    }
    fn wo(&self) -> Vector3f {
        self.get_common().wo
    }

    fn n(&self) -> Normal3f {
        self.get_common().n
    }

    fn is_surface_interaction(&self) -> bool {
        false
    }

    fn is_medium_interaction(&self) -> bool {
        false
    }

    fn get_medium(&self, w: Vector3f) -> Ptr<Medium> {
        let data = self.get_common();
        if !data.medium_interface.inside.is_null() || !data.medium_interface.outside.is_null() {
            if w.dot(data.n.into()) > 0.0 {
                data.medium_interface.outside
            } else {
                data.medium_interface.inside
            }
        } else {
            data.medium
        }
    }

    fn spawn_ray(&self, d: Vector3f) -> Ray {
        let data = self.get_common();
        let mut ray = Ray::spawn(&data.pi, &data.n, data.time, d);

        ray.medium = self.get_medium(d);
        ray
    }

    fn spawn_ray_to_point(&self, p2: Point3f) -> Ray {
        let data = self.get_common();
        let mut ray = Ray::spawn_to_point(&data.pi, &data.n, data.time, p2);
        ray.medium = self.get_medium(ray.d);
        ray
    }

    fn spawn_ray_to_interaction(&self, other: InteractionBase) -> Ray {
        let data = self.get_common();

        let mut ray = Ray::spawn_to_interaction(&data.pi, &data.n, data.time, &other.pi, &other.n);
        ray.medium = self.get_medium(ray.d);
        ray
    }

    fn offset_ray_vector(&self, w: Vector3f) -> Point3f {
        Ray::offset_origin(&self.pi(), &self.n(), &w)
    }

    fn offset_ray_point(&self, pt: Point3f) -> Point3f {
        self.offset_ray_vector(pt - self.p())
    }
}

#[repr(C)]
#[enum_dispatch(InteractionTrait)]
#[derive(Debug, Copy, Clone)]
pub enum Interaction {
    Surface(SurfaceInteraction),
    Medium(MediumInteraction),
    Simple(SimpleInteraction),
}

impl Interaction {
    pub fn set_medium_interface(&mut self, mi: MediumInterface) {
        match self {
            Interaction::Surface(si) => si.common.medium_interface = mi,
            Interaction::Medium(_) => {} // Medium interactions don't usually sit on boundaries
            Interaction::Simple(si) => si.common.medium_interface = mi,
        }
    }
}

#[repr(C)]
#[derive(Debug, Default, Clone, Copy)]
pub struct SimpleInteraction {
    pub common: InteractionBase,
}

impl SimpleInteraction {
    pub fn new(common: InteractionBase) -> Self {
        Self { common }
    }
}

impl InteractionTrait for SimpleInteraction {
    fn get_common(&self) -> &InteractionBase {
        &self.common
    }
    fn get_common_mut(&mut self) -> &mut InteractionBase {
        &mut self.common
    }
    fn is_surface_interaction(&self) -> bool {
        false
    }
    fn is_medium_interaction(&self) -> bool {
        false
    }
}

#[repr(C)]
#[derive(Default, Clone, Copy, Debug)]
pub struct ShadingGeom {
    pub n: Normal3f,
    pub dpdu: Vector3f,
    pub dpdv: Vector3f,
    pub dndu: Normal3f,
    pub dndv: Normal3f,
}

#[repr(C)]
#[derive(Debug, Default, Clone, Copy)]
pub struct SurfaceInteraction {
    pub area_light: Ptr<Light>,
    pub material: Ptr<Material>,
    pub shape: Ptr<Shape>,
    pub common: InteractionBase,
    pub shading: ShadingGeom,
    pub dpdu: Vector3f,
    pub dpdv: Vector3f,
    pub dndu: Normal3f,
    pub dndv: Normal3f,
    pub dpdx: Vector3f,
    pub dpdy: Vector3f,
    pub face_index: i32,
    pub dudx: Float,
    pub dvdx: Float,
    pub dudy: Float,
    pub dvdy: Float,
}

unsafe impl Send for SurfaceInteraction {}
unsafe impl Sync for SurfaceInteraction {}

impl SurfaceInteraction {
    pub fn le(&self, w: Vector3f, lambda: &SampledWavelengths) -> SampledSpectrum {
        if !self.area_light.is_null() {
            self.area_light
                .l(self.p(), self.n(), self.common.uv, w, lambda)
        } else {
            SampledSpectrum::new(0.)
        }
    }

    pub fn compute_differentials(&mut self, r: &Ray, camera: &Camera, samples_per_pixel: i32) {
        let computed = if !r.has_differentials {
            let diff = r.differential;
            let dot_rx = self.common.n.dot(diff.rx_direction.into());
            let dot_ry = self.common.n.dot(diff.ry_direction.into());

            if dot_rx != 0.0 && dot_ry != 0.0 {
                //  Estimate screen-space change in p using ray differentials>
                let p_as_vec = Normal3f::new(self.p().x(), self.p().y(), self.p().z());
                let d = -self.common.n.dot(p_as_vec);

                // Compute t for x-auxiliary ray
                let rx_origin_vec =
                    Normal3f::new(diff.rx_origin.x(), diff.rx_origin.y(), diff.rx_origin.z());
                let tx = (-self.common.n.dot(rx_origin_vec) - d) / dot_rx;

                // Compute intersection point px
                let px = diff.rx_origin + diff.rx_direction * tx;

                // Compute t for y-auxiliary ray
                let ry_origin_vec =
                    Normal3f::new(diff.ry_origin.x(), diff.ry_origin.y(), diff.ry_origin.z());
                let ty = (-self.common.n.dot(ry_origin_vec) - d) / dot_ry;

                let py = diff.ry_origin + diff.ry_direction * ty;

                self.dpdx = px - self.p();
                self.dpdy = py - self.p();

                true
            } else {
                false
            }
        } else {
            false
        };

        if !computed {
            camera.approximate_dp_dxy(
                self.p(),
                self.n(),
                self.time(),
                samples_per_pixel,
                &mut self.dpdx,
                &mut self.dpdy,
            );
        }

        let ata00 = self.dpdu.dot(self.dpdu);
        let ata01 = self.dpdu.dot(self.dpdv);
        let ata11 = self.dpdv.dot(self.dpdv);
        let mut inv_det = 1. / difference_of_products(ata00, ata11, ata01, ata01);
        inv_det = if inv_det.is_finite() { inv_det } else { 0. };
        let atb0x = self.dpdu.dot(self.dpdx);
        let atb1x = self.dpdv.dot(self.dpdx);
        let atb0y = self.dpdu.dot(self.dpdy);
        let atb1y = self.dpdv.dot(self.dpdy);
        // Compute u and v derivatives in x and y
        self.dudx = difference_of_products(ata11, atb0x, ata01, atb1x) * inv_det;
        self.dvdx = difference_of_products(ata00, atb1x, ata01, atb0x) * inv_det;
        self.dudy = difference_of_products(ata11, atb0y, ata01, atb1y) * inv_det;
        self.dvdy = difference_of_products(ata00, atb1y, ata01, atb0y) * inv_det;
        // Clamp derivatives
        self.dudx = if self.dudx.is_finite() {
            clamp(self.dudx, -1e8, 1e8)
        } else {
            0.
        };
        self.dvdx = if self.dvdx.is_finite() {
            clamp(self.dvdx, -1e8, 1e8)
        } else {
            0.
        };
        self.dudy = if self.dudy.is_finite() {
            clamp(self.dudy, -1e8, 1e8)
        } else {
            0.
        };
        self.dvdy = if self.dvdy.is_finite() {
            clamp(self.dvdy, -1e8, 1e8)
        } else {
            0.
        };
    }

    pub fn skip_intersection(&self, ray: &mut Ray, t: Float) {
        let new_ray = Ray::spawn(&self.pi(), &self.n(), ray.time, ray.d);
        ray.o = new_ray.o;
        // Skipping other variables, since they should not change when passing through surface
        if !ray.has_differentials {
            let mut diff = ray.differential;
            diff.rx_origin += diff.rx_direction * t;
            diff.ry_origin += diff.ry_direction * t;
        }
    }

    pub fn compute_bump_geom(
        &mut self,
        tex_eval: &UniversalTextureEvaluator,
        displacement: Ptr<GPUFloatTexture>,
        normal_image: Ptr<Image>,
    ) {
        let ctx = NormalBumpEvalContext::from(&*self);
        let (dpdu, dpdv) = if !displacement.is_null() {
            bump_map(tex_eval, &displacement, &ctx)
        } else if !normal_image.is_null() {
            let map = unsafe { normal_image.as_ref() };
            normal_map(map, &ctx)
        } else {
            (self.shading.dpdu, self.shading.dpdv)
        };

        let mut ns = Normal3f::from(dpdu.cross(dpdv).normalize());
        if ns.dot(self.n()) < 0.0 {
            ns = -ns;
        }

        self.set_shading_geom(ns, dpdu, dpdv, self.shading.dndu, self.shading.dndv, false);
    }

    pub fn spawn_ray_with_differentials(
        &self,
        ray_i: &Ray,
        wi: Vector3f,
        flags: BxDFFlags,
        eta: Float,
    ) -> Ray {
        let mut rd = self.spawn_ray(wi);

        if ray_i.has_differentials {
            let diff_i = ray_i.differential;
            let mut n = self.shading.n;

            let mut dndx = self.shading.dndu * self.dudx + self.shading.dndv * self.dvdx;
            let mut dndy = self.shading.dndu * self.dudy + self.shading.dndv * self.dvdy;

            let dwodx = -diff_i.rx_direction - self.wo();
            let dwody = -diff_i.ry_direction - self.wo();

            let new_diff_rx_origin = self.p() + self.dpdx;
            let new_diff_ry_origin = self.p() + self.dpdy;
            let mut new_diff_rx_dir = Vector3f::default();
            let mut new_diff_ry_dir = Vector3f::default();

            let mut valid_differentials = false;

            if flags.contains(BxDFFlags::SPECULAR_REFLECTION) {
                valid_differentials = true;

                let d_wo_dot_n_dx = dwodx.dot(n.into()) + self.wo().dot(dndx.into());
                let d_wo_dot_n_dy = dwody.dot(n.into()) + self.wo().dot(dndy.into());

                let wo_dot_n = self.wo().dot(n.into());

                new_diff_rx_dir = wi - dwodx
                    + (Vector3f::from(dndx) * wo_dot_n + Vector3f::from(n) * d_wo_dot_n_dx) * 2.0;

                new_diff_ry_dir = wi - dwody
                    + (Vector3f::from(dndy) * wo_dot_n + Vector3f::from(n) * d_wo_dot_n_dy) * 2.0;
            } else if flags.contains(BxDFFlags::SPECULAR_TRANSMISSION) {
                valid_differentials = true;

                if self.wo().dot(n.into()) < 0.0 {
                    n = -n;
                    dndx = -dndx;
                    dndy = -dndy;
                }

                // Compute partial derivatives
                let d_wo_dot_n_dx = dwodx.dot(n.into()) + self.wo().dot(dndx.into());
                let d_wo_dot_n_dy = dwody.dot(n.into()) + self.wo().dot(dndy.into());

                let wo_dot_n = self.wo().dot(n.into());
                let wi_dot_n = wi.dot(Vector3f::from(n));
                let abs_wi_dot_n = wi.abs_dot(n.into());

                let mu = wo_dot_n / eta - abs_wi_dot_n;

                let f_eta = 1.0 / eta;
                let f_eta2 = 1.0 / square(eta);

                let term = f_eta + (f_eta2 * wo_dot_n / wi_dot_n);

                let dmudx = d_wo_dot_n_dx * term;
                let dmudy = d_wo_dot_n_dy * term;

                new_diff_rx_dir =
                    wi - dwodx * eta + (Vector3f::from(dndx) * mu + Vector3f::from(n) * dmudx);
                new_diff_ry_dir =
                    wi - dwody * eta + (Vector3f::from(dndy) * mu + Vector3f::from(n) * dmudy);
            }

            if valid_differentials {
                let threshold = 1e16;

                if new_diff_rx_dir.norm_squared() > threshold
                    || new_diff_ry_dir.norm_squared() > threshold
                    || Vector3f::from(new_diff_rx_origin).norm_squared() > threshold
                    || Vector3f::from(new_diff_ry_origin).norm_squared() > threshold
                {
                    rd.differential = RayDifferential::default();
                } else {
                    rd.differential = RayDifferential {
                        rx_origin: new_diff_rx_origin,
                        ry_origin: new_diff_ry_origin,
                        rx_direction: new_diff_rx_dir,
                        ry_direction: new_diff_ry_dir,
                    };
                }
            }
        }

        rd
    }
}

impl InteractionTrait for SurfaceInteraction {
    fn get_common(&self) -> &InteractionBase {
        &self.common
    }

    fn get_common_mut(&mut self) -> &mut InteractionBase {
        &mut self.common
    }

    fn get_medium(&self, w: Vector3f) -> Ptr<Medium> {
        let interface = self.common.medium_interface;
        if self.n().dot(w.into()) > 0.0 {
            interface.outside
        } else {
            interface.inside
        }
    }

    fn is_surface_interaction(&self) -> bool {
        true
    }
}

impl SurfaceInteraction {
    pub fn new(
        pi: Point3fi,
        uv: Point2f,
        wo: Vector3f,
        dpdu: Vector3f,
        dpdv: Vector3f,
        dndu: Normal3f,
        dndv: Normal3f,
        time: Float,
        flip: bool,
    ) -> Self {
        let mut n = Normal3f::from(dpdu.cross(dpdv).normalize());
        let mut shading_n = n;

        if flip {
            n *= -1.0;
            shading_n *= -1.0;
        }

        Self {
            common: InteractionBase::new_surface_geom(pi, n, uv, wo, time),
            dpdu,
            dpdv,
            dndu,
            dndv,
            shading: ShadingGeom {
                n: shading_n,
                dpdu,
                dpdv,
                dndu,
                dndv,
            },
            material: Ptr::null(),
            face_index: 0,
            area_light: Ptr::null(),
            dpdx: Vector3f::zero(),
            dpdy: Vector3f::zero(),
            dudx: 0.0,
            dudy: 0.0,
            dvdx: 0.0,
            dvdy: 0.0,
            shape: Ptr::null(),
        }
    }

    pub fn new_with_face(
        pi: Point3fi,
        uv: Point2f,
        wo: Vector3f,
        dpdu: Vector3f,
        dpdv: Vector3f,
        dndu: Normal3f,
        dndv: Normal3f,
        time: Float,
        flip: bool,
        face_index: i32,
    ) -> Self {
        let mut si = Self::new(pi, uv, wo, dpdu, dpdv, dndu, dndv, time, flip);
        si.face_index = face_index;
        si
    }

    pub fn set_shading_geom(
        &mut self,
        ns: Normal3f,
        dpdus: Vector3f,
        dpdvs: Vector3f,
        dndus: Normal3f,
        dndvs: Normal3f,
        orientation: bool,
    ) {
        self.shading.n = ns;
        if orientation {
            self.common.n = self.n().face_forward(self.shading.n);
        }
        self.shading.dpdu = dpdus;
        self.shading.dpdv = dpdvs;
        self.shading.dndu = dndus;
        self.shading.dndv = dndvs;
    }

    pub fn new_simple(pi: Point3fi, n: Normal3f, uv: Point2f) -> Self {
        Self {
            common: InteractionBase::new_surface_geom(pi, n, uv, Vector3f::zero(), 0.),
            ..Default::default()
        }
    }

    pub fn new_minimal(pi: Point3fi, uv: Point2f) -> Self {
        Self {
            common: InteractionBase {
                pi,
                uv,
                ..Default::default()
            },
            ..Default::default()
        }
    }

    #[cfg(not(target_os = "cuda"))]
    pub fn set_intersection_properties(
        &mut self,
        mtl: &Material,
        area: &Light,
        ray_medium: &Medium,
        prim_medium_interface: MediumInterface,
    ) {
        self.material = Ptr::from(mtl);
        self.area_light = Ptr::from(area);

        if prim_medium_interface.is_medium_transition() {
            self.common.medium_interface = prim_medium_interface;
        } else {
            self.common.medium = Ptr::from(ray_medium);
        }
    }
}

#[repr(C)]
#[derive(Clone, Copy, Debug)]
pub struct MediumInteraction {
    pub common: InteractionBase,
    pub phase: PhaseFunction,
}

impl MediumInteraction {
    pub fn new(
        p: Point3f,
        wo: Vector3f,
        time: Float,
        medium: Ptr<Medium>,
        phase: PhaseFunction,
    ) -> Self {
        Self {
            common: InteractionBase::new_medium(p, wo, time, medium),
            phase,
        }
    }
}

impl InteractionTrait for MediumInteraction {
    fn is_medium_interaction(&self) -> bool {
        true
    }

    fn get_common(&self) -> &InteractionBase {
        &self.common
    }

    fn get_common_mut(&mut self) -> &mut InteractionBase {
        &mut self.common
    }
}
