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

use std::fs::File;
use std::io::BufWriter;
use std::path::Path;
use std::sync::Arc;
use std::time;
// To use encoder.set()
use png::HasParameters;

//Scene Parameters
use rand::Rng;
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, absorbsion: f64) -> Material<R> {
    Arc::new(
        move |direction: &Vector, normal: &Vector, l: f64, _: f64, rng: &mut R| {
            if rng.gen_range(0.0..=1.0) < absorbsion {
                return None;
            }
            refraction(direction, normal, dispersion_law(l, index, dispersion))
        },
    )
}

fn bubble(scene: &mut Scene, x: f64, y: f64, r: f64, m: &Material<rand_pcg::Pcg64Mcg>) {
    // Create a circle,
    scene.add_object(Segment::curve_from_points(
        (x - r, y - 1.0),
        (x - (r * 0.9375), y + (r * 0.9375)),
        (x + 1.0, y + r),
        m.clone(),
    )); //Bottom Left
    scene.add_object(Segment::curve_from_points(
        (x - 1.0, y + r),
        (x + (r * 0.9375), y + (r * 0.9375)),
        (x + r, y - 1.0),
        m.clone(),
    )); //Bottom Right
    scene.add_object(Segment::curve_from_points(
        (x + r, y + 1.0),
        (x + (r * 0.9375), y - (r * 0.9375)),
        (x - 1.0, y - r),
        m.clone(),
    )); //Rop Right
    scene.add_object(Segment::curve_from_points(
        (x + 1.0, y - r),
        (x - (r * 0.9375), y - (r * 0.9275)),
        (x - r, y + 1.0),
        m.clone(),
    )); //Top Left
}

#[cfg(feature = "gpu")]
fn main() {
    //Initialising scene!
    let now = time::Instant::now();
    let scale = 3.0;
    let width: f64 = 1920.0 * scale;
    let height: f64 = 1080.0 * scale;
    //let rays = 24_000_000; //(width * height).round() as usize;
    let rays = RenderConstraint::Count(80_000_000);
    let threads = 24;
    let mut image = Arc::new(VulkanImage::new(width as usize, height as usize));

    let bubble_m = pbr_transparent_surface(1.03, 6.0 / 100.0, 0.25);

    let mut r = Scene::new(width as usize, height as usize);

    bubble(&mut r, width * 0.05, height, 300.0 * scale, &bubble_m);
    bubble(&mut r, width * 0.1, height * 0.7, 240.0 * scale, &bubble_m);
    bubble(&mut r, width * 0.2, height * 0.9, 220.0 * scale, &bubble_m);
    bubble(&mut r, width * 0.25, height * 0.6, 110.0 * scale, &bubble_m);
    bubble(
        &mut r,
        width * 0.35,
        height * 0.85,
        120.0 * scale,
        &bubble_m,
    );
    bubble(&mut r, width * 0.45, height * 0.5, 150.0 * scale, &bubble_m);
    bubble(&mut r, width * 0.5, height * 0.7, 120.0 * scale, &bubble_m);
    bubble(&mut r, width * 0.6, height * 0.8, 90.0 * scale, &bubble_m);
    bubble(&mut r, width * 0.62, height * 0.65, 90.0 * scale, &bubble_m);
    bubble(&mut r, width * 0.65, height * 0.5, 120.0 * scale, &bubble_m);
    bubble(&mut r, width * 0.7, height * 0.8, 100.0 * scale, &bubble_m);
    bubble(&mut r, width * 0.72, height * 0.6, 90.0 * scale, &bubble_m);
    bubble(&mut r, width * 0.8, height * 0.4, 100.0 * scale, &bubble_m);

    r.add_light(Light {
        power: 50.0.into(),
        location: (
            Sampler::new_gaussian(30.0, 30.0),
            Sampler::new_gaussian(30.0, 30.0),
        )
            .into(),
        polar_distance: 0.0.into(),
        polar_angle: 0.0.into(),
        ray_angle: (180.0, -180.0).into(),
        wavelength: Sampler::new_blackbody(3200.0),
    });
    r.add_light(Light {
        power: 25.0.into(),
        location: (
            Sampler::new_gaussian(width - 100.0, 100.0),
            Sampler::new_gaussian(10.0, 10.0),
        )
            .into(),
        polar_distance: 0.0.into(),
        polar_angle: 0.0.into(),
        ray_angle: (180.0, 0.0).into(),
        wavelength: Sampler::new_blackbody(10000.0),
    });
    r.add_light(Light {
        power: 5.0.into(),
        location: (
            Sampler::new_gaussian(width * 0.6, 60.0),
            Sampler::new_gaussian(height * 0.9, 60.0),
        )
            .into(),
        polar_distance: 0.0.into(),
        polar_angle: 0.0.into(),
        ray_angle: (180.0, -180.0).into(),
        wavelength: Sampler::new_blackbody(2700.0),
    });

    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.7, 1.0 / 2.2); //1.0/2.2
    let downsampling = now.elapsed();

    //Saving!
    let filename = format!("bubbles-gpu-{}t.png", threads);
    let path = Path::new(&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)
    );
}

#[cfg(not(feature = "gpu"))]
fn main() {
    println!("Demo unavailable with GPU feature disabled")
}