/// # Sample
/// Quantum Dynamics
///
/// # Description
/// This example demonstrates quantum dynamics in a style tailored for
/// resource estimation. The sample is specifically the simulation
/// of an Ising model Hamiltonian on an N1xN2 2D lattice using a
/// fourth-order Trotter Suzuki product formula, assuming
/// a 2D qubit architecture with nearest-neighbor connectivity.
/// The is an example of a program that is not amenable to simulating
/// classically, but can be run through resource estimation to determine
/// what size of quantum system would be needed to solve the problem.
namespace QuantumDynamics {

    open Microsoft.Quantum.Math;
    open Microsoft.Quantum.Arrays;


    @EntryPoint()
    operation Main() : Unit {
        // n : Int, m : Int, t: Double, u : Double, tstep : Double

        let n = 10;
        let m = 10;

        let J = 1.0;
        let g = 1.0;

        let totTime = 30.0;
        let dt = 0.9;

        IsingModel2DSim(n, m, J, g, totTime, dt);
    }

    /// # Summary
    /// The function below creates a sequence containing the rotation angles that will be applied with the two operators used in the expansion of the Trotter-Suzuki formula.
    /// # Input
    /// ## p (Double) : Constant used for fourth-order formulas
    ///
    /// ## dt (Double) : Time-step used to compute rotation angles
    ///
    /// ## J (Double) : coefficient for 2-qubit interactions
    ///
    /// ## g (Double) : coefficient for transverse field
    ///
    /// # Output
    /// ## values (Double[]) : The list of rotation angles to be applies in sequence with the corresponding operators
    ///
    function SetAngleSequence(p : Double, dt : Double, J : Double, g : Double) : Double[] {

        let len1 = 3;
        let len2 = 3;
        let valLength = 2 * len1 + len2 + 1;
        mutable values = [0.0, size = valLength];

        let val1 = J * p * dt;
        let val2 = -g * p * dt;
        let val3 = J * (1.0 - 3.0 * p) * dt / 2.0;
        let val4 = g * (1.0 - 4.0 * p) * dt / 2.0;

        for i in 0..len1 {

            if (i % 2 == 0) {
                set values w/= i <- val1;
            } else {
                set values w/= i <- val2;
            }

        }

        for i in len1 + 1..len1 + len2 {
            if (i % 2 == 0) {
                set values w/= i <- val3;
            } else {
                set values w/= i <- val4;
            }
        }

        for i in len1 + len2 + 1..valLength - 1 {
            if (i % 2 == 0) {
                set values w/= i <- val1;
            } else {
                set values w/= i <- val2;
            }
        }
        return values;
    }

    /// # Summary
    /// Applies e^-iX(theta) on all qubits in the 2D lattice as part of simulating the transverse field in the Ising model
    /// # Input
    /// ## n (Int) : Lattice size for an n x n lattice
    ///
    /// ## qArr (Qubit[][]) : Array of qubits representing the lattice
    ///
    /// ## theta (Double) : The angle/time-step for which the unitary simulation is done.
    ///
    operation ApplyAllX(n : Int, qArr : Qubit[][], theta : Double) : Unit {
        // This applies `Rx` with an angle of `2.0 * theta` to all qubits in `qs`
        // using partial application
        for row in 0..n - 1 {
            ApplyToEach(Rx(2.0 * theta, _), qArr[row]);
        }
    }

    /// # Summary
    /// Applies e^-iP(theta) where P = Z o Z as part of the repulsion terms.
    /// # Input
    /// ## n, m (Int, Int) : Lattice sizes for an n x m lattice
    ///
    /// ## qArr (Qubit[]) : Array of qubits representing the lattice
    ///
    /// ## theta (Double) : The angle/time-step for which unitary simulation is done.
    ///
    /// ## dir (Bool) : Direction is true for vertical direction.
    ///
    /// ## grp (Bool) : Group is true for odd starting indices
    ///
    operation ApplyDoubleZ(n : Int, m : Int, qArr : Qubit[][], theta : Double, dir : Bool, grp : Bool) : Unit {
        let start = grp ? 1 | 0;    // Choose either odd or even indices based on group number
        let P_op = [PauliZ, PauliZ];
        let c_end = dir ? m - 1 | m - 2;
        let r_end = dir ? m - 2 | m - 1;

        for row in 0..r_end {
            for col in start..2..c_end {
                // Iterate through even or odd columns based on `grp`

                let row2 = dir ? row + 1 | row;
                let col2 = dir ? col | col + 1;

                Exp(P_op, theta, [qArr[row][col], qArr[row2][col2]]);
            }
        }
    }

    /// # Summary
    /// The main function that takes in various parameters and calls the operations needed to simulate fourth order Trotterizatiuon of the Ising Hamiltonian for a given time-step
    /// # Input
    /// ## N1, N2 (Int, Int) : Lattice sizes for an N1 x N2 lattice
    ///
    /// ## J (Double) : coefficient for 2-qubit interactions
    ///
    /// ## g (Double) : coefficient for transverse field
    ///
    /// ## totTime (Double) : The total time-step for which unitary simulation is done.
    ///
    /// ## dt (Double) : The time the simulation is done for each timestep
    ///
    operation IsingModel2DSim(N1 : Int, N2 : Int, J : Double, g : Double, totTime : Double, dt : Double) : Unit {

        use qs = Qubit[N1 * N2];
        let qubitArray = Chunks(N2, qs); // qubits are re-arranged to be in an N1 x N2 array

        let p = 1.0 / (4.0 - 4.0^(1.0 / 3.0));
        let t = Ceiling(totTime / dt);

        let seqLen = 10 * t + 1;

        let angSeq = SetAngleSequence(p, dt, J, g);

        for i in 0..seqLen - 1 {
            let theta = (i == 0 or i == seqLen - 1) ? J * p * dt / 2.0 | angSeq[i % 10];

            // for even indexes
            if i % 2 == 0 {
                ApplyAllX(N1, qubitArray, theta);
            } else {
                // iterate through all possible combinations for `dir` and `grp`.
                for (dir, grp) in [(true, true), (true, false), (false, true), (false, false)] {
                    ApplyDoubleZ(N1, N2, qubitArray, theta, dir, grp);
                }
            }
        }
    }

}
