#!/bin/bash
#
# Copyright (c) 2026 AMLOGIC
# Neural Network Test Script
# Usage: ./test_nn_en.sh <model_file>
# Example: ./test_nn_en.sh mobilenetv1_subgraph_A9.adla

set -e

# Check parameters
if [ $# -ne 1 ]; then
    echo "Usage: $0 <model_file>"
    echo "Example: $0 mobilenetv1_subgraph_A9.adla"
    exit 1
fi

MODEL_FILE="$1"
INPUT_BIN="224x224x3_fish_uint8.bin"
GOLDEN_FILE="golden.txt"
EXECUTABLE="./nnsample_yocto"

# Check if file exists
check_file() {
    if [ ! -f "$1" ]; then
        echo "Error: File '$1' does not exist"
        exit 1
    fi
}

check_file "$MODEL_FILE"
check_file "$INPUT_BIN"
check_file "$GOLDEN_FILE"
check_file "$EXECUTABLE"

echo "=== Starting Neural Network Test ==="
echo "Model file: $MODEL_FILE"
echo "Input file: $INPUT_BIN"
echo "Golden file: $GOLDEN_FILE"
echo ""

# Parse golden.txt file
parse_golden() {
    local golden_file="$1"
    local test_type="$2"
    
    # Extract labels for corresponding test type from golden.txt
    # Format: [base]---[2 121 120 125 141]
    local line
    line=$(grep "^\[$test_type\]---" "$golden_file" || true)
    
    if [ -n "$line" ]; then
        # Extract labels inside brackets
        local labels
        labels=$(echo "$line" | sed -n 's/.*---\[\([^]]*\)\].*/\1/p')
        echo "$labels"
        return 0
    fi
    
    # Special handling for dynamic_shape
    if [[ "$test_type" == "dynamic_shape" ]]; then
        # Try to find dynamic_shape][shape_0
        line=$(grep "^\[dynamic_shape\]\[shape_0\]---" "$golden_file" || true)
        if [ -n "$line" ]; then
            local labels
            labels=$(echo "$line" | sed -n 's/.*---\[\([^]]*\)\].*/\1/p')
            echo "$labels"
            return 0
        fi
    fi
    
    return 1
}

# Parse specific shape from golden.txt file
parse_golden_shape() {
    local golden_file="$1"
    local shape_name="$2"
    
    # Format: [dynamic_shape][shape_0]---[2 121 120 125 141]
    local line
    line=$(grep "^\[dynamic_shape\]\[$shape_name\]---" "$golden_file" || true)
    
    if [ -n "$line" ]; then
        # Extract labels inside brackets
        local labels
        labels=$(echo "$line" | sed -n 's/.*---\[\([^]]*\)\].*/\1/p')
        echo "$labels"
        return 0
    fi
    
    return 1
}

# Run test command and parse output
run_test() {
    local test_type="$1"
    local expected_labels="$2"
    
    echo "=== Running $test_type test ==="
    
    # Run command
    local output
    output=$("$EXECUTABLE" "$MODEL_FILE" "$INPUT_BIN" "--$test_type" 2>&1)
    local exit_code=$?
    
    if [ $exit_code -ne 0 ]; then
        echo "  Error: Command execution failed (exit code: $exit_code)"
        echo "  Output: $output"
        return 1
    fi
    
    # Special handling for dynamic_shape as it has two shapes
    if [[ "$test_type" == "dynamic_shape" ]]; then
        run_dynamic_shape_test "$output"
        return $?
    fi
    
    # Parse output, extract top5 labels
    # Output format example:
    # [base] output 0 top5:
    #   2: 20.181927
    # 121: 11.818245
    # etc.
    
    local actual_labels=""
    local line
    while IFS= read -r line; do
        # Match format: "  2: 20.181927" or "121: 11.818245"
        if [[ "$line" =~ ^[[:space:]]*([0-9]+):[[:space:]]*[0-9.]+ ]]; then
            local label="${BASH_REMATCH[1]}"
            actual_labels="$actual_labels $label"
            
            # Stop after collecting 5 labels
            if [ $(echo "$actual_labels" | wc -w) -ge 5 ]; then
                break
            fi
        fi
    done <<< "$output"
    
    # Trim leading/trailing spaces
    actual_labels=$(echo "$actual_labels" | xargs)
    
    if [ -z "$actual_labels" ]; then
        echo "  Error: Unable to parse labels from output"
        echo "  Output preview:"
        echo "$output" | head -20
        return 1
    fi
    
    echo "  Actual labels: $actual_labels"
    echo "  Expected labels: $expected_labels"
    
    # Compare labels
    if [ "$actual_labels" = "$expected_labels" ]; then
        echo "  $test_type test pass"
        return 0
    else
        echo "  $test_type test fail"
        echo "  Labels mismatch"
        return 1
    fi
}

# Run dynamic_shape test (special handling, has two shapes)
run_dynamic_shape_test() {
    local output="$1"
    local all_passed=true
    
    echo "  Parsing dynamic_shape output (contains two shapes)..."
    
    # Parse shape_0
    local shape0_labels=""
    local shape1_labels=""
    local current_shape=""
    local line
    
    while IFS= read -r line; do
        # Detect shape start
        if [[ "$line" =~ \[dynamic_shape\]\[shape_0\] ]]; then
            current_shape="shape_0"
            continue
        elif [[ "$line" =~ \[dynamic_shape\]\[shape_1\] ]]; then
            current_shape="shape_1"
            continue
        fi
        
        # If current shape, parse labels
        if [[ "$current_shape" == "shape_0" || "$current_shape" == "shape_1" ]]; then
            # Match format: "  2: 20.181927" or "121: 11.818245"
            if [[ "$line" =~ ^[[:space:]]*([0-9]+):[[:space:]]*[0-9.]+ ]]; then
                local label="${BASH_REMATCH[1]}"
                
                if [[ "$current_shape" == "shape_0" ]]; then
                    shape0_labels="$shape0_labels $label"
                else
                    shape1_labels="$shape1_labels $label"
                fi
                
                # Stop current shape after collecting 5 labels
                if [[ "$current_shape" == "shape_0" && $(echo "$shape0_labels" | wc -w) -ge 5 ]]; then
                    current_shape=""
                elif [[ "$current_shape" == "shape_1" && $(echo "$shape1_labels" | wc -w) -ge 5 ]]; then
                    current_shape=""
                fi
            fi
        fi
    done <<< "$output"
    
    # Trim leading/trailing spaces
    shape0_labels=$(echo "$shape0_labels" | xargs)
    shape1_labels=$(echo "$shape1_labels" | xargs)
    
    # Get expected labels
    local expected_shape0
    expected_shape0=$(parse_golden_shape "$GOLDEN_FILE" "shape_0")
    local expected_shape1
    expected_shape1=$(parse_golden_shape "$GOLDEN_FILE" "shape_1")
    
    # Test shape_0
    if [ -n "$shape0_labels" ] && [ -n "$expected_shape0" ]; then
        echo "  [dynamic_shape][shape_0] Actual labels: $shape0_labels"
        echo "  [dynamic_shape][shape_0] Expected labels: $expected_shape0"
        
        if [ "$shape0_labels" = "$expected_shape0" ]; then
            echo "  [dynamic_shape][shape_0] test pass"
        else
            echo "  [dynamic_shape][shape_0] test fail"
            echo "  Labels mismatch"
            all_passed=false
        fi
    else
        echo "  [dynamic_shape][shape_0] Error: Unable to parse labels or golden data"
        all_passed=false
    fi
    
    echo ""
    
    # Test shape_1
    if [ -n "$shape1_labels" ] && [ -n "$expected_shape1" ]; then
        echo "  [dynamic_shape][shape_1] Actual labels: $shape1_labels"
        echo "  [dynamic_shape][shape_1] Expected labels: $expected_shape1"
        
        if [ "$shape1_labels" = "$expected_shape1" ]; then
            echo "  [dynamic_shape][shape_1] test pass"
        else
            echo "  [dynamic_shape][shape_1] test fail"
            echo "  Labels mismatch"
            all_passed=false
        fi
    else
        echo "  [dynamic_shape][shape_1] Error: Unable to parse labels or golden data"
        all_passed=false
    fi
    
    if $all_passed; then
        return 0
    else
        return 1
    fi
}

# Main test logic
main() {
    local all_passed=true
    
    # Test type list
    local test_types=("base" "dma" "softop_acc" "dynamic_shape")
    
    for test_type in "${test_types[@]}"; do
        # For dynamic_shape, no need to get labels from parse_golden as run_dynamic_shape_test handles it
        if [[ "$test_type" == "dynamic_shape" ]]; then
            # Run test (pass empty labels as dynamic_shape handles it itself)
            if ! run_test "$test_type" ""; then
                all_passed=false
            fi
        else
            # Get expected labels
            local expected_labels
            expected_labels=$(parse_golden "$GOLDEN_FILE" "$test_type")
            
            if [ $? -ne 0 ] || [ -z "$expected_labels" ]; then
                echo "Error: Unable to parse $test_type labels from $GOLDEN_FILE"
                all_passed=false
                continue
            fi
            
            # Run test
            if ! run_test "$test_type" "$expected_labels"; then
                all_passed=false
            fi
        fi
        
        echo ""
    done
    
    # Output summary
    echo "=== Test Summary ==="
    if $all_passed; then
        echo "All tests passed!"
        return 0
    else
        echo "Some tests failed!"
        return 1
    fi
}

# Execute main function
main