import numpy as np
from PIL import Image
import math
import time

EPSILON = 1e-6
PI = math.pi


def normalize(v):
    n = np.linalg.norm(v)
    return v / n if n > 1e-12 else v


def reflect(d, n):
    return d - 2.0 * np.dot(d, n) * n


def refract(incident, normal, eta):
    cos_i = min(-np.dot(incident, normal), 1.0)
    sin_t_sq = eta * eta * (1.0 - cos_i * cos_i)
    if sin_t_sq > 1.0:
        return None
    cos_t = math.sqrt(max(0.0, 1.0 - sin_t_sq))
    return eta * incident + (eta * cos_i - cos_t) * normal


def schlick(cos, ref_idx):
    r0 = (1.0 - ref_idx) / (1.0 + ref_idx)
    r0 *= r0
    return r0 + (1.0 - r0) * math.pow(1.0 - cos, 5.0)


class Ray:
    __slots__ = ('origin', 'direction')

    def __init__(self, origin, direction):
        self.origin = np.asarray(origin, dtype=np.float64)
        self.direction = normalize(np.asarray(direction, dtype=np.float64))


class HitRecord:
    __slots__ = ('t', 'point', 'normal', 'front_face', 'material')

    def __init__(self):
        self.t = float('inf')
        self.point = np.zeros(3)
        self.normal = np.zeros(3)
        self.front_face = True
        self.material = None


class Hittable:
    def hit(self, r, t_min=EPSILON, t_max=float('inf')):
        raise NotImplementedError


class Sphere(Hittable):
    __slots__ = ('center', 'radius', 'material')

    def __init__(self, center, radius, material):
        self.center = np.asarray(center, dtype=np.float64)
        self.radius = float(radius)
        self.material = material

    def hit(self, r, t_min=EPSILON, t_max=float('inf')):
        oc = r.origin - self.center
        a = np.dot(r.direction, r.direction)
        b = np.dot(oc, r.direction)
        c = np.dot(oc, oc) - self.radius * self.radius
        disc = b * b - a * c
        if disc < 0:
            return None
        sq = math.sqrt(disc)
        root = (-b - sq) / a
        if root < t_min or root > t_max:
            root = (-b + sq) / a
            if root < t_min or root > t_max:
                return None
        rec = HitRecord()
        rec.t = root
        rec.point = r.origin + root * r.direction
        out = normalize(rec.point - self.center)
        if np.dot(r.direction, out) < 0:
            rec.normal = out
            rec.front_face = True
        else:
            rec.normal = -out
            rec.front_face = False
        rec.material = self.material
        return rec


class Plane(Hittable):
    __slots__ = ('point', 'normal', 'material')

    def __init__(self, point, normal, material):
        self.point = np.asarray(point, dtype=np.float64)
        self.normal = normalize(np.asarray(normal, dtype=np.float64))
        self.material = material

    def hit(self, r, t_min=EPSILON, t_max=float('inf')):
        denom = np.dot(r.direction, self.normal)
        if abs(denom) < 1e-12:
            return None
        t = np.dot(self.point - r.origin, self.normal) / denom
        if t < t_min or t > t_max:
            return None
        rec = HitRecord()
        rec.t = t
        rec.point = r.origin + t * r.direction
        if np.dot(r.direction, self.normal) < 0:
            rec.normal = self.normal
            rec.front_face = True
        else:
            rec.normal = -self.normal
            rec.front_face = False
        rec.material = self.material
        return rec


class Box(Hittable):
    __slots__ = ('min_pt', 'max_pt', 'material')

    def __init__(self, min_pt, max_pt, material):
        self.min_pt = np.asarray(min_pt, dtype=np.float64)
        self.max_pt = np.asarray(max_pt, dtype=np.float64)
        self.material = material

    def hit(self, r, t_min=EPSILON, t_max=float('inf')):
        t_lo = t_min
        t_hi = t_max
        hit_axis = -1
        inv_ds = [0.0, 0.0, 0.0]
        for axis in range(3):
            inv_d = 1.0 / r.direction[axis] if abs(r.direction[axis]) > 1e-12 else 1e12
            inv_ds[axis] = inv_d
            t0 = (self.min_pt[axis] - r.origin[axis]) * inv_d
            t1 = (self.max_pt[axis] - r.origin[axis]) * inv_d
            if inv_d < 0:
                t0, t1 = t1, t0
            if t0 > t_lo:
                t_lo = t0
                hit_axis = axis
            t_hi = min(t_hi, t1)
            if t_hi < t_lo:
                return None
        if t_lo >= t_max or t_lo <= t_min:
            return None
        rec = HitRecord()
        n = np.zeros(3)
        n[hit_axis] = 1.0 if inv_ds[hit_axis] > 0 else -1.0
        if np.dot(r.direction, n) < 0:
            rec.normal = n
            rec.front_face = True
        else:
            rec.normal = -n
            rec.front_face = False
        rec.t = t_lo
        rec.point = r.origin + t_lo * r.direction
        rec.material = self.material
        return rec


class Scene:
    __slots__ = ('objects', 'lights')

    def __init__(self):
        self.objects = []
        self.lights = []

    def add(self, obj):
        self.objects.append(obj)

    def add_light(self, obj):
        self.lights.append(obj)
        self.objects.append(obj)

    def hit(self, r, t_min=EPSILON, t_max=float('inf')):
        closest = None
        min_t = t_max
        for obj in self.objects:
            rec = obj.hit(r, t_min, min_t)
            if rec is not None and rec.t < min_t:
                closest = rec
                min_t = rec.t
        return closest

    def blocked(self, origin, direction, t_max):
        r = Ray(origin, direction)
        for obj in self.objects:
            if obj.hit(r, EPSILON, t_max) is not None:
                return True
        return False


class Material:
    def scatter(self, ray_in, rec, rng):
        raise NotImplementedError

    def get_emission(self):
        return np.zeros(3)


class Diffuse(Material):
    __slots__ = ('albedo',)

    def __init__(self, albedo):
        self.albedo = np.asarray(albedo, dtype=np.float64)

    def scatter(self, ray_in, rec, rng):
        target = rec.normal + np.random.randn(3) * 0.5
        if np.linalg.norm(target) < 1e-8:
            target = rec.normal.copy()
        return True, self.albedo, Ray(rec.point, normalize(target))


class Mirror(Material):
    __slots__ = ('albedo',)

    def __init__(self, albedo):
        self.albedo = np.asarray(albedo, dtype=np.float64)

    def scatter(self, ray_in, rec, rng):
        reflected = reflect(ray_in.direction, rec.normal)
        if np.dot(reflected, rec.normal) <= 0:
            return False, np.zeros(3), None
        return True, self.albedo, Ray(rec.point, normalize(reflected))


class Dielectric(Material):
    __slots__ = ('ref_idx', 'albedo')

    def __init__(self, ref_idx, albedo=None):
        self.ref_idx = ref_idx
        self.albedo = np.asarray(albedo, dtype=np.float64) if albedo is not None else np.ones(3)

    def scatter(self, ray_in, rec, rng):
        refraction_ratio = 1.0 / self.ref_idx if rec.front_face else self.ref_idx
        unit_direction = ray_in.direction.copy()
        cos_theta = min(-np.dot(unit_direction, rec.normal), 1.0)
        sin_theta = math.sqrt(max(0.0, 1.0 - cos_theta * cos_theta))
        cannot_refract = refraction_ratio * sin_theta > 1.0
        if cannot_refract or schlick(cos_theta, refraction_ratio) > np.random.random():
            direction = reflect(unit_direction, rec.normal)
        else:
            direction = refract(unit_direction, rec.normal, refraction_ratio)
            if direction is None:
                direction = reflect(unit_direction, rec.normal)
        return True, self.albedo, Ray(rec.point, normalize(direction))


class LightSource(Material):
    __slots__ = ('emission',)

    def __init__(self, emission):
        self.emission = np.asarray(emission, dtype=np.float64)

    def get_emission(self):
        return self.emission

    def scatter(self, ray_in, rec, rng):
        return False, np.zeros(3), None


class Camera:
    def __init__(self, lookfrom, lookat, vup, vfov, aspect, aperture=0.0, focus_dist=10.0):
        self.origin = np.asarray(lookfrom, dtype=np.float64)
        half_height = math.tan(math.radians(vfov) / 2.0)
        half_width = aspect * half_height
        w = normalize(self.origin - np.asarray(lookat, dtype=np.float64))
        u = normalize(np.cross(vup, w))
        v = np.cross(w, u)
        self.horizontal = focus_dist * half_width * u
        self.vertical = focus_dist * half_height * v
        self.lower_left_corner = self.origin - self.horizontal / 2.0 - self.vertical / 2.0 - focus_dist * w
        self.aperture = aperture

    def get_ray(self, s, t, rng):
        rd = np.zeros(3)
        if self.aperture > 0:
            rd = (np.random.rand(3) - 0.5) * self.aperture
            rd[2] = 0.0
        direction = self.lower_left_corner + s * self.horizontal + t * self.vertical - self.origin + rd
        return Ray(self.origin, direction)


def cornell_box():
    scene = Scene()

    white = Diffuse(np.array([0.73, 0.73, 0.73]))
    gray = Diffuse(np.array([0.5, 0.5, 0.5]))
    red = Diffuse(np.array([0.65, 0.05, 0.05]))
    green = Diffuse(np.array([0.12, 0.45, 0.15]))
    mirror = Mirror(np.array([0.95, 0.95, 0.95]))
    glass = Dielectric(ref_idx=1.5)
    light = LightSource(np.array([50.0, 50.0, 50.0]))

    scene.add(Plane(np.array([0, 0, 0]), np.array([0, 1, 0]), gray))
    scene.add(Plane(np.array([0, 500, 0]), np.array([0, -1, 0]), white))
    scene.add(Plane(np.array([0, 0, 500]), np.array([0, 0, -1]), white))
    scene.add(Plane(np.array([0, 0, 0]), np.array([1, 0, 0]), red))
    scene.add(Plane(np.array([500, 0, 0]), np.array([-1, 0, 0]), green))

    scene.add(Box(np.array([0, 180, 180]), np.array([25, 220, 320]), light))

    scene.add(Sphere(np.array([250, 180, 250]), 80, glass))
    scene.add(Sphere(np.array([150, 180, 380]), 80, mirror))
    scene.add(Sphere(np.array([350, 160, 300]), 60, red))

    return scene


def trace(scene, ray, depth, max_depth, rng):
    if depth >= max_depth:
        return np.zeros(3)

    rec = scene.hit(ray)
    if rec is None:
        t_val = 0.5 * (ray.direction[1] + 1.0)
        return np.array([1.0, 1.0, 1.0]) * (1.0 - t_val) + np.array([0.5, 0.7, 1.0]) * t_val

    emitted = rec.material.get_emission()
    if isinstance(rec.material, LightSource):
        return emitted

    mat = rec.material
    scattered, attenuation, scattered_ray = mat.scatter(ray, rec, rng)
    if not scattered or scattered_ray is None:
        return emitted

    indirect = trace(scene, scattered_ray, depth + 1, max_depth, rng)
    color = emitted + attenuation * indirect

    if isinstance(mat, Diffuse):
        direct = np.zeros(3)
        for light_obj in scene.lights:
            emission = light_obj.material.get_emission()
            for _ in range(2):
                face_idx = int(rng.random() * 6)
                x = rng.uniform(light_obj.min_pt[0], light_obj.max_pt[0])
                y = rng.uniform(light_obj.min_pt[1], light_obj.max_pt[1])
                z = rng.uniform(light_obj.min_pt[2], light_obj.max_pt[2])
                if face_idx == 0:
                    lp = np.array([light_obj.min_pt[0], y, z])
                elif face_idx == 1:
                    lp = np.array([light_obj.max_pt[0], y, z])
                elif face_idx == 2:
                    lp = np.array([x, light_obj.min_pt[1], z])
                elif face_idx == 3:
                    lp = np.array([x, light_obj.max_pt[1], z])
                elif face_idx == 4:
                    lp = np.array([x, y, light_obj.min_pt[2]])
                else:
                    lp = np.array([x, y, light_obj.max_pt[2]])

                to_light = lp - rec.point
                dist_sq = np.dot(to_light, to_light)
                if dist_sq < 1e-8:
                    continue
                dist = math.sqrt(dist_sq)
                d_to_light = to_light / dist
                cos_at_rec = np.dot(rec.normal, d_to_light)
                if cos_at_rec <= 0:
                    continue

                if scene.blocked(rec.point + EPSILON * rec.normal, d_to_light, dist - EPSILON):
                    continue

                light_area = 6 * (light_obj.max_pt[0] - light_obj.min_pt[0]) * \
                             (light_obj.max_pt[1] - light_obj.min_pt[1]) * \
                             (light_obj.max_pt[2] - light_obj.min_pt[2])
                pdf = light_area / (dist_sq * cos_at_rec)
                if pdf < 1e-12:
                    continue
                brdf = mat.albedo / PI
                direct += emission * cos_at_rec * brdf / pdf
                break

        if np.dot(direct, direct) > 0:
            color = emitted + attenuation * (indirect + direct)

    return color


def render_chunk(args):
    scene, width, height, spp, max_d, \
        lookfrom, lookat, vfov, chunk_start, chunk_end, seed_off = args

    cam = Camera(lookfrom, lookat, np.array([0, 1, 0]), vfov, width / height, 0.0, 1078.0)

    image = np.zeros((chunk_end - chunk_start, width, 3), dtype=np.float64)

    for j_local, j in enumerate(range(chunk_start, chunk_end)):
        for i in range(width):
            color = np.zeros(3)
            rng = np.random.default_rng(seed_off + i * 1000 + j * 7)
            for _ in range(spp):
                u = (i + rng.random()) / (width - 1)
                v = (height - 1 - j + rng.random()) / (height - 1)
                ray = cam.get_ray(u, v, rng)
                color += trace(scene, ray, 0, max_d, rng)
            color /= spp
            color = np.sqrt(np.clip(color, 0, None))
            image[j_local, i] = np.clip(color, 0.0, 1.0)

    return (chunk_start, image)


def render(scene, width=800, height=600, samples_per_pixel=50, max_depth=15,
           lookfrom=None, lookat=None, vfov=40, aperture=0.0, focus_dist=1078.0,
           output="cornell_box.png", workers=16):
    if lookfrom is None:
        lookfrom = np.array([278.0, 278.0, -800.0])
    if lookat is None:
        lookat = np.array([278.0, 278.0, 278.0])

    image = np.zeros((height, width, 3), dtype=np.float64)
    start = time.time()

    chunk = max(1, height // workers)
    ranges = []
    for w in range(workers):
        s = w * chunk
        e = height if w == workers - 1 else (w + 1) * chunk
        ranges.append((scene, width, height, samples_per_pixel,
                       max_depth, lookfrom, lookat, vfov, s, e, w * 100000 + 42))

    from concurrent.futures import ProcessPoolExecutor, as_completed

    from rich.console import Console
    from rich.progress import Progress, SpinnerColumn, BarColumn, TextColumn, TimeElapsedColumn, TimeRemainingColumn
    from rich.table import Table
    from rich.panel import Panel
    from matplotlib import pyplot as plt

    console = Console()

    console.print()
    console.print(Panel.fit(
        "[bold magenta]Cornell Box Raytracer[/bold magenta]\n"
        "[dim]Reflections | Caustics | Multiple Bounces | NEE[/dim]",
        border_style="bright_blue", padding=(0, 2)
    ))

    settings = Table(title="Scene Configuration", show_header=False)
    settings.add_column("Parameter", style="cyan", no_wrap=True)
    settings.add_column("Value", style="white")
    settings.add_row("Resolution", f"[bold]{width}x{height}[/bold]")
    settings.add_row("Samples/pixel", f"[bold yellow]{samples_per_pixel}[/bold yellow]")
    settings.add_row("Max bounces", f"[bold green]{max_depth}[/bold green]")
    settings.add_row("Workers", f"[bold magenta]{workers}[/bold magenta]")
    settings.add_row("Output", f"[bold]{output}[/bold]")
    console.print(settings)
    console.print()

    progress = Progress(
        SpinnerColumn(finished_text="[DONE]", style="green"),
        TextColumn("[progress.description]{task.description}"),
        BarColumn(bar_width=50, style="bright_blue", complete_style="bright_green"),
        TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
        TextColumn("[bold]{task.completed}/{task.total}[/bold] rows"),
        TimeElapsedColumn(),
        TimeRemainingColumn(),
    )

    brightness_history = []
    timestamps = []
    completed_rows = 0
    total_rows = height

    with progress:
        render_task = progress.add_task("Rendering", total=total_rows)
        preview_path = "preview.png"

        with ProcessPoolExecutor(max_workers=workers) as executor:
            futures = {executor.submit(render_chunk, r): r[9] for r in ranges}
            for f in as_completed(futures):
                chunk_start, chunk_img = f.result()
                chunk_end = chunk_start + chunk_img.shape[0]
                image[chunk_start:chunk_end] = chunk_img
                completed_rows += (chunk_end - chunk_start)

                progress.update(render_task, completed=completed_rows)
                progress.refresh()

                img = (np.clip(image, 0.0, 1.0) * 255).astype(np.uint8)
                Image.fromarray(img, 'RGB').save(preview_path)

                elapsed = time.time() - start
                brightness = float(image.mean())
                brightness_history.append(brightness)
                timestamps.append(elapsed)

    elapsed = time.time() - start
    img = (np.clip(image, 0.0, 1.0) * 255).astype(np.uint8)
    Image.fromarray(img, 'RGB').save(output)

    console.print()
    console.print(Panel.fit(
        f"[bold green]Rendered {output}[/bold green]\n"
        f"[dim]Total time: {elapsed:.1f}s | "
        f"Speed: {width * height * samples_per_pixel / elapsed / 1e6:.2f} MPix*samp/s[/dim]",
        border_style="bright_green", padding=(0, 2)
    ))

    region_table = Table(title="Region Brightness Analysis")
    region_table.add_column("Region", style="cyan", no_wrap=True)
    region_table.add_column("Brightness", justify="right", style="white")
    region_table.add_column("Visual", justify="center")
    h_start, h_end = height // 3, 2 * height // 3
    regions = [
        ("Left Wall", image[:, :width // 4]),
        ("Right Wall", image[:, 3 * width // 4:]),
        ("Floor", image[2 * height // 3:, h_start:h_end]),
        ("Ceiling", image[:height // 3, h_start:h_end]),
        ("Center", image[h_start:h_end, h_start:h_end]),
    ]
    for name, region in regions:
        brightness = float(region.mean()) * 255
        bar_len = 20
        filled = int(brightness / 255 * bar_len)
        bar = "#" * filled + "-" * (bar_len - filled)
        color = "green" if brightness > 80 else ("yellow" if brightness > 40 else "red")
        region_table.add_row(name, f"{brightness:.1f}", f"[{color}]{bar}[/{color}]")
    console.print(region_table)
    console.print()

    fig, axes = plt.subplots(1, 2, figsize=(14, 5))
    ax1, ax2 = axes

    t_arr = np.array(timestamps)
    b_arr = np.array(brightness_history)
    ax1.plot(t_arr, b_arr, color='#22cc66', linewidth=2)
    ax1.set_xlabel("Time (s)", fontsize=11)
    ax1.set_ylabel("Mean Brightness", fontsize=11)
    ax1.set_title("Brightness Convergence", fontsize=13, fontweight='bold')
    ax1.grid(True, alpha=0.3)
    ax1.fill_between(t_arr, b_arr, alpha=0.15, color='#22cc66')

    preview = (np.clip(image, 0.0, 1.0) * 255).astype(np.uint8)
    ax2.imshow(preview)
    ax2.set_title("Render Preview", fontsize=13, fontweight='bold')
    ax2.axis('off')

    plt.tight_layout()
    metrics_path = "render_metrics.png"
    plt.savefig(metrics_path, dpi=150, bbox_inches='tight')
    plt.close()
    console.print(f"[dim]Metrics graph saved to [bold]{metrics_path}[/bold][/dim]")

    return image


if __name__ == "__main__":
    import sys

    print("=== Cornell Box Raytracer ===")
    print("Features: reflections, refraction/caustics, multiple bounces, NEE")
    print()

    scene = cornell_box()

    width = int(sys.argv[1]) if len(sys.argv) > 1 else 400
    height = int(sys.argv[2]) if len(sys.argv) > 2 else 300
    spp = int(sys.argv[3]) if len(sys.argv) > 3 else 50
    max_depth = int(sys.argv[4]) if len(sys.argv) > 4 else 15
    output = sys.argv[5] if len(sys.argv) > 5 else "cornell_box.png"

    print(f"Resolution: {width}x{height}")
    print(f"Samples per pixel: {spp}")
    print(f"Max bounces: {max_depth}")
    print(f"Output: {output}")
    print()

    render(scene, width=width, height=height,
           samples_per_pixel=spp, max_depth=max_depth, output=output)