diff --git a/shader/toneshader.cpp b/shader/toneshader.cpp
index d616eda..1991827 100644
--- a/shader/toneshader.cpp
+++ b/shader/toneshader.cpp
@@ -2,112 +2,38 @@
 #include "toneshader.h"
 
 
-ToneShader::ToneShader(const Color &highlightColor, const Color &midtoneColor, const Color &shadowColor,
-                       float edgeNoise)
-        : highlightColor(highlightColor),
+ToneShader::ToneShader(std::shared_ptr<Light> light, const Color &highlightColor, const Color &midtoneColor, const Color &shadowColor,
+                       float edgeFade)
+        : light(light),
+          highlightColor(highlightColor),
           midtoneColor(midtoneColor),
           shadowColor(shadowColor),
-          edgeNoise(edgeNoise)
+          edgeFade(edgeFade)
 {
 
 }
 
 Color ToneShader::shade(const Scene &scene, const Ray &ray) const
 {
-    Color fragmentColor;
     float brightness;
 
-    // Look at first interesting light
-    for (const auto &light: scene.lights())
-    {
-        Light::Illumination illum;
-        illum = light->illuminate(scene, ray);
+    // Look at light
+    Light::Illumination illum;
+    illum = light->illuminate(scene, ray);
 
-        if (illum.distance == 0.0f)
-            continue;   // Skip ambient light
-
-        if (illum.color.r < 0.01f)
-            // Light does not reach the surface
-            brightness = 0;
-        else
-            // Look at angle of light
-            brightness = dotProduct(illum.direction,
-                                    -ray.normal);  // 0 if light is behind surface, 1 if light is in front of surface
-
-        break;  // No other light needed
-    }
-
-    // Accumulate the light over all light sources
-//    for (const auto &light: scene.lights())
-//    {
-//        Light::Illumination illum;
-//        illum = light->illuminate(scene, ray);
-//
-//        float const NdotL = std::max(0.0f, dotProduct(-illum.direction, ray.normal));
-//        if (NdotL <= 0.0f)
-//            continue;
-//
-//        // Diffuse term
-//        Color const diffuse = this->highlightColor / float(PI);
-//        fragmentColor += diffuse * NdotL * illum.color;
-//
-//        // Cook-Torrance term
-//        // half angle vector
-//        Vector3d const H = normalized(-illum.direction - ray.direction);
-//        float const NdotH = std::max(0.0f, dotProduct(ray.normal, H));
-//        float const NdotV = std::max(0.0f, dotProduct(ray.normal, -ray.direction));
-//        float const VdotH = std::max(0.0f, dotProduct(-ray.direction, H));
-//
-//        if (NdotV * NdotL > EPSILON)
-//        {
-//            fragmentColor += NdotL * illum.color;
-//        }
-//    }
-
-    // Accumulate the light over all light sources
-//    for (const auto &light : scene.lights()) {
-//        Light::Illumination const illum = light->illuminate(scene, ray);
-//        // Diffuse term
-//        Color const diffuse = this->highlightColor * std::max(dotProduct(-illum.direction, ray.normal), 0.0f);
-//        fragmentColor += diffuse * illum.color;
-//    }
-
-//    return fragmentColor;
-//    for (const auto &light: scene.lights())
-//        brightness += light->illuminate(scene, ray).color.r > 0.01f;    // Basically count the light sources that hit
-//
-//    brightness /= scene.lights().size();
-
-//    for (const auto &light: scene.lights())
-//        fragmentColor += light->illuminate(scene, ray).color;
-//
-//    brightness = (fragmentColor.r
-//                        + fragmentColor.g
-//                        + fragmentColor.b) / 3;
-
-    // Handle peak brightness
-//    float const PEAK_AREA = 0.0f;
-//    float peak = brightness - 1;
-//    brightness *= 1 - PEAK_AREA;
-//    if (peak > 0)
-//        brightness = 1 - PEAK_AREA + PEAK_AREA * 2 * (1 / (exp(-peak / 10.0f) + 1) - 0.5f);
-
-
-//    float viewFactor = std::acos(dotProduct(ray.direction, -ray.normal)) / float(PI);  // 0 if normal is facing the view, 1 if normal is facing away from the view
-//    viewFactor = pow(viewFactor, 3) * 20;
-//    if (viewFactor > 1)
-//        viewFactor = 1;
-//
-//    brightness = pow(brightness, 1 + 3 * viewFactor);
-//    brightness = pow(brightness, 2);
-
-//    return brightness * Color(1, 1, 1);
+    if (illum.color.r < 0.01f)
+        // Light does not reach the surface
+        brightness = 0;
+    else
+        // Look at angle of light
+        brightness = dotProduct(illum.direction,
+                                -ray.normal);  // 0 if light is behind surface, 1 if light is in front of surface
 
     // Get the tone color base on the brightness
     float const HIGHLIGHT_THRESHOLD = 0.95f;
-    float const MIDTONE_HIGHLIGHT_THRESHOLD = HIGHLIGHT_THRESHOLD - 0.2f * edgeNoise;
+    float const MIDTONE_HIGHLIGHT_THRESHOLD = HIGHLIGHT_THRESHOLD - 0.2f * edgeFade;
     float const SHADOW_MIDTONE_THRESHOLD = 0.01f;
-    float const MIDTONE_THRESHOLD = SHADOW_MIDTONE_THRESHOLD + 0.4f * edgeNoise;
+    float const MIDTONE_THRESHOLD = SHADOW_MIDTONE_THRESHOLD + 0.4f * edgeFade;
 
     // Used for transitions
     std::random_device dev;
@@ -122,10 +48,6 @@ Color ToneShader::shade(const Scene &scene, const Ray &ray) const
         float transition =
                 (brightness - MIDTONE_HIGHLIGHT_THRESHOLD) / (HIGHLIGHT_THRESHOLD - MIDTONE_HIGHLIGHT_THRESHOLD);
         tone = highlightColor * transition + midtoneColor * (1 - transition);
-//        if (randomDistribution(dev) < transition)
-//            tone = highlightColor;
-//        else
-//            tone = midtoneColor;
     } else if (brightness > MIDTONE_THRESHOLD)
     {
         tone = midtoneColor;
@@ -133,10 +55,6 @@ Color ToneShader::shade(const Scene &scene, const Ray &ray) const
     {
         float transition = (brightness - SHADOW_MIDTONE_THRESHOLD) / (MIDTONE_THRESHOLD - SHADOW_MIDTONE_THRESHOLD);
         tone = midtoneColor * transition + shadowColor * (1 - transition);
-//        if (randomDistribution(dev) < transition)
-//            tone = midtoneColor;
-//        else
-//            tone = shadowColor;
     } else
     {
         tone = shadowColor;
diff --git a/shader/toneshader.h b/shader/toneshader.h
index cd0833e..fb769ce 100644
--- a/shader/toneshader.h
+++ b/shader/toneshader.h
@@ -12,10 +12,11 @@ class ToneShader : public Shader
 
 public:
     // Constructor
-    ToneShader(Color const &highlightColor = Color(1, 1, 1),
+    ToneShader(std::shared_ptr<Light> light,
+               Color const &highlightColor = Color(1, 1, 1),
                Color const &midtoneColor = Color(0.5, 0.5, 0.5),
                Color const &shadowColor = Color(0, 0, 0),
-               float edgeNoise = 1.0f);
+               float edgeFade = 1.0f);
 
     // Shader functions
     virtual Color shade(Scene const &scene, Ray const &ray) const;
@@ -25,7 +26,9 @@ protected:
     Color const midtoneColor;
     Color const shadowColor;
 
-    float const edgeNoise;
+    float const edgeFade;
+
+    std::shared_ptr<Light> light;   // Only one light is supported
 };