Untitled

#include <iostream>
#include <fstream>
#include <vector>
#include <string>
#include <sstream>
#include <map>
#include <queue>
#include <cmath>

using namespace std;

// Structure to represent a transition in a Finite State Machine (FSM)
struct FSMTransition
{
    string inputSignal;  // The input signal for the transition
    string currentState; // The current state before the transition
    string nextState;    // The state the machine transitions to
    string outputSignal; // The output signal generated by the transition
};

// Global variables to store FSM and watermark information
extern int numInputSignals;       // Number of input signals
extern int numOutputSignals;      // Number of output signals
extern int numStates;             // Number of states in FSM
extern int numTransitions;        // Number of transitions in FSM
extern int numPossibleInputs;     // Number of possible inputs (2^numInputSignals)
extern int watermarkIndex;        // Index to track the insertion of watermark inputs
extern string initialFSMState;    // The initial state of the FSM
extern string availableState;     // A state available for new transitions
extern map<string, map<string, FSMTransition>> fsmTransitionTable;  // The FSM transition table
extern vector<string> watermarkInputs;   // Vector to store watermark input signals
extern vector<string> watermarkOutputs;  // Vector to store watermark output signals
extern vector<string> transitionPath;    // Path to track transitions during watermark embedding

// Function prototypes
static string toStr(int value);  // Helper function to convert integer to string
inline void createFSMTransition(const string& currentState, const string& inputSignal, const string& outputSignal); // Function to create FSM transitions
void parseFSMFile(const string& filename); // Function to parse the FSM input file
void parseWatermarkFile(const string& filename); // Function to parse the watermark file
void findMissingTransition(const string& state, vector<string> path); // Function to find missing transitions in the FSM
string getStartStateOrInsertTransitions(); // Function to get the start state or insert new transitions
void embedWatermarkToFSM(); // Function to embed the watermark into the FSM
void saveFSMToFile(const string& filename); // Function to save the FSM to a file
void saveLatencyToFile(const string& filename); // Function to save the latency path to a file

// Initialize global variables
int numInputSignals = 0;       // Number of input signals in FSM
int numOutputSignals = 0;      // Number of output signals in FSM
int numStates = 0;             // Total number of states in FSM
int numTransitions = 0;        // Total number of transitions in FSM
int numPossibleInputs = 0;     // 2^numInputSignals possible input combinations
int watermarkIndex = 0;        // Keeps track of where to insert the watermark in the FSM
string initialFSMState;        // The FSM's initial state
string availableState;         // A state available for new transitions

// FSM transition table: Maps current state -> input signal -> FSMTransition
map<string, map<string, FSMTransition>> fsmTransitionTable;

// Vectors for watermark inputs/outputs
vector<string> watermarkInputs;
vector<string> watermarkOutputs;

// Path to track the transition sequence while embedding watermark
vector<string> transitionPath;

// Main function to execute watermark embedding
int main(int argc, char* argv[])
{
    string fsmInputFile = argv[1];   // Input FSM file
    string watermarkFile = argv[2];   // Input watermark file
    string resultFile = argv[3];      // Output FSM file with watermark embedded
    string latencyFile = argv[4];     // Output latency file

    // Parse the FSM and watermark files
    parseFSMFile(fsmInputFile);
    parseWatermarkFile(watermarkFile);

    numPossibleInputs = 1 << numInputSignals;  // Calculate the number of possible inputs

    // Embed the watermark into the FSM
    embedWatermarkToFSM();

    // Save the modified FSM and latency information
    saveFSMToFile(resultFile);
    saveLatencyToFile(latencyFile);

    cout << "Watermark embedding completed successfully!" << endl;
    return 0;
}

// Function to convert an integer to string
static string toStr(int value)
{
    ostringstream oss;
    oss << value;
    return oss.str();
}

// Function to create a new FSM transition with a specified current state, input, and output signals
inline void createFSMTransition(const string& currentState, const string& inputSignal, const string& outputSignal)
{
    FSMTransition transition;
    transition.inputSignal = inputSignal;
    transition.currentState = currentState;
    transition.nextState = "S" + toStr(numStates++);  // Increment and assign a new state name
    transition.outputSignal = outputSignal;
    fsmTransitionTable[currentState][inputSignal] = transition;  // Insert transition into the FSM table
}

// Function to parse the FSM file and populate global FSM variables
void parseFSMFile(const string& filename)
{
    ifstream file(filename.c_str(), ios::in);  // Open the FSM file for reading
    if (!file.is_open())
    {
        cerr << "Error opening FSM file." << endl;
        return;
    }

    string line;
    while (getline(file, line))
    {
        if (line.empty() || line[0] == '#')  // Skip empty lines or comments
        {
            continue;
        }

        if (line[0] == '.')  // Parse FSM configuration settings
        {
            istringstream iss(line);
            if (line[1] == 'i')  // Number of input signals
            {
                iss >> line >> numInputSignals;
            }
            else if (line[1] == 'o')  // Number of output signals
            {
                iss >> line >> numOutputSignals;
            }
            else if (line[1] == 'p')  // Number of transitions
            {
                iss >> line >> numTransitions;
            }
            else if (line[1] == 's')  // Number of states
            {
                iss >> line >> numStates;
            }
            else if (line[1] == 'r')  // Initial state
            {
                iss >> line >> initialFSMState;
            }
            else if (line[1] == 'e')  // End marker
            {
                break;
            }
        }
        else  // Parse FSM transitions
        {
            FSMTransition transition;
            istringstream iss(line);
            iss >> transition.inputSignal >> transition.currentState >> transition.nextState >> transition.outputSignal;
            fsmTransitionTable[transition.currentState][transition.inputSignal] = transition;  // Add transition to FSM table
        }
    }

    file.close();
}

// Function to parse the watermark file and extract input/output signal segments
void parseWatermarkFile(const string& filename)
{
    ifstream file(filename.c_str(), ios::in);  // Open the watermark file for reading
    if (!file.is_open())
    {
        cerr << "Error opening watermark file." << endl;
        return;
    }

    string watermark;
    getline(file, watermark);  // Read the entire watermark into a string
    file.close();

    int segmentLength = numInputSignals + numOutputSignals;
    size_t idx = 0;

    while (idx + segmentLength <= watermark.size())
    {
        // Split the watermark into input/output signal segments
        string inputSegment = watermark.substr(idx, numInputSignals);
        string outputSegment = watermark.substr(idx + numInputSignals, numOutputSignals);

        watermarkInputs.push_back(inputSegment);    // Store input signal
        watermarkOutputs.push_back(outputSegment);  // Store output signal

        idx += segmentLength;
    }
}

// Function to recursively find missing transitions and mark the state path
void findMissingTransition(const string& state, vector<string> path)
{
    if (fsmTransitionTable[state].size() < (size_t)numPossibleInputs)
    {
        // Check if there are missing transitions for the current state
        for (int i = 0; i < numPossibleInputs; i++)
        {
            string binaryInput;
            int temp = i;
            while (temp > 0)
            {
                binaryInput = ((temp % 2) ? "1" : "0") + binaryInput;
                temp /= 2;
            }

            while (binaryInput.size() < (size_t)numInputSignals)
            {
                binaryInput = "0" + binaryInput;  // Pad binary input to match input signal count
            }

            if (fsmTransitionTable[state].find(binaryInput) == fsmTransitionTable[state].end())
            {
                // Record the transition path and available state if transition is missing
                transitionPath = path;
                transitionPath.push_back(binaryInput);
                availableState = state;
                return;
            }
        }
    }
    else
    {
        // Recursively search through the FSM transitions for missing transitions
        for (const auto& pair : fsmTransitionTable[state])
        {
            path.push_back(pair.first);
            findMissingTransition(pair.second.nextState, path);
            path.pop_back();
        }
    }
}

// Function to get the start state for watermark embedding, or insert new transitions if necessary
string getStartStateOrInsertTransitions()
{
    if (fsmTransitionTable[initialFSMState][watermarkInputs[0]].outputSignal == watermarkOutputs[0])
    {
        return initialFSMState;  // If the initial state already matches, return it as the start state
    }
    else
    {
        string startState = "";
        for (; startState == ""; watermarkIndex++)
        {
            // Search through FSM transitions to find matching transitions
            for (const auto& statePair : fsmTransitionTable)
            {
                if (statePair.second.find(watermarkInputs[watermarkIndex]) != statePair.second.end())
                {
                    const FSMTransition& transition = statePair.second.at(watermarkInputs[watermarkIndex]);
                    if (transition.outputSignal == watermarkOutputs[watermarkIndex])
                    {
                        startState = statePair.first;  // Found matching state
                        break;
                    }
                }
            }
        }

        watermarkIndex--;  // Decrease index after finding the start state

        // Find missing transition if any
        findMissingTransition(initialFSMState, vector<string>());

        // Create the missing transition
        createFSMTransition(availableState, transitionPath.back(), watermarkOutputs[0]);
        string newlyCreatedState = "S" + toStr(numStates - 1);

        // Create the remaining watermark transitions
        for (int i = 0; i < watermarkIndex; i++)
        {
            createFSMTransition(newlyCreatedState, watermarkInputs[i], watermarkOutputs[i]);
        }

        // If it's the first watermark input, create a transition back to the start state
        if (watermarkIndex == 0)
        {
            FSMTransition transition;
            transition.inputSignal = watermarkInputs[watermarkIndex];
            transition.currentState = newlyCreatedState;
            transition.nextState = startState;
            transition.outputSignal = watermarkOutputs[watermarkIndex];
            fsmTransitionTable[newlyCreatedState][watermarkInputs[watermarkIndex]] = transition;

            transitionPath.push_back(watermarkInputs[watermarkIndex]);
        }
        else
        {
            // Create transition back to the start state for other watermark inputs
            FSMTransition transition;
            transition.inputSignal = watermarkInputs[watermarkIndex - 1];
            transition.currentState = newlyCreatedState;
            transition.nextState = startState;
            transition.outputSignal = watermarkOutputs[watermarkIndex - 1];
            fsmTransitionTable[newlyCreatedState][watermarkInputs[watermarkIndex - 1]] = transition;

            startState = newlyCreatedState;
        }

        return startState;  // Return the start state after embedding watermark transitions
    }
}

// Function to embed watermark signals into the FSM
void embedWatermarkToFSM()
{
    string currentState = getStartStateOrInsertTransitions();  // Get or insert start state for watermark

    // For each watermark output, embed it into the FSM transitions
    for (size_t i = 0; i < watermarkOutputs.size(); i++)
    {
        if (static_cast<int>(i) < watermarkIndex)
        {
            // Check if the current state already has the transition
            if (fsmTransitionTable[currentState].find(watermarkInputs[i]) == fsmTransitionTable[currentState].end())
            {
                // Create a new transition if it does not exist
                createFSMTransition(currentState, watermarkInputs[i], watermarkOutputs[i]);
                currentState = fsmTransitionTable[currentState][watermarkInputs[i]].nextState;
                numTransitions++;  // Increment transition count
            }
            else if (fsmTransitionTable[currentState][watermarkInputs[i]].outputSignal == watermarkOutputs[i])
            {
                // If the transition already exists and the output signal matches, move to next state
                currentState = fsmTransitionTable[currentState][watermarkInputs[i]].nextState;
            }
            else
            {
                // If the transition output doesn't match, create a new transition
                string lastCreatedState = "S" + toStr(numStates - 1);
                createFSMTransition(lastCreatedState, watermarkInputs[i], watermarkOutputs[i]);
                currentState = lastCreatedState;
                numTransitions++;  // Increment transition count
            }
        }
        else
        {
            // For remaining watermark outputs, always create a new transition
            string lastCreatedState = "S" + toStr(numStates - 1);
            createFSMTransition(lastCreatedState, watermarkInputs[i], watermarkOutputs[i]);
            currentState = lastCreatedState;
            numTransitions++;  // Increment transition count
        }
    }
}

// Function to save the FSM to an output file
void saveFSMToFile(const string& filename)
{
    ofstream file(filename.c_str(), ios::out);  // Open the output file for writing
    if (!file.is_open())
    {
        cerr << "Error." << endl;
        return;
    }

    // Write FSM configuration settings
    file << ".i " << numInputSignals << "\n";
    file << ".o " << numOutputSignals << "\n";
    file << ".p " << numTransitions << "\n";
    file << ".s " << numStates << "\n";
    file << ".r " << initialFSMState << "\n";

    // Write FSM transitions
    for (const auto& statePair : fsmTransitionTable)
    {
        for (const auto& transitionPair : statePair.second)
        {
            const FSMTransition& transition = transitionPair.second;
            file << transition.inputSignal  << " "
                 << transition.currentState << " "
                 << transition.nextState    << " "
                 << transition.outputSignal << "\n";
        }
    }

    file << ".e\n";  // End marker
    file.close();
}

// Function to save the latency path to a file
void saveLatencyToFile(const string& filename)
{
    ofstream file(filename.c_str(), ios::out);  // Open the latency file for writing
    if (!file.is_open())
    {
        cerr << "Error." << endl;
        return;
    }

    file << ".n " << transitionPath.size() << "\n";  // Write the size of the transition path
    for (size_t i = 0; i < transitionPath.size(); i++)
    {
        if (i > 0) file << "_";  // Separator between transition path elements
        file << transitionPath[i];
    }

    file << "\n.e\n";  // End marker
    file.close();
}