#!/bin/bash
# Copyright (c) 2025 Marcus Becker, Uwe Fechner
# SPDX-License-Identifier: BSD-3-Clause

# Analyze memory allocation tracking files (.mem files) and list highest allocations
# This script parses .mem files generated by Julia's --track-allocation flag

# Change to repo root if running from bin directory
if [[ $(basename $(pwd)) == "bin" ]]; then
    cd ..
fi
cd src

# Check if .mem files exist
mem_count=$(find . -name "*.mem" -type f | wc -l)

if [ "$mem_count" -eq 0 ]; then
    echo "No .mem files found."
    echo "Run './bin/test_alloc' first to generate allocation data."
    exit 1
fi

echo "Found $mem_count .mem files"
echo "Analyzing memory allocations..."
echo ""

# Default number of top allocations to show
TOP_N=${1:-20}

echo "=== TOP $TOP_N MEMORY ALLOCATIONS ==="
echo ""

# Parse top allocations safely: filename:line:content where content starts with bytes
find . -name "*.mem" -exec grep -Hn "^[[:space:]]*[0-9]" {} \; | \
    awk -F: '{
        file=$1; line=$2;
        # Rejoin the rest of the fields as content (could contain colons)
        content=$3; for(i=4;i<=NF;i++) content=content":"$i;
        # Extract bytes from beginning of content
        if (match(content, /^[[:space:]]*([0-9]+)/, m)) {
            bytes=m[1]+0;
            if (bytes>0) {
                gsub(/\.mem$/, "", file);
                # Extract source code without the leading byte count
                code=content; sub(/^[[:space:]]*[0-9]+[[:space:]]*/, "", code);
                printf "%d %s:%s %s\n", bytes, file, line, code;
            }
        }
    }' | \
    sort -nr | \
    head -n "$TOP_N" | \
    awk '{
        bytes = $1
        file_line = $2
        # Convert bytes to human readable format
        if (bytes >= 1048576) {
            size = sprintf("%.2f MB", bytes/1048576)
        } else if (bytes >= 1024) {
            size = sprintf("%.2f KB", bytes/1024)
        } else {
            size = sprintf("%s B", bytes)
        }
        printf "%-12s %s", size, file_line
        # Print the rest (source code)
        for(i=3; i<=NF; i++) printf " %s", $i
        printf "\n"
    }'

echo ""
echo "=== ALLOCATION SUMMARY BY FILE ==="
echo ""

# Summary by file (sum leading bytes per line)
find . -name "*.mem" -exec grep -H "^[[:space:]]*[0-9]" {} \; | \
    awk -F: '{
        file=$1;
        content=$2; for(i=3;i<=NF;i++) content=content":"$i;
        if (match(content, /^[[:space:]]*([0-9]+)/, m)) {
            bytes=m[1]+0;
            if (bytes>0) {
                gsub(/\.mem$/, "", file);
                file_total[file] += bytes;
                file_count[file]++;
            }
        }
    }
    END {
        for (file in file_total) {
            printf "%d %d %s\n", file_total[file], file_count[file], file
        }
    }' | sort -nr | \
    awk '{
        bytes = $1
        count = $2
        file = $3
        if (bytes >= 1048576) {
            size = sprintf("%.2f MB", bytes/1048576)
        } else if (bytes >= 1024) {
            size = sprintf("%.2f KB", bytes/1024)
        } else {
            size = sprintf("%s B", bytes)
        }
        printf "%-12s %-4s lines %s\n", size, count, file
    }'

echo ""
echo "=== USAGE ==="
echo "To see more/fewer top allocations: $0 [number]"
echo "Example: $0 50    # Show top 50 allocations"
echo "To clean up .mem files: ./bin/clean_mem_files"
echo ""
echo "=== TIPS ==="
echo "• Focus on allocations > 1KB for optimization"
echo "• Look for allocations in tight loops"
echo "• Check if temporary arrays can be pre-allocated"
echo "• Use @views to avoid array copying"
