// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

export
    PreparePureStateD,
    ApproximatelyPreparePureStateCP,
    PrepareUniformSuperposition;

import
    Std.Diagnostics.Fact,
    Std.Convert.ComplexAsComplexPolar,
    Std.Convert.IntAsDouble,
    Std.Arithmetic.ApplyIfGreaterLE,
    Std.Arithmetic.ApplyIfGreaterL,
    Std.Arithmetic.ReflectAboutInteger,
    Std.Convert.IntAsBigInt,
    Std.Math.*,
    Std.Arrays.*;

/// # Summary
/// Given a set of coefficients and a big-endian quantum register,
/// prepares a state on that register described by the given coefficients.
///
/// # Description
/// This operation prepares an arbitrary quantum
/// state |𝜓⟩ with coefficients 𝑎ⱼ from
/// the n-qubit computational basis state |0...0⟩.
///
/// The action of U on the all-zeros state is given by
/// $$
/// \begin{align}
///     U \ket{0\cdots 0} = \ket{\psi} = \frac{\sum_{j=0}^{2^n-1}\alpha_j \ket{j}}{\sqrt{\sum_{j=0}^{2^n-1}|\alpha_j|^2}}.
/// \end{align}
/// $$
///
/// # Input
/// ## coefficients
/// Array of up to 2ⁿ real coefficients. The j-th coefficient
/// indexes the number state |j⟩ encoded in big-endian format.
///
/// ## qubits
/// Qubit register encoding number states in a big-endian format. This is
/// expected to be initialized in the computational basis state |0...0⟩.
///
/// # Remarks
/// `coefficients` will be normalized and padded with
/// elements 𝑎ⱼ = 0.0 if fewer than 2ⁿ are specified.
///
/// # Example
/// The following snippet prepares the quantum state |𝜓⟩=√(1/8)|0⟩+√(7/8)|2⟩=√(1/8)|00⟩+√(7/8)|10⟩
/// in the qubit register `qubits`.
/// ```qsharp
/// let amplitudes = [Sqrt(0.125), 0.0, Sqrt(0.875), 0.0];
/// use qubits = Qubit[2];
/// PreparePureStateD(amplitudes, qubits);
/// ```
///
/// # References
/// - [arXiv:quant-ph/0406176](https://arxiv.org/abs/quant-ph/0406176)
///   "Synthesis of Quantum Logic Circuits",
///   Vivek V. Shende, Stephen S. Bullock, Igor L. Markov
///
/// # See Also
/// - Std.StatePreparation.ApproximatelyPreparePureStateCP
operation PreparePureStateD(coefficients : Double[], qubits : Qubit[]) : Unit is Adj + Ctl {
    let nQubits = Length(qubits);
    // Pad coefficients at tail length to a power of 2.
    let coefficientsPadded = Padded(-2^nQubits, 0.0, coefficients);
    let idxTarget = 0;

    // Note we use the reversed qubits array to get the endianness ordering that we expect
    // when corresponding qubit state to state vector index.
    let qubits = Reversed(qubits);

    // Since we know the coefficients are real, we can optimize the first round of adjoint approximate unpreparation by directly
    // computing the disentangling angles and the new coefficients on those doubles without producing intermediate complex numbers.

    // For each 2D block, compute disentangling single-qubit rotation parameters
    let (disentanglingY, disentanglingZ, newCoefficients) = StatePreparationSBMComputeCoefficientsD(coefficientsPadded);

    if nQubits == 1 {
        let (abs, arg) = newCoefficients[0]!;
        if (AbsD(arg) > 0.0) {
            Adjoint Exp([PauliI], -1.0 * arg, [qubits[idxTarget]]);
        }
    } elif (Any(c -> AbsComplexPolar(c) > 0.0, newCoefficients)) {
        // Some coefficients are outside tolerance
        let newControl = 2..(nQubits - 1);
        let newTarget = 1;
        Adjoint ApproximatelyUnprepareArbitraryState(0.0, newCoefficients, newControl, newTarget, qubits);
    }

    Adjoint ApproximatelyMultiplexPauli(0.0, disentanglingY, PauliY, qubits[1...], qubits[0]);
    Adjoint ApproximatelyMultiplexPauli(0.0, disentanglingZ, PauliZ, qubits[1...], qubits[0]);
}

/// # Summary
/// Given a set of coefficients and a big-endian quantum register,
/// prepares a state on that register described by the given coefficients,
/// up to a given approximation tolerance.
///
/// # Description
/// This operation prepares an arbitrary quantum
/// state |𝜓⟩ with complex coefficients rⱼ·𝒆^(𝒊·tⱼ) from
/// the n-qubit computational basis state |0...0⟩.
/// In particular, the action of this operation can be simulated by the
/// a unitary transformation U which acts on the all-zeros state as
///
/// $$
/// \begin{align}
///     U\ket{0...0}
///         & = \ket{\psi} \\\\
///         & = \frac{
///                 \sum_{j=0}^{2^n-1} r_j e^{i t_j} \ket{j}
///             }{
///                 \sqrt{\sum_{j=0}^{2^n-1} |r_j|^2}
///             }.
/// \end{align}
/// $$
///
/// # Input
/// ## tolerance
/// The approximation tolerance to be used when preparing the given state.
///
/// ## coefficients
/// Array of up to 2ⁿ complex coefficients represented by their
/// absolute value and phase (rⱼ, tⱼ). The j-th coefficient
/// indexes the number state |j⟩ encoded in a big-endian format.
///
/// ## qubits
/// Qubit register encoding number states in a big-endian format. This is
/// expected to be initialized in the computational basis state
/// |0...0⟩.
///
/// # Remarks
/// `coefficients` will be padded with
/// elements (rⱼ, tⱼ) = (0.0, 0.0) if fewer than 2ⁿ are
/// specified.
///
/// # References
/// - [arXiv:quant-ph/0406176](https://arxiv.org/abs/quant-ph/0406176)
///   "Synthesis of Quantum Logic Circuits",
///   Vivek V. Shende, Stephen S. Bullock, Igor L. Markov
operation ApproximatelyPreparePureStateCP(
    tolerance : Double,
    coefficients : ComplexPolar[],
    qubits : Qubit[]
) : Unit is Adj + Ctl {

    let nQubits = Length(qubits);
    // Pad coefficients at tail length to a power of 2.
    let coefficientsPadded = Padded(-2^nQubits, ComplexPolar(0.0, 0.0), coefficients);
    let idxTarget = 0;
    // Determine what controls to apply
    let rngControl = nQubits > 1 ? (1..(nQubits - 1)) | (1..0);
    // Note we use the reversed qubits array to get the endianness ordering that we expect
    // when corresponding qubit state to state vector index.
    Adjoint ApproximatelyUnprepareArbitraryState(
        tolerance,
        coefficientsPadded,
        rngControl,
        idxTarget,
        Reversed(qubits)
    );
}

/// # Summary
/// Implementation step of arbitrary state preparation procedure.
operation ApproximatelyUnprepareArbitraryState(
    tolerance : Double,
    coefficients : ComplexPolar[],
    rngControl : Range,
    idxTarget : Int,
    register : Qubit[]
) : Unit is Adj + Ctl {

    // For each 2D block, compute disentangling single-qubit rotation parameters
    let (disentanglingY, disentanglingZ, newCoefficients) = StatePreparationSBMComputeCoefficientsCP(coefficients);
    if (AnyOutsideToleranceD(tolerance, disentanglingZ)) {
        ApproximatelyMultiplexPauli(tolerance, disentanglingZ, PauliZ, register[rngControl], register[idxTarget]);
    }
    if (AnyOutsideToleranceD(tolerance, disentanglingY)) {
        ApproximatelyMultiplexPauli(tolerance, disentanglingY, PauliY, register[rngControl], register[idxTarget]);
    }
    // target is now in |0> state up to the phase given by arg of newCoefficients.

    // Continue recursion while there are control qubits.
    if (IsRangeEmpty(rngControl)) {
        let (abs, arg) = newCoefficients[0]!;
        if (AbsD(arg) > tolerance) {
            Exp([PauliI], -1.0 * arg, [register[idxTarget]]);
        }
    } elif (Any(c -> AbsComplexPolar(c) > tolerance, newCoefficients)) {
        // Some coefficients are outside tolerance
        let newControl = (RangeStart(rngControl) + 1)..RangeStep(rngControl)..RangeEnd(rngControl);
        let newTarget = RangeStart(rngControl);
        ApproximatelyUnprepareArbitraryState(tolerance, newCoefficients, newControl, newTarget, register);
    }
}

/// # Summary
/// Applies a Pauli rotation conditioned on an array of qubits, truncating
/// small rotation angles according to a given tolerance.
///
/// # Description
/// This applies a multiply controlled unitary operation that performs
/// rotations by angle $\theta_j$ about single-qubit Pauli operator $P$
/// when controlled by the $n$-qubit number state $\ket{j}$.
/// In particular, the action of this operation is represented by the
/// unitary
///
/// $$
/// \begin{align}
///     U = \sum^{2^n - 1}_{j=0} \ket{j}\bra{j} \otimes e^{i P \theta_j}.
/// \end{align}
/// $$
///
/// # Input
/// ## tolerance
/// A tolerance below which small coefficients are truncated.
///
/// ## coefficients
/// Array of up to $2^n$ coefficients $\theta_j$. The $j$th coefficient
/// indexes the number state $\ket{j}$ encoded in little-endian format.
///
/// ## pauli
/// Pauli operator $P$ that determines axis of rotation.
///
/// ## control
/// $n$-qubit control register that encodes number states $\ket{j}$ in
/// little-endian format.
///
/// ## target
/// Single qubit register that is rotated by $e^{i P \theta_j}$.
///
/// # Remarks
/// `coefficients` will be padded with elements $\theta_j = 0.0$ if
/// fewer than $2^n$ are specified.
operation ApproximatelyMultiplexPauli(
    tolerance : Double,
    coefficients : Double[],
    pauli : Pauli,
    control : Qubit[],
    target : Qubit
) : Unit is Adj + Ctl {
    within {
        MapPauliAxis(PauliZ, pauli, target);
    } apply {
        ApproximatelyMultiplexZ(tolerance, coefficients, control, target);
    }
}

/// # Summary
/// Implementation step of arbitrary state preparation procedure.
/// This version optimized for purely real coefficients represented by an array of doubles.
function StatePreparationSBMComputeCoefficientsD(
    coefficients : Double[]
) : (Double[], Double[], ComplexPolar[]) {
    mutable disentanglingZ = [];
    mutable disentanglingY = [];
    mutable newCoefficients = [];

    for idxCoeff in 0..2..Length(coefficients) - 1 {
        let (rt, phi, theta) = BlockSphereCoordinatesD(coefficients[idxCoeff], coefficients[idxCoeff + 1]);
        set disentanglingZ += [0.5 * phi];
        set disentanglingY += [0.5 * theta];
        set newCoefficients += [rt];
    }

    return (disentanglingY, disentanglingZ, newCoefficients);
}

/// # Summary
/// Implementation step of arbitrary state preparation procedure.
function StatePreparationSBMComputeCoefficientsCP(
    coefficients : ComplexPolar[]
) : (Double[], Double[], ComplexPolar[]) {

    mutable disentanglingZ = [];
    mutable disentanglingY = [];
    mutable newCoefficients = [];

    for idxCoeff in 0..2..Length(coefficients) - 1 {
        let (rt, phi, theta) = BlochSphereCoordinatesCP(coefficients[idxCoeff], coefficients[idxCoeff + 1]);
        set disentanglingZ += [0.5 * phi];
        set disentanglingY += [0.5 * theta];
        set newCoefficients += [rt];
    }

    return (disentanglingY, disentanglingZ, newCoefficients);
}

/// # Summary
/// Computes the Bloch sphere coordinates for a single-qubit state.
///
/// Given two doubles $a0, a1$ that represent the qubit state, computes coordinates
/// on the Bloch sphere such that
/// $a0 \ket{0} + a1 \ket{1} = r e^{it}(e^{-i \phi /2}\cos{(\theta/2)}\ket{0}+e^{i \phi /2}\sin{(\theta/2)}\ket{1})$.
///
/// # Input
/// ## a0
/// Double coefficient of state $\ket{0}$.
/// ## a1
/// Double coefficient of state $\ket{1}$.
///
/// # Output
/// A tuple containing `(ComplexPolar(r, t), phi, theta)`.
function BlockSphereCoordinatesD(
    a0 : Double,
    a1 : Double
) : (ComplexPolar, Double, Double) {
    let abs0 = AbsD(a0);
    let abs1 = AbsD(a1);
    let arg0 = a0 < 0.0 ? PI() | 0.0;
    let arg1 = a1 < 0.0 ? PI() | 0.0;
    let r = Sqrt(abs0 * abs0 + abs1 * abs1);
    let t = 0.5 * (arg0 + arg1);
    let phi = arg1 - arg0;
    let theta = 2.0 * ArcTan2(abs1, abs0);
    return (ComplexPolar(r, t), phi, theta);
}

/// # Summary
/// Computes the Bloch sphere coordinates for a single-qubit state.
///
/// Given two complex numbers $a0, a1$ that represent the qubit state, computes coordinates
/// on the Bloch sphere such that
/// $a0 \ket{0} + a1 \ket{1} = r e^{it}(e^{-i \phi /2}\cos{(\theta/2)}\ket{0}+e^{i \phi /2}\sin{(\theta/2)}\ket{1})$.
///
/// # Input
/// ## a0
/// Complex coefficient of state $\ket{0}$.
/// ## a1
/// Complex coefficient of state $\ket{1}$.
///
/// # Output
/// A tuple containing `(ComplexPolar(r, t), phi, theta)`.
function BlochSphereCoordinatesCP(
    a0 : ComplexPolar,
    a1 : ComplexPolar
) : (ComplexPolar, Double, Double) {

    let abs0 = AbsComplexPolar(a0);
    let abs1 = AbsComplexPolar(a1);
    let arg0 = ArgComplexPolar(a0);
    let arg1 = ArgComplexPolar(a1);
    let r = Sqrt(abs0 * abs0 + abs1 * abs1);
    let t = 0.5 * (arg0 + arg1);
    let phi = arg1 - arg0;
    let theta = 2.0 * ArcTan2(abs1, abs0);
    return (ComplexPolar(r, t), phi, theta);
}

/// # Summary
/// Applies a Pauli Z rotation conditioned on an array of qubits, truncating
/// small rotation angles according to a given tolerance.
///
/// # Description
/// This applies the multiply controlled unitary operation that performs
/// rotations by angle $\theta_j$ about single-qubit Pauli operator $Z$
/// when controlled by the $n$-qubit number state $\ket{j}$.
/// In particular, this operation can be represented by the unitary
///
/// $$
/// \begin{align}
///     U = \sum^{2^n-1}_{j=0} \ket{j}\bra{j} \otimes e^{i Z \theta_j}.
/// \end{align}
/// $$
///
/// # Input
/// ## tolerance
/// A tolerance below which small coefficients are truncated.
///
/// ## coefficients
/// Array of up to $2^n$ coefficients $\theta_j$. The $j$th coefficient
/// indexes the number state $\ket{j}$ encoded in little-endian format.
///
/// ## control
/// $n$-qubit control register that encodes number states $\ket{j}$ in
/// little-endian format.
///
/// ## target
/// Single qubit register that is rotated by $e^{i P \theta_j}$.
///
/// # Remarks
/// `coefficients` will be padded with elements $\theta_j = 0.0$ if
/// fewer than $2^n$ are specified.
///
/// # References
/// - [arXiv:quant-ph/0406176](https://arxiv.org/abs/quant-ph/0406176)
///   "Synthesis of Quantum Logic Circuits",
///   Vivek V. Shende, Stephen S. Bullock, Igor L. Markov
operation ApproximatelyMultiplexZ(
    tolerance : Double,
    coefficients : Double[],
    control : Qubit[],
    target : Qubit
) : Unit is Adj + Ctl {

    body ... {
        let multiplexZParams = GenerateMultiplexZParams(coefficients, control, target);
        ApplyMultiplexZParams(multiplexZParams);
    }

    adjoint ... {
        let adjMultiplexZParams = GenerateAdjMultiplexZParams(coefficients, control, target);
        Adjoint ApplyMultiplexZParams(adjMultiplexZParams);
    }

    controlled (controlRegister, ...) {
        // Pad coefficients length to a power of 2.
        let coefficientsPadded = Padded(2^(Length(control) + 1), 0.0, Padded(-2^Length(control), 0.0, coefficients));
        let (coefficients0, coefficients1) = MultiplexZCoefficients(coefficientsPadded);
        ApproximatelyMultiplexZ(tolerance, coefficients0, control, target);
        if AnyOutsideToleranceD(tolerance, coefficients1) {
            within {
                Controlled X(controlRegister, target);
            } apply {
                ApproximatelyMultiplexZ(tolerance, coefficients1, control, target);
            }
        }
    }

    controlled adjoint invert;
}

operation ApplyMultiplexZParams(params : (Double, Qubit[])[]) : Unit is Adj {
    for (angle, qs) in params {
        if Length(qs) == 2 {
            CNOT(qs[0], qs[1]);
        } elif AbsD(angle) > 0.0 {
            Exp([PauliZ], angle, qs);
        }
    }
}
// Provides the sequence of angles or entangling CNOTs to apply for the multiplex Z step of the state preparation procedure,
// given a set of coefficients and control and target qubits.
// In these sequences, there are always one or two qubits in the qubit array. If there is one qubit, it is the target of a Z rotation by the given angle.
// If there are two qubits, they represent a CNOT with the first qubit as control and the second as target, and the angle is ignored.
function GenerateMultiplexZParams(
    coefficients : Double[],
    control : Qubit[],
    target : Qubit
) : (Double, Qubit[])[] {
    // Pad coefficients length at tail to a power of 2.
    let coefficientsPadded = Padded(-2^Length(control), 0.0, coefficients);

    if Length(coefficientsPadded) == 1 {
        // Termination case
        [(coefficientsPadded[0], [target])]
    } else {
        // Compute new coefficients.
        let (coefficients0, coefficients1) = MultiplexZCoefficients(coefficientsPadded);
        mutable params = GenerateMultiplexZParams(coefficients0, Most(control), target);
        params += [(0.0, [Tail(control), target])] + GenerateMultiplexZParams(coefficients1, Most(control), target);
        params += [(0.0, [Tail(control), target])];
        params
    }
}

// Provides the sequence of angles or entangling CNOTs to apply for the adjoint of the multiplex Z step of the state preparation procedure,
// given a set of coefficients and control and target qubits.
// Note that the adjoint sequence is NOT the reverse of the forward sequence due to the structure of the multiplex Z step, which applies
// the disentangling rotations in between the two recursive calls.
// See `GenerateMultiplexZParams` for more details on the structure of these sequences.
function GenerateAdjMultiplexZParams(
    coefficients : Double[],
    control : Qubit[],
    target : Qubit
) : (Double, Qubit[])[] {
    // Pad coefficients length at tail to a power of 2.
    let coefficientsPadded = Padded(-2^Length(control), 0.0, coefficients);

    if Length(coefficientsPadded) == 1 {
        // Termination case
        [(coefficientsPadded[0], [target])]
    } else {
        // Compute new coefficients.
        let (coefficients0, coefficients1) = MultiplexZCoefficients(coefficientsPadded);
        mutable params = [(0.0, [Tail(control), target])] + GenerateAdjMultiplexZParams(coefficients1, Most(control), target);
        params += [(0.0, [Tail(control), target])];
        params += GenerateAdjMultiplexZParams(coefficients0, Most(control), target);
        params
    }
}


/// # Summary
/// Implementation step of multiply-controlled Z rotations.
function MultiplexZCoefficients(coefficients : Double[]) : (Double[], Double[]) {
    let newCoefficientsLength = Length(coefficients) / 2;
    mutable coefficients0 = [];
    mutable coefficients1 = [];

    for idxCoeff in 0..newCoefficientsLength - 1 {
        set coefficients0 += [0.5 * (coefficients[idxCoeff] + coefficients[idxCoeff + newCoefficientsLength])];
        set coefficients1 += [0.5 * (coefficients[idxCoeff] - coefficients[idxCoeff + newCoefficientsLength])];
    }

    return (coefficients0, coefficients1);
}

function AnyOutsideToleranceD(tolerance : Double, coefficients : Double[]) : Bool {
    // NOTE: This function is not used as the only recursion termination condition
    // only to determine if the multiplex step needs to be applied.
    // For tolerance 0.0 it is always applied due to >= comparison.
    Any(coefficient -> AbsD(coefficient) >= tolerance, coefficients)
}

/// # Summary
/// Prepares a uniform superposition of states that represent integers 0 through
/// `nStates - 1` in a little-endian `qubits` register.
///
/// # Description
/// Given an input state $\ket{0\cdots 0}$ this operation prepares
/// a uniform superposition of all number states $0$ to $M-1$. In other words,
/// $$
/// \begin{align}
///     \ket{0} \mapsto \frac{1}{\sqrt{M}} \sum_{j=0}^{M - 1} \ket{j}
/// \end{align}
/// $$
///
/// The operation is adjointable, but requires that `qubits` register is in a
/// uniform superposition over the first `nStates` basis states in that case.
///
/// # Input
/// ## nStates
/// The number of states in the uniform superposition to be prepared.
/// ## register
/// The little-endian qubit register to store the prepared state.
/// It is assumed to be initialized in the zero state $\ket{0\cdots 0}$.
/// This register must be long enough to store the number $M-1$, meaning that
/// $2^{Length(qubits)} >= M$.
///
/// # Example
/// ```qsharp
///    use qs = Qubit[4];
///    PrepareUniformSuperposition(3, qs);
///    DumpRegister(qs); // The state is (|0000>+|0100>+|1000>)/√3
///    ResetAll(qs);
/// ```
operation PrepareUniformSuperposition(nStates : Int, qubits : Qubit[]) : Unit is Adj + Ctl {
    Fact(nStates > 0, "Number of basis states must be positive.");
    let nQubits = BitSizeI(nStates-1);
    Fact(nQubits <= Length(qubits), $"Qubit register is too short to prepare {nStates} states.");
    let relevantQubits = qubits[...nQubits - 1];
    let nTrailingZeroes = TrailingZeroCountI(nStates);

    ApplyToEachCA(H, relevantQubits);

    if nTrailingZeroes != nQubits {
        // Not a superposition of all relevant states so adjustment is needed.
        use tgt = Qubit();

        let nRelevantStates = 2^nQubits;
        let sqrt = Sqrt(IntAsDouble(nRelevantStates) / IntAsDouble(nStates));
        let angle = 2.0 * ArcSin(0.5 * sqrt);

        ApplyIfGreaterL(Ry(2.0 * angle, _), IntAsBigInt(nStates), relevantQubits, tgt);

        within {
            ApplyToEachA(H, relevantQubits[nTrailingZeroes...]);
        } apply {
            ReflectAboutInteger(0, relevantQubits[nTrailingZeroes...] + [tgt]);
            Ry(-angle, tgt);
        }

        X(tgt);
    }
}
