main.py

"""
main.py
-------
Solar system scene .

Features:
    - Procedural geometry (UV sphere)
    - background rendering
    - Perlin noise procedural texturing (Sun)
    - Image-based texturing + Lambert + Fresnel (Earth)
    - Phong shading (Moon)

"""
import math
from PIL import Image
from pyglm import glm

from rit_window import RitWindow
from transforms import T3
from cgI_engine import (
    CGIengine,
    Texture2D,
    VertexArrayObject,
    VertexAttributeBuffer,
)

# ============================================================
# Geometry: UV Sphere Mesh
# ============================================================
def make_uv_sphere(radius=1.0, segments=48, rings=24):
    """
    Generate a UV-mapped sphere mesh.

    Parameters
    ----------
    radius : float
        Radius of the sphere in object space.
    segments : int
        Number of longitudinal subdivisions (theta).
    rings : int
        Number of latitudinal subdivisions (phi).

    Returns
    -------
    positions : list[glm.vec3]
        Object-space vertex positions.
    normals : list[glm.vec3]
        Unit normals (same as normalized positions).
    uvs : list[glm.vec2]
        UV coordinates in [0,1], with V flipped for image textures.
    indices : list[int]
        Triangle index buffer (two triangles per quad).
    """
    positions, normals, uvs, indices = [], [], [], []

    # Loop over latitude (phi)
    for r in range(rings + 1):
        v = r / rings
        phi = v * math.pi

        # Loop over longitude (theta)
        for s in range(segments + 1):
            u = s / segments
            theta = u * 2.0 * math.pi

            # Spherical to Cartesian conversion
            x = math.sin(phi) * math.cos(theta)
            y = math.cos(phi)
            z = math.sin(phi) * math.sin(theta)

            n = glm.vec3(x, y, z)
            positions.append(radius * n)
            normals.append(glm.normalize(n))
            uvs.append(glm.vec2(u, 1.0 - v))

    # Build triangle indices
    stride = segments + 1
    for r in range(rings):
        for s in range(segments):
            i0 = r * stride + s
            i1 = i0 + 1
            i2 = i0 + stride
            i3 = i2 + 1
            indices += [i0, i2, i1, i1, i2, i3]

    return positions, normals, uvs, indices


# ============================================================
# Geometry: Fullscreen Background Quad)
# ============================================================
def make_background_quad():
    """
    Create a fullscreen quad directly in clip space.

    This quad bypasses model/view/projection transforms and
    fills the screen, making it ideal for skyboxes or backgrounds.

    Returns
    -------
    positions : list[glm.vec3]
        Clip-space positions (z=1 places it at far depth).
    uvs : list[glm.vec2]
        Texture coordinates.
    indices : list[int]
        Two triangles forming a quad.
    """

    positions = [
        glm.vec3(-1, -1, 1),
        glm.vec3( 1, -1, 1),
        glm.vec3( 1,  1, 1),
        glm.vec3(-1,  1, 1),
    ]
    uvs = [
        glm.vec2(0, 0),
        glm.vec2(1, 0),
        glm.vec2(1, 1),
        glm.vec2(0, 1),
    ]
    indices = [0, 1, 2, 0, 2, 3]
    return positions, uvs, indices


# ============================================================
# VAO helper
# ============================================================
def build_vao(positions, normals=None, uvs=None):
    """
    Build a VertexArrayObject from provided attribute lists.

    This mirrors OpenGL-style VAO usage, grouping vertex attributes
    under semantic names.

    Parameters
    ----------
    positions : list[glm.vec3]
    normals : list[glm.vec3] or None
    uvs : list[glm.vec2] or None
    """

    vao = VertexArrayObject()
    vao.add_attribute(VertexAttributeBuffer("position", positions))
    if normals:
        vao.add_attribute(VertexAttributeBuffer("normal", normals))
    if uvs:
        vao.add_attribute(VertexAttributeBuffer("uv", uvs))
    return vao


# ============================================================
# Background shaders
# ============================================================
def vs_background(v, uniforms):
    """
    Vertex shader for background quad.

    Assumes input positions are already in clip space.
    Simply forwards UVs and sets w=1 for proper rasterization.
    """

    v.attach_varying("uv", v.get_varying("uv"))
    v.w = 1.0


def fs_background(p0, p1, p2, a, b, c, uniforms, varyings):
    """
    Fragment shader for background.

    Performs a simple bilinear texture lookup.
    """
    return uniforms["sampler2D"].sample_bilinear(varyings["uv"])


# ============================================================
# Perlin Noise Utilities
# ============================================================
def fade(t):
    # Quintic fade curve
    # Ensures smooth first and second derivatives.
    return t * t * t * (t * (t * 6 - 15) + 10)


def lerp(a, b, t):
    return a + t * (b - a)


def grad(h, x, y):
    """
    Compute dot product between a pseudo-random gradient
    direction (derived from hash h) and offset vector (x,y).
    """
    h = h & 3  # 4 gradient directions
    u = x if h < 2 else y
    v = y if h < 2 else x
    return (u if (h & 1) == 0 else -u) + (v if (h & 2) == 0 else -v)


def hash2(ix, iy):
    """
    Deterministic integer hash for grid corner (ix, iy).
    Returns an int in [0,255].
    """
    return int(math.fmod(math.sin(ix * 127.1 + iy * 311.7) * 43758.5453, 256))


def perlin(uv, base_freq=8.0):
    """
    Classic 2D Perlin gradient noise in [0,1].

    Steps:
        1. Map UV to lattice space
        2. Compute gradients at cell corners
        3. Interpolate using fade curve
    """
    # Scale UV into "noise space" to control feature size
    x, y = uv.x * base_freq, uv.y * base_freq

    # Integer lattice coords
    x0 = int(math.floor(x))
    y0 = int(math.floor(y))

    # Fractional position within lattice cell
    xf = x - x0
    yf = y - y0

    # Mask to keep hash stable
    xi = x0 & 255
    yi = y0 & 255

    # Smooth interpolation weights
    u = fade(xf)
    v = fade(yf)

    # Hash each corner
    aa = hash2(xi,     yi)
    ba = hash2(xi + 1, yi)
    ab = hash2(xi,     yi + 1)
    bb = hash2(xi + 1, yi + 1)

    # Gradient dot products at corners
    g_aa = grad(aa, xf,     yf)
    g_ba = grad(ba, xf - 1, yf)
    g_ab = grad(ab, xf,     yf - 1)
    g_bb = grad(bb, xf - 1, yf - 1)

    # Bilinear interpolation (with fade)
    x1 = lerp(g_aa, g_ba, u)
    x2 = lerp(g_ab, g_bb, u)
    n = lerp(x1, x2, v)

    # Map from approx [-1,1] to [0,1]
    return (n + 1.0) * 0.5


def perlin_fbm(uv, octaves=4, base_freq=6.0):
    """
    Fractal Brownian Motion using Perlin noise.

    Layers multiple frequencies of noise to create
    natural-looking turbulence.
    """
    value = 0.0
    amp = 1.0
    freq = 1.0
    amp_sum = 0.0

    for _ in range(octaves):
        value += amp * perlin(uv * freq, base_freq=base_freq)
        amp_sum += amp
        freq *= 2.0
        amp *= 0.5

    # Normalize to ~[0,1]
    return value / max(1e-8, amp_sum)


# ============================================================
# Vertex Shader for Spheres
# ============================================================
def vs_standard(v, uniforms):
    """
    Standard vertex shader:
        • Object → World → View → Clip
        • Correct normal transformation
        • Passes UVs through unchanged
    """
    pos_obj = glm.vec4(v.x, v.y, v.z, 1.0)
    normal_obj = v.get_varying("normal")
    uv = v.get_varying("uv")

    model = uniforms["modelT"]
    view = uniforms["viewT"]
    proj = uniforms["projectionT"]

    # Transform position through pipeline
    pos_view = view * (model * pos_obj)
    pos_clip = proj * pos_view
    v.x, v.y, v.z, v.w = pos_clip.x, pos_clip.y, pos_clip.z, pos_clip.w

    # Correct normal transform (inverse-transpose of model), then to view space
    n_world = glm.vec3(glm.transpose(glm.inverse(model)) * glm.vec4(normal_obj, 0.0))
    n_view = glm.vec3(view * glm.vec4(n_world, 0.0))

    v.attach_varying("pos_view", glm.vec3(pos_view))
    v.attach_varying("normal_view", glm.normalize(n_view))
    v.attach_varying("uv", uv)


# ============================================================
# Procedural Sun Fragment Shader
# ============================================================
def fs_sun_perlin(p0, p1, p2, a, b, c, uniforms, varyings):
    """
    Fragment shader for the Sun using procedural Perlin noise.

    Features:
        • Perlin fBm surface turbulence
        • Temperature-based color ramp (red → orange → yellow)
        • View-dependent glow halo
        • Slight global darkening for richness
    """

    normal = glm.normalize(varyings["normal_view"])
    pos_view = varyings["pos_view"]
    uv = varyings["uv"]

    # --- Noise in [0,1] ---
    n = perlin_fbm(uv, octaves=4, base_freq=6.0)
    n = max(0.0, min(1.0, n))

    # --- Color ramp: deep red -> orange -> yellow-hot ---
    cool = (0.65, 0.15, 0.05)   # deep red
    warm = (0.95, 0.45, 0.10)   # orange
    hot  = (1.00, 0.80, 0.30)   # yellow-hot

    if n < 0.5:
        t = n * 2.0
        base = (
            cool[0] * (1 - t) + warm[0] * t,
            cool[1] * (1 - t) + warm[1] * t,
            cool[2] * (1 - t) + warm[2] * t,
        )
    else:
        t = (n - 0.5) * 2.0
        base = (
            warm[0] * (1 - t) + hot[0] * t,
            warm[1] * (1 - t) + hot[1] * t,
            warm[2] * (1 - t) + hot[2] * t,
        )

    # Slight global darkening for a richer, less blown-out sun
    brightness = uniforms.get("sun_brightness", 0.85)
    base = (base[0] * brightness, base[1] * brightness, base[2] * brightness)

    # --- Glow halo (view-dependent rim) ---
    view_dir = glm.normalize(-pos_view)
    glow = pow(1.0 - max(0.0, glm.dot(normal, view_dir)), 2.0)

    # Warm, softer glow (less white)
    r = base[0] + glow * 0.9
    g = base[1] + glow * 0.45
    b = base[2] + glow * 0.15

    return (min(1.0, r), min(1.0, g), min(1.0, b))

def fs_earth_texture_fresnel(p0, p1, p2, a, b, c, uniforms, varyings):
    """
    Fragment shader for rendering the Earth using:
        - Image-based texture (albedo)
        - Lambertian diffuse lighting
        - Schlick Fresnel term for atmospheric rim lighting
    """

    # ------------------------------------------------------------
    # 1. Retrieve interpolated varyings
    # ------------------------------------------------------------

    # Surface normal in view space (already perspective-correct)
    normal_view = glm.normalize(varyings["normal_view"])

    # Fragment position in view space (camera at origin)
    pos_view = varyings["pos_view"]

    # Texture coordinates
    uv = varyings["uv"]

    # ------------------------------------------------------------
    # 2. Sample Earth albedo texture
    # ------------------------------------------------------------

    # The texture provides base color (land, ocean, clouds, ice)
    tex = uniforms["sampler2D"]
    base_color = tex.sample_bilinear(uv)

    # ------------------------------------------------------------
    # 3. Construct lighting vectors (view space)
    # ------------------------------------------------------------

    # Light position is provided in view space for convenience
    light_pos_view = uniforms.get("light_pos_view", glm.vec3(0, 0, 5))

    # Direction from surface point toward the light
    light_dir = glm.normalize(light_pos_view - pos_view)

    # Direction from surface point toward the camera (origin)
    view_dir = glm.normalize(-pos_view)

    # ------------------------------------------------------------
    # 4. Lambertian diffuse lighting
    # ------------------------------------------------------------

    # Standard diffuse term: N · L
    ndotl = max(0.0, glm.dot(normal_view, light_dir))

    # Apply diffuse lighting to texture color
    r = base_color[0] * ndotl
    g = base_color[1] * ndotl
    b = base_color[2] * ndotl

    # ------------------------------------------------------------
    # 5. Fresnel term (Schlick approximation)
    # ------------------------------------------------------------

    # Base reflectance at normal incidence:
    # ~0.02 is typical for water / Earth-like surfaces
    F0 = uniforms.get("fresnel_F0", 0.02)

    # Schlick Fresnel approximation:
    # Increases reflectance at grazing angles
    fresnel = F0 + (1.0 - F0) * pow(
        1.0 - max(0.0, glm.dot(normal_view, view_dir)),
        5.0
    )

    # ------------------------------------------------------------
    # 6. Atmospheric rim lighting
    # ------------------------------------------------------------

    # Tint used for atmospheric scattering (bluish by default)
    atm_color = uniforms.get("atm_color", (0.4, 0.6, 1.0))

    # Add Fresnel-scaled atmospheric color
    r += fresnel * atm_color[0]
    g += fresnel * atm_color[1]
    b += fresnel * atm_color[2]

    # ------------------------------------------------------------
    # 7. Clamp final color to valid range
    # ------------------------------------------------------------

    r = max(0.0, min(1.0, r))
    g = max(0.0, min(1.0, g))
    b = max(0.0, min(1.0, b))

    return (r, g, b)



# ============================================================
# Main Program
# ============================================================
def main():
    """
    Entry point: sets up window, engine, geometry, shaders,
    and issues draw calls for a static solar system scene.
    """
    w, h = 800, 800
    win = RitWindow(w, h, "Solar System")
    engine = CGIengine(w, h)

    # Textures
    tex_stars = Texture2D(Image.open("textures/milkyway.jpg").convert("RGB"))
    tex_earth = Texture2D(Image.open("textures/earth.jpg").convert("RGB"))

    # Geometry
    bg_pos, bg_uv, bg_idx = make_background_quad()
    bg_vao = build_vao(bg_pos, uvs=bg_uv)

    sp_pos, sp_nrm, sp_uv, sp_idx = make_uv_sphere()
    sphere_vao = build_vao(sp_pos, sp_nrm, sp_uv)

    # Camera
    eye = glm.vec3(0, 6, 14)
    view = T3.lookAt(eye, glm.vec3(0, 0, 0), glm.vec3(0, 1, 0))
    proj = T3.frustum3D(-1, 1, -1, 1, 1.5, 50)

    light_pos_view = glm.vec3(0, 2, 8)

    def draw(win):
        engine.clear_color(win, (0, 0, 0))
        engine.clearZBuffer()

        # Background
        engine.drawTriangles(
            win, bg_vao, bg_idx,
            vs_background, fs_background,
            {"sampler2D": tex_stars}
        )

        # Sun (procedural Perlin shader)
        engine.drawTriangles(
            win, sphere_vao, sp_idx,
            vs_standard, fs_sun_perlin,
            {
                "modelT": T3.scale3D(1.6, 1.6, 1.6),
                "viewT": view,
                "projectionT": proj,
                "sun_brightness": 0.85,  # tweak darker/lighter here
            }
        )

        # Earth (texture only)
        engine.drawTriangles(
            win, sphere_vao, sp_idx,
            vs_standard, fs_earth_texture_fresnel,
            {
                "modelT": T3.translate3D(5, 0, 0),
                "viewT": view,
                "projectionT": proj,
                "sampler2D": tex_earth,
                "filter": "bilinear",
            }
        )

        # Moon (Phong shading)
        engine.drawTriangles(
            win, sphere_vao, sp_idx,
            CGIengine.vs_phong, CGIengine.fs_phong,
            {
                "modelT": T3.translate3D(6.8, 0.6, 0) * T3.scale3D(0.4, 0.4, 0.4),
                "viewT": view,
                "projectionT": proj,
                "light_pos": light_pos_view,
                "light_color": (1.0, 1.0, 1.0),
                "kd": 0.9,
                "ks": 0.3,
                "shininess": 16,
                "ambient": 0.08,
            }
        )

    win.doRun(draw)


if __name__ == "__main__":
    main()