import { readFileSync, writeFileSync } from "node:fs";
import { createHash } from "node:crypto";
import { fileURLToPath } from "node:url";
import { inferRedshift } from "../../../src/scripts/physics-lab/experiments/galaxy-photometry/model/redshift-posterior.ts";
import type { FluxObservation } from "../../../src/scripts/physics-lab/experiments/galaxy-photometry/model/photometry-model.ts";

const root = new URL("../../../", import.meta.url);
const sourceUrl = new URL("docs/astronomy/ac04/sources/sdss-dr17.csv", root);
const sourceBytes = readFileSync(sourceUrl);
const lines = sourceBytes
  .toString("utf8")
  .trim()
  .split(/\r?\n/)
  .filter((line) => !line.startsWith("#"));
const columns = lines.shift()!.split(",");
const rawRows = lines.map((line) =>
  Object.fromEntries(
    line.split(",").map((value, index) => [columns[index], value]),
  ),
);
const bands = ["u", "g", "r", "i", "z"];
// These approximate offsets are frozen before inspecting spectroscopic validation.
const magnitudeAbOffsets = [-0.04, 0, 0, 0, 0.02];
const calibrationFraction = 0.03;

function observationFromRow(
  row: Record<string, string>,
  useAbOffsets: boolean,
): FluxObservation {
  const flux: number[] = [];
  const sigma: number[] = [];
  bands.forEach((band, index) => {
    const correction =
      10 **
      (0.4 * Number(row["extinction_" + band]) -
        0.4 * (useAbOffsets ? magnitudeAbOffsets[index] : 0));
    const correctedFlux = Number(row["modelFlux_" + band]) * correction;
    const measurementSigma =
      correction / Math.sqrt(Number(row["modelFluxIvar_" + band]));
    flux.push(correctedFlux);
    sigma.push(
      Math.hypot(
        measurementSigma,
        calibrationFraction * Math.abs(correctedFlux),
      ),
    );
  });
  return {
    flux,
    sigma,
    included: [true, true, true, true, true],
    kind: "SDSS DR17",
  };
}

const results = rawRows.map((row, index) => {
  const observation = observationFromRow(row, true);
  const posterior = inferRedshift(observation);
  const sensitivity = inferRedshift(observationFromRow(row, false));
  // The holdout is first accessed here, after the inference and fixed sensitivity run.
  const spectroscopicRedshift = Number(row.z);
  const normalizedError =
    (posterior.mode - spectroscopicRedshift) / (1 + spectroscopicRedshift);
  return {
    row: index + 1,
    objectId: row.objID,
    spectrumId: row.specObjID,
    raDegrees: Number(row.ra),
    decDegrees: Number(row.dec),
    rawFluxNanomaggies: bands.map((band) => Number(row["modelFlux_" + band])),
    rawInverseVariance: bands.map((band) =>
      Number(row["modelFluxIvar_" + band]),
    ),
    foregroundExtinctionMagnitudes: bands.map((band) =>
      Number(row["extinction_" + band]),
    ),
    observation,
    spectroscopicRedshift,
    spectroscopicRedshiftError: Number(row.zErr),
    posterior,
    posteriorWithoutApproximateAbOffsets: sensitivity,
    standardizedResiduals: observation.flux.map(
      (value, band) =>
        (value - posterior.predictedFlux[band]) / observation.sigma[band],
    ),
    normalizedError,
    catastrophic: Math.abs(normalizedError) > 0.15,
    coveredByCentral68:
      spectroscopicRedshift >= posterior.interval68[0] &&
      spectroscopicRedshift <= posterior.interval68[1],
    modeWithoutApproximateAbOffsets: sensitivity.mode,
  };
});
const median = (values: number[]) => {
  const sorted = [...values].sort((a, b) => a - b);
  return (
    (sorted[Math.floor((sorted.length - 1) / 2)] +
      sorted[Math.ceil((sorted.length - 1) / 2)]) /
    2
  );
};
const errors = results.map((row) => row.normalizedError);
const medianError = median(errors);
const summary = {
  count: results.length,
  negativeFluxMeasurements: results.reduce(
    (total, row) =>
      total + row.rawFluxNanomaggies.filter((value) => value < 0).length,
    0,
  ),
  catastrophicCount: results.filter((row) => row.catastrophic).length,
  central68CoverageCount: results.filter((row) => row.coveredByCentral68)
    .length,
  medianNormalizedError: medianError,
  normalizedMedianAbsoluteDeviation:
    1.48 * median(errors.map((value) => Math.abs(value - medianError))),
  maximumAbOffsetModeShift: Math.max(
    ...results.map((row) =>
      Math.abs(row.posterior.mode - row.modeWithoutApproximateAbOffsets),
    ),
  ),
};
const document = {
  source: {
    release: "SDSS DR17",
    file: "sdss-dr17.csv",
    bytes: sourceBytes.length,
    sha256: createHash("sha256").update(sourceBytes).digest("hex"),
    selection:
      "TOP 24 sciencePrimary GALAXY, zWarning=0, clean photometry, 0.02<z<0.8, ORDER BY specObjID; no photo-z selection",
    query: readFileSync(
      new URL("docs/astronomy/ac04/sources/sdss-dr17-query.sql", root),
      "utf8",
    ).trim(),
    rights:
      "SDSS public data-release policy; source and collaboration acknowledgement retained",
  },
  model: {
    name: "Limited BC03 SSP + Calzetti Bayesian teaching model; independent of EAZY",
    amplitude:
      "Observed r-band mean fnu in nanomaggies, half-normal prior scale 10000; analytically marginalized",
    redshiftPrior:
      "Uniform density from 0 to 1.2, spacing 0.005; trapezoidal normalization",
    nuisancePriors:
      "Equal weights for seven frozen SSP ages allowed by flat H0=70, Om=0.3 cosmic age; six continuum E(B-V) values 0..0.5, conditional on z",
    noise:
      "Independent Gaussian flux errors; quoted sigma plus 3% of absolute corrected observed flux in quadrature, held fixed across models",
    magnitudeAbOffsets,
    foregroundCorrection:
      "Flux and statistical sigma both multiplied by 10^(0.4*extinction)",
    omissions:
      "No nebular emission, composite star-formation histories, metallicity variation, cross-band covariance, IGM absorption or learned magnitude prior",
  },
  summary,
  results,
};
const destination = new URL("docs/astronomy/ac04/sdss-results.json", root);
if (process.argv.includes("--check")) {
  const previous = readFileSync(destination, "utf8");
  if (previous !== JSON.stringify(document, null, 2) + "\n")
    throw new Error("Frozen SDSS results differ; review numerical change");
} else writeFileSync(destination, JSON.stringify(document, null, 2) + "\n");
console.log(
  JSON.stringify({ output: fileURLToPath(destination), summary }, null, 2),
);