#!/bin/bash
# compare-3x.sh — Run A/B/C comparison 3 times for statistical validity
#
# Runs the full comparison (Vanilla Qwen, Qwen + Claude Style, Claude Code)
# three separate times and averages results.
#
# Usage: ./scripts/compare-3x.sh

set -e

# Colors
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
RED='\033[0;31m'
BLUE='\033[0;34m'
NC='\033[0m'

PROMPT="Create a Python file 'user_database.py' with a class that manages a user database. Support adding users, finding by email, updating fields, and deleting users. Store in a dict."
OUTPUT_DIR="/tmp/claude-style-compare-3x"
SCRIPT_PATH="$(readlink -f "${BASH_SOURCE[0]}")"
SCRIPT_DIR="$(cd "$(dirname "$SCRIPT_PATH")" && pwd)"
PROJECT_DIR="$(cd "$SCRIPT_DIR/.." && pwd)"
WORK_DIR=""
mkdir -p "$OUTPUT_DIR"

cleanup_work_dir() {
    if [ -n "$WORK_DIR" ] && [ -d "$WORK_DIR" ]; then
        rm -rf "$WORK_DIR"
    fi
}
trap cleanup_work_dir EXIT

echo -e "${BLUE}╔══════════════════════════════════════════════════════╗${NC}"
echo -e "${BLUE}║   3x A/B/C Comparison: Qwen vs Style vs Claude      ║${NC}"
echo -e "${BLUE}╚══════════════════════════════════════════════════════╝${NC}"
echo ""

declare -a VANILLA_CODES VANILLA_WORDS STYLE_CODES STYLE_WORDS NATIVE_CODES NATIVE_WORDS

# Prerequisite: check for Claude Code
if ! command -v claude &> /dev/null; then
    echo -e "${YELLOW}⚠ Claude Code not found — native runs will be skipped${NC}"
    echo ""
fi

for run in 1 2 3; do
    echo -e "${BLUE}════════════════════════════════════════════════════${NC}"
    echo -e "${BLUE}  RUN $run of 3${NC}"
    echo -e "${BLUE}════════════════════════════════════════════════════${NC}"
    echo ""

    RUN_DIR="$OUTPUT_DIR/run-$run"
    mkdir -p "$RUN_DIR"

    # ---- 1. Vanilla Qwen ----
    echo -e "${YELLOW}[1/3] Qwen WITHOUT Claude Style...${NC}"
    "$PROJECT_DIR/scripts/uninstall.sh" 2>&1 | grep -E "^(✓|✗|⚠)" || true
    WORK_DIR=$(mktemp -d)
    pushd "$WORK_DIR" > /dev/null
    qwen --prompt "$PROMPT" --output-format text -y 2>&1 | tee "$RUN_DIR/vanilla.txt" > /dev/null 2>&1
    if [ -f "$WORK_DIR/user_database.py" ]; then
        cp "$WORK_DIR/user_database.py" "$RUN_DIR/vanilla_code.py"
        VANILLA_CODE=$(wc -l < "$RUN_DIR/vanilla_code.py")
    else
        VANILLA_CODE=0
    fi
    popd > /dev/null
    VANILLA_CODES[$run]=$VANILLA_CODE
    VANILLA_WORDS[$run]=$(wc -w < "$RUN_DIR/vanilla.txt")
    echo -e "${GREEN}  ✓ Code: $VANILLA_CODE lines | Text: ${VANILLA_WORDS[$run]} words${NC}"
    echo ""

    # ---- 2. Qwen + Claude Style ----
    echo -e "${YELLOW}[2/3] Qwen WITH Claude Style...${NC}"
    "$PROJECT_DIR/scripts/setup.sh" 2>&1 | grep -E "^(✓|✗|⚠)" || true
    WORK_DIR=$(mktemp -d)
    pushd "$WORK_DIR" > /dev/null
    qwen --prompt "$PROMPT" --output-format text -y 2>&1 | tee "$RUN_DIR/claude-style.txt" > /dev/null 2>&1
    if [ -f "$WORK_DIR/user_database.py" ]; then
        cp "$WORK_DIR/user_database.py" "$RUN_DIR/claude-style_code.py"
        STYLE_CODE=$(wc -l < "$RUN_DIR/claude-style_code.py")
    else
        STYLE_CODE=0
    fi
    popd > /dev/null
    WORK_DIR=""
    STYLE_CODES[$run]=$STYLE_CODE
    STYLE_WORDS[$run]=$(wc -w < "$RUN_DIR/claude-style.txt")
    echo -e "${GREEN}  ✓ Code: $STYLE_CODE lines | Text: ${STYLE_WORDS[$run]} words${NC}"
    echo ""

    # ---- 3. Claude Code Native ----
    echo -e "${YELLOW}[3/3] Claude Code (native)...${NC}"
    if command -v claude &> /dev/null; then
        WORK_DIR=$(mktemp -d)
        pushd "$WORK_DIR" > /dev/null
        claude -p "$PROMPT" --dangerously-skip-permissions 2>&1 | tee "$RUN_DIR/claude-native.txt" > /dev/null 2>&1
        if [ -f "$WORK_DIR/user_database.py" ]; then
            cp "$WORK_DIR/user_database.py" "$RUN_DIR/claude-native_code.py"
            NATIVE_CODE=$(wc -l < "$RUN_DIR/claude-native_code.py")
        else
            NATIVE_CODE=0
        fi
        popd > /dev/null
        WORK_DIR=""
        NATIVE_CODES[$run]=$NATIVE_CODE
        NATIVE_WORDS[$run]=$(wc -w < "$RUN_DIR/claude-native.txt")
        echo -e "${GREEN}  ✓ Code: $NATIVE_CODE lines | Text: ${NATIVE_WORDS[$run]} words${NC}"
    else
        NATIVE_CODES[$run]=0
        NATIVE_WORDS[$run]=0
        echo -e "${RED}  ✗ Skipped — Claude Code not installed${NC}"
    fi
    echo ""
done

# ============================================================
# AVERAGES
# ============================================================
echo -e "${BLUE}╔══════════════════════════════════════════════════════╗${NC}"
echo -e "${BLUE}║                  AGGREGATE RESULTS                    ║${NC}"
echo -e "${BLUE}╚══════════════════════════════════════════════════════╝${NC}"
echo ""

VANILLA_CODE_AVG=$(( (${VANILLA_CODES[1]:-0} + ${VANILLA_CODES[2]:-0} + ${VANILLA_CODES[3]:-0}) / 3 ))
STYLE_CODE_AVG=$(( (${STYLE_CODES[1]:-0} + ${STYLE_CODES[2]:-0} + ${STYLE_CODES[3]:-0}) / 3 ))
NATIVE_CODE_AVG=$(( (${NATIVE_CODES[1]:-0} + ${NATIVE_CODES[2]:-0} + ${NATIVE_CODES[3]:-0}) / 3 ))

VANILLA_WORD_AVG=$(( (${VANILLA_WORDS[1]:-0} + ${VANILLA_WORDS[2]:-0} + ${VANILLA_WORDS[3]:-0}) / 3 ))
STYLE_WORD_AVG=$(( (${STYLE_WORDS[1]:-0} + ${STYLE_WORDS[2]:-0} + ${STYLE_WORDS[3]:-0}) / 3 ))
NATIVE_WORD_AVG=$(( (${NATIVE_WORDS[1]:-0} + ${NATIVE_WORDS[2]:-0} + ${NATIVE_WORDS[3]:-0}) / 3 ))

echo -e "${YELLOW}─── Code Lines (lower is better) ───${NC}"
printf "  %-25s %6s %6s %6s %8s\n" "" "Run 1" "Run 2" "Run 3" "AVG"
printf "  %-25s %6s %6s %6s %8s\n" "---" "---" "---" "---" "---"
printf "  %-25s %6d %6d %6d %8d\n" "Qwen Vanilla" ${VANILLA_CODES[1]:-0} ${VANILLA_CODES[2]:-0} ${VANILLA_CODES[3]:-0} $VANILLA_CODE_AVG
printf "  %-25s %6d %6d %6d %8d\n" "Qwen + Claude Style" ${STYLE_CODES[1]:-0} ${STYLE_CODES[2]:-0} ${STYLE_CODES[3]:-0} $STYLE_CODE_AVG
printf "  %-25s %6d %6d %6d %8d\n" "Claude Code Native" ${NATIVE_CODES[1]:-0} ${NATIVE_CODES[2]:-0} ${NATIVE_CODES[3]:-0} $NATIVE_CODE_AVG
echo ""

echo -e "${YELLOW}─── Explanation Words (lower is better) ───${NC}"
printf "  %-25s %6s %6s %6s %8s\n" "" "Run 1" "Run 2" "Run 3" "AVG"
printf "  %-25s %6s %6s %6s %8s\n" "---" "---" "---" "---" "---"
printf "  %-25s %6d %6d %6d %8d\n" "Qwen Vanilla" ${VANILLA_WORDS[1]:-0} ${VANILLA_WORDS[2]:-0} ${VANILLA_WORDS[3]:-0} $VANILLA_WORD_AVG
printf "  %-25s %6d %6d %6d %8d\n" "Qwen + Claude Style" ${STYLE_WORDS[1]:-0} ${STYLE_WORDS[2]:-0} ${STYLE_WORDS[3]:-0} $STYLE_WORD_AVG
printf "  %-25s %6d %6d %6d %8d\n" "Claude Code Native" ${NATIVE_WORDS[1]:-0} ${NATIVE_WORDS[2]:-0} ${NATIVE_WORDS[3]:-0} $NATIVE_WORD_AVG
echo ""

# Warning if any run produced zero lines
if [ ${VANILLA_CODES[1]:-0} -eq 0 ] || [ ${VANILLA_CODES[2]:-0} -eq 0 ] || [ ${VANILLA_CODES[3]:-0} -eq 0 ]; then
    echo -e "${YELLOW}⚠ One or more vanilla runs produced no code — averages may be misleading${NC}"
    echo ""
fi
if [ ${STYLE_CODES[1]:-0} -eq 0 ] || [ ${STYLE_CODES[2]:-0} -eq 0 ] || [ ${STYLE_CODES[3]:-0} -eq 0 ]; then
    echo -e "${YELLOW}⚠ One or more Claude Style runs produced no code — averages may be misleading${NC}"
    echo ""
fi

# Docstring counts (from run 3 as representative)
RUN3_DIR="$OUTPUT_DIR/run-3"
if [ -f "$RUN3_DIR/vanilla_code.py" ]; then
    VANILLA_DOCS=$(grep -c '"""' "$RUN3_DIR/vanilla_code.py" 2>/dev/null || echo "0")
else
    VANILLA_DOCS=0
fi
if [ -f "$RUN3_DIR/claude-style_code.py" ]; then
    STYLE_DOCS=$(grep -c '"""' "$RUN3_DIR/claude-style_code.py" 2>/dev/null || echo "0")
else
    STYLE_DOCS=0
fi
if [ -f "$RUN3_DIR/claude-native_code.py" ]; then
    NATIVE_DOCS=$(grep -c '"""' "$RUN3_DIR/claude-native_code.py" 2>/dev/null || echo "0")
else
    NATIVE_DOCS=0
fi

echo -e "${YELLOW}─── Docstrings (from Run 3) ───${NC}"
printf "  %-25s %d\n" "Qwen Vanilla" $VANILLA_DOCS
printf "  %-25s %d\n" "Qwen + Claude Style" $STYLE_DOCS
printf "  %-25s %d\n" "Claude Code Native" $NATIVE_DOCS
echo ""

# Summary
echo -e "${BLUE}─── Summary ───${NC}"
if [ $STYLE_CODE_AVG -lt $VANILLA_CODE_AVG ] && [ $VANILLA_CODE_AVG -gt 0 ]; then
    echo -e "${GREEN}✓ Claude Style reduces code by $(( VANILLA_CODE_AVG - STYLE_CODE_AVG )) lines avg ($(( (VANILLA_CODE_AVG - STYLE_CODE_AVG) * 100 / VANILLA_CODE_AVG ))%) vs Vanilla Qwen${NC}"
fi
if [ $NATIVE_CODE_AVG -gt 0 ] && [ $STYLE_CODE_AVG -gt $NATIVE_CODE_AVG ]; then
    echo -e "${YELLOW}⚠ Claude Style is $(( STYLE_CODE_AVG - NATIVE_CODE_AVG )) lines more than Claude Native ($(( (STYLE_CODE_AVG - NATIVE_CODE_AVG) * 100 / NATIVE_CODE_AVG ))% more)${NC}"
elif [ $NATIVE_CODE_AVG -gt 0 ] && [ $STYLE_CODE_AVG -eq $NATIVE_CODE_AVG ]; then
    echo -e "${GREEN}✓ Claude Style matches Claude Native code length${NC}"
fi

echo ""
echo -e "${BLUE}────────────────────────────────────────────────────────${NC}"
echo ""
echo -e "${YELLOW}Full outputs saved to:${NC} $OUTPUT_DIR/"
echo -e "${YELLOW}  Run 1:${NC} $OUTPUT_DIR/run-1/"
echo -e "${YELLOW}  Run 2:${NC} $OUTPUT_DIR/run-2/"
echo -e "${YELLOW}  Run 3:${NC} $OUTPUT_DIR/run-3/"