use std::clone::Clone;
use std::f64;
use std::fmt;
use std::fs::File;
use std::io::Write;
use std::ops::Add;
use std::ops::Mul;
use std::ops::Sub;

use log::info;
use rand::Rng;
use structopt::StructOpt;

fn vdot(v0: &Vector, v1: &Vector) -> f64 {
    v0.x * v1.x + v0.y * v1.y + v0.z * v1.z
}

fn vcross(v0: &Vector, v1: &Vector) -> Vector {
    Vector {
        x: v0.y * v1.z - v0.z * v1.y,
        y: v0.z * v1.x - v0.x * v1.z,
        z: v0.x * v1.y - v0.y * v1.x,
    }
}

#[derive(Debug, Copy, Clone)]
struct Vector {
    x: f64,
    y: f64,
    z: f64,
}

impl Vector {
    fn zero() -> Vector {
        Vector {
            x: 0.,
            y: 0.,
            z: 0.,
        }
    }
    fn new(x: f64, y: f64, z: f64) -> Vector {
        Vector { x, y, z }
    }

    fn normalize(&mut self) {
        let length = vdot(self, self).sqrt();
        if length.abs() < 1.0e-17 {
            return;
        }

        self.x = self.x / length;
        self.y = self.y / length;
        self.z = self.z / length;
    }
}

impl fmt::Display for Vector {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "({:.4}, {:.4}, {:.4})", self.x, self.y, self.z)
    }
}

impl Add for Vector {
    type Output = Vector;

    fn add(self, other: Vector) -> Vector {
        Vector {
            x: self.x + other.x,
            y: self.y + other.y,
            z: self.z + other.z,
        }
    }
}

impl Mul for Vector {
    type Output = Vector;

    fn mul(self, other: Vector) -> Vector {
        Vector {
            x: self.x * other.x,
            y: self.y * other.y,
            z: self.z * other.z,
        }
    }
}

impl Mul<f64> for Vector {
    type Output = Vector;

    fn mul(self, other: f64) -> Vector {
        Vector {
            x: self.x * other,
            y: self.y * other,
            z: self.z * other,
        }
    }
}

impl Sub for Vector {
    type Output = Vector;

    fn sub(self, other: Vector) -> Vector {
        Vector {
            x: self.x - other.x,
            y: self.y - other.y,
            z: self.z - other.z,
        }
    }
}

#[derive(Copy, Clone)]
struct ISect {
    t: f64,
    p: Vector,
    n: Vector,
    hit: bool,
}

struct Sphere {
    center: Vector,
    radius: f64,
}

struct Plane {
    p: Vector,
    n: Vector,
}

#[derive(Debug)]
struct Ray {
    org: Vector,
    dir: Vector,
}

impl fmt::Display for Ray {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "o: {} d: {}", self.org, self.dir)
    }
}

fn ray_sphere_intersect(mut isect: ISect, ray: &Ray, sphere: &Sphere) -> ISect {
    let rs = ray.org - sphere.center;

    let b = vdot(&rs, &ray.dir);
    let c = vdot(&rs, &rs) - sphere.radius * sphere.radius;
    let d = b * b - c;
    if d > 0.0 {
        let t = -b - d.sqrt();
        if t > 0.0 && t < isect.t {
            isect.t = t;
            isect.hit = true;

            isect.p = ray.org + ray.dir * t;

            isect.n = isect.p - sphere.center;
            isect.n.normalize();
        }
    }

    isect
}

fn ray_plane_intersect(mut isect: ISect, ray: &Ray, plane: &Plane) -> ISect {
    let d = -vdot(&plane.p, &plane.n);
    let v = vdot(&ray.dir, &plane.n);

    if v.abs() < 1.0e-17 {
        return isect;
    }

    let t = -(vdot(&ray.org, &plane.n) + d) / v;

    if (t > 0.0) && (t < isect.t) {
        isect.t = t;
        isect.hit = true;

        isect.p = ray.org + ray.dir * t;
        isect.n = plane.n;
    }
    isect
}

struct World {
    spheres: Vec<Sphere>,
    plane: Plane,

    width: usize,
    height: usize,
    n_subsamples: usize,
    n_aosamples: usize,
}

fn ortho_basis(n: &Vector) -> [Vector; 3] {
    let mut basis: [Vector; 3] = [Vector::zero(), Vector::zero(), n.clone()];

    if (n.x < 0.6) && (n.x > -0.6) {
        basis[1].x = 1.0;
    } else if (n.y < 0.6) && (n.y > -0.6) {
        basis[1].y = 1.0;
    } else if (n.z < 0.6) && (n.z > -0.6) {
        basis[1].z = 1.0;
    } else {
        basis[1].x = 1.0;
    }

    basis[0] = vcross(&basis[1], &basis[2]);
    basis[0].normalize();

    basis[1] = vcross(&basis[2], &basis[0]);
    basis[1].normalize();

    basis
}

impl World {
    fn ambient_occlusion(&self, isect: ISect) -> Vector {
        let n_theta: usize = self.n_aosamples;
        let n_phi: usize = self.n_aosamples;
        let eps: f64 = 0.0001;

        let p = isect.p + isect.n * eps;
        let basis = ortho_basis(&isect.n);
        let mut occlusion: f64 = 0.0;
        let mut rng = rand::thread_rng();
        for _ in 0..n_theta {
            for _ in 0..n_phi {
                let theta: f64 = rng.gen::<f64>().sqrt();
                let phi: f64 = rng.gen::<f64>() * 2.0 * f64::consts::PI;

                let x = phi.cos() * theta;
                let y = phi.sin() * theta;
                let z = (1.0 - theta * theta).sqrt();

                let rx = x * basis[0].x + y * basis[1].x + z * basis[2].x;
                let ry = x * basis[0].y + y * basis[1].y + z * basis[2].y;
                let rz = x * basis[0].z + y * basis[1].z + z * basis[2].z;

                let ray = Ray {
                    org: p.clone(),
                    dir: Vector::new(rx, ry, rz),
                };

                let mut occ_isect = ISect {
                    t: 1.0e+17,
                    hit: false,
                    p: Vector::zero(),
                    n: Vector::zero(),
                };

                occ_isect = ray_sphere_intersect(occ_isect, &ray, &self.spheres[0]);
                occ_isect = ray_sphere_intersect(occ_isect, &ray, &self.spheres[1]);
                occ_isect = ray_sphere_intersect(occ_isect, &ray, &self.spheres[2]);
                occ_isect = ray_plane_intersect(occ_isect, &ray, &self.plane);

                if occ_isect.hit {
                    occlusion += 1.0;
                }
            }
        }
        let n_theta = n_theta as f64;
        let n_phi = n_phi as f64;
        occlusion = (n_theta * n_phi - occlusion) / (n_theta * n_phi);
        Vector::new(occlusion, occlusion, occlusion)
    }
}

fn clamp(f: f64) -> u8 {
    let i = (f * 255.5) as i32;

    if i < 0 {
        return 0;
    }
    if i > 255 {
        return 255;
    }
    i as u8
}

impl World {
    fn render(&mut self) -> Vec<u8> {
        let w = self.width;
        let h = self.height;
        info!(
            "Rendering world {}x{} w/ {} subsampling and {} ambient occlusion sampling",
            w, h, self.n_subsamples, self.n_aosamples
        );
        let mut img = vec![0; w * h * 3];
        let mut fimg = vec![0.0; w * h * 3];
        for y in 0..h as usize {
            for x in 0..w as usize {
                for v in 0..self.n_subsamples {
                    for u in 0..self.n_subsamples {
                        let fx = x as f64;
                        let fy = y as f64;
                        let fw = w as f64;
                        let fh = h as f64;
                        let fu = u as f64;
                        let fv = v as f64;
                        let n_subsamples = self.n_subsamples as f64;
                        let px = (fx + (fu / n_subsamples) - (fw / 2.0)) / (fw / 2.);
                        let py = -(fy + (fv / n_subsamples) - (fh / 2.0)) / (fh / 2.);

                        let mut dir = Vector::new(px, py, -1.0);
                        dir.normalize();

                        let ray = Ray {
                            org: Vector::zero(),
                            dir,
                        };
                        let mut isect = ISect {
                            t: 1.0e+17,
                            hit: false,
                            p: Vector::zero(),
                            n: Vector::zero(),
                        };

                        isect = ray_sphere_intersect(isect, &ray, &self.spheres[0]);

                        isect = ray_sphere_intersect(isect, &ray, &self.spheres[1]);
                        isect = ray_sphere_intersect(isect, &ray, &self.spheres[2]);
                        isect = ray_plane_intersect(isect, &ray, &self.plane);

                        if isect.hit {
                            let col = self.ambient_occlusion(isect);
                            fimg[(x + y * w) * 3 + 0] += col.x;
                            fimg[(x + y * w) * 3 + 1] += col.y;
                            fimg[(x + y * w) * 3 + 2] += col.z;
                        }
                    }
                }
                let n2 = (self.n_subsamples * self.n_subsamples) as f64;
                fimg[3 * (y * w + x) + 0] /= n2;
                fimg[3 * (y * w + x) + 1] /= n2;
                fimg[3 * (y * w + x) + 2] /= n2;

                img[3 * (y * w + x) + 0] = clamp(fimg[3 * (y * w + x) + 0]);
                img[3 * (y * w + x) + 1] = clamp(fimg[3 * (y * w + x) + 1]);
                img[3 * (y * w + x) + 2] = clamp(fimg[3 * (y * w + x) + 2]);
            }
        }
        img
    }
}

fn save(width: usize, height: usize, img: &[u8], path: &str) -> std::io::Result<()> {
    info!("Saving scene to {}", path);
    let mut buf = File::create(path)?;
    writeln!(&buf, "P6\n{} {}\n255", width, height)?;
    buf.write_all(img)
}

fn init_scene(width: usize, height: usize, n_subsamples: usize, n_aosamples: usize) -> World {
    World {
        spheres: vec![
            Sphere {
                center: Vector::new(-2.0, 0.0, -3.5),
                radius: 0.5,
            },
            Sphere {
                center: Vector::new(-0.5, 0.0, -3.0),
                radius: 0.5,
            },
            Sphere {
                center: Vector::new(1.0, 0.0, -2.2),
                radius: 0.5,
            },
        ],
        plane: Plane {
            p: Vector::new(0.0, -0.5, 0.0),
            n: Vector::new(0.0, 1.0, 0.0),
        },
        width,
        height,
        n_subsamples,
        n_aosamples,
    }
}

#[derive(StructOpt, Debug)]
#[structopt()]
struct Opt {
    #[structopt(short = "w", long = "width", default_value = "512")]
    width: usize,
    #[structopt(short = "h", long = "height", default_value = "512")]
    height: usize,
    #[structopt(short = "s", long = "subsamples", default_value = "2")]
    n_subsamples: usize,
    #[structopt(short = "a", long = "aosamples", default_value = "8")]
    n_aosamples: usize,
}

fn main() -> std::io::Result<()> {
    let opt = Opt::from_args();
    stderrlog::new()
        .verbosity(3)
        .timestamp(stderrlog::Timestamp::Millisecond)
        .init()
        .unwrap();

    let mut w = init_scene(opt.width, opt.height, opt.n_subsamples, opt.n_aosamples);
    let img = w.render();
    save(opt.width, opt.height, &img, "ao.ppm")
}
