#!/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 with normal Yosys AIG flow"
    echo -e "${GREEN}2)${RESET} Synthesize to NEM and perform MUXIG optimisaiton run"
    echo -e "${GREEN}3)${RESET} Synthesize to NEM and perform MIG optimisation run"
    echo -e "${GREEN}4)${RESET} Print initial design"
    echo -e "${GREEN}5)${RESET} Print out NEM optimized design"
    echo -e "${GREEN}6)${RESET} Perform SAT comparison"
    echo -e "${GREEN}7)${RESET} Export FSM as KISS2 format"
    echo -e "${GREEN}8)${RESET} Start shell with modules"
    echo -e "${GREEN}9)${RESET} Switch from normal 3T gate library to new 4T"
    echo -e "${GREEN}10)${RESET} Run series of test and create comparison report"
    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() {
    local override="$1"
    if [[ ("$LIBERTY_FILE" == "nem_basic_yosys.lib" && "$override" != "3T") || "$override" == "4T" ]]; then
        LIBERTY_FILE="nem_basic_yosys_extended.lib"
        LIBERTY_USED="4T"
        echo "Now using extended (4T devices) library"
    elif [[ "$LIBERTY_FILE" == "nem_basic_yosys_extended.lib" || "$override" == "3T" ]]; then
        LIBERTY_FILE="nem_basic_yosys.lib"
        LIBERTY_USED="3T"
        echo "Now using normal library"
    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%"
    } > "./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,Mux,Ratio Mux,and_3T,nand_3T,or_3T,nor_3T,xor_3T,xnor_3T,inv_3T,mux_3T,muxig_inv_3T,mux_4T,device_count_3T,device_count_muxig,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
            area3=$(grep "Chip area for module" "$file" | sed -n '3s/.*: //p')  # Area MUX

            # Extract sections
            normal_3T_section=$(awk '/------------- normal 3T ---------------/,/------------- 4T ---------------/' "$file")
            muxig_section=$(awk '/------------- Muxig ---------------/,/------------- Stats ---------------/' "$file")

            # Extract gate counts for "normal 3T"
            and_3T=$(echo "$normal_3T_section" | awk '/ and_3T/ {print $2}' || echo 0)
            nand_3T=$(echo "$normal_3T_section" | awk '/ nand_3T/ {print $2}' || echo 0)
            or_3T=$(echo "$normal_3T_section" | awk '/ or_3T/ {print $2}' || echo 0)
            nor_3T=$(echo "$normal_3T_section" | awk '/ nor_3T/ {print $2}' || echo 0)
            xor_3T=$(echo "$normal_3T_section" | awk '/ xor_3T/ {print $2}' || echo 0)
            xnor_3T=$(echo "$normal_3T_section" | awk '/ xnor_3T/ {print $2}' || echo 0)
            inv_3T=$(echo "$normal_3T_section" | awk '/ inv_3T/ {print $2}' || echo 0)
            mux_3T=$(echo "$normal_3T_section" | awk '/ mux_3T/ {print $2}' || echo 0)

            # Extract gate counts for "Muxig"
            muxig_inv_3T=$(echo "$muxig_section" | awk '/ inv_3T/ {print $2}' || echo 0)
            mux_4T=$(echo "$muxig_section" | awk '/ mux_4T/ {print $2}' || echo 0)

            # Calculate device counts
            device_count_3T=$((and_3T * 6 + nand_3T * 4 + or_3T * 6 + nor_3T * 4 + xor_3T * 12 + xnor_3T * 12 + mux_3T * 12 + inv_3T * 2))
            device_count_muxig=$((muxig_inv_3T * 2 + mux_4T * 2))

            # Extract the ratio
            ratio_4T=$(grep -oP '(?<=Ratio 4T->3T: )[\d.]+' "$file" || echo 0)  # Extract the ratio
            ratio_MUX=$(grep -oP '(?<=Ratio MUX->3T: )[\d.]+' "$file" || echo 0)  # Extract the ratio
            ratio_4T=$(printf "%.5f" "$ratio_4T")
            ratio_MUX=$(printf "%.5f" "$ratio_MUX")

            ratio_device_count=$(echo "($device_count_muxig / $device_count_3T)" | bc -l)

            # Append the data to the CSV file
            echo "$module_name,$area1,$area3,$ratio_MUX,$and_3T,$nand_3T,$or_3T,$nor_3T,$xor_3T,$xnor_3T,$inv_3T,$mux_3T,$muxig_inv_3T,$mux_4T,$device_count_3T,$device_count_muxig,$ratio_device_count" >> "$csv_output"

            # Print the results in the table format
            printf "%-20s %-20s %-20s %-20s %-20s %-20s\n" "$module_name" "$area1" "$area3" "$ratio_4T" "$area3" "$ratio_MUX"
        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")
                
                # 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"
                    switch_liberty "3T"
                    run_yosys_file "synth_nem" 0
                    #run_yosys_file "sat_test" 0
                    switch_liberty "4T"
                    run_yosys_file "synth_nem" 0
                    #run_yosys_file "sat_test" 0
                    switch_liberty "4T"
                    run_yosys_file "bruteforce" 0
                    switch_liberty "3T"
                    compare_area
                    
                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
                compare_area
                create_report
                exit 0
            else
                echo "File does not exist: $file"
                exit 1
            fi
        else
            echo "Missing module (-m) for the file (-f)."
        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 with normal Yosys flow"
        run_yosys_file "synth_nem" 0
        ;;
        2)
        echo "exporting BLIf and running in mockturtle for muxig"
        run_yosys_file "bruteforce" 0
        ;;
        3)
        echo "exporting BLIF and running in mockturtle for MIG"
        run_yosys_file "bruteforce" 1
        ;;
        4)
        echo "Plotting the initial design with $FILE and $MODULE"
        run_yosys_file "visual" 0
        ;;
        5)
        echo "Plotting the NEM design with $FILE and $MODULE"
        run_yosys_file "visual" 1
        ;;
        6)
        echo "Performing SAT test on $FILE and $MODULE"
        run_yosys_file "sat_test" 0
        ;;
        7)
        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
        ;;
        8)
        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
        ;;
        9)
        echo "Switching libary"
        switch_liberty
        ;;
        10)
        echo "running sequence of test commands"
        switch_liberty "3T"
        run_yosys_file "synth_nem" 0
        #run_yosys_file "visual" 1
        switch_liberty "4T"
        run_yosys_file "synth_nem" 0
        #run_yosys_file "visual" 1
        run_yosys_file "bruteforce" 0
        compare_area
        ;;
        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