package main

import "fmt" // for IO and standard library
import "os" // for handling the progress bar
import "math" // for maths
import "unsafe" // for fast inverse square pointers
import "math/rand" // for random
import "time" // for random


func init() {
    rand.Seed(time.Now().UnixNano())
}

// ================ VEC3 CLASS =====================

type Vec3 struct {
	E [3]float32
}

//Basic vector functions

func NewVec3(e0, e1, e2 float32) Vec3 {
	return Vec3{E: [3]float32{e0, e1, e2}}
}

func (v Vec3) X() float32 {
	return v.E[0]
}

// can be executed as v.X() in main

func (v Vec3) Y() float32 {
	return v.E[1]
}

func (v Vec3) Z() float32 {
	return v.E[2]
}

func (v Vec3) Neg() Vec3 {
	return NewVec3(-v.E[0], -v.E[1], -v.E[2])
}

func (v Vec3) Get(i int) float32 {
	return v.E[i]
}

func (v *Vec3) Set(i int, val float32) {
	v.E[i] = val
}

func (v Vec3) Add(v2 Vec3) Vec3 {
	return NewVec3(v.E[0]+v2.E[0], v.E[1]+v2.E[1], v.E[2]+v2.E[2])
}

func (v Vec3) Mult(t float32) Vec3 {
	return NewVec3(v.E[0]*t, v.E[1]*t, v.E[2]*t)
}

func (v Vec3) Div(t float32) Vec3 {
	if t != 0 {
		return NewVec3(v.E[0] / t, v.E[1] / t, v.E[2] / t)
	}
	return v
}


func (v Vec3) Length() float32 {
	return float32(math.Sqrt(float64(v.E[0]*v.E[0] + v.E[1]*v.E[1] + v.E[2]*v.E[2])))
}

func (v Vec3) Length_squared() float32 {
	return v.E[0]*v.E[0] + v.E[1]*v.E[1] + v.E[2]*v.E[2]
}

func (v Vec3) Near_zero() bool {
	var s float64 = 1e-8
	return (math.Abs(float64(v.E[0])) < s && math.Abs(float64(v.E[1])) < s && math.Abs(float64(v.E[2])) < s)
}

// Vector utility functions

func (v Vec3) String() string {
	return fmt.Sprintf("%v %v %v", v.E[0], v.E[1], v.E[2])
}

func (v Vec3) Sub(v2 Vec3) Vec3 {
	return NewVec3(v.E[0]-v2.E[0], v.E[1]-v2.E[1], v.E[2]-v2.E[2])
}

func (v Vec3) MultVec(v2 Vec3) Vec3 {
	return NewVec3(v.E[0]*v2.E[0], v.E[1]*v2.E[1], v.E[2]*v2.E[2])
}

func (v Vec3) DivVec(v2 Vec3) Vec3 {
	if v2.E[0] != 0 && v2.E[1] != 0 && v2.E[2] != 0 {
		return NewVec3(v.E[0]/v2.E[0], v.E[1]/v2.E[1], v.E[2]/v2.E[2])
	}
	return v
}

func Dot(v1 Vec3, v2 Vec3) float32 {
	return (v1.E[0]*v2.E[0] + v1.E[1]*v2.E[1] + v1.E[2]*v2.E[2])
}

func Cross(v1 Vec3, v2 Vec3) Vec3 {
	return NewVec3(v1.E[1]*v2.E[2] - v1.E[2]*v2.E[1],
		v1.E[2]*v2.E[0] - v1.E[0]*v2.E[2],
		v1.E[0]*v2.E[1] - v1.E[1]*v2.E[0])
}

func q_rsqrt(v Vec3) float32 {
	var x float32 = v.E[0]*v.E[0] + v.E[1]*v.E[1] + v.E[2]*v.E[2]
	i := *(*int32)(unsafe.Pointer(&x)) // evil floating point bit level hacking
    i = 0x5f3759df - (i >> 1)          // what the fuck?
    y := *(*float32)(unsafe.Pointer(&i))
    return y
}

func Unit_vector(v Vec3) Vec3 {
	//new_v := v.Mult(q_rsqrt(v))
	new_v := v.Mult(1.0/v.Length())
    return new_v
}

func Random_in_unit_disk() Vec3 {
	for true {
		p := NewVec3(RandomDoubleInRange(-1,1), RandomDoubleInRange(-1,1), 0)
		if p.Length_squared() < 1 {
			return p
		}
	}
	return NewVec3(0,0,0)
}

func RandomInUnitSphere() Vec3 {
	for true {
		p := RandomVec3(-1, 1)
		if (p.Length_squared() < 1) {
			return p
		}
	}
	return NewVec3(0,0,0)
}

func RandomUnitVector() Vec3 {
	return Unit_vector(RandomInUnitSphere())
}

func RandomOnHemisphere(normal Vec3) Vec3 {
	var on_unit_sphere Vec3 = RandomUnitVector()
	if (Dot(on_unit_sphere, normal) > 0.0) {
		return on_unit_sphere
	} else {
		return on_unit_sphere.Neg()
	}
}

func Reflect(v Vec3, n Vec3) Vec3 {
	return v.Sub(n.Mult(2*Dot(v,n)))
}

func Refract(uv, n Vec3, etai_over_etat float32) Vec3 {
    cos_theta := float32(math.Min(float64(Dot(uv.Neg(), n)), 1.0))
    r_out_perp := uv.Add(n.Mult(cos_theta)).Mult(etai_over_etat)
    r_out_parallel := n.Mult(float32(math.Sqrt(math.Abs(1.0 - float64(r_out_perp.Length_squared())))))
    return r_out_perp.Add(r_out_parallel)
}

const pi float32 = 3.1415926535897932385

func Degrees_to_radians(degrees float32) float64 {
	return float64((degrees * pi) / 180.0)
}

func Linear_to_gamma(linear_component float32) float32 {
	return float32(math.Sqrt(float64(linear_component)))
}

func RandomDouble() float32 {
    // Returns a random real in [0,1).
    return rand.Float32()
}

func RandomDoubleInRange(min, max float32) float32 {
    // Returns a random real in [min,max).
    return min + (max-min)*RandomDouble()
}

func RandomVec3(borders ...float32) Vec3 {
	if (len(borders) == 2) {
		return NewVec3(RandomDoubleInRange(borders[0], borders[1]),RandomDoubleInRange(borders[0], borders[1]),RandomDoubleInRange(borders[0], borders[1]))
	} else {
		return NewVec3(RandomDouble(),RandomDouble(),RandomDouble())
	}
}
// ============== COLOUR CLASS ==============

func Write_color(v Vec3, samples_per_pixel int) {
	// Averaging

	r := v.X() / float32(samples_per_pixel)
	g := v.Y() / float32(samples_per_pixel)
	b := v.Z() / float32(samples_per_pixel)

	r = Linear_to_gamma(r)
	g = Linear_to_gamma(g)
	b = Linear_to_gamma(b)

	// Write the translate [0, 255] value of each color component
	intensity := NewInterval(0.000, 0.999)
	fmt.Println(int(intensity.Clamp(r)* 256.0), int(intensity.Clamp(g)* 256.0), int(intensity.Clamp(b)* 256.0))
}

func NewColor(e0, e1, e2 float32) Vec3 {
	return Vec3{E: [3]float32{e0, e1, e2}}
}

// ============== RAY CLASS =================

func NewPoint3(e0, e1, e2 float32) Vec3 {
	return Vec3{E: [3]float32{e0, e1, e2}}
}
type Ray struct {
	Orig Vec3
	Dir Vec3
}

func NewRay(orig Vec3, dir Vec3) *Ray {
    return &Ray{Orig: orig, Dir: dir}
}

func (r *Ray) Origin() Vec3 {
	return r.Orig
}

func (r *Ray) Direction() Vec3 {
	return r.Dir
}

func (r *Ray) At(t float32) Vec3 {
	return Vec3{
		E: [3]float32{
			r.Orig.X() + t*r.Dir.X(),
			r.Orig.Y() + t*r.Dir.Y(),
			r.Orig.Z() + t*r.Dir.Z(),
		},
	}
}

// =============== INTERVAL =================

type Interval struct {
	min, max float32
}

func NewInterval(borders ...float32) *Interval {
	if (len(borders) == 2) {
		return &Interval{
			min: borders[0],
			max: borders[1],
		}
	} else {
		return &Interval{
			min: float32(math.Inf(1)),
			max: float32(math.Inf(-1)),
		}
	}
}

func (i *Interval) Contains(x float32) bool {
	return i.min <= x && x <= i.max
}

func (i *Interval) Surrounds(x float32) bool {
	return i.min < x && x < i.max
}

func (i *Interval) Clamp(x float32) float32 {
	if (x < i.min) {
		return i.min
	} else if (x > i.max) {
		return i.max
	}
	return x
}

var Empty *Interval = NewInterval()
var Universe *Interval = NewInterval(float32(math.Inf(-1)), float32(math.Inf(1)))


// =============== HIT ======================

type Hit_record struct {
	p Vec3
	normal Vec3
	t float32
	front_face bool
	mat Material
}

func (rec *Hit_record) Set_face_normal(r *Ray, outward_normal Vec3) {
	if (Dot(r.Direction(), outward_normal) < 0) {
		rec.front_face = true
	}
	if (rec.front_face) {
		rec.normal = outward_normal
	} else {
		rec.normal = outward_normal.Neg()
	}
}

// =============== SPHERE ===================

type Sphere struct {
	center Vec3
	radius float32
	mat Material
}

func (s Sphere) Hit_sphere(r *Ray, ray_t *Interval, rec *Hit_record) bool {
	oc := r.Origin().Sub(s.center)
	a := r.Direction().Length_squared()
	half_b := Dot(oc, r.Direction())
	c := oc.Length_squared() - s.radius*s.radius

	discriminant := half_b*half_b - a*c
	if (discriminant < 0) {
		return false
	}
	sqrtd := float32(math.Sqrt(float64(discriminant)))

	// find the nearest root that lies in the acceptable range
	root := (-half_b - sqrtd) / a
	if (!ray_t.Surrounds(root)) {
		root = (sqrtd-half_b) / a
		if (!ray_t.Surrounds(root)) {
			return false
		}
	}
	rec.t = root
	rec.p = r.At(rec.t)
	outward_normal := rec.p.Sub(s.center).Div(s.radius)
	rec.Set_face_normal(r, outward_normal)
	rec.mat = s.mat
	return true
}

func NewSphere(center Vec3, radius float32, mat Material) Sphere {
	return Sphere{
		center: center,
		radius: radius,
		mat:    mat,
	}
}


// =============== HITTABLE ==================
type Hittable struct {
	spheres []Sphere
}

func NewHittable() *Hittable {
	return &Hittable{}
}

func (hl *Hittable) Add(s Sphere) {
	hl.spheres = append(hl.spheres, s)
}

func (hl *Hittable) Clear() {
	hl.spheres = []Sphere{}
}

func (hl Hittable) Hit(r *Ray, ray_t Interval, rec *Hit_record) bool {
	var temp_rec Hit_record
	var hit_anything bool = false
	var closest_so_far float32 = ray_t.max

	for _, sphere := range hl.spheres {
		if (sphere.Hit_sphere(r, NewInterval(ray_t.min, closest_so_far), &temp_rec)) {
			hit_anything = true
			closest_so_far = temp_rec.t
			*rec = temp_rec
		}
	}
	return hit_anything
}

// =============== MATERIAL =================
type Material struct {
	material int
	albedo Vec3
	fuzz float32
	ir float32
}

func NewMaterial(m int, a Vec3, f float32, i float32) Material {
	if f >= 1.0 {
		return Material{
			material: m,
			albedo : a,
			fuzz: 1.0,
			ir: i,
		}
	} else {
		return Material{
			material: m,
			albedo : a,
			fuzz: f,
			ir: i,
		}
	}

}

func Reflectance(cosine float32, ref_idx float32) float32 {
	// Use Schlick's approximation for reflectance
	r0 := (1.0-ref_idx) / (1.0+ref_idx)
	r0 = r0*r0
	return r0 + (1.0-r0)*float32(math.Pow(1.0 - float64(cosine), 5))
}

func (mat Material) Scatter(r_in *Ray, rec *Hit_record, attenuation *Vec3, scattered *Ray) bool {
	if (mat.material == 0) {	//Lambertian
		scatter_direction := rec.normal.Add(RandomUnitVector())
		// Catch degenerate scatter direction
		if scatter_direction.Near_zero() {
			scatter_direction = rec.normal
		}

		*scattered = *NewRay(rec.p, scatter_direction)
		*attenuation = mat.albedo
	} else if (mat.material == 1) {   //Metall
		var reflected Vec3 = Reflect(Unit_vector(r_in.Direction()), rec.normal)
		*scattered = *NewRay(rec.p, reflected.Add(RandomUnitVector().Mult(mat.fuzz)))
		*attenuation = mat.albedo
		if !(Dot(scattered.Direction(), rec.normal) > 0) {
			return false
		}
	} else if (mat.material == 2) { //Dielectric
		*attenuation = NewColor(1.0, 1.0, 1.0)
		var refraction_ration float32 = 1.0
		if rec.front_face {
			refraction_ration /= mat.ir
		} else {
			refraction_ration = mat.ir
		}
		unit_direction := Unit_vector(r_in.Direction())

		cos_theta := float32(math.Min(float64(Dot(unit_direction.Neg(), rec.normal)), 1.0))
		sin_theta := float32(math.Sqrt(float64(1.0 - cos_theta*cos_theta)))
		var direction Vec3
		cannot_refract := refraction_ration*sin_theta > 1.0
		
		if (cannot_refract || Reflectance(cos_theta, refraction_ration) > RandomDouble()) {
			direction = Reflect(unit_direction, rec.normal)
		} else {
			direction = Refract(unit_direction, rec.normal, refraction_ration)
		}
		
		*scattered = *NewRay(rec.p, direction)
	}
	return true
}


// =============== CAMERA ===================

type Camera struct {
	aspect_ratio float32
	image_width int
	image_height int
	center Vec3
	pixel00_loc Vec3
	pixel_delta_u Vec3
	pixel_delta_v Vec3
	samples_per_pixel int
	max_depth int
	vfov float32
	lookfrom Vec3
	lookat Vec3
	vup Vec3
	u Vec3
	v Vec3
	w Vec3
	defocus_angle float32
	focus_dist float32
	defocus_disk_u Vec3
	defocus_disk_v Vec3
} 

func NewCamera() *Camera {
	return &Camera{}
}

func Ray_color(r Ray, depth int, world *Hittable) Vec3 {
	var rec Hit_record

	if (depth <= 0) {
		return NewColor(0,0,0)
	}
	if (world.Hit(&r, *NewInterval(0.001, float32(math.Inf(1))), &rec)) {
		var scattered Ray
		var attenuation Vec3
		if ((rec.mat).Scatter(&r, &rec, &attenuation, &scattered)) {
			return attenuation.MultVec(Ray_color(scattered, depth-1, world))
		}
		return NewColor(0.0,0.0,0.0)
	}
	unit_direction := Unit_vector(r.Direction())
	a := (unit_direction.Y() + 1.0)*0.5
	return NewColor(1.0,1.0,1.0).Mult(float32(1.0-a)).Add(NewColor(0.5,0.7,1.0).Mult(a))
}

func (cam *Camera) Initialize() {
    // Calculate the image height
    cam.image_height = int(float32(cam.image_width) / cam.aspect_ratio)
    if cam.image_height < 1 {
        cam.image_height = 1
    }
	cam.center = cam.lookfrom


    // Viewport
	theta := Degrees_to_radians(cam.vfov)
	h := float32(math.Tan(theta/2.0))
    var viewport_height float32 = 2.0 * h * cam.focus_dist
    var viewport_width float32 = viewport_height * float32(cam.image_width)/float32(cam.image_height)

	// Calculate the u,v,w unit basis vectors for the camera coordinate frame
	cam.w = Unit_vector(cam.lookfrom.Sub(cam.lookat))
	cam.u = Unit_vector(Cross(cam.vup, cam.w))
	cam.v = Cross(cam.w, cam.u)

    // Calculate the vectors across the horizontal and down the vertical viewport edges
    viewport_u := cam.u.Mult(viewport_width)
    viewport_v := cam.v.Mult(-viewport_height)

    //Calculate the horizontal and vertical delta vectors from pixel to pixel
    cam.pixel_delta_u = viewport_u.Div(float32(cam.image_width))
    cam.pixel_delta_v = viewport_v.Div(float32(cam.image_height))

    // Calculate the location of the upper left pixel
    viewport_upper_left := cam.center.Sub(cam.w.Mult(cam.focus_dist)).Sub(viewport_u.Div(2)).Sub(viewport_v.Div(2))
    cam.pixel00_loc = viewport_upper_left.Add((cam.pixel_delta_u.Add(cam.pixel_delta_v)).Div(2))

	// Calculate the camera defocus disk basis vectors
	defocus_radius := cam.focus_dist*float32(math.Tan(Degrees_to_radians(cam.defocus_angle / 2)))
	cam.defocus_disk_u = cam.u.Mult(defocus_radius)
	cam.defocus_disk_v = cam.v.Mult(defocus_radius)
}

func (cam *Camera) Render(world *Hittable) {
    cam.Initialize()

    // Rendering
    fmt.Println("P3")
    fmt.Println(cam.image_width, " ", cam.image_height, "\n255")
    for j := 0; j < cam.image_height; j++ {
        fmt.Fprintf(os.Stderr, "\rScanlines remaining: %d ", cam.image_height-j)
        for i := 0; i < cam.image_width; i++ {
			pixel_color := NewColor(0,0,0)
			for sample := 0; sample < cam.samples_per_pixel; sample++ {
				r := cam.GetRay(i, j)
				pixel_color = pixel_color.Add(Ray_color(*r, cam.max_depth, world))
			}
			Write_color(pixel_color, cam.samples_per_pixel)
        }
    }
}


func (cam *Camera) Defocus_disk_sample() Vec3 {
	// Returns a random point in the camera defocus disk
	p := Random_in_unit_disk()
	return cam.center.Add(cam.defocus_disk_u.Mult(p.E[0])).Add(cam.defocus_disk_v.Mult(p.E[1]))
}

func (cam *Camera) GetRay(i, j int) *Ray {
	// Get a randomly-sampled camera ray for the pixel at location i,j originating from the camera defocus disk
	pixel_center := cam.pixel00_loc.Add(cam.pixel_delta_u.Mult(float32(i))).Add(cam.pixel_delta_v.Mult(float32(j)))
	pixel_sample := pixel_center.Add(cam.Pixel_sample_square())

	var ray_origin Vec3
	if cam.defocus_angle <= 0 {
		ray_origin = cam.center
	} else {
		ray_origin = cam.Defocus_disk_sample()
	}
	ray_direction := pixel_sample.Sub(ray_origin)
	return NewRay(ray_origin , ray_direction)
}

func (cam *Camera) Pixel_sample_square() Vec3 {
	// Returns a random point in the square surrounding a pixel at the origin
	px := -0.5 + RandomDouble()
	py := -0.5 + RandomDouble()
	return (cam.pixel_delta_u.Mult(px).Add(cam.pixel_delta_v.Mult(py)))
}


// =============== MAIN =====================
func main() {
	world := NewHittable()

	groundMaterial := NewMaterial(0, NewColor(0.5, 0.5, 0.5), 0.0, 0.0)
	world.Add(Sphere{
		center: NewVec3(0, -1000, 0),
		radius: 1000,
		mat:    groundMaterial,
	})

	for a := -11; a < 11; a++ {
		for b := -11; b < 11; b++ {
			chooseMat := RandomDouble()
			center := NewVec3(float32(a)+0.9*RandomDouble(), 0.2, float32(b)+0.9*RandomDouble())

			if (center.Sub(NewVec3(4, 0.2, 0)).Length() > 0.9) {
				var sphereMaterial Material

				if chooseMat < 0.8 {
					// diffuse
					albedo := RandomVec3().MultVec(RandomVec3())
					sphereMaterial = NewMaterial(0, albedo, 0.0, 0.0)
				} else if chooseMat < 0.95 {
					// metal
					albedo := RandomVec3(0.5, 1)
					fuzz := RandomDoubleInRange(0, 0.5)
					sphereMaterial = NewMaterial(1, albedo, fuzz, 0.0)
				} else {
					// glass
					sphereMaterial = NewMaterial(2, NewColor(1.0, 1.0, 1.0), 0.0, 1.5)
				}
				world.Add(Sphere{
					center: center,
					radius: 0.2,
					mat:    sphereMaterial,
				})
			}
		}
	}

	material1 := NewMaterial(2, NewColor(1.0, 1.0, 1.0), 0.0, 1.5)
	world.Add(Sphere{
		center: NewVec3(0, 1, 0),
		radius: 1.0,
		mat:    material1,
	})

	material2 := NewMaterial(0, NewColor(0.4, 0.2, 0.1), 0.0, 0.0)
	world.Add(Sphere{
		center: NewVec3(-4, 1, 0),
		radius: 1.0,
		mat:    material2,
	})

	material3 := NewMaterial(1, NewColor(0.7, 0.6, 0.5), 0.0, 0.0)
	world.Add(Sphere{
		center: NewVec3(4, 1, 0),
		radius: 1.0,
		mat:    material3,
	})

	cam := NewCamera()
	cam.aspect_ratio = 16.0 / 9.0
	cam.image_width = 1200
	cam.samples_per_pixel = 200
	cam.max_depth = 45
	cam.vfov = 20
	cam.lookfrom = NewVec3(13, 2, 3)
	cam.lookat = NewVec3(0, 0, 0)
	cam.vup = NewVec3(0, 1, 0)
	cam.defocus_angle = 0.6
	cam.focus_dist = 10.0

	cam.Render(world)
}