// Demo code and the images generated by it (c) by SEGFAULT
//
// This code and all images generated by it is licensed under a
// Creative Commons Attribution-ShareAlike 4.0 International License.
//
// You should have received a copy of the license along with this
// work. If not, see <http://creativecommons.org/licenses/by-sa/4.0/>.

extern crate png;
extern crate rand;
extern crate rustic_zen;

use std::env;
use std::ffi::OsString;
use std::fs::File;
use std::io::BufWriter;
use std::sync::Arc;
use std::time;

// To use encoder.set()
use png::HasParameters;

use rand::prelude::*;
use rustic_zen::geom::Vector;
use rustic_zen::material::Material;
use rustic_zen::prelude::*;

fn refraction(direction: &Vector, normal: &Vector, ior: f64) -> Option<Vector> {
    let tangent = Vector {
        x: normal.y,
        y: -normal.x,
    };
    let backface = normal.dot(direction);
    if backface <= 0.0 {
        let dti = direction.dot(&tangent) / ior;
        Some((tangent * (dti)) - (*normal * f64::sqrt(1.0 - dti.powi(2))))
    } else {
        let dti = direction.dot(&tangent) * ior;
        if dti > 1.0 {
            Some(direction.reflect(normal))
        } else {
            Some((tangent * (dti)) + (*normal * f64::sqrt(1.0 - dti.powi(2))))
        }
    }
}

fn dispersion_law(ln: f64, b: f64, c: f64) -> f64 {
    let lm2 = (ln / 1000.0).powi(2);
    f64::sqrt(1.0 + ((b * lm2) / (lm2 - c)))
}

fn pbr_transparent_surface<R: Rng>(index: f64, dispersion: f64) -> Material<R> {
    Arc::new(
        move |direction: &Vector, normal: &Vector, l: f64, _: f64, _: &mut R| {
            refraction(direction, normal, dispersion_law(l, index, dispersion))
        },
    )
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let (quality, filename) = if args.len() < 2 {
        (1, "prism.png".to_owned())
    } else {
        let quality = usize::from_str_radix(&args[1], 10).expect("Usage <script> [threads]");
        (quality, format!("prism-quality-{}.png", quality))
    };

    //Initialising scene!
    let now = time::Instant::now();
    let width: f64 = 2048.0;
    let height: f64 = 2048.0;
    let rays = 1_000_000 * quality;
    let threads = num_cpus::get();
    let angle = -25.0;

    let m = pbr_transparent_surface(1.03, 6.0 / 80.0);

    //let m = hqz_legacy(0.0, 1.0, 0.0);

    let points: Vec<(f64, f64)> = [(0.8, 0.8), (0.5, 0.2), (0.2, 0.8), (0.8, 0.8)]
        .iter()
        .map(|(x, y)| (x * width, y * height))
        .collect();

    let l1 = Light {
        power: Sampler::new_const(50.0),
        location: (
            Sampler::new_gaussian(width * 0.5, width * 0.25),
            Sampler::new_gaussian(height * 0.8, height * 0.1),
        )
            .into(),
        polar_angle: Sampler::new_const(0.0),
        polar_distance: Sampler::new_const(0.0),
        ray_angle: Sampler::new_range(180.0, -180.0),
        wavelength: Sampler::new_blackbody(5600.0),
    };

    let l2 = Light {
        power: Sampler::new_const(1.0),
        location: (
            Sampler::new_const(0.0),
            Sampler::new_gaussian(height * 0.666666666, 0.3),
        )
            .into(),
        polar_angle: Sampler::new_const(0.0),
        polar_distance: Sampler::new_const(0.0),
        ray_angle: Sampler::new_gaussian(angle, 0.7),
        wavelength: Sampler::new_blackbody(4500.0),
    };

    let mut r = Scene::new(width as usize, height as usize)
        .with_light(l1)
        .with_light(l2);

    for i in 0..points.len() - 1 {
        r.add_object(Segment::line_from_points(
            points[i],
            points[i + 1],
            m.clone(),
        ));
    }

    let mut image = Arc::new(Image::new(width as usize, height as usize));

    let setup = now.elapsed();

    println!("Tracing {} rays, with {} threads!\n", rays, threads);
    let now = time::Instant::now();
    let rays = r.render(rays, threads, &mut image);
    let tracing = now.elapsed();

    //Downsampling Image!
    let now: time::Instant = time::Instant::now();
    let data = image.to_rgba8(rays, 0.4, 2.2);
    let downsampling = now.elapsed();

    //Saving!
    let path = OsString::from(filename);
    let file = File::create(path).unwrap();
    let ref mut w = BufWriter::new(file);

    let mut encoder = png::Encoder::new(w, width as u32, height as u32);
    encoder.set(png::ColorType::RGBA).set(png::BitDepth::Eight);
    let mut writer = encoder.write_header().unwrap();
    writer.write_image_data(&data).unwrap();

    println!("\nTiming Summary:");
    println!("  Setup:        {}ms", setup.as_millis());
    println!(
        "  Tracing:      {}ms ({}cpu ms)",
        tracing.as_millis(),
        tracing.as_millis() * threads as u128
    );
    println!("  Downsampling: {}ms\n", downsampling.as_millis());

    println!(
        "{} rays in {}ms: {}rays/s",
        rays,
        tracing.as_millis(),
        rays as f32 / (tracing.as_millis() as f32 / 1000.0)
    );
}