#!/bin/bash
#
# run_optuna.sh - Wrapper script for Optuna-based PPA optimization
#
# This script sets up the environment and runs the Optuna optimizer
# to explore synthesis parameter combinations for timing, area, and power.
#
# Usage:
#   ./run_optuna.sh <sv_file> [options]
#
# Examples:
#   ./run_optuna.sh sv_files/ast_module_00001.sv
#   ./run_optuna.sh sv_files/ast_module_00001.sv --n-trials 200
#   ./run_optuna.sh sv_files/ast_module_00001.sv --sampler random --n-trials 100
#

set -e

# Script directory
SCRIPT_DIR="$(dirname $(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd))"

# Display help
show_help() {
    cat << EOF
Usage: $(basename "$0") <sv_file> [options]

Run Optuna-based PPA optimization for synthesis parameter exploration.
Optimizes for timing (TNS), area, and power.

Prerequisites:
    Docker container 'OR-PPA' must be running with the project directory
    mounted at the SAME PATH inside the container. Example:

        docker run -d --name OR-PPA \\
          -v /path/to/PPA_training:/path/to/PPA_training \\
          <image_name> tail -f /dev/null

    For non-rootless Docker (e.g., AWS), use sudo and set DOCKER_SUDO:
        export DOCKER_SUDO=sudo

Arguments:
    sv_file             Path to the SystemVerilog source file

Environment Variables:
    DOCKER_SUDO         Set to "sudo" for non-rootless Docker systems (e.g., AWS)
                        Example: export DOCKER_SUDO=sudo
    DOCKER_CONTAINER    Container name (default: OR-PPA)

Options:
    --n-trials N        Number of optimization trials (default: 100)
    --sampler TYPE      Sampler type: tpe, random, cmaes (default: tpe)
    --study-name NAME   Name for the optimization study
    --output-dir DIR    Directory for output files
    --timeout SECS      Timeout per EDA run in seconds (default: 600)
    --n-jobs N          Parallel jobs (default: 80 for 96-core machine)
    --resume            Resume a previous study (requires --study-name)
    --verbose           Enable verbose logging

    Objective Weights (for multi-objective optimization):
    --weight-tns W      Weight for TNS in objective (default: 1.0)
    --weight-area W     Weight for area in objective (default: 0.001)
    --weight-power W    Weight for power in objective (default: 1000.0)

    -h, --help          Display this help message

Examples:
    # Basic optimization with 100 trials
    $(basename "$0") sv_files/ast_module_00001.sv

    # More trials with TPE sampler
    $(basename "$0") sv_files/ast_module_00001.sv --n-trials 200 --sampler tpe

    # Resume a previous study
    $(basename "$0") sv_files/ast_module_00001.sv --resume --study-name my_study

    # Custom objective weights (prioritize timing)
    $(basename "$0") sv_files/ast_module_00001.sv --weight-tns 2.0 --weight-area 0.0005

Parameter Search Space:
    TOP_MODULE_PREFIX : "logic_then_pipe" or "pipe_then_logic" (pipeline style)
    PIPE_STAGES       : 1-5  (Pipeline depth)
    OPT_LEVEL         : 0-4  (Optimization aggressiveness)
    SHARE_LEVEL       : 0-3  (Resource sharing level)
    ABC_SCRIPT        : 0-5  (ABC optimization strategy)
    RETIME_MODE       : 0-2  (Retiming mode)

    Total search space: 2 * 5 * 5 * 4 * 6 * 3 = 3,600 combinations

Metrics Optimized:
    TNS   - Total Negative Slack (from OpenSTA)
    Area  - Chip area (from Yosys synthesis)
    Power - Total power in Watts (from OpenSTA report_power)

EOF
    exit 0
}

# Check for help flag
if [[ "$1" == "-h" || "$1" == "--help" || -z "$1" ]]; then
    show_help
fi

# Auto-detect if sudo is needed for Docker (if DOCKER_SUDO not already set)
if [[ -z "${DOCKER_SUDO:-}" ]]; then
    if docker info >/dev/null 2>&1; then
        export DOCKER_SUDO=""
    elif sudo docker info >/dev/null 2>&1; then
        export DOCKER_SUDO="sudo"
        echo "Note: Using sudo for Docker (non-rootless Docker detected)"
    else
        echo "Error: Docker is not available. Please start the Docker container."
        exit 1
    fi
fi

# Check if SV file argument is provided
SV_FILE="$1"
shift

if [[ ! -f "$SV_FILE" ]]; then
    echo "Error: SV file not found: $SV_FILE"
    exit 1
fi

# Check for Python and required packages
check_python() {
    if ! command -v python3 &> /dev/null; then
        echo "Error: python3 not found. Please install Python 3."
        exit 1
    fi

    # Check for optuna
    if ! python3 -c "import optuna" 2>/dev/null; then
        echo "Warning: optuna package not found."
        echo "Installing required packages..."
        pip3 install optuna
    fi
}

# Find or create virtual environment
setup_venv() {
    # Check for existing virtual environments
    if [[ -d "${SCRIPT_DIR}/.venv_new" ]]; then
        source "${SCRIPT_DIR}/.venv_new/bin/activate"
        echo "Using virtual environment: .venv_new"
    elif [[ -d "${SCRIPT_DIR}/.venv" ]]; then
        source "${SCRIPT_DIR}/.venv/bin/activate"
        echo "Using virtual environment: .venv"
    else
        echo "No virtual environment found, using system Python"
    fi
}

# Get SV file basename for display
SV_BASENAME=$(basename "$SV_FILE")

# Main execution (minimal header - Python script provides details)
echo ""
echo "══════════════════════════════════════════════════════════════════════"
echo "  Optuna PPA Optimizer: $SV_BASENAME"
echo "══════════════════════════════════════════════════════════════════════"

# Setup environment
setup_venv
check_python

# Create output directory for this run
TIMESTAMP=$(date +"%Y%m%d_%H%M%S")
SV_NAME=$(basename "$SV_FILE" .sv)
OUTPUT_DIR="${SCRIPT_DIR}/optuna_results/${SV_NAME}_${TIMESTAMP}"
mkdir -p "$OUTPUT_DIR"

# Ensure tmp_eda exists
mkdir -p "${SCRIPT_DIR}/tmp_eda"
mkdir -p "${SCRIPT_DIR}/tmp_eda/temp_sv"

# Run the optimizer (it prints its own formatted output)

python3 "${SCRIPT_DIR}/scripts/optuna_ppa_optimizer.py" \
    --sv-file "$SV_FILE" \
    --project-root "$SCRIPT_DIR" \
    --output-dir "$OUTPUT_DIR" \
    "$@"

EXIT_CODE=$?

# Reset terminal to sane state after parallel Docker execution
# This fixes terminal corruption that can occur with parallel subprocess output
stty sane 2>/dev/null || true

if [[ $EXIT_CODE -eq 0 ]]; then
    echo ""
    echo "Output files:"
    echo "  • optimization_report.json  (detailed results)"
    echo "  • optimization_results.csv  (spreadsheet format)"
    echo "  • optuna_study.db           (Optuna database)"
    echo ""
    echo "Visualize: python3 scripts/visualize_optuna_results.py --results-dir $OUTPUT_DIR"
    echo "══════════════════════════════════════════════════════════════════════"
else
    echo ""
    echo "Optimization FAILED (exit code: $EXIT_CODE)"
    echo "Check logs in: $OUTPUT_DIR"
    echo "══════════════════════════════════════════════════════════════════════"
fi

# Check for failed trials and provide debug hints
DEBUG_DIR="${SCRIPT_DIR}/tmp_eda/debug"
if [[ -d "$DEBUG_DIR" ]] && [[ -n "$(ls -A "$DEBUG_DIR" 2>/dev/null)" ]]; then
    FAIL_COUNT=$(ls -1 "$DEBUG_DIR"/*.log 2>/dev/null | wc -l)
    if [[ $FAIL_COUNT -gt 0 ]]; then
        echo ""
        echo "Debug: $FAIL_COUNT trial failure log(s) saved to: $DEBUG_DIR"
        echo "View first failure: cat $DEBUG_DIR/$(ls -1 "$DEBUG_DIR" | head -1)"
        echo ""
        echo "Common issues:"
        echo "  • Docker container not running: docker ps | grep OR-PPA"
        echo "  • Docker needs sudo: export DOCKER_SUDO=sudo"
        echo "  • Permission issues: check tmp_eda/ directory permissions"
        echo "══════════════════════════════════════════════════════════════════════"
    fi
fi

exit $EXIT_CODE
