import functools
import numpy as np
from PIL import Image, ImageTk
import tkinter as tk

def extract_elements(condition, arr):

    a = np.asarray(arr)
    if a.ndim == 0:
        # broadcast scalar to mask shape, then select
        return np.full(condition.shape, a)[condition]
    return a[condition]


class clVector3d:
    def __init__(self, x, y, z):
        # x, y, z can be scalars or numpy arrays of the same shape
        self.x = np.array(x)
        self.y = np.array(y)
        self.z = np.array(z)

    # element-wise ops
    def __mul__(self, other):
        if isinstance(other, clVector3d):
            return clVector3d(self.x * other.x, self.y * other.y, self.z * other.z)
        else:
            return clVector3d(self.x * other, self.y * other, self.z * other)

    __rmul__ = __mul__

    def __truediv__(self, other):
        return clVector3d(self.x / other, self.y / other, self.z / other)

    def __add__(self, o):
        return clVector3d(self.x + o.x, self.y + o.y, self.z + o.z)

    def __sub__(self, o):
        return clVector3d(self.x - o.x, self.y - o.y, self.z - o.z)

    def dot(self, o):
        return (self.x * o.x) + (self.y * o.y) + (self.z * o.z)

    def mag(self):
        return np.sqrt(self.dot(self))

    def normalize(self):
        m = np.sqrt(self.dot(self))
        m_safe = np.where(m == 0, 1.0, m)
        return self * (1.0 / m_safe)

    def get_components(self):
        return (self.x, self.y, self.z)

    def get_np_array(self):
        return np.array((self.x, self.y, self.z))

    def cross(self, o):
        a = self.get_np_array()
        b = o.get_np_array()
        c = np.cross(a, b, axis=0) if a.ndim == 2 or b.ndim == 2 else np.cross(a, b)
        return clVector3d(c[0], c[1], c[2])

    def extract_elements(self, condition):

        return clVector3d(
            extract_elements(condition, self.x),
            extract_elements(condition, self.y),
            extract_elements(condition, self.z),
        )

    def replace_elements(self, condition):

        n_true = int(np.count_nonzero(condition))

        def place(comp):
            a = np.asarray(comp)

            # Output buffer (full-size)
            out = np.zeros(condition.shape, dtype=a.dtype if a.ndim > 0 else float)

            if a.ndim == 0:
                # broadcast scalar to True positions
                out[condition] = a
                return out


            if a.shape == condition.shape:
                out = np.where(condition, a, 0)
                return out


            flat = a.ravel()
            if flat.size < n_true:

                flat = np.pad(flat, (0, n_true - flat.size))
            elif flat.size > n_true:
                # trim excess
                flat = flat[:n_true]
            out[condition] = flat
            return out

        return clVector3d(place(self.x), place(self.y), place(self.z))


class clPoint3d(clVector3d):
    # Same representation as vector; arithmetic semantics are what matter
    pass


class Sphere:
    def __init__(self, center, radius, diffuse, specular_coef=0.5, texture=None):
        self.center = center
        self.radius = radius
        self.diffuse = diffuse
        self.specular_coef = specular_coef
        self.texture = texture

    def intersect(self, ray_starting_point, normalized_ray_direction):
        eye_to_center = ray_starting_point - self.center
        b = 2 * normalized_ray_direction.dot(eye_to_center)
        c = eye_to_center.dot(eye_to_center) - (self.radius * self.radius)
        discriminant = (b ** 2) - (4 * c)
        discriminant_root = np.sqrt(np.maximum(0, discriminant))
        root1 = (-b - discriminant_root) / 2.0
        root2 = (-b + discriminant_root) / 2.0
        selected_root = np.where((root1 > 0) & (root1 < root2), root1, root2)
        return np.where((discriminant > 0) & (selected_root > 0), selected_root, infinity_plus)

    def diffuse_color(self, M=None):
        if self.texture is not None:
            # Simple checker using x and z
            checker = ((M.x * 2).astype(int) % 2) == ((M.z * 2).astype(int) % 2)
            return self.diffuse * checker
        else:
            return self.diffuse

    def calculate_intensity(self, ray_starting_point, normalized_ray_direction, intersection, objects_list,
                            number_of_bounces):
        # Intersection point and normal
        M = (ray_starting_point + normalized_ray_direction * intersection)
        N = (M - self.center) * (1.0 / self.radius)

        # Lighting vectors
        toL = (light_position - M).normalize()
        toO = (eye_position - M).normalize()

        # Nudge to avoid self-hit
        nudged = M + N * 0.0001

        # Shadow test: can M see the light?
        light_distances = [obj.intersect(nudged, toL) for obj in objects_list]
        light_nearest = functools.reduce(np.minimum, light_distances)
        seelight = light_distances[objects_list.index(self)] == light_nearest

        # Ambient
        color = rgb(0.05, 0.05, 0.05)

        # Lambert (diffuse)
        lv = np.maximum(N.dot(toL), 0)
        color += self.diffuse_color(M) * lv * seelight


        if number_of_bounces < 5 and self.specular_coef > 0:
            rayD = (normalized_ray_direction - N * 2 * normalized_ray_direction.dot(N)).normalize()
            color += raytrace(nudged, rayD, objects_list, number_of_bounces + 1) * self.specular_coef

        # Blinn-Phong (specular highlight)
        phong = N.dot((toL + toO).normalize())
        color += rgb(1, 1, 1) * np.power(np.clip(phong, 0, 1), 50) * seelight
        return color


class clCamera:
    def __init__(self, eye, look_at, up, field_of_view, resolution_width, resolution_height):
        self.eye = eye
        self.look_at = look_at
        self.up = up
        self.fov = field_of_view
        self.resolution_width = resolution_width
        self.resolution_height = resolution_height

        self.N = (self.look_at - self.eye).normalize()
        self.U = self.up.cross(self.N).normalize()
        self.V = self.N.cross(self.U).normalize()

        self.aspect_ratio = self.resolution_width / self.resolution_height
        self.fov_rad = np.radians(field_of_view)
        self.half_width = np.tan(self.fov_rad / 2)
        self.half_height = self.half_width / self.aspect_ratio

        # grid coordinates in camera plane
        x = np.tile(np.linspace(-self.half_width, self.half_width, self.resolution_width),
                    self.resolution_height)
        y = np.repeat(np.linspace(self.half_height, -self.half_height, self.resolution_height),
                      self.resolution_width)

        image_center = self.eye + self.N  # 1 unit away
        self.pixels_world_coordinates = (
            image_center.get_np_array().reshape(-1, 1) +
            x * self.U.get_np_array().reshape(-1, 1) +
            y * self.V.get_np_array().reshape(-1, 1)
        )

def rgb(r, g=None, b=None):
    """Convenience: rgb(1,1,1) or rgb(clVector3d)."""
    if isinstance(r, clVector3d) and g is None and b is None:
        return r
    return clVector3d(np.array(r), np.array(g), np.array(b))

def raytrace(ray_starting_point, normalized_ray_directions, objects_list, number_of_bounces=1):
    # Intersections to all objects
    all_objects_intersections = [obj.intersect(ray_starting_point, normalized_ray_directions)
                                 for obj in objects_list]
    # nearest t per ray
    nearest = functools.reduce(np.minimum, all_objects_intersections)

    color = rgb(0, 0, 0)
    for (obj, single_object_intersections) in zip(objects_list, all_objects_intersections):
        hit = (nearest != infinity_plus) & (single_object_intersections == nearest)
        if np.any(hit):
            dc = extract_elements(hit, single_object_intersections)
            Oc = ray_starting_point.extract_elements(hit)
            Dc = normalized_ray_directions.extract_elements(hit)
            cc = obj.calculate_intensity(Oc, Dc, dc, objects_list, number_of_bounces)
            color += cc.replace_elements(hit)
    return color


class TkIncrementalRenderer:
    def __init__(self, width, height, title="Ray tracing… (incremental)"):
        self.root = tk.Tk()
        self.root.title(title)
        self.width = width
        self.height = height

        # UI
        self.canvas = tk.Canvas(self.root, width=width, height=height, highlightthickness=0)
        self.canvas.pack()
        self.status = tk.Label(self.root, text="Starting…", anchor="w")
        self.status.pack(fill="x")

        # image buffer & Tk image
        self.img_array = np.zeros((height, width, 3), dtype=np.uint8)
        self.pil_img = Image.fromarray(self.img_array, "RGB")
        self.tk_img = ImageTk.PhotoImage(self.pil_img)
        self.image_id = self.canvas.create_image(0, 0, image=self.tk_img, anchor="nw")

        # keep a ref so PhotoImage isn't garbage collected
        self.canvas.image = self.tk_img

    def update_rows(self, y0, y1):
        # recreate the PhotoImage (simple and reliable)
        self.pil_img = Image.fromarray(self.img_array, "RGB")
        self.tk_img = ImageTk.PhotoImage(self.pil_img)
        self.canvas.itemconfig(self.image_id, image=self.tk_img)
        self.canvas.image = self.tk_img

    def set_status(self, text):
        self.status.config(text=text)

    def step(self, pause_ms=0):
        # let Tk process pending events
        self.root.update_idletasks()
        self.root.update()
        if pause_ms and pause_ms > 0:
            self.root.after(pause_ms)

    def finish(self, save_path=None):
        self.set_status("Done.")
        self.step(0)
        if save_path:
            Image.fromarray(self.img_array, "RGB").save(save_path)


def render_incremental_tk(camera, eye_position, objects_list,
                          chunk_rows=32, pause_ms=10, save_final="raytrace.png", window_title=None):

    H, W = camera.resolution_height, camera.resolution_width
    viewer = TkIncrementalRenderer(W, H, title=(window_title or "Ray tracing… (incremental)"))

    # Prepare world coords (flattened) once
    P = camera.pixels_world_coordinates
    pixels_world = clPoint3d(P[0, :], P[1, :], P[2, :])

    # Loop over strips
    for y0 in range(0, H, chunk_rows):
        y1 = min(y0 + chunk_rows, H)
        idx = np.arange(y0 * W, y1 * W)

        # Directions for current strip
        strip_points = clPoint3d(pixels_world.x[idx], pixels_world.y[idx], pixels_world.z[idx])
        dirs = (strip_points - eye_position).normalize()

        # Trace current strip
        color = raytrace(eye_position, dirs, objects_list)
        R, G, B = color.get_components()

        # --- Final safety: pad/trim and reshape for the strip ---
        expected = (y1 - y0) * W

        def prep(chan):
            a = np.asarray(chan).ravel()
            if a.size < expected:
                a = np.pad(a, (0, expected - a.size))
            elif a.size > expected:
                a = a[:expected]
            a = (255 * np.clip(a, 0, 1)).astype(np.uint8)
            return a.reshape(y1 - y0, W)

        R = prep(R); G = prep(G); B = prep(B)

        viewer.img_array[y0:y1, :, 0] = R
        viewer.img_array[y0:y1, :, 1] = G
        viewer.img_array[y0:y1, :, 2] = B

        # Refresh window
        percent = int(100 * y1 / H)
        viewer.set_status(f"Rendering… {percent}%")
        viewer.update_rows(y0, y1)
        viewer.step(pause_ms)

    # Done
    viewer.finish(save_path=save_final)


if __name__ == "__main__":
    image_resolution_width, image_resolution_height = 1000, 800

    # Camera
    eye_position = clPoint3d(1.0, 1.0, -1.0)
    look_at_point = clPoint3d(0.0, 0.0, 0.0)
    up = clVector3d(0.0, 1.0, 0.0)
    field_of_view = 90  # degrees

    # Light & sentinel (globals used in Sphere.calculate_intensity)
    light_position = clPoint3d(5.0, 5.0, 5.0)
    infinity_plus = 1.0e39

    # Scene objects
    objects_list = [
        Sphere(clPoint3d(0.0, 0.0, 0.0), 0.1, rgb(1, 0, 0)),
        Sphere(clPoint3d(0.0, 0.10, 0.0), 0.1, rgb(0, 1, 0)),
        Sphere(clPoint3d(0.0, 0.0, 1.0), 0.1, rgb(1, 1, 0)),
        Sphere(clPoint3d(-0.4, 0.2, 1.0), 0.1, rgb(0.5, 0.223, 0.5)),
        Sphere(clPoint3d(0.0, 0.3, 1.0), 0.2, rgb(0.5, 0.223, 0.5)),
        Sphere(clPoint3d(0.4, 0.4, 1.0), 0.3, rgb(0.5, 0.223, 0.5)),
        Sphere(clPoint3d(0, -99999.5, 0), 99999, rgb(0.75, 0.75, 0.0), 0.7, texture=1),
    ]

    # Camera instance
    camera = clCamera(eye_position, look_at_point, up, field_of_view,
                      image_resolution_width, image_resolution_height)

    # Incremental render — tweak to taste
    render_incremental_tk(
        camera,
        eye_position,
        objects_list,
        chunk_rows=8,     # smaller = smoother UI; larger = fewer updates
        pause_ms=1,       # 0 for fastest; >0 so you can watch it paint
        save_final="raytrace.png",
        window_title="Ray Tracer (Tkinter, incremental)",
    )