/* Resolver-only benchmark for the fast-mode FFTW global convolution.
 *
 * Fills one deterministic impulse buffer, then times the spatial reference and
 * the FFTW production resolver on that identical buffer.  First-use plan/setup
 * is reported separately from steady-state execution.  No catalog or geodesic
 * work is involved. */
#ifndef FAST_PSF_FFTW
#error "benchmark_fast_psf_fftw requires -DFAST_PSF_FFTW (PSF_BACKEND=cpu)"
#endif

#include "fast_psf_fftw.h"
#include "optics.h"

#include <math.h>
#include <omp.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/resource.h>

typedef struct {
  int width, height, supersample, repeats;
  int run_spatial;
  int measure;
  FastPsfDeposit deposit;
} Options;

static int parse_int(const char *text, int *out)
{
  char *end = NULL;
  const long value = strtol(text, &end, 10);
  if (end == text || *end != '\0' || value <= 0 || value > 1000000)
    return -1;
  *out = (int)value;
  return 0;
}

static int parse_options(int argc, char **argv, Options *options)
{
  *options = (Options){.width = 320,
                       .height = 240,
                       .supersample = 2,
                       .repeats = 5,
                       .run_spatial = 0,
                       .measure = 0,
                       .deposit = FAST_PSF_DEPOSIT_NEAREST};
  for (int i = 1; i < argc; ++i) {
    if (!strcmp(argv[i], "--spatial")) {
      options->run_spatial = 1;
    } else if (!strcmp(argv[i], "--measure")) {
      options->measure = 1;
    } else if (!strcmp(argv[i], "--bilinear")) {
      options->deposit = FAST_PSF_DEPOSIT_BILINEAR;
    } else if (!strcmp(argv[i], "--width") && i + 1 < argc) {
      if (parse_int(argv[++i], &options->width))
        return -1;
    } else if (!strcmp(argv[i], "--height") && i + 1 < argc) {
      if (parse_int(argv[++i], &options->height))
        return -1;
    } else if (!strcmp(argv[i], "--supersample") && i + 1 < argc) {
      if (parse_int(argv[++i], &options->supersample))
        return -1;
    } else if (!strcmp(argv[i], "--repeats") && i + 1 < argc) {
      if (parse_int(argv[++i], &options->repeats))
        return -1;
    } else {
      fprintf(stderr, "unknown argument '%s'\n", argv[i]);
      return -1;
    }
  }
  return 0;
}

static uint64_t hash_bytes(const double *values, size_t count)
{
  const unsigned char *bytes = (const unsigned char *)values;
  const size_t nbytes = count * sizeof *values;
  uint64_t hash = 1469598103934665603ULL;
  for (size_t i = 0; i < nbytes; ++i) {
    hash ^= bytes[i];
    hash *= 1099511628211ULL;
  }
  return hash;
}

static double wall_seconds(void)
{
  return omp_get_wtime();
}

static int compare_double(const void *a, const void *b)
{
  const double x = *(const double *)a;
  const double y = *(const double *)b;
  return x < y ? -1 : x > y ? 1 : 0;
}

int main(int argc, char **argv)
{
  Options options;
  if (parse_options(argc, argv, &options)) {
    fprintf(stderr,
            "usage: %s [--width W] [--height H] [--supersample N] "
            "[--repeats R] [--bilinear] [--spatial] [--measure]\n",
            argv[0]);
    return 2;
  }
  const PointSpreadFunction psf = {.fwhm_pixels = 2.7, .moffat_beta = 4.5};
  const double relative_tail = 1e-8;
  if (options.repeats > 128)
    options.repeats = 128;

  printf("impulse: %dx%d supersample=%d deposit=%s repeats=%d spatial=%d\n",
         options.width, options.height, options.supersample,
         options.deposit == FAST_PSF_DEPOSIT_NEAREST ? "nearest" : "bilinear",
         options.repeats, options.run_spatial);
  printf("plan_mode=%s\n", options.measure ? "measure" : "estimate");
  fast_psf_fftw_set_plan_mode(options.measure);

  FastPsfAccumulator acc = {0};
  const double init_start = wall_seconds();
  if (fast_psf_accumulator_init(&acc, options.width, options.height,
                                options.supersample, options.deposit, &psf,
                                relative_tail, 0.0, 1)) {
    fputs("accumulator initialization failed\n", stderr);
    return 1;
  }
  const double init_seconds = wall_seconds() - init_start;
  printf("setup: init=%.6f s fftw_plan=%.6f s kernel_fft=%.6f s\n",
         init_seconds, acc.fftw_setup_seconds, acc.fftw_kernel_seconds);
  fast_psf_fftw_report(acc.fftw, stdout);

  const size_t ss_count = (size_t)acc.supersampled_width *
                          acc.supersampled_height;
  uint64_t state = 0x9e3779b97f4a7c15ULL;
  for (size_t i = 0; i < ss_count; ++i) {
    state = state * 6364136223846793005ULL + 1442695040888963407ULL;
    if ((state >> 58) != 0)
      continue;
    state = state * 6364136223846793005ULL + 1442695040888963407ULL;
    acc.buffer[3 * i] = (double)(state >> 40) / (double)(1ULL << 24);
    state = state * 6364136223846793005ULL + 1442695040888963407ULL;
    acc.buffer[3 * i + 1] = (double)(state >> 40) / (double)(1ULL << 24);
    state = state * 6364136223846793005ULL + 1442695040888963407ULL;
    acc.buffer[3 * i + 2] = (double)(state >> 40) / (double)(1ULL << 24);
  }
  const size_t hdr_count = (size_t)options.width * options.height * 3;
  const size_t buffer_count = ss_count * 3;
  const uint64_t impulse_hash = hash_bytes(acc.buffer, buffer_count);
  printf("impulse_hash=%016llx\n", (unsigned long long)impulse_hash);

  double *fftw_hdr = calloc(hdr_count, sizeof *fftw_hdr);
  double *spatial_hdr = calloc(hdr_count, sizeof *spatial_hdr);
  if (fftw_hdr == NULL || spatial_hdr == NULL) {
    fputs("HDR allocation failed\n", stderr);
    free(fftw_hdr);
    free(spatial_hdr);
    fast_psf_accumulator_destroy(&acc);
    return 1;
  }

  /* Warm-up: first execution includes any lazy per-frame allocation. */
  if (fast_psf_accumulator_resolve(&acc, fftw_hdr, omp_get_max_threads())) {
    fputs("FFTW warm-up resolve failed\n", stderr);
    return 1;
  }
  double *samples = calloc((size_t)options.repeats, sizeof *samples);
  if (samples == NULL) {
    fputs("sample allocation failed\n", stderr);
    return 1;
  }
  for (int i = 0; i < options.repeats; ++i) {
    memset(fftw_hdr, 0, hdr_count * sizeof *fftw_hdr);
    const double start = wall_seconds();
    if (fast_psf_accumulator_resolve(&acc, fftw_hdr, omp_get_max_threads())) {
      fputs("FFTW resolve failed\n", stderr);
      return 1;
    }
    samples[i] = wall_seconds() - start;
  }
  const uint64_t fftw_hash = hash_bytes(fftw_hdr, hdr_count);
  double sorted[128];
  memcpy(sorted, samples, (size_t)options.repeats * sizeof *sorted);
  qsort(sorted, (size_t)options.repeats, sizeof *sorted, compare_double);
  printf("fftw: min=%.6f s median=%.6f s max=%.6f s\n", sorted[0],
         sorted[options.repeats / 2], sorted[options.repeats - 1]);
  printf("fftw samples:");
  for (int i = 0; i < options.repeats; ++i)
    printf(" %.6f", samples[i]);
  printf("\nfftw_hdr_hash=%016llx\n", (unsigned long long)fftw_hash);

  if (options.run_spatial) {
    memset(spatial_hdr, 0, hdr_count * sizeof *spatial_hdr);
    const double start = wall_seconds();
    if (fast_psf_accumulator_resolve_spatial_reference(
            &acc, spatial_hdr, omp_get_max_threads())) {
      fputs("spatial resolve failed\n", stderr);
      return 1;
    }
    const double spatial_seconds = wall_seconds() - start;
    printf("spatial: %.6f s\n", spatial_seconds);
    printf("spatial_hdr_hash=%016llx\n",
           (unsigned long long)hash_bytes(spatial_hdr, hdr_count));
    double max_abs = 0.0, peak = 0.0;
    for (size_t i = 0; i < hdr_count; ++i) {
      peak = fmax(peak, fabs(spatial_hdr[i]));
      max_abs = fmax(max_abs, fabs(fftw_hdr[i] - spatial_hdr[i]));
    }
    printf("difference: max_abs=%.3g peak=%.3g speedup=%.3fx\n", max_abs, peak,
           spatial_seconds / sorted[options.repeats / 2]);
  }

  struct rusage usage;
  if (getrusage(RUSAGE_SELF, &usage) == 0)
    printf("rss_max_kb=%ld\n", usage.ru_maxrss);

  free(samples);
  free(fftw_hdr);
  free(spatial_hdr);
  fast_psf_accumulator_destroy(&acc);
  return 0;
}