#!/bin/bash

FILE=""
FILE_BASENAME=""
MODULE=""
LIBERTY_FILE="nem_basic_yosys.lib"
LIBERTY_USED="3T"
visualize=0

# Function to display the menu and get user input
show_menu() {
    

    # Define color codes
    GREEN='\033[0;32m'
    YELLOW='\033[1;33m'
    CYAN='\033[0;36m'
    RESET='\033[0m'

    echo "--------------------------------------------------------------"
    echo -e "${CYAN}Current file: $FILE with module: $MODULE${RESET}"
    echo -e "${YELLOW}Please select your options (you can choose multiple options):${RESET}"
    echo
    echo -e "${GREEN}1)${RESET} Synthesize to NEM technology"
    echo -e "${GREEN}2)${RESET} Print initial design"
    echo -e "${GREEN}3)${RESET} Print out NEM optimized design"
    echo -e "${GREEN}4)${RESET} Perform SAT comparison"
    echo -e "${GREEN}5)${RESET} Export FSM as KISS2 format"
    echo -e "${GREEN}6)${RESET} Start shell with modules"
    echo -e "${GREEN}7)${RESET} Switch from normal 3T gate library to new 4T"
    echo -e "${GREEN}8)${RESET} Run series of test and create comparison report"
    echo -e "${GREEN}9)${RESET} Perform MUXIG synthesis run"
    echo -e "${GREEN}10)${RESET} export truth table and MIG run"
    echo -e "${GREEN}11)${RESET} Select a new Verilog file"
    echo -e "${GREEN}0)${RESET} Exit the program"
    echo "--------------------------------------------------------------"
}

# Request the file to process
request_data(){
    echo "-:- Enter the file to map to NEM"
    read -e -p "What is the file name?: " FILE
    read -p "What is the name of the top module? (press ENTER for the same as the file name): " MODULE

    if [ ! -f "$FILE" ]; then
    echo "File not found"
    request_data
    fi

    FILE_BASENAME=$(basename "$FILE" | cut -d. -f1)
    #echo $FILE_BASENAME

    if [ -z "$MODULE" ]; then
    #echo "setting name equal to that of the file"
    MODULE=$FILE_BASENAME
    fi
}

#run a yosys file specified to the function
run_yosys_file() {
    local yosys_file="$1"
    local depth="$2"
    local additional_yosys_args="$3"
    
    # Start with basic sed commands
    sed_command=$(sed -e "s|{{FILE}}|$FILE|g" \
                      -e "s|{{FILE_BASENAME}}|$FILE_BASENAME|g" \
                      -e "s|{{MODULE}}|$MODULE|g" \
                      -e "s|{{LIBERTY_FILE}}|$LIBERTY_FILE|g" \
                      -e "s|{{LIBERTY_USED}}|$LIBERTY_USED|g"\
                      "./yosys/${yosys_file}.ys")

    # Apply additional sed expressions based on DEPTH value
    if [[ $depth -eq 0 ]]; then
        sed_command=$(echo "$sed_command" | sed -e "/#IF {{DEPTH}}==0/d" \
                                                -e "/#ELSE/,/#END/d")
    elif [[ $depth -eq 1 ]]; then
        sed_command=$(echo "$sed_command" | sed -e "/#IF {{DEPTH}}==0/,/#ELSE/d" \
                                                -e "/#END/d")
    fi

    # Write the result to a temp file and run yosys
    echo "$sed_command" > "./temp/${yosys_file}_temp.ys"
    yosys $additional_yosys_args "./temp/${yosys_file}_temp.ys"
}

#Switch between 3T and 4T pass through gates
switch_liberty() {
    if [ "$LIBERTY_FILE" == "nem_basic_yosys.lib" ]; then
        LIBERTY_FILE="nem_basic_yosys_extended.lib"
        LIBERTY_USED="4T"
        echo "Now using extended (4T devices) libary"
    elif [ "$LIBERTY_FILE" == "nem_basic_yosys_extended.lib" ]; then
        LIBERTY_FILE="nem_basic_yosys.lib"
        LIBERTY_USED="3T"
        echo "Now using normal libary"
    else
        echo "Unknown LIBERTY_FILE value: $LIBERTY_FILE"
    fi
}

compare_area() {
    # Extract area values from .stat files
    local area_3T=$(grep "Chip area for module" "./temp/${FILE_BASENAME}_3T.stat" | awk '{print $6}')
    local area_4T=$(grep "Chip area for module" "./temp/${FILE_BASENAME}_4T.stat" | awk '{print $6}')
    local area_MUX=$(grep "Chip area for module" "./temp/${FILE_BASENAME}_MUX.stat" | awk '{print $6}')

    # Calculate ratio as (area_3T / area_4T) * 100
    local ratio_4T=$(echo "($area_4T / $area_3T)" | bc -l)
    local ratio_MUX=$(echo "($area_MUX / $area_3T)" | bc -l)

        {
        echo "------------- normal 3T ---------------"
        cat "./temp/${FILE_BASENAME}_3T.stat"
        echo "------------- 4T ---------------"
        cat "./temp/${FILE_BASENAME}_4T.stat"
        echo "------------- Muxig ---------------"
        cat "./temp/${FILE_BASENAME}_MUX.stat"
        echo "------------- Stats ---------------"
        echo "Area 3T: $area_3T"
        echo "Area 4T: $area_4T"
        echo "Ratio 4T->3T: $ratio_4T%"
        echo "Area MUX: $area_MUX"
        echo "Ratio MUX->3T: $ratio_MUX%"
        
        ec
    } > "./output/${FILE_BASENAME}.ratio"

    # Output the areas and the ratio
    echo "Area 3T: $area_3T, Area 4T: $area_4T, Ratio 4T->3T: $ratio_4T, Area MUX: $area_MUX, Ratio MUX->3T: $ratio_MUX%" 
}

create_report() {

# Output CSV file name
csv_output="./output/output_report.csv"

# Clear the CSV file by redirecting an empty string to it
> "$csv_output"

# Write the CSV header
echo "Module,3T,4T,Ratio" > "$csv_output"
# Print the header of the table
printf "%-20s %-20s %-20s %-20s %-20s %-20s\n" "Module" "3T" "4T" "Ratio 4T" "MUX" "Ratio MUx"
printf "%-20s %-20s %-20s %-20s %-20s %-20s\n" "-------" "------" "------" "-----" "-----" "-----"

# Loop through each .ratio file in the directory
for file in ./output/*.ratio; do
    # Check if the file exists
    if [[ -f "$file" ]]; then
        # Extract the module name
        module_name=$(grep -m 1 -oP '(?<==== ).*(?= ===)' "$file")  # Extract the module name
        
        # Extract areas using grep and sed
        area1=$(grep "Chip area for module" "$file" | sed -n '1s/.*: //p')  # Area 3T
        area2=$(grep "Chip area for module" "$file" | sed -n '2s/.*: //p')  # Area 4T
        area3=$(grep "Chip area for module" "$file" | sed -n '3s/.*: //p')  # Area MUX

        # Extract the ratio
        ratio_4T=$(grep -oP '(?<=Ratio 4T->3T: )[\d.]+' "$file")  # Extract the ratio
        ratio_MUX=$(grep -oP '(?<=Ratio MUX->3T: )[\d.]+' "$file")  # Extract the ratio
        # Append the data to the CSV file
        echo "$module_name,$area1,$area2,$ratio_4T,$area3,$ratio_MUX" >> "$csv_output"
        
        # Print the results in the table format
        printf "%-20s %-20s %-20s %-20s\n" "$module_name" "$area1" "$area2" "$ratio"
    fi
done
}


#START ACTUAL EXECUTION

#Check if in menu mode or in CLI mode
if [ -z "$1" ]; then
    # in menu mode
    request_data
else
    
    #in cli mode. Filter through all the parameters
    while getopts ":d:f:m:v:x:r:" opt; do
        case $opt in
            d)  # -d option for directory
                file_directory="$OPTARG"
                ;;
            f)  # -f option for file
                FILE="$OPTARG"
                ;;
            m)  # -m option for module (requires -f to be set)
                MODULE="$OPTARG"
                ;;
            v) # -v visualize before and after synthesis
                echo "found visualize"
                visualize=1
                ;;
            x) # -x switch to extended nem liberty file
                echo "switching to 4T libert file"
                switch_liberty
                ;;
            r) # -r generate report of output
                echo "generating report"
                create_report
                ;;
            \?) # Invalid option
                echo "Invalid option: -$OPTARG" >&2
                usage
                ;;
            :)  # Missing argument for an option
                echo "Option -$OPTARG requires an argument." >&2
                usage
                ;;
        esac
    done

    #running synthesis on al lthe files in the directory
    if [ -n "$file_directory" ]; then
        if [ -d "$file_directory" ]; then
            echo "Directory exists: $file_directory"

            for file in "$file_directory"/*.v; do
            # Check if it's a regular file
            if [ -f "$file" ]; then
                # Use grep to find the line that starts with 'module' and extract the module name
                module_name=$(grep -m 1 -oP '^module\s+\K\w+' "$file")
                #module_name=
                # If the module name is found, print the file path and the module name
                if [ -n "$module_name" ]; then
                    echo "File: $file"
                    echo "Module: $module_name"
                    echo

                    FILE=$file
                    FILE_BASENAME=$(basename "$FILE" | cut -d. -f1)
                    #MODULE=$module_name
                    MODULE=$FILE_BASENAME

                    #synthesise the file
                    echo "running sequence of test commands"
                    run_yosys_file "synth_nem" 0
                    #run_yosys_file "sat_test" 0
                    switch_liberty
                    run_yosys_file "synth_nem" 0
                    compare_area
                    switch_liberty
                else
                    echo "No module found in file: $file"
                    echo
                fi
            fi
        done

        #done with synthesis
        create_report
        exit 0
        
        else
            echo "Directory does not exist: $file_directory"
            exit 1
        fi
    fi

    #running synthesis on the file requested
    if [ -n "$FILE" ]; then
        if [ -n "$MODULE" ]; then
            if [ -f "$FILE" ]; then
                echo "File exists: $file"
                echo "Module: $module"
                FILE_BASENAME=$(basename "$FILE" | cut -d. -f1)
                run_yosys_file "synth_nem" 0
                if [ "$visualize" -eq 1 ]; then
                run_yosys_file "visual" 0
                run_yosys_file "visual" 1
                else
                echo "no visualize set"
                fi
                exit 0
            else
                echo "File does not exist: $file"
                exit 1
            fi
        else
            echo "Missing module (-m) for the file (-f)."
            usage
        fi
    fi
    exit 1
 
fi


# Loop to allow multiple selections
while true; do
  show_menu
  read -p "Enter your choices (e.g., 1 2 3, or 0 to finish): " -a choices

  for choice in "${choices[@]}"; do
    case $choice in
        1)
        echo "performing synthesis"
        run_yosys_file "synth_nem" 0
        ;;
        2)
        echo "Plotting the initial design with $FILE and $MODULE"
        run_yosys_file "visual" 0
        ;;
        3)
        echo "Plotting the NEM design with $FILE and $MODULE"
        run_yosys_file "visual" 1
        ;;
        4)
        echo "Performing SAT test on $FILE and $MODULE"
        run_yosys_file "sat_test" 0
        ;;
        5)
        echo "Exporting FSM overview of the design"
        make clean #to make sure no previous .kiss2 file remains
        run_yosys_file "fsm_export" 0

        if [ -f "./temp/${FILE_BASENAME}.kiss2" ]; then
            # If the file exists, run the python script and xdot
            python3 ./yosys/kiss2dot.py ./temp/${FILE_BASENAME}.kiss2 > ./temp/${FILE_BASENAME}.dot
            xdot ./temp/${FILE_BASENAME}.dot
        else
            # If the file doesn't exist, print a message
            echo "Could not detect an FSM in ${MODULE}"
        fi
        ;;
        6)
        echo "Plotting the initial design with $FILE and $MODULE"
        make clean #Clean directories
        run_yosys_file "synth_nem" 0
        make all #build plugins
        ls ./plugins/*.so
        run_yosys_file "start_shell" 0 "$(for so_file in ./plugins/*.so; do echo -m "$so_file"; done)" #create a list of all plugins to load
        ;;
        7)
        echo "Switching libary"
        switch_liberty
        ;;
        8)
        echo "running sequence of test commands"
        run_yosys_file "synth_nem" 0
        #run_yosys_file "visual" 1
        switch_liberty
        run_yosys_file "synth_nem" 0
        #run_yosys_file "visual" 1
        run_yosys_file "bruteforce" 0
        compare_area
        ;;
        9)
        echo "exporting truth table and running in mockturtle for muxig"
        run_yosys_file "bruteforce" 0
        ;;
        10)
        echo "exporting truth table and running in mockturtle for MIG"
        run_yosys_file "bruteforce" 1
        ;;
        11)
        echo "requesting new module"
        request_data
        ;;
        0)
        echo "exiting"
        break 2
        ;;
        *)
        echo "Invalid choice. Please select a number between 1 and 6."
        ;;
    esac
  done

  echo
done
