// galaxycube.cc - GalaxyCube renderer for ArtCube
// Rotating spiral galaxy projected onto a 6-face 64x64 cube strip.
//
// Face order matches globe:
// 0=TOP(+Y), 1=RIGHT(+X), 2=REAR(-Z), 3=LEFT(-X), 4=FRONT(+Z), 5=BOTTOM(-Y)
//
// Build: link with cube_orientation.cc and rgbmatrix.

#include <algorithm>
#include <chrono>
#include <cmath>
#include <csignal>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <random>
#include <thread>
#include <vector>

#include "led-matrix.h"
#include "cube_orientation.h"

using rgb_matrix::Canvas;
using rgb_matrix::FrameCanvas;
using rgb_matrix::RGBMatrix;

static volatile std::sig_atomic_t g_running = 1;
static void InterruptHandler(int) { g_running = 0; }

static inline float clampf(float v, float lo, float hi) {
  return v < lo ? lo : (v > hi ? hi : v);
}

static inline uint8_t clamp_u8(int v) {
  if (v < 0) return 0;
  if (v > 255) return 255;
  return (uint8_t)v;
}

static inline uint8_t scale_u8(uint8_t v, float s) {
  return clamp_u8((int)std::lround((float)v * s));
}

struct Vec3 { float x, y, z; };

static inline Vec3 vadd(Vec3 a, Vec3 b) { return {a.x + b.x, a.y + b.y, a.z + b.z}; }
static inline Vec3 vmul(Vec3 a, float s) { return {a.x * s, a.y * s, a.z * s}; }
static inline float vdot(Vec3 a, Vec3 b) { return a.x*b.x + a.y*b.y + a.z*b.z; }
static inline float vlen(Vec3 a) { return std::sqrt(vdot(a,a)); }
static inline Vec3 vnorm(Vec3 a) {
  float L = vlen(a);
  if (L <= 1e-6f) return {0,0,1};
  return {a.x/L, a.y/L, a.z/L};
}
static inline Vec3 rotate_y(Vec3 v, float a) {
  float c = std::cos(a), s = std::sin(a);
  return {v.x*c + v.z*s, v.y, -v.x*s + v.z*c};
}
static inline Vec3 rotate_x(Vec3 v, float a) {
  float c = std::cos(a), s = std::sin(a);
  return {v.x, v.y*c - v.z*s, v.y*s + v.z*c};
}
static inline Vec3 rotate_z(Vec3 v, float a) {
  float c = std::cos(a), s = std::sin(a);
  return {v.x*c - v.y*s, v.x*s + v.y*c, v.z};
}

struct FaceXY { int face, x, y; };

// 3D direction -> cube face + N×N pixel coords.
// IMPORTANT: This mapping matches the final SolarCube orientation fixes so edge transitions feel continuous.
static FaceXY DirToFaceXY(Vec3 d, int N) {
  d = vnorm(d);
  float ax = std::fabs(d.x), ay = std::fabs(d.y), az = std::fabs(d.z);

  int face = 4;
  float u = 0.0f, v = 0.0f;

  if (ay >= ax && ay >= az) {
    if (d.y >= 0) {
      face = 0;              // +Y TOP
      u = -d.x / ay;
      v = -d.z / ay;
    } else {
      face = 5;              // -Y BOTTOM
      u = d.x / ay;
      v = -d.z / ay;
    }
  } else if (ax >= ay && ax >= az) {
    if (d.x >= 0) {
      face = 1;              // +X RIGHT
      u = -d.z / ax;
      v = -d.y / ax;
    } else {
      face = 3;              // -X LEFT
      u = d.z / ax;
      v = -d.y / ax;
    }
  } else {
    if (d.z >= 0) {
      face = 4;              // +Z FRONT
      u = d.x / az;
      v = -d.y / az;
    } else {
      face = 2;              // -Z REAR
      u = -d.x / az;
      v = -d.y / az;
    }
  }

  int px = (int)std::lround((u + 1.0f) * 0.5f * (float)(N - 1));
  int py = (int)std::lround((v + 1.0f) * 0.5f * (float)(N - 1));
  if (px < 0) px = 0; else if (px >= N) px = N - 1;
  if (py < 0) py = 0; else if (py >= N) py = N - 1;
  return {face, px, py};
}

static inline void PutPixel(Canvas* draw, int face, int x, int y,
                            uint8_t r, uint8_t g, uint8_t b,
                            float brightness_scale, int N) {
  draw->SetPixel(face * N + x, y,
                 scale_u8(r, brightness_scale),
                 scale_u8(g, brightness_scale),
                 scale_u8(b, brightness_scale));
}

struct Star {
  Vec3 dir0;          // base direction
  float omega;        // rotation speed factor
  float phase;        // twinkle phase
  uint8_t r, g, b;    // base color
  uint8_t amp;        // twinkle amplitude
};

static inline float frand(std::mt19937& rng, float a, float b) {
  std::uniform_real_distribution<float> d(a, b);
  return d(rng);
}
static inline float gauss(std::mt19937& rng, float mean, float sigma) {
  std::normal_distribution<float> d(mean, sigma);
  return d(rng);
}

static Star MakeBackgroundStar(std::mt19937& rng) {
  float z = frand(rng, -1.0f, 1.0f);
  float t = frand(rng, 0.0f, 2.0f * 3.1415926f);
  float r = std::sqrt(std::max(0.0f, 1.0f - z*z));
  Vec3 d{r*std::cos(t), z, r*std::sin(t)};
  uint8_t v = (uint8_t)frand(rng, 8.0f, 26.0f);
  uint8_t tint = (uint8_t)frand(rng, 0.0f, 18.0f);
  return {d, frand(rng, 0.02f, 0.05f), frand(rng, 0.0f, 6.28f),
          (uint8_t)clamp_u8((int)v + (int)tint), v, (uint8_t)clamp_u8((int)v - (int)tint),
          (uint8_t)frand(rng, 6.0f, 14.0f)};
}

static Star MakeGalaxyStar(std::mt19937& rng, int arm_idx) {
  float u = frand(rng, 0.0f, 1.0f);
  float r = 0.25f + 2.75f * (u*u);  // center-biased
  float arm_offset = float(arm_idx) * (2.0f * 3.1415926f / 3.0f);

  float tight = 2.1f;
  float theta = arm_offset + tight * std::log(r + 0.15f) + gauss(rng, 0.0f, 0.18f);

  float spread = 0.07f + 0.06f * r;
  float rr = r + gauss(rng, 0.0f, spread);

  float x = rr * std::cos(theta);
  float z = rr * std::sin(theta);
  float y = gauss(rng, 0.0f, 0.12f + 0.05f * rr);

  Vec3 pos{x, y + 0.10f, z + 0.05f};
  Vec3 d = vnorm(pos);

  float cool = frand(rng, 0.0f, 1.0f);
  uint8_t base = (uint8_t)frand(rng, 18.0f, 70.0f);
  uint8_t r8 = (uint8_t)clamp_u8((int)std::lround(base * (0.85f + 0.30f * (1.0f - cool))));
  uint8_t g8 = (uint8_t)clamp_u8((int)std::lround(base * (0.95f + 0.20f * (1.0f - cool))));
  uint8_t b8 = (uint8_t)clamp_u8((int)std::lround(base * (1.05f + 0.35f * cool)));

  float omega = 0.07f + 0.06f * (1.0f - clampf(rr / 3.0f, 0.0f, 1.0f));
  uint8_t amp = (uint8_t)frand(rng, 10.0f, 26.0f);

  return {d, omega, frand(rng, 0.0f, 6.28f), r8, g8, b8, amp};
}

static Star MakeCoreStar(std::mt19937& rng) {
  float rr = std::fabs(gauss(rng, 0.0f, 0.35f));
  float theta = frand(rng, 0.0f, 2.0f * 3.1415926f);
  float x = rr * std::cos(theta);
  float z = rr * std::sin(theta);
  float y = gauss(rng, 0.0f, 0.22f);

  Vec3 pos{x, y + 0.08f, z + 0.05f};
  Vec3 d = vnorm(pos);

  uint8_t base = (uint8_t)frand(rng, 40.0f, 120.0f);
  uint8_t r8 = (uint8_t)clamp_u8((int)base + (int)frand(rng, 10.0f, 40.0f));
  uint8_t g8 = (uint8_t)clamp_u8((int)base + (int)frand(rng, 0.0f, 20.0f));
  uint8_t b8 = (uint8_t)clamp_u8((int)std::lround(base * frand(rng, 0.60f, 0.85f)));

  float omega = 0.14f + frand(rng, 0.0f, 0.10f);
  uint8_t amp = (uint8_t)frand(rng, 12.0f, 30.0f);

  return {d, omega, frand(rng, 0.0f, 6.28f), r8, g8, b8, amp};
}

static Star MakeGiant(std::mt19937& rng) {
  float z = frand(rng, -1.0f, 1.0f);
  float t = frand(rng, 0.0f, 2.0f * 3.1415926f);
  float r = std::sqrt(std::max(0.0f, 1.0f - z*z));
  Vec3 d{r*std::cos(t), z, r*std::sin(t)};
  uint8_t base = (uint8_t)frand(rng, 120.0f, 220.0f);
  uint8_t r8 = (uint8_t)clamp_u8((int)std::lround(base * 0.80f));
  uint8_t g8 = (uint8_t)clamp_u8((int)std::lround(base * 0.90f));
  uint8_t b8 = (uint8_t)clamp_u8((int)std::lround(base * 1.05f));
  return {d, frand(rng, 0.01f, 0.03f), frand(rng, 0.0f, 6.28f), r8, g8, b8, (uint8_t)frand(rng, 8.0f, 18.0f)};
}

int main(int argc, char** argv) {
  std::signal(SIGINT, InterruptHandler);
  std::signal(SIGTERM, InterruptHandler);

  RGBMatrix::Options options;
  rgb_matrix::RuntimeOptions runtime_opt;
  options.rows = 64;
  options.cols = 64;
  options.chain_length = 6;
  options.parallel = 1;
  options.hardware_mapping = "adafruit-hat";
  options.brightness = 65;

  rgb_matrix::ParseOptionsFromFlags(&argc, &argv, &options, &runtime_opt);

  // GalaxyCube custom args (leftover after rgb-matrix flag parsing):
  //   --yaw=<deg>                rotate scene around +Y (0/90/180/270 handy)
  //   --brightness-scale=<0..1>  multiply rendered colors (museum dimmer)
  double yaw_deg = 0.0;
  float brightness_scale = 1.0f;

  for (int i = 1; i < argc; i++) {
    if (std::strncmp(argv[i], "--yaw=", 6) == 0) {
      yaw_deg = std::atof(argv[i] + 6);
    } else if (std::strncmp(argv[i], "--brightness-scale=", 19) == 0) {
      brightness_scale = (float)std::atof(argv[i] + 19);
    }
  }
  brightness_scale = clampf(brightness_scale, 0.05f, 1.0f);
  const float yaw_rad = (float)(yaw_deg * 3.1415926535 / 180.0);

  RGBMatrix* matrix = RGBMatrix::CreateFromOptions(options, runtime_opt);
  if (matrix == nullptr) return 1;

  matrix->SetBrightness((uint8_t)(options.brightness ? options.brightness : 65));

  FrameCanvas* canvas = matrix->CreateFrameCanvas();
  cube::CubeCanvas cube_canvas(canvas, 64, 64);
  Canvas* draw = &cube_canvas;

  const int N = 64;

  std::mt19937 rng((unsigned)std::chrono::high_resolution_clock::now().time_since_epoch().count());

  std::vector<Star> bg;
  std::vector<Star> arms;
  std::vector<Star> core;
  std::vector<Star> giants;

  bg.reserve(700);
  for (int i = 0; i < 700; i++) bg.push_back(MakeBackgroundStar(rng));

  arms.reserve(1800);
  for (int i = 0; i < 1800; i++) arms.push_back(MakeGalaxyStar(rng, i % 3));

  core.reserve(900);
  for (int i = 0; i < 900; i++) core.push_back(MakeCoreStar(rng));

  giants.reserve(90);
  for (int i = 0; i < 90; i++) giants.push_back(MakeGiant(rng));

  const float incl_x = 0.55f;
  const float incl_z = 0.20f;

  auto t0 = std::chrono::steady_clock::now();
  const int target_ms = 33;

  while (g_running) {
    auto frame_start = std::chrono::steady_clock::now();
    float t = std::chrono::duration<float>(frame_start - t0).count();

    draw->Clear();

    auto draw_star_list = [&](const std::vector<Star>& stars, float tw_speed, float base_spin) {
      for (const auto& s0 : stars) {
        float a = t * (base_spin + s0.omega);
        Vec3 d = s0.dir0;
        d = rotate_y(d, a);
        d = rotate_x(d, incl_x);
        d = rotate_z(d, incl_z);
        d = rotate_y(d, yaw_rad);

        float tw = 0.70f + 0.30f * std::sin(t * tw_speed + s0.phase);
        int rr = (int)s0.r + (int)std::lround((float)s0.amp * (tw - 0.70f));
        int gg = (int)s0.g + (int)std::lround((float)s0.amp * (tw - 0.70f));
        int bb = (int)s0.b + (int)std::lround((float)s0.amp * (tw - 0.70f));

        FaceXY p = DirToFaceXY(d, N);
        PutPixel(draw, p.face, p.x, p.y, clamp_u8(rr), clamp_u8(gg), clamp_u8(bb), brightness_scale, N);
      }
    };

    draw_star_list(bg,     1.05f, 0.008f);
    draw_star_list(arms,   1.35f, 0.013f);
    draw_star_list(core,   1.15f, 0.018f);
    draw_star_list(giants, 0.80f, 0.006f);

    canvas = matrix->SwapOnVSync(canvas);
    cube_canvas.SetBase(canvas);

    auto frame_end = std::chrono::steady_clock::now();
    int elapsed_ms = (int)std::chrono::duration_cast<std::chrono::milliseconds>(frame_end - frame_start).count();
    int sleep_ms = target_ms - elapsed_ms;
    if (sleep_ms > 0) std::this_thread::sleep_for(std::chrono::milliseconds(sleep_ms));
  }

  draw->Clear();
  canvas = matrix->SwapOnVSync(canvas);
  cube_canvas.SetBase(canvas);

  delete matrix;
  return 0;
}
