SHELL := /bin/bash
PYTHON ?= python
ATLAS ?= schaefer200
ATLAS_SOURCE ?= auto
MAX_CANDIDATES ?= 200
MIN_VOXELS ?= 20
PROMPT_TOP_K ?= 30
SELECTION_BUDGET ?= 20
SEED ?= 42
EPISODES ?= 8
PARCEL_MANIFEST ?= configs/parcel_candidates.json
GRPO_OUTPUT_DIR ?= outputs/grpo_brainrl
GRPO_MODEL ?= Qwen/Qwen2.5-0.5B-Instruct
GRPO_MAX_STEPS ?= 50
GRPO_BATCH_SIZE ?= 4
GRPO_NUM_GENERATIONS ?= 4
GRPO_MAX_PROMPT_LENGTH ?= 1792
GRPO_MAX_COMPLETION_LENGTH ?= 64
GRPO_TRAIN_PROMPT_TOP_K ?= 16
GRPO_TEMPERATURE ?= 1.2
GRPO_TOP_P ?= 0.95
GRPO_DUPLICATE_ACTION_PENALTY ?= 0.1
GRPO_REWARD_DELTA_SCALE ?= 5.0
CUDA_VISIBLE_DEVICES ?= 0

CONDITION ?= single_m
PARTICIPANT_INFO ?= configs/participant_run_info.json
TRAIN_SUBJECTS ?= sub-01:sub-20
TEST_SUBJECTS ?= sub-21:sub-26
EXCLUDE_SUBJECTS ?= sub-03,sub-18
STIMULUS_DIR ?= /home/mohith/ds005345/data/annotation
STIMULUS_WINDOW_SIZE ?= 30
STIMULUS_TOP_WORDS ?= 10
EVAL_OUT ?= outputs/eval/$(CONDITION)
TRAIN_OUT ?= outputs/grpo_$(CONDITION)
HF_REPO_ID ?= you/brainrl-region-selection
HF_VISIBILITY ?= --public
HF_DATA_REPO ?= you/brainrl-config-data
HF_MODEL_REPO ?= you/brainrl-grpo-single-m
HF_DATA_REVISION ?= main
HF_JOB_FLAVOR ?= a10g-small
HF_ENV_URL ?=

.PHONY: all prepare prepare-easy dryrun dryrun-single inference inference-llm \
        inference-single eval eval-llm eval-single-train eval-single-test \
        train-single grpo grpo-easy server clean help deploy deploy-dry \
        hf-data-export hf-data-upload hf-data-download hf-job-command hf-job-launch

all: prepare dryrun-single eval-single-test

help:
	@echo "Common targets (CONDITION=$(CONDITION), TRAIN=$(TRAIN_SUBJECTS), TEST=$(TEST_SUBJECTS), EXCLUDE=$(EXCLUDE_SUBJECTS), STIMULUS_DIR=$(STIMULUS_DIR), WINDOW_SIZE=$(STIMULUS_WINDOW_SIZE)):"
	@echo "  make prepare              - Build configs/parcel_candidates.json (Schaefer-200, $(MAX_CANDIDATES) parcels)"
	@echo "  make dryrun               - Plain GRPO verifier dry-run"
	@echo "  make dryrun-single        - Dry-run scoped to CONDITION + TRAIN_SUBJECTS"
	@echo "  make inference            - Static prompt-policy rollout (any condition)"
	@echo "  make inference-llm        - LLM prompt-policy rollout (requires HF_TOKEN)"
	@echo "  make inference-single     - Rollout one episode of CONDITION on test subject"
	@echo "  make eval                 - Compare baselines across all conditions"
	@echo "  make eval-llm             - Add LLM prompt policy to baseline eval"
	@echo "  make eval-single-train    - Compare baselines on CONDITION/train split + plots"
	@echo "  make eval-single-test     - Compare baselines on CONDITION/test split + plots"
	@echo "  make train-single         - GRPO training on CONDITION using TRAIN_SUBJECTS + plots"
	@echo "  make grpo                 - Tiny GRPO smoke run on the full subset"
	@echo "  make server               - Boot the OpenEnv FastAPI server on :8000"
	@echo "  make deploy-dry           - Dry-run HF Spaces deploy plan"
	@echo "  make deploy               - Deploy server.app to HF Spaces ($(HF_REPO_ID))"
	@echo "  make hf-data-upload       - Upload config/stimulus bundle to HF Dataset ($(HF_DATA_REPO))"
	@echo "  make hf-job-command       - Print HF Jobs training command"
	@echo "  make clean                - Drop generated manifests and outputs"

prepare:
	$(PYTHON) prepare_parcels.py \
	  --atlas $(ATLAS) \
	  --atlas-source $(ATLAS_SOURCE) \
	  --max-candidates $(MAX_CANDIDATES) \
	  --min-voxels $(MIN_VOXELS) \
	  --prompt-top-k $(PROMPT_TOP_K) \
	  --selection-budget $(SELECTION_BUDGET) \
	  --seed $(SEED) \
	  --output $(PARCEL_MANIFEST)

prepare-easy:
	@echo "Easy curriculum uses the legacy 7-ROI config. Edit configs/subset_config.yaml:"
	@echo "  candidate_mode: roi_priors"
	@echo "  selection_budget: 5"
	@echo "Then run 'make eval' / 'make dryrun' as usual (no manifest needed)."

dryrun:
	$(PYTHON) train_grpo.py --dry-run

dryrun-single:
	BRAINRL_STIMULUS_DIR=$(STIMULUS_DIR) \
	BRAINRL_STIMULUS_WINDOW_SIZE=$(STIMULUS_WINDOW_SIZE) \
	BRAINRL_STIMULUS_TOP_WORDS=$(STIMULUS_TOP_WORDS) \
	$(PYTHON) train_grpo.py --dry-run \
	  --condition $(CONDITION) \
	  --participant-info $(PARTICIPANT_INFO) \
	  --train-subjects $(TRAIN_SUBJECTS) \
	  --test-subjects $(TEST_SUBJECTS) \
	  --exclude-subjects $(EXCLUDE_SUBJECTS)

inference:
	$(PYTHON) inference.py --show-prompts

inference-llm:
	$(PYTHON) inference.py --use-llm --show-prompts

inference-single:
	BRAINRL_STIMULUS_DIR=$(STIMULUS_DIR) \
	BRAINRL_STIMULUS_WINDOW_SIZE=$(STIMULUS_WINDOW_SIZE) \
	BRAINRL_STIMULUS_TOP_WORDS=$(STIMULUS_TOP_WORDS) \
	$(PYTHON) inference.py \
	  --condition $(CONDITION) \
	  --participant-info $(PARTICIPANT_INFO) \
	  --train-subjects $(TRAIN_SUBJECTS) \
	  --test-subjects $(TEST_SUBJECTS) \
	  --exclude-subjects $(EXCLUDE_SUBJECTS) \
	  --split test \
	  --episodes 1

eval:
	$(PYTHON) evaluate.py --episodes $(EPISODES) --seed $(SEED)

eval-llm:
	$(PYTHON) evaluate.py --episodes $(EPISODES) --seed $(SEED) --use-llm

eval-single-train:
	BRAINRL_STIMULUS_DIR=$(STIMULUS_DIR) \
	BRAINRL_STIMULUS_WINDOW_SIZE=$(STIMULUS_WINDOW_SIZE) \
	BRAINRL_STIMULUS_TOP_WORDS=$(STIMULUS_TOP_WORDS) \
	$(PYTHON) evaluate.py \
	  --episodes $(EPISODES) \
	  --seed $(SEED) \
	  --condition $(CONDITION) \
	  --participant-info $(PARTICIPANT_INFO) \
	  --train-subjects $(TRAIN_SUBJECTS) \
	  --test-subjects $(TEST_SUBJECTS) \
	  --exclude-subjects $(EXCLUDE_SUBJECTS) \
	  --split train \
	  --output-csv $(EVAL_OUT)/baselines_train.csv \
	  --plot-dir $(EVAL_OUT)/plots_train

eval-single-test:
	BRAINRL_STIMULUS_DIR=$(STIMULUS_DIR) \
	BRAINRL_STIMULUS_WINDOW_SIZE=$(STIMULUS_WINDOW_SIZE) \
	BRAINRL_STIMULUS_TOP_WORDS=$(STIMULUS_TOP_WORDS) \
	$(PYTHON) evaluate.py \
	  --episodes $(EPISODES) \
	  --seed $(SEED) \
	  --condition $(CONDITION) \
	  --participant-info $(PARTICIPANT_INFO) \
	  --train-subjects $(TRAIN_SUBJECTS) \
	  --test-subjects $(TEST_SUBJECTS) \
	  --exclude-subjects $(EXCLUDE_SUBJECTS) \
	  --split test \
	  --output-csv $(EVAL_OUT)/baselines_test.csv \
	  --plot-dir $(EVAL_OUT)/plots_test

train-single:
	PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True CUDA_VISIBLE_DEVICES=$(CUDA_VISIBLE_DEVICES) \
	BRAINRL_STIMULUS_DIR=$(STIMULUS_DIR) \
	BRAINRL_STIMULUS_WINDOW_SIZE=$(STIMULUS_WINDOW_SIZE) \
	BRAINRL_STIMULUS_TOP_WORDS=$(STIMULUS_TOP_WORDS) \
	$(PYTHON) train_grpo.py \
	  --condition $(CONDITION) \
	  --participant-info $(PARTICIPANT_INFO) \
	  --train-subjects $(TRAIN_SUBJECTS) \
	  --test-subjects $(TEST_SUBJECTS) \
	  --exclude-subjects $(EXCLUDE_SUBJECTS) \
	  --model-name $(GRPO_MODEL) \
	  --max-steps $(GRPO_MAX_STEPS) \
	  --train-prompt-top-k $(GRPO_TRAIN_PROMPT_TOP_K) \
	  --per-device-train-batch-size $(GRPO_BATCH_SIZE) \
	  --num-generations $(GRPO_NUM_GENERATIONS) \
	  --max-prompt-length $(GRPO_MAX_PROMPT_LENGTH) \
	  --max-completion-length $(GRPO_MAX_COMPLETION_LENGTH) \
	  --temperature $(GRPO_TEMPERATURE) \
	  --top-p $(GRPO_TOP_P) \
	  --duplicate-action-penalty $(GRPO_DUPLICATE_ACTION_PENALTY) \
	  --reward-delta-scale $(GRPO_REWARD_DELTA_SCALE) \
	  --output-dir $(TRAIN_OUT) \
	  --plot-dir $(TRAIN_OUT)/plots

grpo:
	$(PYTHON) train_grpo.py \
	  --model-name $(GRPO_MODEL) \
	  --max-steps $(GRPO_MAX_STEPS) \
	  --output-dir $(GRPO_OUTPUT_DIR)

server:
	uvicorn server.app:app --host 0.0.0.0 --port 8000

deploy:
	$(PYTHON) deploy_to_hf.py --repo-id $(HF_REPO_ID) $(HF_VISIBILITY) --include-manifest

deploy-dry:
	$(PYTHON) deploy_to_hf.py --repo-id $(HF_REPO_ID) --dry-run --include-manifest

hf-data-export:
	$(PYTHON) hf_data.py export --output-dir hf_data_bundle --annotation-dir $(STIMULUS_DIR)

hf-data-upload:
	$(PYTHON) hf_data.py upload --repo-id $(HF_DATA_REPO) --public --annotation-dir $(STIMULUS_DIR)

hf-data-download:
	$(PYTHON) hf_data.py download --repo-id $(HF_DATA_REPO) --revision $(HF_DATA_REVISION)

hf-job-command:
	$(PYTHON) hf_jobs.py \
	  --data-repo $(HF_DATA_REPO) \
	  --model-repo $(HF_MODEL_REPO) \
	  --space-repo $(HF_REPO_ID) \
	  --data-revision $(HF_DATA_REVISION) \
	  --model-name $(GRPO_MODEL) \
	  --condition $(CONDITION) \
	  --train-subjects $(TRAIN_SUBJECTS) \
	  --test-subjects $(TEST_SUBJECTS) \
	  --exclude-subjects $(EXCLUDE_SUBJECTS) \
	  --max-steps $(GRPO_MAX_STEPS) \
	  --flavor $(HF_JOB_FLAVOR) \
	  $(if $(HF_ENV_URL),--env-url $(HF_ENV_URL),) \
	  --secret-hf-token

hf-job-launch:
	$(PYTHON) hf_jobs.py \
	  --data-repo $(HF_DATA_REPO) \
	  --model-repo $(HF_MODEL_REPO) \
	  --space-repo $(HF_REPO_ID) \
	  --data-revision $(HF_DATA_REVISION) \
	  --model-name $(GRPO_MODEL) \
	  --condition $(CONDITION) \
	  --train-subjects $(TRAIN_SUBJECTS) \
	  --test-subjects $(TEST_SUBJECTS) \
	  --exclude-subjects $(EXCLUDE_SUBJECTS) \
	  --max-steps $(GRPO_MAX_STEPS) \
	  --flavor $(HF_JOB_FLAVOR) \
	  $(if $(HF_ENV_URL),--env-url $(HF_ENV_URL),) \
	  --secret-hf-token \
	  --launch

clean:
	rm -rf $(PARCEL_MANIFEST) outputs/ __pycache__ */__pycache__