#include "spacetime.h"

#include <math.h>
#include <stdlib.h>

typedef struct {
  double mass;
  double escape_radius;
  double capture_radius;
} SchwarzschildKsContext;

/* Schwarzschild in ingoing Cartesian Kerr--Schild coordinates:
 * g_mu_nu = eta_mu_nu + (2 M / r) l_mu l_nu, l_mu = (1, x_i / r).
 * These slices are regular at r = 2 M; only the physical r = 0 singularity
 * is excluded by the conservative capture cutoff. */
static int schwarzschild_ks_eval(const SpacetimeSource *source, double t,
                                 const double x[3], MetricData *metric) {
  const SchwarzschildKsContext *context = source->context;
  double r2 = 0.0;
  (void)t;
  for (int i = 0; i < 3; ++i)
    r2 += x[i] * x[i];
  if (!isfinite(r2) || r2 <= 0.0)
    return -1;
  const double r = sqrt(r2);
  const double m = context->mass;
  const double f = 2.0 * m / r;
  const double alpha = 1.0 / sqrt(1.0 + f);
  const double beta_scale = 2.0 * m / (r * (r + 2.0 * m));
  const double dbeta_scale_dr =
      -4.0 * m * (r + m) / (r2 * (r + 2.0 * m) * (r + 2.0 * m));
  *metric = (MetricData){.alpha = alpha};
  for (int i = 0; i < 3; ++i) {
    metric->beta[i] = beta_scale * x[i];
    metric->d_alpha[i] = m * alpha * alpha * alpha * x[i] / (r2 * r);
    for (int j = 0; j < 3; ++j) {
      metric->gamma[i][j] = (i == j ? 1.0 : 0.0) +
                             2.0 * m * x[i] * x[j] / (r2 * r);
      metric->d_beta[i][j] = beta_scale * (i == j ? 1.0 : 0.0) +
                              dbeta_scale_dr * x[i] * x[j] / r;
      for (int k = 0; k < 3; ++k)
        metric->d_gamma[i][j][k] =
            2.0 * m * (((i == j ? 1.0 : 0.0) * x[k] +
                        (i == k ? 1.0 : 0.0) * x[j]) /
                           (r2 * r) -
                       3.0 * x[i] * x[j] * x[k] / (r2 * r2 * r));
    }
  }
  /* K_ij = (D_i beta_j + D_j beta_i)/(2 alpha) for these stationary slices.
   * Build the connection from the analytic spatial-metric derivatives above;
   * this avoids finite differences in the geodesic RHS. */
  for (int i = 0; i < 3; ++i)
    for (int j = 0; j < 3; ++j) {
      double d_beta_cov_i_j =
          2.0 * m * ((i == j ? 1.0 : 0.0) / r2 -
                     2.0 * x[i] * x[j] / (r2 * r2));
      double d_beta_cov_j_i = d_beta_cov_i_j;
      double connection_term_ij = 0.0;
      double connection_term_ji = 0.0;
      for (int ell = 0; ell < 3; ++ell) {
        double gamma_inverse_ell_k[3];
        for (int k = 0; k < 3; ++k)
          gamma_inverse_ell_k[k] = (ell == k ? 1.0 : 0.0) -
                                   f / (1.0 + f) * x[ell] * x[k] / r2;
        double gamma_ell_ij = 0.0, gamma_ell_ji = 0.0;
        for (int k = 0; k < 3; ++k) {
          gamma_ell_ij += 0.5 * gamma_inverse_ell_k[k] *
                          (metric->d_gamma[i][k][j] +
                           metric->d_gamma[j][k][i] -
                           metric->d_gamma[k][i][j]);
          gamma_ell_ji += 0.5 * gamma_inverse_ell_k[k] *
                          (metric->d_gamma[j][k][i] +
                           metric->d_gamma[i][k][j] -
                           metric->d_gamma[k][j][i]);
        }
        const double beta_cov_ell = 2.0 * m * x[ell] / r2;
        connection_term_ij += gamma_ell_ij * beta_cov_ell;
        connection_term_ji += gamma_ell_ji * beta_cov_ell;
      }
      metric->K[i][j] = (d_beta_cov_i_j - connection_term_ij +
                         d_beta_cov_j_i - connection_term_ji) /
                        (2.0 * alpha);
    }
  return 0;
}

static SpacetimeRayStatus schwarzschild_ks_classify(
    const SpacetimeSource *source, double t, const double x[3]) {
  const SchwarzschildKsContext *context = source->context;
  double r2 = 0.0;
  (void)t;
  for (int i = 0; i < 3; ++i)
    r2 += x[i] * x[i];
  if (!isfinite(r2) || r2 <= context->capture_radius * context->capture_radius)
    return SPACETIME_RAY_CAPTURED;
  return r2 >= context->escape_radius * context->escape_radius
             ? SPACETIME_RAY_ESCAPED
             : SPACETIME_RAY_ACTIVE;
}

static void schwarzschild_ks_destroy(SpacetimeSource *source) {
  free(source->context);
  source->context = NULL;
  source->ops = NULL;
}

static const SpacetimeOps schwarzschild_ks_ops = {
    .eval = schwarzschild_ks_eval,
    .classify = schwarzschild_ks_classify,
    .destroy = schwarzschild_ks_destroy,
};

int spacetime_create_schwarzschild_ks(SpacetimeSource *source, double mass,
                                      double escape_radius,
                                      double capture_radius) {
  if (source == NULL || mass <= 0.0 || escape_radius <= 2.0 * mass ||
      capture_radius <= 0.0 || capture_radius >= 2.0 * mass ||
      capture_radius >= escape_radius)
    return -1;
  SchwarzschildKsContext *context = malloc(sizeof *context);
  if (context == NULL)
    return -1;
  *context = (SchwarzschildKsContext){mass, escape_radius, capture_radius};
  source->ops = &schwarzschild_ks_ops;
  source->context = context;
  return 0;
}

int spacetime_create_default(SpacetimeSource *source) {
  return spacetime_create_schwarzschild_ks(source, 1.0, 256.0, 1.5);
}