//
// Copyright (c) 2026 BEL ESPRIT D ACCORD TRUST HOLDINGS INC
// All rights reserved.

// BlackHoleEngine/QuantumPrimitives.qs
namespace BlackHoleEngine.Primitives {
    open Microsoft.Quantum.Intrinsic;
    open Microsoft.Quantum.Canon;
    open Microsoft.Quantum.Arithmetic;
    open Microsoft.Quantum.Measurement;
    open Microsoft.Quantum.Arrays;
    open Microsoft.Quantum.Convert;
    open Microsoft.Quantum.Math;

    newtype ViolatingAssignment = (Bool, Bool, Bool);

    operation ApplyClauseProjector(
        qubits : Qubit[],
        viol : ViolatingAssignment
    ) : Unit is Adj + Ctl {
        let (va, vb, vc) = viol;
        within {
            if va { X(qubits[0]); }
            if vb { X(qubits[1]); }
            if vc { X(qubits[2]); }
        } apply {
            // Projector |000><000| applied via controlled rotation on ancilla
        }
    }

    operation HadamardTestClause(
        ancilla : Qubit,
        qubits : Qubit[],
        viol : ViolatingAssignment
    ) : (Double, Double) {
        let (va, vb, vc) = viol;

        within {
            H(ancilla);
        } apply {
            Controlled ApplyClauseProjector([ancilla], (qubits, viol));
        }

        H(ancilla);
        let real = M(ancilla) == Zero ? 1.0 | -1.0;
        Reset(ancilla);

        within {
            H(ancilla);
            S(ancilla);
        } apply {
            Controlled ApplyClauseProjector([ancilla], (qubits, viol));
        }

        H(ancilla);
        let imag = M(ancilla) == Zero ? 1.0 | -1.0;
        Reset(ancilla);
        return (real, imag);
    }

    operation LindbladStepClause(
        ancilla : Qubit,
        qubits : Qubit[],
        viol : ViolatingAssignment,
        gamma_dt : Double
    ) : Unit is Adj + Ctl {
        let sqrt_gamma = Sqrt(gamma_dt);

        within {
            Ry(2.0 * Arcsin(sqrt_gamma), ancilla);
        } apply {
            Controlled ApplyClauseProjector([ancilla], (qubits, viol));
        }

        let outcome = M(ancilla);
        Reset(ancilla);
    }

    operation HaydenPreskillScrambler(
        qubits : Qubit[],
        scrambleTime : Double,
        couplings : Double[][],
        fields : Double[]
    ) : Unit is Adj + Ctl {
        let n = Length(qubits);
        let dt = scrambleTime / 10.0;
        let steps = 10;

        for _ in 1..steps {
            for i in 0..n-1 {
                Rz(2.0 * fields[i] * dt, qubits[i]);
            }

            for layer in GetXXLayers(n, couplings) {
                for (i, j) in layer {
                    let J = couplings[i][j];
                    if (AbsD(J) > 1e-10) {
                        CNOT(qubits[i], qubits[j]);
                        Rz(2.0 * J * dt, qubits[j]);
                        CNOT(qubits[i], qubits[j]);
                    }
                }
            }
        }
    }

    function GetXXLayers(n : Int, couplings : Double[][]) : (Int, Int)[][] {
        mutable layers = [];
        for i in 0..n-1 {
            for j in i+1..n-1 {
                if (AbsD(couplings[i][j]) > 1e-10) {
                    mutable placed = false;
                    for layerIdx in 0..Length(layers)-1 {
                        mutable conflict = false;
                        for (a, b) in layers[layerIdx] {
                            if (a == i or a == j or b == i or b == j) {
                                set conflict = true;
                            }
                        }
                        if (not conflict) {
                            set layers[layerIdx] += [(i, j)];
                            set placed = true;
                        }
                    }
                    if (not placed) {
                        set layers += [[(i, j)]];
                    }
                }
            }
        }
        return layers;
    }

    operation MeasureEnergy(
        qubits : Qubit[],
        clauses : (Int, Int, Int)[],
        violations : ViolatingAssignment[]
    ) : Double {
        mutable total = 0.0;
        use ancilla = Qubit();

        for idx in 0..Length(clauses)-1 {
            let (a, b, c) = clauses[idx];
            let viol = violations[idx];
            let (re, im) = HadamardTestClause(ancilla, [qubits[a], qubits[b], qubits[c]], viol);
            set total += re;
        }

        return total;
    }
}