// QuantumCircuit.hpp — State-vector quantum circuit simulator
// SPDX-License-Identifier: LGPL-3.0-or-later
// Copyright (c) 2026 calx contributors

#ifndef CALX_QUANTUM_QUANTUMCIRCUIT_HPP
#define CALX_QUANTUM_QUANTUMCIRCUIT_HPP

#include <vector>
#include <array>
#include <complex>
#include <cmath>
#include <cstddef>
#include <stdexcept>
#include <random>
#include <numbers>

namespace calx::quantum {

using Complex = std::complex<double>;
using StateVector = std::vector<Complex>;

// ---------------------------------------------------------------
// 単一量子ビットゲートの適用
// gate は列優先 2x2: { g00, g10, g01, g11 }
// ただし慣例として行優先 { g00, g01, g10, g11 } を採用
// ---------------------------------------------------------------
inline StateVector applyGate(const StateVector& state, size_t qubit,
                              const std::array<Complex, 4>& gate, size_t nQubits) {
    const size_t dim = size_t{1} << nQubits;
    if (state.size() != dim)
        throw std::invalid_argument("applyGate: state size mismatch");
    if (qubit >= nQubits)
        throw std::out_of_range("applyGate: qubit index out of range");

    StateVector result = state;
    const size_t bit = size_t{1} << qubit;

    for (size_t i = 0; i < dim; ++i) {
        if (i & bit) continue; // 下位ペアのみ処理
        size_t j = i | bit;
        Complex a = state[i];
        Complex b = state[j];
        result[i] = gate[0] * a + gate[1] * b;
        result[j] = gate[2] * a + gate[3] * b;
    }
    return result;
}

// ---------------------------------------------------------------
// 定数ゲート行列 (行優先 { g00, g01, g10, g11 })
// ---------------------------------------------------------------
namespace gates {
    inline constexpr double INV_SQRT2 = 0.7071067811865475244;

    inline const std::array<Complex, 4> H = {
        Complex{INV_SQRT2}, Complex{INV_SQRT2},
        Complex{INV_SQRT2}, Complex{-INV_SQRT2}
    };
    inline const std::array<Complex, 4> X = {
        Complex{0}, Complex{1},
        Complex{1}, Complex{0}
    };
    inline const std::array<Complex, 4> Y = {
        Complex{0},            Complex{0, -1},
        Complex{0, 1},         Complex{0}
    };
    inline const std::array<Complex, 4> Z = {
        Complex{1},  Complex{0},
        Complex{0},  Complex{-1}
    };
    inline const std::array<Complex, 4> S = {
        Complex{1},  Complex{0},
        Complex{0},  Complex{0, 1}
    };
    inline const std::array<Complex, 4> T = {
        Complex{1},  Complex{0},
        Complex{0},  std::polar(1.0, std::numbers::pi / 4.0)
    };

    inline std::array<Complex, 4> Rx(double theta) {
        double c = std::cos(theta / 2.0);
        double s = std::sin(theta / 2.0);
        return { Complex{c}, Complex{0, -s},
                 Complex{0, -s}, Complex{c} };
    }
    inline std::array<Complex, 4> Ry(double theta) {
        double c = std::cos(theta / 2.0);
        double s = std::sin(theta / 2.0);
        return { Complex{c}, Complex{-s},
                 Complex{s}, Complex{c} };
    }
    inline std::array<Complex, 4> Rz(double theta) {
        return { std::polar(1.0, -theta / 2.0), Complex{0},
                 Complex{0}, std::polar(1.0, theta / 2.0) };
    }
} // namespace gates

// ---------------------------------------------------------------
// QuantumCircuit クラス
// ---------------------------------------------------------------
class QuantumCircuit {
    size_t nQubits_;
    StateVector state_;

    // 単一量子ビットゲートをインプレース適用
    void applySingle(size_t q, const std::array<Complex, 4>& gate) {
        if (q >= nQubits_)
            throw std::out_of_range("QuantumCircuit: qubit index out of range");
        const size_t dim = size_t{1} << nQubits_;
        const size_t bit = size_t{1} << q;
        for (size_t i = 0; i < dim; ++i) {
            if (i & bit) continue;
            size_t j = i | bit;
            Complex a = state_[i];
            Complex b = state_[j];
            state_[i] = gate[0] * a + gate[1] * b;
            state_[j] = gate[2] * a + gate[3] * b;
        }
    }

public:
    explicit QuantumCircuit(size_t nQubits)
        : nQubits_(nQubits), state_(size_t{1} << nQubits, Complex{0}) {
        if (nQubits == 0 || nQubits > 24)
            throw std::invalid_argument("QuantumCircuit: nQubits must be 1..24");
        state_[0] = Complex{1}; // |00...0>
    }

    // --- 単一量子ビットゲート ---
    void h(size_t q)  { applySingle(q, gates::H); }
    void x(size_t q)  { applySingle(q, gates::X); }
    void y(size_t q)  { applySingle(q, gates::Y); }
    void z(size_t q)  { applySingle(q, gates::Z); }
    void s(size_t q)  { applySingle(q, gates::S); }
    void t(size_t q)  { applySingle(q, gates::T); }

    void rx(size_t q, double theta) { applySingle(q, gates::Rx(theta)); }
    void ry(size_t q, double theta) { applySingle(q, gates::Ry(theta)); }
    void rz(size_t q, double theta) { applySingle(q, gates::Rz(theta)); }

    // --- CNOT (controlled-X) ---
    void cnot(size_t control, size_t target) {
        if (control >= nQubits_ || target >= nQubits_)
            throw std::out_of_range("QuantumCircuit::cnot: qubit index out of range");
        if (control == target)
            throw std::invalid_argument("QuantumCircuit::cnot: control == target");

        const size_t dim = size_t{1} << nQubits_;
        const size_t cBit = size_t{1} << control;
        const size_t tBit = size_t{1} << target;

        for (size_t i = 0; i < dim; ++i) {
            // control ビットが 1 かつ target ビットが 0 のペアのみ処理
            if ((i & cBit) && !(i & tBit)) {
                size_t j = i | tBit;
                std::swap(state_[i], state_[j]);
            }
        }
    }

    // --- SWAP ---
    void swap(size_t q1, size_t q2) {
        // SWAP = 3 CNOT
        cnot(q1, q2);
        cnot(q2, q1);
        cnot(q1, q2);
    }

    // --- 測定 ---
    int measure(size_t q, unsigned seed = 42) {
        if (q >= nQubits_)
            throw std::out_of_range("QuantumCircuit::measure: qubit index out of range");

        const size_t dim = size_t{1} << nQubits_;
        const size_t bit = size_t{1} << q;

        // |1> の確率を計算
        double prob1 = 0.0;
        for (size_t i = 0; i < dim; ++i) {
            if (i & bit)
                prob1 += std::norm(state_[i]);
        }

        // 乱数で測定結果を決定
        std::mt19937 rng(seed);
        std::uniform_real_distribution<double> dist(0.0, 1.0);
        int result = (dist(rng) < prob1) ? 1 : 0;

        // 状態の崩壊 (射影 + 正規化)
        double normFactor = 0.0;
        for (size_t i = 0; i < dim; ++i) {
            bool hasBit = (i & bit) != 0;
            if ((result == 1 && !hasBit) || (result == 0 && hasBit)) {
                state_[i] = Complex{0};
            } else {
                normFactor += std::norm(state_[i]);
            }
        }
        if (normFactor > 0.0) {
            double invNorm = 1.0 / std::sqrt(normFactor);
            for (size_t i = 0; i < dim; ++i)
                state_[i] *= invNorm;
        }

        return result;
    }

    // --- 確率分布 ---
    std::vector<double> probabilities() const {
        std::vector<double> probs(state_.size());
        for (size_t i = 0; i < state_.size(); ++i)
            probs[i] = std::norm(state_[i]);
        return probs;
    }

    const StateVector& state() const { return state_; }

    void reset() {
        std::fill(state_.begin(), state_.end(), Complex{0});
        state_[0] = Complex{1};
    }
};

} // namespace calx::quantum

#endif // CALX_QUANTUM_QUANTUMCIRCUIT_HPP
