#!/bin/bash

# 分析accept length的脚本

echo "=========================================="
echo "Accept Length Analysis"
echo "=========================================="
echo ""

# 检查results目录
if [ ! -d "results" ]; then
    echo "Error: results directory not found"
    exit 1
fi

# 查找所有结果文件
result_files=$(ls results/*.jsonl 2>/dev/null)

if [ -z "$result_files" ]; then
    echo "No result files found in results/ directory"
    echo ""
    echo "Please run the benchmark first:"
    echo "  python benchmarks/bench_eagle3.py ..."
    exit 1
fi

echo "Found result files:"
ls -lh results/*.jsonl
echo ""
echo "=========================================="
echo ""

# 分析每个结果文件
for file in $result_files; do
    filename=$(basename "$file")
    echo "File: $filename"
    echo "----------------------------------------"

    # 检查文件是否包含mtbench结果
    if grep -q "mtbench" "$file"; then
        # 提取accept_length
        echo "Accept lengths:"
        cat "$file" | jq -r '.mtbench[0].metrics[] | "  Sample \(.sample_id): accept_length=\(.accept_length // "N/A"), output_tokens=\(.output_tokens // "N/A")"' 2>/dev/null

        echo ""
        echo "Statistics:"
        # 计算平均值
        avg_accept=$(cat "$file" | jq -r '.mtbench[0].metrics[] | .accept_length' 2>/dev/null | awk '{sum+=$1; count++} END {if(count>0) printf "  Average accept_length: %.4f\n", sum/count; else print "  No data"}')
        echo "$avg_accept"

        # 计算最小值和最大值
        min_accept=$(cat "$file" | jq -r '.mtbench[0].metrics[] | .accept_length' 2>/dev/null | sort -n | head -1)
        max_accept=$(cat "$file" | jq -r '.mtbench[0].metrics[] | .accept_length' 2>/dev/null | sort -n | tail -1)
        echo "  Min accept_length: $min_accept"
        echo "  Max accept_length: $max_accept"

        # 样本数量
        sample_count=$(cat "$file" | jq -r '.mtbench[0].metrics | length' 2>/dev/null)
        echo "  Total samples: $sample_count"
    else
        echo "  No mtbench results found in this file"
    fi

    echo ""
    echo "=========================================="
    echo ""
done

# 如果有baseline和trained的结果，进行对比
baseline_file=$(ls results/baseline*.jsonl 2>/dev/null | head -1)
trained_file=$(ls results/trained*.jsonl 2>/dev/null | head -1)

if [ -n "$baseline_file" ] && [ -n "$trained_file" ]; then
    echo "Comparison: Baseline vs Trained"
    echo "----------------------------------------"

    baseline_avg=$(cat "$baseline_file" | jq -r '.mtbench[0].metrics[] | .accept_length' 2>/dev/null | awk '{sum+=$1; count++} END {if(count>0) print sum/count}')
    trained_avg=$(cat "$trained_file" | jq -r '.mtbench[0].metrics[] | .accept_length' 2>/dev/null | awk '{sum+=$1; count++} END {if(count>0) print sum/count}')

    if [ -n "$baseline_avg" ] && [ -n "$trained_avg" ]; then
        echo "Baseline average: $baseline_avg"
        echo "Trained average:  $trained_avg"

        # 计算提升百分比
        improvement=$(echo "$baseline_avg $trained_avg" | awk '{printf "%.2f%%", ($2-$1)/$1*100}')
        echo "Improvement:      $improvement"
    fi
    echo ""
fi

echo "Done!"