extern crate image;

mod camera;
mod color;
mod hittable;
mod hittable_list;
mod material;
mod ray;
mod utility;
mod vec3;

use camera::Camera;
use hittable::{HitRecord, Hittable, Sphere, Triangle};
use hittable_list::HittableList;
use image::{Rgb, RgbImage};
use indicatif::ProgressBar;
use material::{Dielectric, Lambertian, Material, Metal, Rainbow};
use ray::Ray;
use rayon::prelude::*;
use std::env;
use std::sync::Arc;
use tobj;
use vec3::{Color, Point3, Vec3};

/* Gets the pixel color for the passed ray */
fn ray_color(r: &Ray, world: &HittableList, depth: u32) -> Color {
    let mut rec = HitRecord::empty();

    if depth <= 0 {
        return Color::null();
    }

    if world.hit(r, 0.001, f64::INFINITY, &mut rec) {
        let mut scattered = Ray::new(Point3::null(), Vec3::null());
        let mut attenuation = Color::null();
        if rec
            .mat_ptr
            .scatter(r, &rec, &mut attenuation, &mut scattered)
        {
            return attenuation * ray_color(&scattered, world, depth - 1);
        }
        return Color::null();
    }

    let unit_direction = r.direction();
    let t = 0.5 * (unit_direction.y() + 1.0);
    return (1.0 - t) * Color::new(1.0, 1.0, 1.0) + t * Color::new(0.5, 0.7, 1.0);
}

// Current world view:
//
//             I y
//             I
//             I
//             I
//           /   \
//         /       \
//       /           \
//     / z             \ x

/* 
* Main function that builds everything and runs the raytracing
*/
fn main() {
    // File
    let mut default_file = "image.ppm";

    // Image
    let aspect_ratio = 10.0 / 7.5; //16.0 / 9.0;
    let image_width = 300;
    let image_height = (image_width as f64 / aspect_ratio) as u32;
    let samples_per_pixel = 1_u32;
    let max_depth = 50;
    let antialiasing_threshold = 0.2; // at what diff between two colors will a pixel be antialiased

    let vfov = 43.0;
    let lookfrom = Point3::new(2.0, 1.0, 1.0);
    let lookat = Point3::new(0.0, 0.2, 0.0);
    let vup = Vec3::new(0.0, 1.0, 0.0);
    let dist_to_focus = 1.0;
    let aperture = 0.0; // disable depth of field

    // limit rayon multithreading thread count
    let thread_count = 0; // if 0, for each logical cpu core a thread wil be created
    if thread_count > 0 {
        env::set_var("RAYON_NUM_THREADS", thread_count.to_string());
    }

    // World
    eprintln!("[1/4] Loading meshes from file...");
    let world = from_obj("obj/viking_room.obj");
    // let world = random_world();

    // Camera
    let cam = Camera::new(
        lookfrom,
        lookat,
        vup,
        vfov,
        aspect_ratio,
        aperture,
        dist_to_focus,
    );

    // Render
    let args: Vec<String> = env::args().collect();
    if args.len() > 1 && args[1] != "" {
        default_file = &args[1];
    }

    eprintln!("[2/4] Generating image...");
    let bar = ProgressBar::new(image_height as u64);

    let mut image = RgbImage::new(image_width, image_height);

    let color_lines: Vec<_> = (0..image_height)
        .into_par_iter() // threadded/parallel variant
        //.into_iter()  // iterative variant
        .rev()
        .map(|j| {
            bar.inc(1);

            let mut colors = Vec::new();
            for i in 0..image_width {
                let pixel_color = ray_color(
                    &cam.get_ray(
                        i as f64 / (image_width - 1) as f64,
                        j as f64 / (image_height - 1) as f64,
                    ),
                    &world,
                    max_depth,
                );
                colors.push(pixel_color);
            }
            return colors;
        })
        .collect();
    bar.finish_and_clear();

    // no antialiasing
    if samples_per_pixel == 1_u32 {
        eprintln!("[4/4] Exporting image to disk...");  
        (0..image_height).into_iter().rev().for_each(|j| {
            (0..image_width).into_iter().for_each(|i| {
                color::put_color(
                    &mut image,
                    &color_lines[(image_height - j - 1) as usize][i as usize],
                    i,
                    image_height - j - 1,
                    1,
                );
            })
        });

        image.save(default_file).unwrap();

        return;
    }

    eprintln!("[3/4] Antialiasing image...");
    let mut antialiasing: Vec<Vec<bool>> = Vec::new();
    let mut antialiasing_counter = 0;
    let mut antialiasing_col_counter = 0;
    let mut antialiases_per_line: Vec<u64> = Vec::new();
    (0..image_height).into_iter().for_each(|j| {
        antialiasing.push(Vec::new());
        (0..image_width).into_iter().for_each(|i| {
            if j != 0
                && Color::diff(
                    &color_lines[(j - 1) as usize][i as usize],
                    &color_lines[j as usize][i as usize],
                ) >= antialiasing_threshold
            {
                antialiasing[j as usize].push(true);
                antialiasing_col_counter += 1;
            } else if j != image_height - 1
                && Color::diff(
                    &color_lines[(j + 1) as usize][i as usize],
                    &color_lines[j as usize][i as usize],
                ) >= antialiasing_threshold
            {
                antialiasing[j as usize].push(true);
                antialiasing_col_counter += 1;
            } else if i != 0
                && Color::diff(
                    &color_lines[j as usize][(i - 1) as usize],
                    &color_lines[j as usize][i as usize],
                ) >= antialiasing_threshold
            {
                antialiasing[j as usize].push(true);
                antialiasing_col_counter += 1;
            } else if i != image_width - 1
                && Color::diff(
                    &color_lines[j as usize][(i + 1) as usize],
                    &color_lines[j as usize][i as usize],
                ) >= antialiasing_threshold
            {
                antialiasing[j as usize].push(true);
                antialiasing_col_counter += 1;
            }
            if antialiasing[j as usize].len() < (i + 1) as usize {
                antialiasing[j as usize].push(false);
            }
        });
        antialiases_per_line.push(antialiasing_col_counter);
        antialiasing_counter += antialiasing_col_counter;
        antialiasing_col_counter = 0;
    });

    let bar = ProgressBar::new(antialiasing_counter as u64);
    let color_lines_antialiased: Vec<_> = (0..image_height)
        .into_par_iter() // threadded/parallel variant
        //.into_iter()  // iterative variant
        .map(|j| {
            let mut colors = Vec::new();
            for i in 0..image_width {
                if samples_per_pixel > 0 && antialiasing[j as usize][i as usize] {
                    let mut pixel_color = Color::null();
                    
                    for _ in 0..samples_per_pixel {
                        let u = (i as f64 + utility::random_f64()) / (image_width - 1) as f64;
                        let v = (image_height as f64- 1.0 -(j as f64 + utility::random_f64())) / (image_height - 1) as f64;
                        let new_pixel_color = ray_color(&cam.get_ray(u, v), &world, max_depth);
                        pixel_color += new_pixel_color;
                    }
                    // Correct antialiasing gamma
                    let fin_color = (1.0/samples_per_pixel as f64) * pixel_color;
                    //println!("x: {}, y: {}, z: {}", pixel_color.x(), pixel_color.y(), pixel_color.z());
                    
                    colors.push(fin_color);
                } else {
                    colors.push(color_lines[j as usize][i as usize]);
                }
            }
            bar.inc(antialiases_per_line[j as usize]);
            return colors;
        })
        .collect();
    bar.finish_and_clear();

    eprintln!("[4/4] Exporting image to disk...");  
    (0..image_height).into_iter().rev().for_each(|j| {
        (0..image_width).into_iter().for_each(|i| {
            color::put_color(
                &mut image,
                &color_lines_antialiased[(image_height - j - 1) as usize][i as usize],
                i,
                image_height - j - 1,
                1,
            );
        })
    });

    image.save(default_file).unwrap();
}


/********************
*  WORLD GENERATION *
********************/


/*
* Generates world based on .obj mesh file passed by path.
*
* Currently only works for .obj files which contain face normals.
*/
fn from_obj(path: &str) -> HittableList {
    let mut world = HittableList::new();

    /*
    let material_ground = Arc::new(Lambertian::new(&Color::new(
        29.0 / 255.0,
        71.0 / 255.0,
        14.0 / 255.0,
    )));
    world.add(Box::<Sphere>::new(Sphere::new(
        Point3::new(-500.0, -5005.0, -500.0),
        5000.0,
        material_ground.clone(),
    )));
    */

    //let material = Arc::new(Lambertian::new(&Color::new(
    //    77.0 / 255.0,
    //    77.0 / 255.0,
    //    118.0 / 255.0,
    //)));
    //let material = Arc::new(Dielectric::new(2.0));
    let material = Arc::new(Metal::new(&Color::new(0.7, 0.6, 0.5), 0.0));
    //let material = Arc::new(Rainbow::new());

    let cornell_box = tobj::load_obj(path, &tobj::OFFLINE_RENDERING_LOAD_OPTIONS);
    let (models, materials) = cornell_box.expect("Failed to load OBJ file");
    let materials = materials.expect("Failed to load MTL file");

    for (i, m) in models.iter().enumerate() {
        let mesh = &m.mesh;

        /* Mesh vector lengths
        println!("positions: {}", mesh.positions.len());
        println!("vertex_color: {}", mesh.vertex_color.len());
        println!("normals: {}", mesh.normals.len());
        println!("texcoords: {}", mesh.texcoords.len());
        println!("indices: {}", mesh.indices.len());
        println!("face_arities: {}", mesh.face_arities.len());
        println!("texcoord_indices: {}", mesh.texcoord_indices.len());
        println!("normal_indices: {}", mesh.normal_indices.len());
        */

        let mut next_face = 0;
        for f in 0..mesh.face_arities.len() {
            let end = next_face + mesh.face_arities[f] as usize;
            let face_indices: Vec<_> = mesh.indices[next_face..end].iter().collect();
            println!("    face[{}] = {:?}", f, face_indices);
            next_face = end;
        }

        for v in 0..mesh.indices.len() / 3 {
            let index_a = mesh.indices[3 * v] as usize;
            let index_b = mesh.indices[3 * v + 2] as usize;
            let index_c = mesh.indices[3 * v + 1] as usize;
            let index_normal_a = mesh.normal_indices[3 * v] as usize;
            let index_normal_b = mesh.normal_indices[3 * v + 1] as usize;
            let index_normal_c = mesh.normal_indices[3 * v + 2] as usize;

            let normal_avg = Vec3::unit_vector(
                Vec3::new(
                    mesh.positions[3 * index_normal_a] as f64,
                    mesh.positions[3 * index_normal_a + 2] as f64,
                    mesh.positions[3 * index_normal_a + 1] as f64,
                ) + Vec3::new(
                    mesh.positions[3 * index_normal_b] as f64,
                    mesh.positions[3 * index_normal_b + 2] as f64,
                    mesh.positions[3 * index_normal_b + 1] as f64,
                ) + Vec3::new(
                    mesh.positions[3 * index_normal_c] as f64,
                    mesh.positions[3 * index_normal_c + 2] as f64,
                    mesh.positions[3 * index_normal_c + 1] as f64,
                ),
            );

            /*
            println!("a:{},{},{}; b:{},{},{}; c:{},{},{}",
            mesh.normals[3*index_normal_a],
            mesh.normals[3*index_normal_a+1],
            mesh.normals[3*index_normal_a+2],
            mesh.normals[3*index_normal_b],
            mesh.normals[3*index_normal_b+1],
            mesh.normals[3*index_normal_b+2],
            mesh.normals[3*index_normal_c],
            mesh.normals[3*index_normal_c+1],
            mesh.normals[3*index_normal_c+2]);
            */

            world.add(Box::<Triangle>::new(Triangle::new(
                Point3::new(
                    mesh.positions[3 * index_a] as f64,
                    mesh.positions[3 * index_a + 2] as f64,
                    mesh.positions[3 * index_a + 1] as f64,
                ),
                Point3::new(
                    mesh.positions[3 * index_b] as f64,
                    mesh.positions[3 * index_b + 2] as f64,
                    mesh.positions[3 * index_b + 1] as f64,
                ),
                Point3::new(
                    mesh.positions[3 * index_c] as f64,
                    mesh.positions[3 * index_c + 2] as f64,
                    mesh.positions[3 * index_c + 1] as f64,
                ),
                normal_avg,
                material.clone(),
            )));
        }
    }
    return world;
}

/*
* Generates a world with a bunch of spheres
*/
fn random_world() -> HittableList {
    let mut world = HittableList::new();

    let material_ground = Arc::new(Lambertian::new(&Color::new(0.05, 0.05, 0.05)));
    world.add(Box::<Sphere>::new(Sphere::new(
        Point3::new(0.0, -50000.0, 0.0),
        50000.0,
        material_ground.clone(),
    )));

    (-6..5).into_iter().for_each(|a| {
        (-6..5).into_iter().for_each(|b| {
            let choose_mat = utility::random_f64();
            let rad = utility::random_rng(0.1, 0.5);
            let center = Point3::new(
                1.5 * a as f64 + 1.3 * utility::random_f64(),
                rad,
                1.5 * b as f64 + 1.3 * utility::random_f64(),
            );
            if (center - Point3::new(4.0, rad, 0.0)).length() > 0.9 {
                if choose_mat < 0.8 {
                    // diffuse
                    let sphere_material = Arc::new(Lambertian::new(
                        &(Color::random_f64() * Color::random_f64()),
                    ));
                    world.add(Box::<Sphere>::new(Sphere::new(
                        center,
                        rad,
                        sphere_material.clone(),
                    )));
                } else if choose_mat < 0.95 {
                    // metal
                    let sphere_material = Arc::new(Metal::new(
                        &Color::random_rng(0.5, 1.0),
                        utility::random_rng(0.0, 0.5),
                    ));
                    world.add(Box::<Sphere>::new(Sphere::new(
                        center,
                        0.2,
                        sphere_material.clone(),
                    )));
                } else {
                    // glass
                    let sphere_material = Arc::new(Dielectric::new(1.5));
                    world.add(Box::<Sphere>::new(Sphere::new(
                        center,
                        0.2,
                        sphere_material.clone(),
                    )));
                }
            }
        });
    });

    let material1 = Arc::new(Dielectric::new(1.5));
    let material2 = Arc::new(Lambertian::new(&Color::new(0.4, 0.2, 0.1)));
    let material3 = Arc::new(Metal::new(&Color::new(0.7, 0.6, 0.5), 0.0));
    let material4 = Arc::new(Dielectric::new(2.0));
    let material5 = Arc::new(Metal::new(&Color::new(0.9, 0.9, 0.7), 0.0));

    world.add(Box::<Sphere>::new(Sphere::new(
        Point3::new(0.0, 1.0, 0.0),
        1.0,
        material2.clone(),
    )));

    world.add(Box::<Sphere>::new(Sphere::new(
        Point3::new(-4.0, 1.0, 0.0),
        1.0,
        material1.clone(),
    )));

    world.add(Box::<Sphere>::new(Sphere::new(
        Point3::new(4.0, 1.0, 0.0),
        1.0,
        material3.clone(),
    )));

    world.add(Box::<Sphere>::new(Sphere::new(
        Point3::new(-2.0, 2.0, -5.0),
        2.0,
        material4.clone(),
    )));

    world.add(Box::<Sphere>::new(Sphere::new(
        Point3::new(-3.6, 2.0, -2.0),
        0.6,
        material5.clone(),
    )));

    /*
    world.add(Box::<Triangle>::new(Triangle::new(
        Point3::new(0.0, 1.0, 5.0),
        Point3::new(3.0, 2.0, 0.0),
        Point3::new(0.0, 4.0, 0.0),
        material2.clone(),
    )));

    world.add(Box::<Triangle>::new(Triangle::new(
        Point3::new(5.0, 1.0, -6.0),
        Point3::new(1.0, 3.0, -5.0),
        Point3::new(6.0, 4.0, -6.0),
        material5.clone(),
    )));
    world.add(Box::<Triangle>::new(Triangle::new(
        Point3::new(5.0, 1.0, -6.0),
        Point3::new(8.0, 1.0, -7.0),
        Point3::new(6.0, 4.0, -6.0),
        material5.clone(),
    )));
    world.add(Box::<Triangle>::new(Triangle::new(
        Point3::new(8.0, 4.0, -5.0),
        Point3::new(8.0, 1.0, -7.0),
        Point3::new(6.0, 4.0, -6.0),
        material5.clone(),
    )));
    */

    return world;
}