main.cpp

#include <cmath>
#include <iostream>
#include <random>
#include <simd/simd.h>
#include <vector>

#include "MetalRenderer.hpp"
#include "SharedTypes.hpp"

// ─── Random helpers (CPU side, for scene generation) ─────────────────

static float random_double() {
  thread_local std::mt19937 gen;
  std::uniform_real_distribution<float> dist(0.0f, 1.0f);
  return dist(gen);
}

static float random_double(float min, float max) {
  thread_local std::mt19937 gen;
  std::uniform_real_distribution<float> dist(min, max);
  return dist(gen);
}

static simd_float3 random_color() {
  return simd_make_float3(random_double(), random_double(), random_double());
}

static simd_float3 random_color(float min, float max) {
  return simd_make_float3(random_double(min, max), random_double(min, max),
                          random_double(min, max));
}

static float vec_length(simd_float3 v) {
  return std::sqrt(v.x * v.x + v.y * v.y + v.z * v.z);
}

// ─── Factory helpers ─────────────────────────────────────────────────

static GPUMaterial lambertian(simd_float3 albedo) {
  GPUMaterial m{};
  m.type = MATERIAL_LAMBERTIAN;
  m.albedo =
      simd_make_float4(albedo.x, albedo.y, albedo.z, 0.0f); // w = fuzz (unused)
  return m;
}

static GPUMaterial metal(simd_float3 albedo, float fuzz) {
  GPUMaterial m{};
  m.type = MATERIAL_METAL;
  m.albedo = simd_make_float4(albedo.x, albedo.y, albedo.z, fuzz); // w = fuzz
  return m;
}

static GPUMaterial dielectric(float refraction_index) {
  GPUMaterial m{};
  m.type = MATERIAL_DIELECTRIC;
  m.albedo = simd_make_float4(1, 1, 1, 0);
  m.refractionIndex = refraction_index;
  return m;
}

static GPUHittable makeSphere(simd_float3 center, simd_float3 center2,
                              float radius, GPUMaterial mat) {
  GPUHittable h{};
  h.type = HITTABLE_SPHERE;
  h.center1 =
      simd_make_float4(center.x, center.y, center.z, radius); // w = radius
  h.center2 = simd_make_float4(center2.x, center2.y, center2.z, 0);
  h.material = mat;
  return h;
}

static GPUHittable makeStaticSphere(simd_float3 center, float radius,
                                    GPUMaterial mat) {
  return makeSphere(center, center, radius, mat);
}

// ─── Main ────────────────────────────────────────────────────────────

int main() {
  std::vector<GPUHittable> world;

  // ─── Scene setup ──────────────────────────────────────────────────

  // Ground
  world.push_back(
      makeStaticSphere(simd_make_float3(0, -1000, 0), 1000,
                       lambertian(simd_make_float3(0.5f, 0.5f, 0.5f))));

  // Random small spheres
  for (int a = -11; a < 11; a++) {
    for (int b = -11; b < 11; b++) {
      float choose_mat = random_double();
      simd_float3 center = simd_make_float3(a + 0.9f * random_double(), 0.2f,
                                            b + 0.9f * random_double());

      if (vec_length(center - simd_make_float3(4, 0.2f, 0)) > 0.9f) {
        if (choose_mat < 0.8f) {
          simd_float3 albedo = random_color() * random_color();
          simd_float3 center2 =
              center + simd_make_float3(0, random_double(0, 0.5), 0);
          world.push_back(
              makeSphere(center, center2, 0.2f, lambertian(albedo)));
        } else if (choose_mat < 0.95f) {
          simd_float3 albedo = random_color(0.5f, 1.0f);
          float fuzz = random_double(0, 0.5f);
          world.push_back(makeStaticSphere(center, 0.2f, metal(albedo, fuzz)));
        } else {
          world.push_back(makeStaticSphere(center, 0.2f, dielectric(1.5f)));
        }
      }
    }
  }

  // Three hero spheres
  world.push_back(
      makeStaticSphere(simd_make_float3(0, 1, 0), 1.0f, dielectric(1.5f)));
  world.push_back(
      makeStaticSphere(simd_make_float3(-4, 1, 0), 1.0f,
                       lambertian(simd_make_float3(0.4f, 0.2f, 0.1f))));
  world.push_back(
      makeStaticSphere(simd_make_float3(4, 1, 0), 1.0f,
                       metal(simd_make_float3(0.7f, 0.6f, 0.5f), 0.0f)));

  std::clog << "Scene: " << world.size() << " hittables" << std::endl;

  // ─── Camera setup ────────────────────────────────────────────────

  constexpr int imageWidth = 400;
  constexpr float aspect_ratio = 16.0f / 9.0f;
  int imageHeight = static_cast<int>(imageWidth / aspect_ratio);
  if (imageHeight < 1)
    imageHeight = 1;

  constexpr int samplesPerPixel = 100;
  constexpr int maxDepth = 50;
  constexpr int vfov = 20;
  float defocusAngle = 0.6f;
  float focusDistance = 10.0f;

  simd_float3 lookFrom = simd_make_float3(13, 2, 3);
  simd_float3 lookAt = simd_make_float3(0, 0, 0);
  simd_float3 upVector = simd_make_float3(0, 1, 0);

  // Compute camera basis
  simd_float3 w =
      simd_normalize(lookFrom - lookAt); // Negative camera look direction
  simd_float3 u = simd_normalize(simd_cross(upVector, w)); // Camera right
  simd_float3 v = simd_normalize(simd_cross(w, u));

  constexpr float pi = 3.1415926535897932385f;
  float theta = vfov * pi / 180.0f;
  float h = std::tan(theta / 2.0f);
  float viewportHeight = 2.0f * h * focusDistance;
  float viewportWidth = viewportHeight * (float(imageWidth) / imageHeight);

  simd_float3 viewport_u = viewportWidth * u;
  simd_float3 viewport_v = viewportHeight * (-v);

  simd_float3 pixel_delta_u = viewport_u / (float)imageWidth;
  simd_float3 pixel_delta_v = viewport_v / (float)imageHeight;

  simd_float3 viewport_upper_left =
      lookFrom - (focusDistance * w) - viewport_u / 2.0f - viewport_v / 2.0f;
  simd_float3 pixel00_loc =
      viewport_upper_left + 0.5f * (pixel_delta_u + pixel_delta_v);

  float defocusRadius =
      focusDistance * std::tan(defocusAngle / 2.0f * pi / 180.0f);
  simd_float3 defocusDiscU = defocusRadius * u;
  simd_float3 defocusDiscV = defocusRadius * v;

  // Pack into GPUCameraParams using float4
  GPUCameraParams camera{};
  camera.center =
      simd_make_float4(lookFrom.x, lookFrom.y, lookFrom.z, defocusAngle);
  camera.pixel00_loc =
      simd_make_float4(pixel00_loc.x, pixel00_loc.y, pixel00_loc.z, 0);
  camera.pixel_delta_u =
      simd_make_float4(pixel_delta_u.x, pixel_delta_u.y, pixel_delta_u.z, 0);
  camera.pixel_delta_v =
      simd_make_float4(pixel_delta_v.x, pixel_delta_v.y, pixel_delta_v.z, 0);
  camera.defocusDiscU =
      simd_make_float4(defocusDiscU.x, defocusDiscU.y, defocusDiscU.z, 0);
  camera.defocusDiscV =
      simd_make_float4(defocusDiscV.x, defocusDiscV.y, defocusDiscV.z, 0);
  camera.imageWidth = imageWidth;
  camera.imageHeight = imageHeight;
  camera.samplesPerPixel = samplesPerPixel;
  camera.maxDepth = maxDepth;

  // ─── Render on GPU ───────────────────────────────────────────────

  MetalRenderer renderer;
  std::vector<float> pixels = renderer.render(world, camera);

  // ─── Write PPM to stdout ─────────────────────────────────────────

  std::cout << "P3\n" << imageWidth << ' ' << imageHeight << "\n255\n";

  for (int i = 0; i < imageHeight; i++) {
    for (int j = 0; j < imageWidth; j++) {
      int idx = (i * imageWidth + j) * 3;
      int r = int(256 * pixels[idx + 0]);
      int g = int(256 * pixels[idx + 1]);
      int b = int(256 * pixels[idx + 2]);
      r = std::min(r, 255);
      g = std::min(g, 255);
      b = std::min(b, 255);
      std::cout << r << ' ' << g << ' ' << b << '\n';
    }
  }

  std::clog << "Done.\n";
  return 0;
}