#!/usr/bin/env python3

import math
import time
import random

def create_image(width, height):
    #Create an empty image (list of lists) initialized to black
    return [[(0.0, 0.0, 0.0) for _ in range(width)] for _ in range(height)]

def put_pixel(img, x, y, color):
    #Safely place a pixel into the image buffer, if in range.
    if 0 <= x < len(img[0]) and 0 <= y < len(img):
        img[y][x] = color

def clamp_color(c):
    #Clamp an (r, g, b) to integer 0..255
    return (
        max(0, min(255, int(c[0]))),
        max(0, min(255, int(c[1]))),
        max(0, min(255, int(c[2]))),
    )

def normalize(v):
    #Normalize a 3D vector.
    length = math.sqrt(v[0]*v[0] + v[1]*v[1] + v[2]*v[2])
    if length == 0:
        return (0, 0, 0)
    return (v[0]/length, v[1]/length, v[2]/length)

def vector_sub(a, b):
    #Subtract 3D vectors.
    return (a[0]-b[0], a[1]-b[1], a[2]-b[2])

def vector_add(a, b):
    #Add 3D vectors.
    return (a[0]+b[0], a[1]+b[1], a[2]+b[2])

def dot(a, b):
    #Dot product.
    return a[0]*b[0] + a[1]*b[1] + a[2]*b[2]

def mul_scalar(v, s):
    #Multiply a vector by a scalar.
    return (v[0]*s, v[1]*s, v[2]*s)


# 1. Procedural Terrain (Static + Dynamic Updates)
def generate_terrain(width, height, scale):
    #Create a 2D heightmap using a simple wave function.
    terrain = [[0.0 for _ in range(width)] for _ in range(height)]
    for y in range(height):
        for x in range(width):
            nx = x / scale
            ny = y / scale
            val = math.sin(nx) * math.cos(ny) * 10.0
            terrain[y][x] = val
    return terrain

def update_terrain(terrain, particles, erosion_rate):
    #Erode the terrain based on particle positions.
    #If a particle is over a certain cell, lower the height a bit.
    h = len(terrain)
    w = len(terrain[0]) if h > 0 else 0
    for p in particles:
        px = int(p[0] % w)
        py = int(p[1] % h)
        old_val = terrain[py][px]
        new_val = old_val - erosion_rate
        if new_val < 0:
            new_val = 0
        terrain[py][px] = new_val


# 2. Particle Simulation (with Terrain Collision)
def simulate_particles(num_particles, bounds, terrain, steps):
    #Particles bounce off terrain and walls.
    #Positions also used later for lighting/fog calculation.
    particles = []
    w, h, z_bound = bounds

    # Init
    for _ in range(num_particles):
        x = random.uniform(0, w)
        y = random.uniform(0, h)
        z = random.uniform(-z_bound, z_bound)
        vx = random.uniform(-1, 1)
        vy = random.uniform(-1, 1)
        vz = random.uniform(-1, 1)
        particles.append([x, y, z, vx, vy, vz])

    terrain_h = len(terrain)
    terrain_w = len(terrain[0]) if terrain_h > 0 else 0

    for _ in range(steps):
        for p in particles:
            p[0] += p[3]
            p[1] += p[4]
            p[2] += p[5]

            # Terrain collision
            tx = int(p[0] % terrain_w)
            ty = int(p[1] % terrain_h)
            th = terrain[ty][tx]
            if p[2] <= th:
                p[2] = th
                p[5] *= -1  # Bounce in z direction

            # Reflect off bounding walls
            # X bound
            if p[0] < 0:
                p[0] = 0
                p[3] *= -1
            elif p[0] > w:
                p[0] = w
                p[3] *= -1
            # Y bound
            if p[1] < 0:
                p[1] = 0
                p[4] *= -1
            elif p[1] > h:
                p[1] = h
                p[4] *= -1
            # Z bound
            if p[2] < -z_bound:
                p[2] = -z_bound
                p[5] *= -1
            elif p[2] > z_bound:
                p[2] = z_bound
                p[5] *= -1

    return particles


# 3. Cloud Map
def generate_cloud_map(width, height, swirl_scale):
    #Generate a simple swirl-based 'cloud' overlay in [0..1].
    #Each pixel has a swirl_value that we can use to blend in final color.

    cloud_data = [[0.0 for _ in range(width)] for _ in range(height)]
    for j in range(height):
        for i in range(width):
            x = i / swirl_scale
            y = j / swirl_scale
            swirl_val = math.sin(x + math.sin(y)) + math.cos(y - 0.5*math.sin(x))
            # Normalize swirl_val from about -2..2 to [0..1]
            normalized = (swirl_val + 2.0) / 4.0
            cloud_data[j][i] = normalized
    return cloud_data


# 4. Mandelbrot Fractal (for overlay)
def mandelbrot_value(x, y, max_iter):
    #Compute a Mandelbrot iteration ratio at point (x, y).
    zr, zi = 0.0, 0.0
    cr, ci = x, y
    iter_count = 0
    while (zr*zr + zi*zi) <= 4.0 and iter_count < max_iter:
        tmp = zr*zr - zi*zi + cr
        zi = 2.0*zr*zi + ci
        zr = tmp
        iter_count += 1
    return iter_count / max_iter

def generate_fractal_overlay(width, height, max_iter):
    #Generate a 2D float array in [0..1], storing the Mandelbrot ratio at each pixel.
    fractal_data = [[0.0 for _ in range(width)] for _ in range(height)]
    for j in range(height):
        for i in range(width):
            # Map (i, j) to a region in the complex plane
            xx = (i/width)*3.0 - 2.0
            yy = (j/height)*2.0 - 1.0
            val = mandelbrot_value(xx, yy, max_iter)
            fractal_data[j][i] = val
    return fractal_data


# 5. Ray Tracer with Reflections, Fog, Particle Lighting,
#    Terrain in background, Fractal overlay, Cloud overlay
def intersect_sphere(ray_orig, ray_dir, center, radius):
    #Return the distance t along ray_dir where the ray hits the sphere, or None.
    ox, oy, oz = ray_orig
    cx, cy, cz = center
    oc = (ox-cx, oy-cy, oz-cz)
    a = dot(ray_dir, ray_dir)
    b = 2.0 * dot(oc, ray_dir)
    c = dot(oc, oc) - radius*radius
    disc = b*b - 4*a*c
    if disc < 0:
        return None
    sqrt_d = math.sqrt(disc)
    t1 = (-b - sqrt_d) / (2*a)
    t2 = (-b + sqrt_d) / (2*a)
    if t1 > 0 and (t1 < t2 or t2 < 0):
        return t1
    if t2 > 0 and (t2 < t1 or t1 < 0):
        return t2
    return None

def compute_reflection(ray_dir, normal):
    #Reflect the ray_dir around the normal: R = D - 2(D·N)N
    dot_dn = dot(ray_dir, normal)
    return vector_sub(ray_dir, mul_scalar(normal, 2.0 * dot_dn))

def compute_diffuse_light(point, normal, particles):
    #Sum diffuse contributions from each particle acting as a 'light source'.
    color_sum = (0.0, 0.0, 0.0)
    for p in particles:
        # Light color depends on particle velocity, for variety
        speed = math.sqrt(p[3]*p[3] + p[4]*p[4] + p[5]*p[5])
        # For instance: (r ~ speed, g ~ inversely ~ speed, b ~ 0.5)
        light_col = (min(1, speed/3.0), max(0, 1.0 - speed/3.0), 0.5)

        light_dir = normalize((p[0] - point[0], p[1] - point[1], p[2] - point[2]))
        diff = max(0.0, dot(light_dir, normal))
        color_sum = (
            color_sum[0] + light_col[0]*diff,
            color_sum[1] + light_col[1]*diff,
            color_sum[2] + light_col[2]*diff,
        )
    return color_sum


def compute_fog(color, dist, fog_density):
    #Exponential fog factor based on travel distance.
    fog_factor = math.exp(-fog_density*dist)
    # Fog color is a simple gray
    fog_col = (0.5, 0.5, 0.5)
    return (
        color[0]*fog_factor + fog_col[0]*(1 - fog_factor),
        color[1]*fog_factor + fog_col[1]*(1 - fog_factor),
        color[2]*fog_factor + fog_col[2]*(1 - fog_factor),
    )

def trace_ray(ray_orig, ray_dir, spheres, particles, depth=0):
    #Trace a single ray with optional reflection bounce.

    closest_t = float("inf")
    hit_sphere = None
    for center, radius, color in spheres:
        t = intersect_sphere(ray_orig, ray_dir, center, radius)
        if t is not None and t < closest_t:
            closest_t = t
            hit_sphere = (center, radius, color)

    if hit_sphere is None:
        # No intersection => background
        return (0.0, 0.0, 0.0), None

    hit_point = vector_add(ray_orig, mul_scalar(ray_dir, closest_t))
    normal = normalize(vector_sub(hit_point, hit_sphere[0]))

    # Basic diffuse from particle lights
    diffuse = compute_diffuse_light(hit_point, normal, particles)
    base_color = (
        hit_sphere[2][0] * diffuse[0],
        hit_sphere[2][1] * diffuse[1],
        hit_sphere[2][2] * diffuse[2],
    )

    # Reflection
    if depth < 1:  # Single reflection bounce
        reflect_dir = normalize(compute_reflection(ray_dir, normal))
        reflected_col, _ = trace_ray(hit_point, reflect_dir, spheres, particles, depth+1)
        # Blend reflection
        reflection_factor = 0.3
        base_color = (
            base_color[0]*(1 - reflection_factor) + reflected_col[0]*reflection_factor,
            base_color[1]*(1 - reflection_factor) + reflected_col[1]*reflection_factor,
            base_color[2]*(1 - reflection_factor) + reflected_col[2]*reflection_factor,
        )

    return base_color, closest_t


# 6. Depth of Field (Naive Implementation)
def camera_rays_with_dof(i, j, width, height, fov, aspect_ratio, focal_length, aperture, samples):

    #Generate multiple rays for naive Depth of Field effect.
    #Each sample uses a random offset in camera-aperture space.

    # Basic pinhole without DoF: direction = normalize((x, y, -1))
    # We'll gather multiple rays with random offsets in [aperture].
    angle = math.tan(math.pi * 0.5 * fov / 180.0)

    # Pixel center in camera space
    px = (2.0*((i+0.5)/float(width)) - 1.0)*angle*aspect_ratio
    py = (1.0 - 2.0*((j+0.5)/float(height)))*angle

    # We'll treat the camera as located at (0,0,0) looking down -Z.
    # The focal plane is at z = -focal_length.
    # We'll randomize the origin in a small lens circle for each sample.
    # Then we define the direction to the focal point, which is (px, py, -1)*focal_length.

    rays = []
    for _ in range(samples):
        # Random offset in lens space
        lens_x = (random.random() - 0.5) * aperture
        lens_y = (random.random() - 0.5) * aperture

        # Focal point
        focal_point = (px * focal_length, py * focal_length, -focal_length)

        # Ray origin offset
        origin = (lens_x, lens_y, 0.0)
        dir_approx = vector_sub(focal_point, origin)
        dir_norm = normalize(dir_approx)
        rays.append((origin, dir_norm))

    return rays


# 7. Final Raytrace Scene with DoF
def raytrace_scene_with_dof(
    img, terrain, particles,
    fractal_img, cloud_map,
    width, height
):
    #Perform ray tracing with Depth of Field, fractal overlay, cloud overlay, etc.
    # Hardcode some spheres
    spheres = [
        ((0, 0, -20), 4, (1, 0, 0)),     # Red
        ((5, -1, -15), 2, (0, 1, 0)),    # Green
        ((-4, 2, -18), 3, (0.2, 0.6, 1)) # Blue-ish
    ]

    fov = 60.0
    aspect_ratio = width / float(height)
    fog_density = 0.02

    # Depth of Field parameters
    focal_length = 1.5
    aperture = 0.1
    samples_per_pixel = 4  # More samples => heavier computation

    for j in range(height):
        for i in range(width):
            final_col = (0.0, 0.0, 0.0)
            # Generate multiple rays for DoF
            rays = camera_rays_with_dof(i, j, width, height, fov, aspect_ratio, focal_length, aperture, samples_per_pixel)
            for origin, direction in rays:
                color, dist = trace_ray(origin, direction, spheres, particles, depth=0)
                if dist is None:
                    dist = 9999.0

                # Fog
                color = compute_fog(color, dist, fog_density)

                # Fractal overlay
                fractal_c = fractal_img[j][i]  # 0..1
                color = (
                    color[0]*fractal_c,
                    color[1]*fractal_c,
                    color[2]*fractal_c
                )

                # Cloud overlay
                cloud_c = cloud_map[j][i]  # 0..1
                color = (
                    color[0]*(1 - cloud_c) + 1.0*cloud_c,
                    color[1]*(1 - cloud_c) + 1.0*cloud_c,
                    color[2]*(1 - cloud_c) + 1.0*cloud_c
                )

                # Accumulate
                final_col = (
                    final_col[0] + color[0],
                    final_col[1] + color[1],
                    final_col[2] + color[2],
                )
            # Average
            final_col = (
                final_col[0]/samples_per_pixel,
                final_col[1]/samples_per_pixel,
                final_col[2]/samples_per_pixel,
            )
            put_pixel(img, i, j, final_col)

def main():
    width, height = 300, 300

    # Terrain + Particle sim bounds
    terrain_scale = 50
    erosion_rate = 0.01
    num_particles = 100
    bounds = (width, height, 20)
    particle_steps = 200

    # For fractal and clouds
    fractal_iter = 100
    swirl_scale = 80.0

    random.seed(42)  # For reproducibility

    start_all = time.time()

    # 1) Generate terrain
    start_terrain = time.time()
    terrain = generate_terrain(width, height, terrain_scale)
    end_terrain = time.time()

    # 2) Simulate particles
    start_particles = time.time()
    particles = simulate_particles(num_particles, bounds, terrain, particle_steps)
    end_particles = time.time()

    # 2.5) Erode terrain after simulation
    update_terrain(terrain, particles, erosion_rate)

    # 3) Generate cloud map
    start_clouds = time.time()
    cloud_map = generate_cloud_map(width, height, swirl_scale)
    end_clouds = time.time()

    # 4) Generate fractal overlay
    start_fractal = time.time()
    fractal_img = generate_fractal_overlay(width, height, fractal_iter)
    end_fractal = time.time()

    # 5) Ray Trace scene (with Depth of Field)
    start_ray = time.time()
    img = create_image(width, height)
    raytrace_scene_with_dof(
        img, terrain, particles,
        fractal_img, cloud_map,
        width, height
    )
    end_ray = time.time()

    end_all = time.time()

    # Print times
    print("=== TIMING ===")
    print(f"Terrain generation: {end_terrain - start_terrain:.2f} s")
    print(f"Particle simulation: {end_particles - start_particles:.2f} s")
    print(f"Cloud generation: {end_clouds - start_clouds:.2f} s")
    print(f"Fractal generation: {end_fractal - start_fractal:.2f} s")
    print(f"Ray tracing (with DoF): {end_ray - start_ray:.2f} s")
    print(f"TOTAL: {end_all - start_all:.2f} s")

    # Save final output
    print("Saving final image to 'final_output.ppm'...")
    with open("final_output.ppm", "w") as f:
        f.write(f"P3\n{width} {height}\n255\n")
        for row in img:
            for (r_f, g_f, b_f) in row:
                (r, g, b) = clamp_color((r_f*255, g_f*255, b_f*255))
                f.write(f"{r} {g} {b} ")
            f.write("\n")

    print("Done!")

if __name__ == "__main__":
    main()