#include <iostream>
#include <iomanip>
#include <atomic>
#include <memory>
#include <cstdio>
#include <utility>

#define LODEPNG_NO_COMPILE_DECODER
#include "lodepng.h"

#include <boost/lockfree/queue.hpp>

#include "color.h"
#include "vec3.h"
#include "ray.h"
#include "util.h"
#include "hittable_list.h"
#include "sphere.h"
#include "moving_sphere.h"
#include "camera.h"
#include "aarect.h"
#include "box.h"

Color ray_color(const Ray& r, const Color& background, const Hittable& world, int depth) {
    hit_record rec;

    if (depth <= 0)
        return Color(0, 0, 0);

    if (!world.hit(r, 0.001, infinity, rec))
        return background;

    Ray scattered;
    Color attenuation;
    Color emitted = rec.mat_ptr->emitted(rec.u, rec.v, rec.p);

    if (!rec.mat_ptr->scatter(r, rec, attenuation, scattered))
        return emitted;

    return emitted + attenuation * ray_color(scattered, background, world, depth - 1);

    // fading background
    /*Vec3 unit_direction = unit_vector(r.direction());
    auto t = 0.5 * (unit_direction.y() + 1.0);
    return (1.0 - t) * Color(1.0, 1.0, 1.0) + t * Color(0.5, 0.7, 1.0);*/
}

std::pair<Bvh_node, Camera> setup_random_scene(const int sph_i, const double aspect_ratio) {
    Hittable_list world;

    //auto ground_material = std::make_shared<Lambertian>(std::make_shared<Solid_color>(Color(0.5, 0.5, 0.5)));
    auto checker = std::make_shared<Checker_texture>(
            std::make_shared<Solid_color>(0.2, 0.3, 0.1),
            std::make_shared<Solid_color>(0.9, 0.9, 0.9)
    );
    auto ground_material = std::make_shared<Lambertian>(checker);
    world.add(std::make_shared<Sphere>(Point3(0, -1000, 0), 1000, ground_material));
    for (int a = -sph_i; a<sph_i; a++) {
        for (int b = -sph_i; b<sph_i; b++) {
            auto choose_mat = random_double();
            Point3 center(a + 0.9 * (11/sph_i) * random_double(), 0.2, b + 0.9 * (11/sph_i) * random_double());

            if ((center - Point3(4, 0.2, 0)).length() > 0.9) {
                std::shared_ptr<Material> sphere_material;
                if (choose_mat < 0.2) {
                    // diffuse moving
                    auto albedo = Color::random() * Color::random();
                    sphere_material = std::make_shared<Lambertian>(std::make_shared<Solid_color>(albedo));
                    auto center2 = center + Vec3(0, random_double(0, 0.1), 0);
                    world.add(std::make_shared<Moving_sphere>(center, center2, 0.0, 1.0, 0.2, sphere_material));
                }
                else if (choose_mat < 0.7) {
                    // diffuse
                    auto albedo = Color::random() * Color::random();
                    sphere_material = std::make_shared<Lambertian>(std::make_shared<Solid_color>(albedo));
                    world.add(std::make_shared<Sphere>(center, 0.2, sphere_material));
                } else if (choose_mat < 0.95) {
                    // metal
                    auto albedo = Color::random(0.5, 1);
                    auto fuzz = random_double(0, 0.5);

                    sphere_material = std::make_shared<Metal>(albedo, fuzz);
                    world.add(std::make_shared<Sphere>(center, 0.2, sphere_material));
                } else {
                    // glass
                    sphere_material = std::make_shared<Dielectric>(1.45);
                    world.add(std::make_shared<Sphere>(center, 0.2, sphere_material));
                }
            }
        }
    }

    auto material1 = std::make_shared<Dielectric>(1.45);
    world.add(std::make_shared<Sphere>(Point3(0, 1, 0), 1.0, material1));

    auto material2 = std::make_shared<Lambertian>(std::make_shared<Solid_color>(Color(0.4, 0.2, 0.1)));
    world.add(std::make_shared<Sphere>(Point3(-4, 1, 0), 1.0, material2));

    auto material3 = std::make_shared<Metal>(Color(0.7, 0.6, 0.5), 0.0);
    world.add(std::make_shared<Sphere>(Point3(4, 1, 0), 1.0, material3));

    Point3 lookfrom(13, 2, 3);
    Point3 lookat(0, 0, 0);
    Vec3 vup(0, 1, 0);
    auto dist_to_focus = (lookfrom - lookat).length();
    //auto dist_to_focus = 10;
    auto aperture = 0.1;
    auto vfov = 20;

    Camera cam(lookfrom, lookat, vup, vfov, aspect_ratio, aperture, dist_to_focus, 0.0, 1.0);

    return std::make_pair(world.generate_bvh(0, 1), cam);
}

Hittable_list two_spheres() {
    Hittable_list objects;

    auto checker = std::make_shared<Checker_texture>(
            std::make_shared<Solid_color>(0.2, 0.3, 0.1),
            std::make_shared<Solid_color>(0.9, 0.9, 0.9)
    );

    auto mat = std::make_shared<Lambertian>(checker);

    objects.add(std::make_shared<Sphere>(Point3(0, -10, 0), 10, mat));
    objects.add(std::make_shared<Sphere>(Point3(0,  10, 0), 10, mat));

    return objects;
}

Hittable_list simple_light() {
    Hittable_list objects;

    auto checker = std::make_shared<Checker_texture>(
            std::make_shared<Solid_color>(0.2, 0.3, 0.1),
            std::make_shared<Solid_color>(0.9, 0.9, 0.9)
    );
    auto material = std::make_shared<Lambertian>(checker);

    objects.add(std::make_shared<Sphere>(Point3(0, -1000, 0), 1000, material));
    objects.add(std::make_shared<Sphere>(Point3(0, 2, 0), 2, material));

    auto difflight = std::make_shared<Diffuse_light>(std::make_shared<Solid_color>(4, 4, 4));
    objects.add(std::make_shared<Sphere>(Point3(0, 7, 0), 2, difflight));
    objects.add(std::make_shared<Xy_rect>(3, 5, 1, 3, -2, difflight));

    return objects;
}

std::pair<Hittable_list, Camera> cornell_box(const double aspect_ratio) {
    Hittable_list objects;

    auto red   = std::make_shared<Lambertian>(std::make_shared<Solid_color>(0.65, 0.05, 0.05));
    auto white = std::make_shared<Lambertian>(std::make_shared<Solid_color>(0.73, 0.73, 0.73));
    auto green = std::make_shared<Lambertian>(std::make_shared<Solid_color>(0.12, 0.45, 0.15));

    //auto light = std::make_shared<Diffuse_light>(std::make_shared<Solid_color>(15, 15, 15));
    auto light = std::make_shared<Diffuse_light>(std::make_shared<Solid_color>(30, 30, 30));
    //auto light = std::make_shared<Diffuse_light>(std::make_shared<Solid_color>(100, 100, 100));

    objects.add(std::make_shared<Flip_face>(std::make_shared<Yz_rect>(  0, 555,   0, 555, 555, green)));
    objects.add(std::make_shared<Yz_rect>(  0, 555,   0, 555,   0, red  ));
    objects.add(std::make_shared<Xz_rect>(213, 343, 227, 332, 554, light));
    objects.add(std::make_shared<Flip_face>(std::make_shared<Xz_rect>(  0, 555,   0, 555,   0, white)));
    objects.add(std::make_shared<Xz_rect>(  0, 555,   0, 555, 555, white));
    objects.add(std::make_shared<Flip_face>(std::make_shared<Xy_rect>(  0, 555,   0, 555, 555, white)));

    std::shared_ptr<Hittable> box1 = std::make_shared<Box>(Point3(0, 0,  0), Point3(165, 300, 165), white);
    box1 = std::make_shared<Rotate_y>(box1, 15);
    box1 = std::make_shared<Translate>(box1, Vec3(265, 0, 295));
    objects.add(box1);

    std::shared_ptr<Hittable> box2 = std::make_shared<Box>(Point3(0, 0, 0), Point3(165, 165, 165), white);
    box2 = std::make_shared<Rotate_y>(box2, -18);
    box2 = std::make_shared<Translate>(box2, Vec3(130, 0, 65));
    objects.add(box2);

    Point3 lookfrom(278, 278, -800);
    Point3 lookat(278, 278, 0);
    Vec3 vup(0, 1, 0);
    //auto dist_to_focus = (lookfrom - lookat).length();
    auto dist_to_focus = 10;
    auto aperture = 0.0;
    auto vfov = 40;

    Camera cam(lookfrom, lookat, vup, vfov, aspect_ratio, aperture, dist_to_focus, 0.0, 1.0);

    return std::make_pair(objects, cam);
}

struct render_tile {
    int start_x;
    int start_y;
    int end_x;
    int end_y;
};

int main() {
    const auto aspect_ratio = 16.0 / 9.0;
    //const int image_width = 1920;
    //const int image_width = 1280;
    const int image_width = 768;
    //const int image_width = 384;
    const int image_height = static_cast<int>(image_width / aspect_ratio);
    //const int samples_per_pixel = 10000;
    const int samples_per_pixel = 5000;
    //const int samples_per_pixel = 2000;
    //const int samples_per_pixel = 1000;
    //const int samples_per_pixel = 400;
    //const int samples_per_pixel = 50;
    const int max_depth = 50;

    //const int tile_size = 128;
    const int tile_size = 64;

    //const int sph_i = 3;
    //const int sph_i = 5;
    //const int sph_i = 7;
    //const int sph_i = 11;

    //auto [world, camera] = setup_random_scene(sph_i, aspect_ratio);
    //auto [world, camera] = two_spheres(aspect_ratio);
    //auto [world, camera] = simple_light(aspect_ratio);
    auto [world, camera] = cornell_box(aspect_ratio);

    const Color background(0, 0, 0);

    std::atomic<int> queue_counter;
    boost::lockfree::queue<render_tile> queue(0);

    for (int i = 0; i<std::ceil(double(image_height)/double(tile_size)); i++) {
        for (int j = 0; j<std::ceil(double(image_width)/double(tile_size)); j++) {
            render_tile rt = {
                i*tile_size,
                j*tile_size,
                std::min((i+1)*tile_size, image_height),
                std::min((j+1)*tile_size, image_width)
            };
            queue.push(rt);
            ++queue_counter;
        }
    }

    std::vector<Color> image(image_width * image_height);

    #pragma omp parallel
    {
        render_tile rt;
        while (queue.pop(rt)) {
            --queue_counter;
            int qc = queue_counter;
            printf("\rTiles remaining: %4d", qc);
            std::cout << std::flush;
            for (int i = rt.start_x; i<rt.end_x; ++i) {
                for (int j = rt.start_y; j<rt.end_y; ++j) {
                    Color pixel_color(0, 0, 0);
                    for (int s = 0; s<samples_per_pixel; ++s) {
                        auto u = double(j + random_double(-0.5, 0.5)) / (image_width - 1);
                        auto v = double(i + random_double(-0.5, 0.5)) / (image_height - 1);
                        Ray r = camera.get_ray(u, v);
                        pixel_color += ray_color(r, background, world, max_depth);
                    }
                    image[i*image_width+j] = pixel_color;
                }
            }
        }
    }

    std::cerr << "\nAssembling image.\n";

    std::vector<unsigned char> img_lode(image_width * image_height * 3);
    for (int i = 0; i<image_height; ++i) {
        for (int j = 0; j<image_width; ++j) {
            write_color_vec(img_lode, i*image_width*3+j*3, image[(image_height-1-i)*image_width+j], samples_per_pixel);
        }
    }
    std::cerr << "Writing file \"image.png\".\n";

    unsigned error = lodepng::encode("image.png", img_lode, image_width, image_height, LCT_RGB);

    if (error)
        std::cerr << "lodepng encoder error " << error << ": "<< lodepng_error_text(error) << std::endl;

    std::cerr << "Done.\n";
}