# p_astro production pipeline for the O4 bank
#
#   make               everything: component models -> combined + stamped model
#                      -> plots -> reference-candidate check -> provenance record
#   make model         pastro-model.h5 (build components, combine, stamp)
#   make stamp         (re)stamp the runtime normalization on the existing model,
#                      refitting chi on $(INJECTIONS) (HDF5 attrs only -- seconds)
#   make plot          decision-surface map at FAR
#   make plot-far      p_astro vs FAR operating curves, one template per region
#   make check         the reference SGNL candidate through BOTH consumers
#                      (sgn-manifold CLI and sgnl-pastro's coinc script); they
#                      must agree bit-for-bit
#   make provenance    versions.txt: package versions + model attrs, for the release
#   make mark-current  declare the existing component/model files up to date
#                      (touch, in build order) -- use after a cosmetic yaml edit,
#                      since make rebuilds a component whenever its yaml is newer
#                      than its h5 (component builds take HOURS on the O4 bank)
#   make clean         remove generated files EXCEPT the component models
#   make clean-all     remove generated files INCLUDING the component models
#                      (an hours-long rebuild)
#
# Inputs:  {bns,nsbh,bbh,pop}.yaml (mass-model specs; the three regions
#          carry cuts, pop has none), $(INJECTIONS) (an SGNL injection campaign,
#          stillsuit SQLite: the rho(FAR) calibration AND the chi fit), $(COINC)
#          (a GraceDB coinc document for the check).
# Outputs: pastro-model.h5 -- ONE self-contained file: the validated combined model
#          with the rho(FAR) calibration and the runtime normalization stamped on
#          it, consumed by
#            sgn-manifold-cbc-p-astro --model pastro-model.h5 ...
#            sgnl-pastro-calculate-fgmc-coinc --model pastro-model.h5 coinc.xml
#            SGNL's sgnl-ll-pastro-uploader (pastro.load; `read-only: true`)
#
# Normalization (stamped; also the defaults for every tool below): 150 detected
# events/yr anchored at a 160 Mpc fiducial-BNS range, counting threshold 1/yr
# (3.17e-8 Hz), live range = the anchor range.  chi is FIT on $(INJECTIONS) by
# ensemble self-consistency (sgn-manifold-cbc-p-astro-configure
# --fit-far-calibration; the campaign path is stamped as provenance) -- its
# precision is set by the recoveries in the last few e-folds of FAR below the
# counting threshold, and the tool prints the last-e-fold fraction as a
# cross-check.  Set FAR_CALIBRATION to a number to override the fit (e.g. the
# reference-campaign value 0.037).  Refit whenever the ranking statistic or its
# FAR normalization changes.  See docs/user-guide/pastro.md.

BANK = H1L1V1-O4_MANIFOLD_BANK-0-2000000000.h5
BUILD_FLAGS = --processes 32 --window 50000 --match-dispersion 0.1 --verbose

#
# INJECTION DERIVED QUANTITIES
#

# SGNL injection campaign (stillsuit SQLite results database, .sqlite or
# .sqlite.gz): the rho(FAR) calibration and the chi fit.  A legacy FAR,SNR CSV
# (e.g. far_snr.txt) is accepted in its place.
INJECTIONS = far_snr.txt 
# INJECTIONS = results_sgnl_final.sqlite.gz
# recoveries a FAR decade needs to enter the rho(FAR) fit (300 suits ~1e5-1e6
# recoveries; this campaign has ~2500 over ~85 decades)
FARSNR_MIN_PER_DECADE = 100


COMPONENTS = bns.h5 nsbh.h5 bbh.h5 pop.h5
MODEL = pastro-model.h5

# runtime normalization
EVENT_RATE = 150 # events per year
EVENT_RATE_RANGE = 160 # Mpc
EVENT_RATE_FAR = 3.17e-8 # 1 / year in Hz (approx)
BNS_RANGE = 160 # expected range of the run that we are modeling

# empty: fit chi on $(INJECTIONS); a number: use it instead
FAR_CALIBRATION =
STAMP_FLAGS = --event-rate $(EVENT_RATE) --event-rate-range-mpc $(EVENT_RATE_RANGE) \
	--event-rate-far $(EVENT_RATE_FAR) --bns-range-mpc $(BNS_RANGE) \
	$(if $(FAR_CALIBRATION),--far-calibration $(FAR_CALIBRATION),--fit-far-calibration $(INJECTIONS))

# reference SGNL candidate (GraceDB coinc document) for `make check`
COINC = sgnl-coinc.xml
PLOT_PREFIX = o4

all : bank model plot plot-far check provenance

$(BANK) :
	wget https://dcc.ligo.org/public/0184/T2200343/003/H1L1V1-O4_MANIFOLD_BANK-0-2000000000.h5 

bank : $(BANK)

model : $(MODEL)

%.h5 : %.yaml $(BANK)
	sgn-manifold-cbc-bank-mass-model --bank $(BANK) --yaml $< --output-h5 $@ $(BUILD_FLAGS)

# combine + validate, then stamp.  Written under a temporary name so an
# interrupted run never leaves a half-built pastro-model.h5 behind.
$(MODEL) : $(COMPONENTS) $(INJECTIONS)
	sgn-manifold-cbc-p-astro-combine \
		--bns bns.h5 --nsbh nsbh.h5 --bbh bbh.h5 --pop pop.h5 \
		--far-snr $(INJECTIONS) --far-snr-min-per-decade $(FARSNR_MIN_PER_DECADE) \
		--output-h5 $@.tmp --verbose
	sgn-manifold-cbc-p-astro-configure --model $@.tmp $(STAMP_FLAGS)
	mv $@.tmp $@

# (re)stamp the normalization -- and refit chi on $(INJECTIONS) -- in place
stamp : $(MODEL) $(INJECTIONS)
	sgn-manifold-cbc-p-astro-configure --model $(MODEL) $(STAMP_FLAGS)

plot : $(MODEL)
	sgn-manifold-cbc-p-astro-plot --model $(MODEL) --far $(EVENT_RATE_FAR) --output-prefix $(PLOT_PREFIX)

# p_astro vs FAR operating curves for one representative template per region
plot-far : $(MODEL)
	for sel in "1.4 1.4 bns" "10 1.4 nsbh" "35 30 bbh"; do \
		set -- $$sel; \
		sgn-manifold-cbc-p-astro-plot --model $(MODEL) --m1 $$1 --m2 $$2 \
			--output-prefix $(PLOT_PREFIX)-$$3; \
	done

# The reference candidate through both consumers of the model file.
check : $(MODEL) $(COINC)
	sgn-manifold-cbc-p-astro --model $(MODEL) --coinc $(COINC) --output check-manifold.json
	sgnl-pastro-calculate-fgmc-coinc --model $(MODEL) $(COINC) > check-sgnl-pastro.jsonl
	python -c "import json; \
	    a = json.load(open('check-manifold.json')); \
	    b = [json.loads(l) for l in open('check-sgnl-pastro.jsonl')]; \
	    assert len(a) == len(b) > 0; \
	    keys = (('p_bns','BNS'),('p_nsbh','NSBH'),('p_bbh','BBH'),('p_noise','Terrestrial')); \
	    assert all(x[k] == y['p_astro'][n] for x, y in zip(a, b) for k, n in keys), 'MISMATCH'; \
	    assert all(sum(y['p_astro'].values()) == 1.0 for y in b); \
	    [print('  tid %d  FAR %.3g  ->  %s' % (x['diagnostics']['template_id'], x['combined_far'], y['p_astro'])) for x, y in zip(a, b)]; \
	    print('check: %d event(s), sgn-manifold == sgnl-pastro bit-for-bit, sums exactly 1' % len(a))"

# Package versions and the model's stored attributes, for the release record.
provenance : $(MODEL) $(BANK)
	{ \
	  echo "generated: $$(date -u +%Y-%m-%dT%H:%M:%SZ) on $$(hostname)"; \
	  echo "bank: $(BANK)"; \
	  echo "build flags: $(BUILD_FLAGS)"; \
	  echo "stamp flags: $(STAMP_FLAGS)"; \
	  echo; \
	  echo "## packages"; \
	  pip list 2>/dev/null | grep -iE "^(sgn-manifold|sgnl-pastro|sgnl|sgn|sgn-ts|sgn-ligo|gw-strike|lalsuite|igwn-ligolw|numpy|scipy|h5py) "; \
	  echo; \
	  echo "## model attributes"; \
	  python -c "import h5py, json; f = h5py.File('$(MODEL)', 'r'); \
	    [print('  %s = %s' % (k, json.loads(v) if k == 'warnings' else v)) for k, v in sorted(f.attrs.items())]; \
	    print('  templates =', len(f['template_id']))"; \
	} > versions.txt
	@cat versions.txt

# Remove generated files.  `clean` keeps the component models (they take HOURS
# to rebuild on the O4 bank); `clean-all` removes those too.
clean :
	rm -f $(MODEL) $(MODEL).tmp
	rm -f $(PLOT_PREFIX)_pastro_far-*.png
	rm -f $(PLOT_PREFIX)-bns_pastro-vs-far_*.png \
		$(PLOT_PREFIX)-nsbh_pastro-vs-far_*.png \
		$(PLOT_PREFIX)-bbh_pastro-vs-far_*.png
	rm -f check-manifold.json check-sgnl-pastro.jsonl
	rm -f versions.txt

clean-all : clean
	rm -f $(COMPONENTS)
	rm -f H1L1V1-O4_MANIFOLD_BANK-0-2000000000.h5

# Declare the existing files up to date (touch in dependency order): after a
# cosmetic yaml edit, or when the components are known to be current, so make
# does not launch an hours-long rebuild.  Verify first that each component's
# stored spec matches its yaml (mass_model_spec attribute).
mark-current :
	touch $(COMPONENTS)
	touch $(MODEL)

.PHONY : all model stamp plot plot-far check provenance mark-current clean clean-all