    // The contents of this file are in the public domain. See LICENSE_FOR_EXAMPLE_PROGRAMS.txt
    /*
        This is an example illustrating the use of the quantum computing
        simulation classes from the dlib C++ Library.

        This example assumes you are familiar with quantum computing and
        Grover's search algorithm and Shor's 9 bit error correcting code
        in particular.   The example shows how to simulate both of these
        algorithms.

        The code to simulate Grover's algorithm is primarily here to show
        you how to make custom quantum gate objects.  The Shor ECC example
        is simpler and uses just the default gates that come with the
        library.

    */

    #include <iostream>
    #include <complex>
    #include <ctime>
    #include <dlib/quantum_computing.h>
    #include <dlib/string.h>

    using namespace std;
    using namespace dlib;

    // ----------------------------------------------------------------------------------------
    // ----------------------------------------------------------------------------------------
    // ----------------------------------------------------------------------------------------

    // This declares a random number generator that we will be using below
    dlib::rand rnd;

    // ----------------------------------------------------------------------------------------

    void shor_encode (
        quantum_register& reg
    )
    /*!
        requires
            - reg.num_bits() == 1
        ensures
            - #reg.num_bits() == 9
            - #reg == the Shor error coding of the input register
    !*/
    {
        DLIB_CASSERT(reg.num_bits() == 1,"");

        quantum_register zeros;
        zeros.set_num_bits(8);
        reg.append(zeros);

        using namespace dlib::quantum_gates;
        const gate<1> h = hadamard();
        const gate<1> i = noop();

        // Note that the expression (h,i) represents the tensor product of the 1 qubit
        // h gate with the 1 qubit i gate and larger versions of this expression
        // represent even bigger tensor products.  So as you see below, we make gates
        // big enough to apply to our quantum register by listing out all the gates we
        // want to go into the tensor product and then we just apply the resulting gate
        // to the quantum register.

        // Now apply the gates that constitute Shor's encoding to the input register.
        (cnot<3,0>(),i,i,i,i,i).apply_gate_to(reg);
        (cnot<6,0>(),i,i).apply_gate_to(reg);
        (h,i,i,h,i,i,h,i,i).apply_gate_to(reg);
        (cnot<1,0>(),i,cnot<1,0>(),i,cnot<1,0>(),i).apply_gate_to(reg);
        (cnot<2,0>(),cnot<2,0>(),cnot<2,0>()).apply_gate_to(reg);
    }

    // ----------------------------------------------------------------------------------------

    void shor_decode (
        quantum_register& reg
    )
    /*!
        requires
            - reg.num_bits() == 9
        ensures
            - #reg.num_bits() == 1
            - #reg == the decoded qubit that was in the given input register
    !*/
    {
        DLIB_CASSERT(reg.num_bits() == 9,"");

        using namespace dlib::quantum_gates;
        const gate<1> h = hadamard();
        const gate<1> i = noop();

        // Now apply the gates that constitute Shor's decoding to the input register

        (cnot<2,0>(),cnot<2,0>(),cnot<2,0>()).apply_gate_to(reg);
        (cnot<1,0>(),i,cnot<1,0>(),i,cnot<1,0>(),i).apply_gate_to(reg);

        (toffoli<0,1,2>(),toffoli<0,1,2>(),toffoli<0,1,2>()).apply_gate_to(reg);

        (h,i,i,h,i,i,h,i,i).apply_gate_to(reg);

        (cnot<6,0>(),i,i).apply_gate_to(reg);
        (cnot<3,0>(),i,i,i,i,i).apply_gate_to(reg);
        (toffoli<0,3,6>(),i,i).apply_gate_to(reg);

        // Now that we have decoded the value we don't need the extra 8 bits any more so
        // remove them from the register.
        for (int i = 0; i < 8; ++i)
            reg.measure_and_remove_bit(0,rnd);
    }

    // ----------------------------------------------------------------------------------------
    // ----------------------------------------------------------------------------------------
    // ----------------------------------------------------------------------------------------

    // This is the function we will use in Grover's search algorithm.  In this
    // case the value we are searching for is 257.
    bool is_key (unsigned long n)
    {
        return n == 257;
    }

    // ----------------------------------------------------------------------------------------

    template <int bits>
    class uf_gate;

    namespace dlib {
    template <int bits>
    struct gate_traits<uf_gate<bits> >
    {
        static const long num_bits = bits;
        static const long dims = dlib::qc_helpers::exp_2_n<num_bits>::value;
    };}

    template <int bits>
    class uf_gate : public gate_exp<uf_gate<bits> >
    {
        /*!
            This gate represents the black box function in Grover's search algorithm.
            That is, it is the gate defined as follows:
                Uf|x>|y> = |x>|y XOR is_key(x)>

            See the documentation for the gate_exp object for the details regarding
            the compute_state_element() and operator() functions defined below.
        !*/
    public:
        uf_gate() : gate_exp<uf_gate>(*this) {}

        static const long num_bits = gate_traits<uf_gate>::num_bits;
        static const long dims = gate_traits<uf_gate>::dims;

        const qc_scalar_type operator() (long r, long c) const
        {
            unsigned long output = c;
            // if the input control bit is set
            if (is_key(output>>1))
            {
                output = output^0x1;
            }

            if ((unsigned long)r == output)
                return 1;
            else
                return 0;
        }

        template <typename exp>
        qc_scalar_type compute_state_element (
            const matrix_exp<exp>& reg,
            long row_idx
        ) const
        {
            unsigned long output = row_idx;
            // if the input control bit is set
            if (is_key(output>>1))
            {
                output = output^0x1;
            }

            return reg(output);
        }
    };

    // ----------------------------------------------------------------------------------------

    template <int bits>
    class w_gate;

    namespace dlib {
    template <int bits>
    struct gate_traits<w_gate<bits> >
    {
        static const long num_bits = bits;
        static const long dims = dlib::qc_helpers::exp_2_n<num_bits>::value;
    }; }

    template <int bits>
    class w_gate : public gate_exp<w_gate<bits> >
    {
        /*!
            This is the W gate from the Grover algorithm
        !*/
    public:

        w_gate() : gate_exp<w_gate>(*this) {}

        static const long num_bits = gate_traits<w_gate>::num_bits;
        static const long dims = gate_traits<w_gate>::dims;

        const qc_scalar_type operator() (long r, long c) const
        {
            qc_scalar_type res = 2.0/dims;
            if (r != c)
                return res;
            else
                return res - 1.0;
        }

        template <typename exp>
        qc_scalar_type compute_state_element (
            const matrix_exp<exp>& reg,
            long row_idx
        ) const
        {
            qc_scalar_type temp = sum(reg)*2.0/dims;
            // compute this value: temp = temp - reg(row_idx)*2.0/dims + reg(row_idx)*(2.0/dims - 1.0)
            temp = temp - reg(row_idx);

            return temp;
        }
    };

    // ----------------------------------------------------------------------------------------
    // ----------------------------------------------------------------------------------------
    // ----------------------------------------------------------------------------------------

    int main()
    {
        // seed the random number generator
        rnd.set_seed(cast_to_string(time(0)));

        // Pick out some of the gates we will be using below
        using namespace dlib::quantum_gates;
        const gate<1> h = quantum_gates::hadamard();
        const gate<1> z = quantum_gates::z();
        const gate<1> x = quantum_gates::x();
        const gate<1> i = quantum_gates::noop();

        quantum_register reg;

        // We will be doing the 12 qubit version of Grover's search algorithm.
        const int bits=12;
        reg.set_num_bits(bits);

        // set the quantum register to its initial state
        (i,i, i,i,i,i,i, i,i,i,i,x).apply_gate_to(reg);

        // Print out the starting bits
        cout << "starting bits: ";
        for (int i = reg.num_bits()-1; i >= 0; --i)
            cout << reg.probability_of_bit(i);
        cout << endl;

        // Now apply the Hadamard gate to all the input bits
        (h,h, h,h,h,h,h, h,h,h,h,h).apply_gate_to(reg);

        // Here we do the grover iteration
        for (int j = 0; j < 35; ++j)
        {
            (uf_gate<bits>()).apply_gate_to(reg);
            (w_gate<bits-1>(),i).apply_gate_to(reg);

            cout << j << " probability: bit 1 = " << reg.probability_of_bit(1) << ", bit 9 = " << reg.probability_of_bit(9) << endl;
        }

        cout << endl;

        // Print out the final probability of measuring a 1 for each of the bits
        for (int i = reg.num_bits()-1; i >= 1; --i)
            cout << "probability for bit " << i << " = " << reg.probability_of_bit(i) << endl;
        cout << endl;

        cout << "The value we want grover's search to find is 257 which means we should measure a bit pattern of 00100000001" << endl;
        cout << "Measured bits: ";
        // finally, measure all the bits and print out what they are.
        for (int i = reg.num_bits()-1; i >= 1; --i)
            cout << reg.measure_bit(i,rnd);
        cout << endl;

        // Now let's test out the Shor 9 bit encoding
        cout << "\n\n\n\nNow let's try playing around with Shor's 9bit error correcting code" << endl;

        // Reset the quantum register to contain a single bit
        reg.set_num_bits(1);
        // Set the state of this single qubit to some random mixture of the two computational bases
        reg.state_vector()(0) = qc_scalar_type(rnd.get_random_double(),rnd.get_random_double());
        reg.state_vector()(1) = qc_scalar_type(rnd.get_random_double(),rnd.get_random_double());
        // Make sure the state of the quantum register is a unit vector
        reg.state_vector() /= sqrt(sum(norm(reg.state_vector())));

        cout << "state: " << trans(reg.state_vector());

        shor_encode(reg);
        cout << "x bit corruption on bit 8" << endl;
        (x,i,i,i,i,i,i,i,i).apply_gate_to(reg); // mess up the high order bit
        shor_decode(reg); // try to decode the register

        cout << "state: " << trans(reg.state_vector());

        shor_encode(reg);
        cout << "x bit corruption on bit 1" << endl;
        (i,i,i,i,i,i,i,x,i).apply_gate_to(reg);
        shor_decode(reg);

        cout << "state: " << trans(reg.state_vector());

        shor_encode(reg);
        cout << "z bit corruption on bit 8" << endl;
        (z,i,i,i,i,i,i,i,i).apply_gate_to(reg);
        shor_decode(reg);

        cout << "state: " << trans(reg.state_vector());

        cout << "\nThe state of the input qubit survived all the corruptions in tact so the code works." << endl;

    }