#!/bin/bash
# Download and extract HEVC conformance bitstreams from ITU
# Usage: ./conformance/fetch.sh [VECTOR_NAME...]
# If no names given, downloads a core inter-prediction subset.
#
# Files are cached in conformance/vectors/ — only downloads once.
# Each vector extracts to a directory containing the .bin bitstream
# and a reference .yuv (generated by dec265 if available).

set -euo pipefail

VECTORS_DIR="$(dirname "$0")/vectors"
ITU_BASE="https://www.itu.int/wftp3/av-arch/jctvc-site/bitstream_exchange/draft_conformance/HEVC_v1"
DEC265="/home/lilith/work/heic/libde265-src/build/dec265/dec265"

mkdir -p "$VECTORS_DIR"

# Core inter-prediction conformance vectors
CORE_INTER=(
    # AMVP (Advanced Motion Vector Prediction)
    AMVP_A_MTK_4
    AMVP_B_MTK_4
    AMVP_C_Samsung_7
    # Merge mode
    MERGE_A_TI_3
    MERGE_B_TI_3
    MERGE_C_TI_3
    MERGE_D_TI_3
    MERGE_E_TI_3
    MERGE_F_MTK_4
    MERGE_G_HHI_4
    # Temporal MVP
    TMVP_A_MS_3
    # AMP (Asymmetric Motion Partition)
    AMP_A_Samsung_7
    AMP_B_Samsung_7
    AMP_D_Hisilicon_3
    AMP_E_Hisilicon_3
    AMP_F_Hisilicon_3
    # MV clipping/edge
    MVCLIP_A_qualcomm_3
    MVDL1ZERO_A_docomo_4
    MVEDGE_A_qualcomm_3
    # Prediction merge
    PMERGE_A_TI_3
    PMERGE_B_TI_3
    PMERGE_C_TI_3
    PMERGE_D_TI_3
    PMERGE_E_TI_3
    # Weighted prediction
    WP_A_Toshiba_3
    WP_B_Toshiba_3
    # Reference picture sets
    RPS_A_docomo_5
    RPS_B_qualcomm_5
    RPS_C_ericsson_5
    RPS_D_ericsson_6
    RPS_E_qualcomm_5
    RPS_F_docomo_5
    # Random access point
    RAP_A_docomo_6
    RAP_B_Bossen_2
    # Deblocking
    DBLK_A_SONY_3
    DBLK_B_SONY_3
    DBLK_C_SONY_3
    DBLK_D_VIXS_2
    DBLK_E_VIXS_2
    DBLK_F_VIXS_2
    DBLK_G_VIXS_2
    # SAO
    SAO_A_MediaTek_4
    SAO_B_MediaTek_5
    SAO_C_Samsung_5
    SAO_D_Samsung_5
    SAO_E_Canon_4
    SAO_F_Canon_3
    SAO_G_Canon_3
    SAO_H_Parabola_1
    # Intra prediction (baseline)
    IPRED_A_Qualcomm_3
    IPRED_B_Nokia_3
    IPRED_C_HHI_3
    # Transform
    RQT_A_HHI_4
    RQT_B_HHI_4
    RQT_C_HHI_4
    RQT_D_HHI_4
    RQT_E_HHI_4
    RQT_F_HHI_4
    RQT_G_HHI_4
    # Slices
    SLICES_A_Rovi_3
    # Sign data hiding
    SDH_A_Orange_4
    # Scaling lists
    SLIST_A_Sony_5
    SLIST_B_Sony_9
    SLIST_C_Sony_4
    SLIST_D_Sony_9
    # Tiles
    TILES_A_Cisco_2
    TILES_B_Cisco_1
    # CIP, conformance window, etc.
    CIP_A_Panasonic_3
    CONFWIN_A_Sony_1
    # Struct
    STRUCT_A_Samsung_7
    STRUCT_B_Samsung_7
    # POC
    POC_A_Bossen_3
    # Delta QP
    DELTAQP_A_BRCM_4
    DELTAQP_B_SONY_3
    DELTAQP_C_SONY_3
    # WPP
    WPP_A_ericsson_MAIN_2
    WPP_B_ericsson_MAIN_2
    WPP_C_ericsson_MAIN_2
    WPP_D_ericsson_MAIN_2
    WPP_E_ericsson_MAIN_2
    WPP_F_ericsson_MAIN_2
    # Misc
    NUT_A_ericsson_5
    PICSIZE_A_Bossen_1
    PICSIZE_B_Bossen_1
    PICSIZE_C_Bossen_1
    PICSIZE_D_Bossen_1
    BUMPING_A_ericsson_1
    CAINIT_A_SHARP_4
    CAINIT_B_SHARP_4
    CAINIT_C_SHARP_3
    CAINIT_D_SHARP_3
    CAINIT_E_SHARP_3
    CAINIT_F_SHARP_3
    CAINIT_G_SHARP_3
    CAINIT_H_SHARP_3
    TUSIZE_A_Samsung_1
    TSKIP_A_MS_3
)

# Use command-line args if given, otherwise use CORE_INTER
if [ $# -gt 0 ]; then
    VECTORS=("$@")
else
    VECTORS=("${CORE_INTER[@]}")
fi

download_one() {
    local name="$1"
    local dir="$VECTORS_DIR/$name"

    if [ -d "$dir" ] && ls "$dir"/*.bin "$dir"/*.265 2>/dev/null | head -1 >/dev/null 2>&1; then
        return 0  # Already downloaded
    fi

    echo "Downloading $name..."
    local zip="/tmp/${name}.zip"
    curl -sSfL "${ITU_BASE}/${name}.zip" -o "$zip" 2>/dev/null || {
        echo "  FAILED to download $name"
        return 1
    }

    mkdir -p "$dir"
    unzip -o -q "$zip" -d "$dir" 2>/dev/null || {
        echo "  FAILED to extract $name"
        rm -f "$zip"
        return 1
    }
    rm -f "$zip"

    # Find the bitstream file (could be .bin, .265, or .bit)
    local bitstream
    bitstream=$(find "$dir" -name "*.bit" -o -name "*.bin" -o -name "*.265" | head -1)
    if [ -z "$bitstream" ]; then
        echo "  WARNING: no bitstream found in $name"
        return 1
    fi

    # Generate reference YUV using dec265 if available
    if [ -x "$DEC265" ] && [ ! -f "$dir/reference.yuv" ]; then
        echo "  Generating reference YUV with dec265..."
        "$DEC265" "$bitstream" -o "$dir/reference.yuv" --quiet 2>/dev/null || {
            echo "  WARNING: dec265 failed for $name (may use unsupported features)"
        }
    fi

    # Also get frame info via ffprobe
    if command -v ffprobe &>/dev/null && [ ! -f "$dir/info.txt" ]; then
        ffprobe -v error -select_streams v:0 \
            -show_entries stream=width,height,pix_fmt,nb_frames \
            -show_entries frame=pict_type \
            -of csv=p=0 "$bitstream" > "$dir/info.txt" 2>/dev/null || true
    fi

    echo "  OK: $(basename "$bitstream")"
}

FAILED=0
TOTAL=0
for name in "${VECTORS[@]}"; do
    TOTAL=$((TOTAL + 1))
    download_one "$name" || FAILED=$((FAILED + 1))
done

echo ""
echo "Downloaded $((TOTAL - FAILED))/$TOTAL vectors to $VECTORS_DIR"
if [ $FAILED -gt 0 ]; then
    echo "($FAILED failed)"
fi
