#include <iostream>
#include <vector>
#include <queue>
#include <map>
#include <set>
#include <algorithm>
#include <random>
#include <ctime>

using namespace std;

//=========================================
// CONFIGURATION
//=========================================

const int NUM_PACKETS = 100;
const int REORDER_WINDOW = 32;
const int XOR_KEY = 0xAA;

double PACKET_LOSS_RATE = 0.10;      //10%
double OUT_OF_ORDER_RATE = 0.20;     //20%

//=========================================
// PACKET STRUCTURE
//=========================================
struct PDCPPacket
{
    int SN;
    int originalSize;
    int compressedSize;

    vector<unsigned char> payload;

    unsigned int integrity;
    bool duplicate;

    int txOrder;
    int rxOrder;
};

//=========================================
// RANDOM NUMBER GENERATOR
//=========================================
random_device rd;
mt19937 gen(rd());

//=========================================
// GENERATE RANDOM PAYLOAD
//=========================================
vector<unsigned char> generatePayload(int size)
{
    vector<unsigned char> data(size);

    uniform_int_distribution<> dis(0,255);

    for(int i=0;i<size;i++)
    {
        data[i] = dis(gen);
    }

    return data;
}

//=========================================
// HEADER COMPRESSION
// Original IP Header = 40 Bytes
// Compressed Header = 5 Bytes
//=========================================
int compressHeader(int packetSize)
{
    int originalHeader = 40;
    int compressedHeader = 5;

    return packetSize - originalHeader + compressedHeader;
}

//=========================================
// CHECKSUM FOR INTEGRITY
//=========================================
unsigned int calculateIntegrity(const vector<unsigned char>& data)
{
    unsigned int sum = 0;

    for(auto b : data)
    {
        sum += b;
    }

    return sum;
}

//=========================================
// XOR CIPHER
//=========================================
void xorCipher(vector<unsigned char>& data)
{
    for(auto &b : data)
    {
        b ^= XOR_KEY;
    }
}

//=========================================
// CREATE PDCP PACKET
//=========================================
PDCPPacket createPacket(int sn,int txOrder)
{
    uniform_int_distribution<> sizeDis(100,1500);

    int size = sizeDis(gen);

    PDCPPacket pkt;

    pkt.SN = sn;
    pkt.originalSize = size;
    pkt.compressedSize = compressHeader(size);

    pkt.payload = generatePayload(pkt.compressedSize);

    pkt.integrity = calculateIntegrity(pkt.payload);

    xorCipher(pkt.payload);

    pkt.txOrder = txOrder;
    pkt.rxOrder = -1;

    pkt.duplicate = false;

    return pkt;
}

//=========================================
// CHANNEL SIMULATION
//=========================================
vector<PDCPPacket> transmitChannel(vector<PDCPPacket> txPackets)
{
    vector<PDCPPacket> rxPackets;

    uniform_real_distribution<> lossDis(0,1);

    for(auto pkt : txPackets)
    {
        double loss = lossDis(gen);

        if(loss < PACKET_LOSS_RATE)
        {
            continue;
        }

        rxPackets.push_back(pkt);

        // Create duplicate packet
        if(lossDis(gen) < 0.05)
        {
            rxPackets.push_back(pkt);
        }
    }

    // Out of order delivery
    int swaps = rxPackets.size()*OUT_OF_ORDER_RATE;

    uniform_int_distribution<> indexDis(0, rxPackets.size()-1);

    for(int i=0;i<swaps;i++)
    {
        int a=indexDis(gen);
        int b=indexDis(gen);

        swap(rxPackets[a],rxPackets[b]);
    }

    return rxPackets;
}

//=========================================
// RECEIVER PROCESSING
//=========================================
void receiverProcessing(vector<PDCPPacket>& rxPackets)
{
    map<int,PDCPPacket> reorderBuffer;

    set<int> receivedSN;

    int expectedSN = 0;

    int delivered = 0;
    int duplicates = 0;
    int integrityFailures = 0;

    long totalDelay = 0;

    cout<<"\n====================================";
    cout<<"\nRECEIVER PROCESSING";
    cout<<"\n====================================\n";

    for(int i=0;i<rxPackets.size();i++)
    {
        rxPackets[i].rxOrder = i;

        PDCPPacket pkt = rxPackets[i];

        if(receivedSN.count(pkt.SN))
        {
            duplicates++;
            continue;
        }

        vector<unsigned char> copyPayload = pkt.payload;

        xorCipher(copyPayload);

        unsigned int recvIntegrity = calculateIntegrity(copyPayload);

        if(recvIntegrity != pkt.integrity)
        {
            integrityFailures++;
            continue;
        }

        receivedSN.insert(pkt.SN);

        reorderBuffer[pkt.SN] = pkt;

        while(reorderBuffer.find(expectedSN)!=reorderBuffer.end())
        {
            PDCPPacket deliverPkt = reorderBuffer[expectedSN];

            delivered++;

            totalDelay +=
            abs(deliverPkt.rxOrder - deliverPkt.txOrder);

            reorderBuffer.erase(expectedSN);

            expectedSN++;
        }
    }

    double avgDelay = 0;

    if(delivered>0)
        avgDelay = (double)totalDelay/delivered;

    cout<<"Packets Delivered       : "<<delivered<<endl;
    cout<<"Duplicate Packets       : "<<duplicates<<endl;
    cout<<"Integrity Failures      : "<<integrityFailures<<endl;
    cout<<"Average Reorder Delay   : "<<avgDelay<<endl;
    cout<<"Remaining In Buffer     : "<<reorderBuffer.size()<<endl;
}

//=========================================
// MAIN
//=========================================
int main()
{
    cout<<"=====================================\n";
    cout<<" PDCP LAYER PROCESSING SIMULATOR\n";
    cout<<" LTE / 5G NR\n";
    cout<<"=====================================\n";

    vector<PDCPPacket> txPackets;

    long totalOriginal = 0;
    long totalCompressed = 0;

    //--------------------------------------
    // Select SN Size
    //--------------------------------------
    int SN_BITS;

    cout<<"Enter PDCP SN Length (12 or 18): ";
    cin>>SN_BITS;

    int MAX_SN = (1<<SN_BITS);

    //--------------------------------------
    // TRANSMITTER
    //--------------------------------------
    for(int i=0;i<NUM_PACKETS;i++)
    {
        PDCPPacket pkt =
        createPacket(i%MAX_SN,i);

        txPackets.push_back(pkt);

        totalOriginal += pkt.originalSize;
        totalCompressed += pkt.compressedSize;
    }

    //--------------------------------------
    // Compression Statistics
    //--------------------------------------
    double efficiency =
    100.0*(totalOriginal-totalCompressed)
    /totalOriginal;

    cout<<"\nTRANSMITTER STATISTICS\n";
    cout<<"---------------------------\n";
    cout<<"Packets Generated      : "
        <<NUM_PACKETS<<endl;

    cout<<"Original Bytes         : "
        <<totalOriginal<<endl;

    cout<<"Compressed Bytes       : "
        <<totalCompressed<<endl;

    cout<<"Compression Efficiency : "
        <<efficiency<<" %"<<endl;

    //--------------------------------------
    // Channel
    //--------------------------------------
    vector<PDCPPacket> rxPackets =
    transmitChannel(txPackets);

    cout<<"\nCHANNEL STATISTICS\n";
    cout<<"---------------------------\n";
    cout<<"Packets After Channel : "
        <<rxPackets.size()<<endl;

    //--------------------------------------
    // Receiver
    //--------------------------------------
    receiverProcessing(rxPackets);

    cout<<"\n=====================================\n";
    cout<<"SIMULATION COMPLETED\n";
    cout<<"=====================================\n";

    return 0;
}