1 files changed, 23 insertions, 12 deletions

diff --git a/src/trace/pathtrace.rs b/src/trace/pathtrace.rs
index 2d15976..f53e433 100644
--- a/src/trace/pathtrace.rs
+++ b/src/trace/pathtrace.rs
@@ -1,23 +1,23 @@
 use crate::world::{Hittable, Scene};
 use crate::core::{Ray, Spectrum};
+use crate::material::Material;
 use crate::sample::Sampler;
-use crate::material::{Lambertian, Material};
 use super::Tracer;
 
 pub struct PathTracer<'a> {
     depth: i32,
     scn: &'a Scene,
-    default_mat: Box<dyn Material>,
+    background: Option<Box<dyn Material>>,
 }
 
 impl PathTracer<'_> {
-    pub fn new(scn: &Scene, depth: Option<i32>) -> PathTracer {
+    pub fn new(scn: &Scene, depth: Option<i32>, background: Option<Box<dyn Material>>) -> PathTracer {
         let depth = depth.unwrap_or(-1);
 
         PathTracer {
             depth,
             scn,
-            default_mat: Box::new(Lambertian::new(Spectrum::ZERO)),
+            background,
         }
     }
 
@@ -28,17 +28,28 @@ impl PathTracer<'_> {
         }
 
         if let Some(i) = self.scn.intersect(ray) {
-            if let Some((scalar, nray)) = i.m.unwrap_or(self.default_mat.as_ref()).scatter(ray, &i, sampler) {
-                return self.trace_recur(sampler, &nray, depth-1) * scalar;
-            } else {
-                return Spectrum::ZERO;
+            // Extract material or default
+            if let Some(mat) = i.m {
+                let mut col = Spectrum::ZERO;
+
+                if let Some((scalar, nray)) = mat.scatter(ray, &i, sampler) {
+                    col += &(self.trace_recur(sampler, &nray, depth-1) * scalar);
+                }
+
+                if let Some(c) = mat.emitted(ray) {
+                    col += &c;
+                }
+
+                return col;
             }
         }
 
-        // Simulates a sky
-        let t = (ray.direction.norm().y + 1.0) * 0.5;
-        Spectrum::new_rgb(1.0, 1.0, 1.0) * (1.0-t) + Spectrum::new_rgb(0.5, 0.7, 1.0) * t
-
+        // If no color return background
+        if let Some(back) = &self.background {
+            back.emitted(ray).unwrap_or(Spectrum::ZERO)
+        } else {
+            Spectrum::ZERO
+        }
     }
 }