#include <iostream>
#include <fstream>
#include <stdlib.h>
#include <math.h>
#include <stdio.h>
#include <vector>
#include <sstream>
#include "BasicFunctions.h"
#include "MatchingDecoder.h"
#include "DistanceFunctions.hpp"

using namespace std;

int main()
{

    // * Initialise DataOutput file *cd
    ofstream DataOutput; // Name variable for text file output
    // Name textfile
    DataOutput.open("QSUBFILENAME.txt");

    for (int BASHSYSSIZE = 3; BASHSYSSIZE < 7; BASHSYSSIZE += 2)
    {

        for (int BASHpINT = 0; BASHpINT < 20; BASHpINT += 2)
        {

            // * Set seed *
            srand((unsigned int)time(NULL) + BASHSYSSIZE + BASHpINT);

            // *** Initialise system ***

            // * Initialise simulation parameters *
            int L = BASHSYSSIZE; // System size
            cout << "L =  " << L << endl;

            int NoOfVertices = 3 * L * (L + 1) / 2; // Number of stabilizers in triangle
            int NoOfQubits = 1 + 3 * L * (L + 1);   // Number of qubits in triangle

            int NoOfSamples = 10000; // Vary this parameter accordingly

            // * Initalise lattice *
            bool QubitsX[NoOfQubits];
            bool StabilizersZ[NoOfVertices];

            // * Stabilizer X and Y coordinates in non-orthgonal lattice space are defined and stored in StabXPos and StabYPos*
            int StabXPos[NoOfVertices];
            int StabYPos[NoOfVertices];

            // * Qubit X and Y coordinates in non-orthgonal lattice space are defined and stored in QubXPos and QubYPos*
            int QubXPos[NoOfQubits];
            int QubYPos[NoOfQubits];

            // fill above intiialized arrays with the correct values
            XYPosFiller(StabXPos, StabYPos, QubXPos, QubYPos, L, NoOfVertices);

            // * Stabilizer coordinates for a particular lattice size have been stored *

            // ****  creating lookup tables for
            // (1) shortest distances between every point on the mobius strip, and
            // (2) which creases are crossed when travelling along these distances  ****

            int AllRLVertices = 3 * NoOfVertices;            // total no. of stabilizers on Mobius Strip
            int totalpoints = AllRLVertices * AllRLVertices; // total number of computable distances on MS

            vector<int> WrappedCostMatrix = GetAdjacencyMatrix(L, StabXPos, StabYPos); // returns the adjacency matrix graph representation for the MS (mobius strip)
            vector<int> DoubleWrap = Floyds(WrappedCostMatrix, NoOfVertices);          // returns container vector for true shortest distances betweeb all points, unwrapped below

            vector<int> WrappedDistances; // matrix to store shortest distance between every point in MS
            vector<int> WrappedFlips;     // matrix to store which creases are 'flipped' while travelling along the corresponding shortest paths in WrappedDistances

            for (int i = 0; i < AllRLVertices; i++)
            {
                for (int j = 0; j < AllRLVertices; j++)
                {
                    WrappedDistances.push_back(DoubleWrap[i * AllRLVertices + j]);
                    WrappedFlips.push_back(DoubleWrap[totalpoints + i * AllRLVertices + j]); // unwrapping vector returned from Floyds
                }
            }

            // the below alternate forces a matching that does NOT cross the red edge
            vector<int> AltDummyDoubleWrap = GetAlternateDummies(WrappedCostMatrix, NoOfVertices, StabXPos, StabYPos);
            vector<int> WrappedAltDistances;
            vector<int> WrappedAltFlips;
            vector<int> WrappedAltPreds;

            for (int i = 0; i < AllRLVertices; i++)
            {
                for (int j = 0; j < AllRLVertices; j++)
                {
                    WrappedAltDistances.push_back(AltDummyDoubleWrap[i * AllRLVertices + j]);
                    WrappedAltFlips.push_back(AltDummyDoubleWrap[totalpoints + i * AllRLVertices + j]);
                    WrappedAltPreds.push_back(AltDummyDoubleWrap[2 * totalpoints + i * AllRLVertices + j]); // unwrapping vector
                }
            }

            vector<int> DistancesRedEdgeRB = GetRPerpDistances(L, StabXPos, StabYPos);         // get distances of stabilizers on mobius strip from red edge going towards RB lattice
            vector<int> DistancesRedEdgeRG = GetRPerpDistancesOtherDir(L, StabXPos, StabYPos); // get distances of stabilizers on mobius strip from red edge going towards RG lattice

            // ****  lookuptables are now created ****

            double p = 0.078 + 0.001 * (double)(BASHpINT); // physical error rate
            cout << "p =  " << p << endl;

            int FailCount = 0; // count failures

            for (int samples = 0; samples < NoOfSamples; samples++)
            {

                for (int j = 0; j < NoOfVertices; j++) // initilize system with 0 errors
                {
                    QubitsX[2 * j] = 0;
                    QubitsX[2 * j + 1] = 0;

                    StabilizersZ[j] = 0;
                }
                QubitsX[NoOfQubits - 1] = 0; // noofqubits = noofstabilizers*2 + 1 !

                // Add an error to each qubit with probability 'p', initialised above
                for (int j = 0; j < NoOfQubits; j++)
                {
                    if (((double)rand() / (double)RAND_MAX) < p)
                    {
                        QubitsX[j] = 1;
                    }
                }

                // *** Begin simulation ***

                // * Syndrome evaluation *
                // We make a list of defects ONLY on the triangular lattice
                // Defects are marked by boolian variables in the '1' state on an array
                for (int j = 0; j < NoOfVertices; j++)
                {
                    // Function 'StabilizerMeasurement' evaluates hexagonal measurements ONLY on the triangular lattice
                    StabilizersZ[j] = StabilizerMeasurement(QubitsX, j, L, StabXPos, StabYPos);
                }

                // expand StabilizersZ on triangle to Mobius Strip
                bool MobiusStabs[NoOfVertices * 3];
                ExpandToMobius(MobiusStabs, StabXPos, StabYPos, StabilizersZ, NoOfVertices);

                // take the locations of the even number of defects on the MS + the lookup tables on the MS and returns a MWPM
                vector<int> DecoderPairs = MatchingDecoder(MobiusStabs, L, WrappedDistances, WrappedFlips, StabXPos, StabYPos);
                int ModDecoderEstimate = DecoderOutcome(DecoderPairs, L); // outputs a value indicating which creases on the MS were crossed an odd number of times
                int RedCrossed = ModDecoderEstimate / 8;
                int DecoderEstimate = ModDecoderEstimate % 8;

                // now we prepare an alternate correction
                int DummyDecoderEstimate = 0;
                vector<int> DummyPairs;
                if (RedCrossed % 2 == 1) // red edge has been crossed an odd no of times; now dummy must not cross red
                {
                    DummyPairs = MatchingDecoder(MobiusStabs, L, WrappedAltDistances, WrappedAltFlips, StabXPos, StabYPos);
                    DummyDecoderEstimate = DecoderOutcome(DummyPairs, L) % 8;
                }
                else // now dummy must cross red edge
                {
                    DummyPairs = DummyMatchingDecoder(MobiusStabs, L, WrappedAltDistances, WrappedAltFlips, DistancesRedEdgeRB, DistancesRedEdgeRG, DecoderPairs, StabXPos, StabYPos);
                    DummyDecoderEstimate = DummyOutcome(DummyPairs, L);
                }

                int SelectedEstimate = DecoderEstimate;

                vector<int> w1w2 = GetMatchingWeights(DecoderPairs, DummyPairs); // calculate orginal and alternate matching weights. return {original, alternate}

                if (abs(w1w2[1] - w1w2[0]) == 1 && (3 * L - w1w2[0]) % 2 == 1) // if two conditions are fulfilled then we use the alternate correction instead of the original
                {
                    SelectedEstimate = DummyDecoderEstimate; // COMMENT OUT THIS LINE TO ONLY USE THE ORIGINAL DECODER
                }

                // Measurement returns a value indicating which boundaries on the RGB triangle have an odd number of qubits in state '1' on them
                int Measurement = LogicalMeasurement(QubitsX, L);

                if (SelectedEstimate != Measurement)
                {
                    FailCount++;
                }
            }

            DataOutput << L << " " << p << " " << FailCount << " " << NoOfSamples << endl;
        }
    }

    // The text output file is closed
    DataOutput.close();
    return 0;
}