#!/bin/bash

set -euxo pipefail

# Computes pair questions using the RGB embedding.
PATCH_EMBEDDING_PBTXT="$PWD/experimental/james/embedding/patch_emb_1.pbtxt"
DATA_SOURCE_TR_PBTXT="$PWD/experimental/james/embedding/data_source_tr_1.pbtxt"
DATA_SOURCE_TE_PBTXT="$PWD/experimental/james/embedding/data_source_te_1.pbtxt"
CLASSIFIER_PBTXT="s3cache://zoox-testdata/vision/models/multiclass_v14_1_v4_trt.pbtxt"
DETECTIONS_TR_LMDB="/tmp/det_tr_1.lmdb"
DETECTIONS_TE_LMDB="/tmp/det_te_1.lmdb"
EMBEDDINGS_TR_LMDB="/tmp/emb_tr_1.lmdb"
EMBEDDINGS_TR_HTML="/tmp/emb_tr_1.html"
EMBEDDINGS_TE_LMDB="/tmp/emb_te_1.lmdb"
EMBEDDINGS_TE_HTML="/tmp/emb_te_1.html"
PAIRS_TR_LMDB="/tmp/pairs_tr_1.lmdb"
PAIRS_TE_LMDB="/tmp/pairs_te_1.lmdb"
PAIR_QUS_TR_LMDB="/tmp/pair_qus_tr_1.lmdb"
PAIR_QUS_TE_LMDB="/tmp/pair_qus_te_1.lmdb"
PAIR_QUS_TR_HTML="/tmp/pair_qus_tr_1.html"
PAIR_QUS_TE_HTML="/tmp/pair_qus_te_1.html"
# Labels to generate pairs for. 1003 = PED, 2000 = BIKE, 2001 = VEHICLE.
LABEL_IDS="1003,2000,2001"
# Re-balance the labels to get a good distribution of classes. Otherwise
# vehicles dominate.
LABEL_TR_NUM_PAIRS="1003:7500,2000:5000,2001:7500"
LABEL_TE_NUM_PAIRS="1003:750,2000:500,2001:750"

# Compute the detections.
#bazel run -c opt --config=cuda vision/detection/tools:run_detector_main -- \
#  --classifier="$CLASSIFIER_PBTXT" \
#  --data_source="$DATA_SOURCE_TR_PBTXT" \
#  --output="$DETECTIONS_TR_LMDB" \
#  --alsologtostderr
#
#bazel run -c opt --config=cuda vision/detection/tools:run_detector_main -- \
#  --classifier="$CLASSIFIER_PBTXT" \
#  --data_source="$DATA_SOURCE_TE_PBTXT" \
#  --output="$DETECTIONS_TE_LMDB" \
#  --alsologtostderr

# Compute the embeddings for each set.
#bazel run -c opt --config=cuda vision/embedding/tools:compute_embeddings -- \
#  --patch_embedding_pbtxt="$PATCH_EMBEDDING_PBTXT" \
#  --detection_lmdb="$DETECTIONS_TR_LMDB" \
#  --html="$EMBEDDINGS_TR_HTML" \
#  --lmdb="$EMBEDDINGS_TR_LMDB" \
#  --label_ids="$LABEL_IDS" \
#  --alsologtostderr
#
#bazel run -c opt --config=cuda vision/embedding/tools:compute_embeddings -- \
#  --patch_embedding_pbtxt="$PATCH_EMBEDDING_PBTXT" \
#  --detection_lmdb="$DETECTIONS_TE_LMDB" \
#  --html="$EMBEDDINGS_TE_HTML" \
#  --lmdb="$EMBEDDINGS_TE_LMDB" \
#  --label_ids="$LABEL_IDS" \
#  --alsologtostderr

# Compute the pairs for each set.
#bazel run -c opt --config=cuda vision/embedding/tools:compute_pairs -- \
#  --embedding_lmdb="$EMBEDDINGS_TR_LMDB" \
#  --lmdb="$PAIRS_TR_LMDB" \
#  --label_num_pairs="$LABEL_TR_NUM_PAIRS" \
#  --alsologtostderr
#
#bazel run -c opt --config=cuda vision/embedding/tools:compute_pairs -- \
#  --embedding_lmdb="$EMBEDDINGS_TE_LMDB" \
#  --lmdb="$PAIRS_TE_LMDB" \
#  --label_num_pairs="$LABEL_TE_NUM_PAIRS" \
#  --alsologtostderr

# Compute pair questions from the generated pairs.
bazel run -c opt --config=cuda vision/embedding/tools:compute_pair_questions -- \
  --pair_lmdb="$PAIRS_TR_LMDB" \
  --detection_lmdb="$DETECTIONS_TR_LMDB" \
  --lmdb="$PAIR_QUS_TR_LMDB" \
  --html="$PAIR_QUS_TR_HTML" \
  --alsologtostderr

bazel run -c opt --config=cuda vision/embedding/tools:compute_pair_questions -- \
  --pair_lmdb="$PAIRS_TE_LMDB" \
  --detection_lmdb="$DETECTIONS_TE_LMDB" \
  --lmdb="$PAIR_QUS_TE_LMDB" \
  --html="$PAIR_QUS_TE_HTML" \
  --alsologtostderr
