SAMPLE_DIR := samples
BW_DIR := $(SAMPLE_DIR)/trajectory/backward
LYAP_DIR := $(SAMPLE_DIR)/lyap-exp
LVEC_DIR := $(SAMPLE_DIR)/lyap-vec

BASE_FILES := main.py simulator.py models/base_sampler.py

PERTURBS = 0.1 0.5 1.0

SINGULAR_ID := singular-blob
SINGULAR_PT := \
	$(BW_DIR)/$(SINGULAR_ID).pt \
	$(LYAP_DIR)/$(SINGULAR_ID).pt \
	$(LVEC_DIR)/$(SINGULAR_ID).pt \

SINGULAR_PERTURBED_PT := \
	$(foreach p,$(PERTURBS),\
        $(BW_DIR)/$(SINGULAR_ID)-$(p).pt \
		$(LYAP_DIR)/$(SINGULAR_ID)-$(p).pt \
		$(LVEC_DIR)/$(SINGULAR_ID)-$(p).pt \
    )

HALFMOONS_ID := halfmoons
HALFMOONS_PT := \
	$(BW_DIR)/$(HALFMOONS_ID).pt \
	$(LYAP_DIR)/$(HALFMOONS_ID).pt \
	$(LVEC_DIR)/$(HALFMOONS_ID).pt \

HALFMOONS_PERTURBED_PT := \
	$(foreach p,$(PERTURBS),\
        $(BW_DIR)/$(HALFMOONS_ID)-$(p).pt \
		$(LYAP_DIR)/$(HALFMOONS_ID)-$(p).pt \
		$(LVEC_DIR)/$(HALFMOONS_ID)-$(p).pt \
    )

MNIST_ID := mnist
MNIST_LYAP := \
	$(BW_DIR)/$(MNIST_ID).pt \
	$(LYAP_DIR)/$(MNIST_ID).pt \
	$(LVEC_DIR)/$(MNIST_ID).pt \

MNIST_SAMPS := $(BW_DIR)/$(MNIST_ID)-samples.pt

all: singular halfmoons mnist

singular: $(SINGULAR_PT) $(SINGULAR_PERTURBED_PT) plt-singular-blob.py
	@echo "Plotting blob"
	python3 plt-singular-blob.py

$(SINGULAR_PT): $(BASE_FILES) models/singular_2d.py 
	@echo "Generating backward trajectory with Lyapunov Spectrum\n"
	python3 main.py --model Singular2D --number-of-noise-realizations 3000 \
	--output $(SINGULAR_ID) --T 0.9 --n-grid 1000 --conf-file singular_blob.json \
	--noise-schedule cosine

$(SINGULAR_PERTURBED_PT): $(BASE_FILES) models/singular_2d.py
	@echo "Processing multiple perturbation values...\n"
	for p in $(PERTURBS); do \
		echo "Perturbation strength: $$p"; \
		python3 main.py --model Singular2D --number-of-noise-realizations 3000 \
		--output $(SINGULAR_ID)-$$p --T 0.9 --n-grid 1000 --conf-file singular_blob.json \
		--noise-schedule cosine --perturb-size $$p;\
	done
	@echo "Finished processing the list of perturbation values.\n"

halfmoons: $(HALFMOONS_PT) $(HALFMOONS_PERTURBED_PT) plt-halfmoons.py
	@echo "Plotting halfmoons"
	python3 plt-halfmoons.py

$(HALFMOONS_PT): $(BASE_FILES) models/half_moons.py 
	@echo "Generating backward trajectory with Lyapunov Spectrum\n"
	python3 main.py --model HalfMoons --number-of-noise-realizations 3000 \
	--output $(HALFMOONS_ID) --T 0.9 --n-grid 1000 \
	--noise-schedule cosine

$(HALFMOONS_PERTURBED_PT): $(BASE_FILES) models/half_moons.py
	@echo "Processing multiple perturbation values...\n"
	for p in $(PERTURBS); do \
		echo "Perturbation strength: $$p"; \
		python3 main.py --model HalfMoons --number-of-noise-realizations 3000 \
		--output $(HALFMOONS_ID)-$$p --T 0.9 --n-grid 1000 \
		--noise-schedule cosine --perturb-size $$p;\
	done
	@echo "Finished processing the list of perturbation values.\n"

mnist: $(MNIST_LYAP) $(MNIST_SAMPS) plt-mnist.py
	@echo "Plotting mnist"
	python3 plt-mnist.py

$(MNIST_LYAP): $(BASE_FILES) models/trained_images.py models/base_trained.py models/unet.py models/attend.py models/weights/ScoreModelImage.mnist_digits.pt
	@echo "Generating backward trajectory with Lyapunov Spectrum"
	@echo "WARNING: This takes a while\n"
	python3 main.py --model TrainedImage --number-of-noise-realizations 1 \
	--output $(MNIST_ID) --T 0.9 --n-grid 4000 --conf-file mnist.json \
	--noise-schedule cosine

$(MNIST_SAMPS): $(BASE_FILES) models/trained_images.py models/base_trained.py models/unet.py models/attend.py models/weights/ScoreModelImage.mnist_digits.pt
	@echo "Generating a few samples from MNIST"
	python3 main.py --model TrainedImage --number-of-noise-realizations 20 \
	--output $(MNIST_ID)-samples --T 0.9 --n-grid 4000 --conf-file mnist.json \
	--noise-schedule cosine --no-spectrum

models/weights/ScoreModelImage.mnist_digits.pt: trainer.py
	python3 trainer.py