#!/usr/bin/env bash
#
# Permission to use, copy, modify, and/or distribute this software for
# any purpose with or without fee is hereby granted.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

# sweep-klein-mask.sh -- disocclusion-mask A/B for klein dual-ref reprojection.
#
# Context (NOTES.md, "Klein dual-ref GPU validation", 2026-09-13): under
# dual-ref WITHOUT reprojection the loop freezes into a static consensus --
# the flow warp is the only motion carrier, so it is load-bearing. The cost
# is warp quality: grayscale Farneback + BORDER_REPLICATE smears disoccluded
# regions, and klein faithfully preserves those smears as "appearance".
#
# This driver A/Bs --reproject-mask (imaging.make_reprojector's forward-
# backward consistency masking): pixels whose round-trip flow disagrees
# (> 1.5 px) or that sample outside the state frame are patched from the
# FRESH FRAME instead of being replicate-smeared.
#
#   leg fbmask : --reproject --reproject-mask   (run tag ..._dualref-fbmask)
#   leg plain  : --reproject                    (control, ..._dualref)
#
# Both legs: --mode stateful --conditioning dual-ref --reproject, identical
# prompt/steps/seed/geometry. The fbmask leg runs FIRST so the new variant
# is on disk early for mid-sweep inspection. Legs land in separate subtrees
# (untracked/output/<model>/[/prefix/]mask-ab/{fbmask,plain}/<stem>_dualref[-fbmask]_...)
# so nothing collides.
#
# Usage (from any directory inside the repo):
#   scripts/sweep-klein-mask.sh
#   DRY_RUN=1 scripts/sweep-klein-mask.sh
#   MODEL=flux2-klein-4b scripts/sweep-klein-mask.sh     # ungated, faster
#   PROMPT="" scripts/sweep-klein-mask.sh                # neutral leg
#
# Environment overrides:
#   MODEL          klein model key      (default flux2-klein-9b; 4b ungated)
#   CLIP           source video         (default: untracked/input/inputs.env if
#                                        present, else input_example/
#                                        video_cropped.mp4; scripts/inputs.sh)
#   PROMPT         edit prompt (default: role-naming enhance-slight; EMPTY
#                  string = neutral/no-prompt leg). Indices follow REF_ORDER.
#   REF_ORDER      state-first | frame-first (default state-first)
#   STEPS          --num-inference-steps (default 4, the klein card default)
#   MAX_FRAMES     source frames per run (default 300; empty = whole clip)
#   TAIL_FRAMES    frames per tail       (default 60; 0 = no tails)
#   TAIL_MODES     tail scenario list    (default freeze)
#   SAVE_EVERY     save every Nth frame  (default 10)
#   OUT_PREFIX     extra path component before mask-ab/ (default none)
#   EXTRA_ARGS     extra flags, word-split, appended to every run
#   DRY_RUN=1      print commands without executing anything
#   SKIP_GPU_CHECK=1  bypass the accelerator preflight
#
# Log: untracked/output/sweep_klein_mask_<UTC timestamp>.log
# Stops at the first failing run (set -euo pipefail).

set -euo pipefail

cd "$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"

# Input files: env > untracked/input/inputs.env (user) > input_example/inputs.env.
source scripts/inputs.sh

MODEL="${MODEL:-flux2-klein-9b}"
CLIP="${CLIP:-input_example/video_cropped.mp4}"
REF_ORDER="${REF_ORDER:-state-first}"
STEPS="${STEPS:-4}"
MAX_FRAMES="${MAX_FRAMES-300}"
TAIL_FRAMES="${TAIL_FRAMES:-60}"
TAIL_MODES="${TAIL_MODES:-freeze}"
SAVE_EVERY="${SAVE_EVERY:-10}"
OUT_PREFIX="${OUT_PREFIX:-}"
EXTRA_ARGS="${EXTRA_ARGS:-}"
DRY_RUN="${DRY_RUN:-0}"

case "$REF_ORDER" in
    state-first|frame-first) ;;
    *) echo "sweep-klein-mask: REF_ORDER must be state-first or frame-first (got '$REF_ORDER')" >&2; exit 1 ;;
esac

# Role-naming indices mirror scripts/sweep-klein.sh (REF_ORDER is known not
# to affect motion -- kept configurable for texture-side comparability).
if [[ "$REF_ORDER" == "state-first" ]]; then
    STATE_IMG=1
    FRAME_IMG=2
else
    STATE_IMG=2
    FRAME_IMG=1
fi
DEFAULT_PROMPT="image ${FRAME_IMG} is the current frame; keep the appearance of image ${STATE_IMG}, slightly enhancing fine details"
# Unset -> default prompt; explicitly empty -> neutral (no --prompt flag).
PROMPT="${PROMPT-${DEFAULT_PROMPT}}"

OUT_PREFIX="${OUT_PREFIX#/}"; OUT_PREFIX="${OUT_PREFIX%/}"
OUT_ROOT="untracked/output/$MODEL${OUT_PREFIX:+/$OUT_PREFIX}/mask-ab"

# -------------------------------------------------------------- preflight ----
if [[ "$DRY_RUN" != "1" ]]; then
    [[ -f "$CLIP" ]] || { echo "sweep-klein-mask: clip not found: $CLIP" >&2; exit 1; }
    command -v uv >/dev/null || { echo "sweep-klein-mask: 'uv' not on PATH" >&2; exit 1; }

    if [[ "$MODEL" == "flux2-klein-9b" && -z "${HF_TOKEN:-}" ]]; then
        echo "sweep-klein-mask: HF_TOKEN is not set (required for flux2-klein-9b)." >&2
        echo "             Accept the license at https://huggingface.co/black-forest-labs/FLUX.2-klein-9B" >&2
        echo "             then: export HF_TOKEN=hf_... and re-run." >&2
        exit 1
    fi

    if [[ "${SKIP_GPU_CHECK:-0}" != "1" ]]; then
        if ! uv run python -c 'import sys, torch; sys.exit(0 if (torch.cuda.is_available() or torch.backends.mps.is_available()) else 1)'; then
            echo "sweep-klein-mask: torch reports no CUDA or MPS device -- run this on a GPU pod or Apple Silicon" >&2
            echo "             (SKIP_GPU_CHECK=1 to bypass, DRY_RUN=1 to preview)." >&2
            exit 1
        fi
    fi
fi

mkdir -p untracked/output
LOG="untracked/output/sweep_klein_mask_$(date -u +%Y%m%d-%H%M%S).log"

common=(--model "$MODEL" --input "$CLIP" --save-every "$SAVE_EVERY"
        --mode stateful --conditioning dual-ref --ref-order "$REF_ORDER"
        --reproject --num-inference-steps "$STEPS")
if [[ -n "$MAX_FRAMES" ]]; then common+=(--max-frames "$MAX_FRAMES"); fi
if [[ "$TAIL_FRAMES" -gt 0 ]]; then
    common+=(--tail-frames "$TAIL_FRAMES" --tail-modes "$TAIL_MODES")
fi

log() { echo "$@" | tee -a "$LOG"; }

run() {
    log ""
    log "== $(date -u +%Y-%m-%dT%H:%M:%SZ) dltb-klein $*"
    if [[ "$DRY_RUN" == "1" ]]; then
        log "DRY RUN"
        return 0
    fi
    uv run dltb-klein "$@" 2>&1 | tee -a "$LOG"
}

log "sweep-klein-mask: model=$MODEL clip=$CLIP ref_order=$REF_ORDER steps=$STEPS"
log "sweep-klein-mask: prompt='${PROMPT:-<empty>}' max_frames=${MAX_FRAMES:-<all>} tail=${TAIL_FRAMES}x${TAIL_MODES} out_root=$OUT_ROOT"
log "sweep-klein-mask: A/B = --reproject-mask on (fbmask) vs off (plain); reproject always on"
log "sweep-klein-mask: log=$LOG"

# --------------------------------------------------------------- A/B legs ----
for MASK in 1 0; do
    if [[ "$MASK" == "1" ]]; then leg="fbmask"; else leg="plain"; fi
    args=(${common[@]+"${common[@]}"} --output-dir "$OUT_ROOT/$leg")
    if [[ "$MASK" == "1" ]]; then args+=(--reproject-mask); fi
    if [[ -n "$PROMPT" ]]; then args+=(--prompt "$PROMPT"); fi

    log ""
    log "########## leg '$leg' [reproject-mask=$MASK]: ${PROMPT:-<empty>} ##########"
    run "${args[@]}" ${EXTRA_ARGS:+$EXTRA_ARGS}
    log "sweep-klein-mask: leg done: $OUT_ROOT/$leg (safe to copy off mid-sweep)"
done

log ""
log "sweep-klein-mask: done -- log: $LOG"