#!/bin/bash
# File Size Comparison Script for Quantized Checkpoints
#
# Compares FP32 vs INT8 checkpoint sizes across all models.
#
# Usage:
#   ./ml/scripts/compare_checkpoint_sizes.sh

set -e

CHECKPOINT_DIR="ml/trained_models"

echo "=== Checkpoint Size Comparison (FP32 vs INT8) ==="
echo ""

# Function to get file size in MB
get_size_mb() {
    local file="$1"
    if [ -f "$file" ]; then
        du -h "$file" | awk '{print $1}'
    else
        echo "N/A"
    fi
}

# Function to calculate percentage reduction
calc_reduction() {
    local fp32_size="$1"
    local int8_size="$2"

    if [ "$fp32_size" != "N/A" ] && [ "$int8_size" != "N/A" ]; then
        # Strip units and calculate
        fp32_bytes=$(du -b "$fp32_size" 2>/dev/null | awk '{print $1}')
        int8_bytes=$(du -b "$int8_size" 2>/dev/null | awk '{print $1}')

        if [ "$fp32_bytes" -gt 0 ]; then
            reduction=$(echo "scale=1; 100 * (1 - $int8_bytes / $fp32_bytes)" | bc)
            echo "${reduction}%"
        else
            echo "N/A"
        fi
    else
        echo "N/A"
    fi
}

echo "Model          | FP32 Size | INT8 Size | Reduction"
echo "---------------|-----------|-----------|----------"

# DQN checkpoints
DQN_FP32="${CHECKPOINT_DIR}/dqn_final_epoch3.safetensors"
DQN_INT8="${CHECKPOINT_DIR}/dqn_quantized_demo.safetensors"

if [ -f "$DQN_FP32" ] || [ -f "$DQN_INT8" ]; then
    fp32_size=$(get_size_mb "$DQN_FP32")
    int8_size=$(get_size_mb "$DQN_INT8")
    echo "DQN            | $fp32_size    | $int8_size    | $(calc_reduction "$DQN_FP32" "$DQN_INT8")"
fi

# PPO checkpoints
PPO_FP32="${CHECKPOINT_DIR}/ppo_checkpoint_epoch_3.safetensors"
PPO_INT8="${CHECKPOINT_DIR}/ppo_quantized_epoch_3.safetensors"

if [ -f "$PPO_FP32" ] || [ -f "$PPO_INT8" ]; then
    fp32_size=$(get_size_mb "$PPO_FP32")
    int8_size=$(get_size_mb "$PPO_INT8")
    echo "PPO            | $fp32_size    | $int8_size    | $(calc_reduction "$PPO_FP32" "$PPO_INT8")"
fi

# MAMBA-2 checkpoints
MAMBA_FP32="${CHECKPOINT_DIR}/mamba2_real_data/mamba2_epoch_24.safetensors"
MAMBA_INT8="${CHECKPOINT_DIR}/mamba2_real_data/mamba2_quantized_epoch_24.safetensors"

if [ -f "$MAMBA_FP32" ] || [ -f "$MAMBA_INT8" ]; then
    fp32_size=$(get_size_mb "$MAMBA_FP32")
    int8_size=$(get_size_mb "$MAMBA_INT8")
    echo "MAMBA-2        | $fp32_size    | $int8_size    | $(calc_reduction "$MAMBA_FP32" "$MAMBA_INT8")"
fi

# TFT checkpoints
TFT_FP32="${CHECKPOINT_DIR}/tft_225_epoch_0.safetensors"
TFT_INT8="${CHECKPOINT_DIR}/tft_quantized_225_epoch_0.safetensors"

if [ -f "$TFT_FP32" ] || [ -f "$TFT_INT8" ]; then
    fp32_size=$(get_size_mb "$TFT_FP32")
    int8_size=$(get_size_mb "$TFT_INT8")
    echo "TFT            | $fp32_size    | $int8_size    | $(calc_reduction "$TFT_FP32" "$TFT_INT8")"
fi

echo ""
echo "Legend:"
echo "  FP32 Size: Float32 checkpoint size (original)"
echo "  INT8 Size: Quantized INT8 checkpoint size"
echo "  Reduction: Percentage size reduction (target: ~75%)"
echo ""
echo "Target file sizes:"
echo "  DQN:    < 50 MB  (INT8)"
echo "  PPO:    < 75 MB  (INT8)"
echo "  MAMBA:  < 100 MB (INT8)"
echo "  TFT:    < 100 MB (INT8)"