// Copyright (C) 2026 Kiyotsugu Arai
// SPDX-License-Identifier: LGPL-3.0-or-later

// QuantumECC.hpp — 量子誤り訂正 (QC-4)
// Bit-flip code (3 qubit), Phase-flip code, Shor code (9 qubit)

#ifndef CALX_QUANTUM_QUANTUMECC_HPP
#define CALX_QUANTUM_QUANTUMECC_HPP

#include "QuantumCircuit.hpp"

namespace calx::quantum {

    // ---------------------------------------------------------------
    // Bit-flip code (3 qubit)
    // 論理量子ビット |ψ> = α|0> + β|1> を
    //   α|000> + β|111> にエンコード
    // ---------------------------------------------------------------

    // エンコード: 1 qubit → 3 qubit
    // 入力回路は qubit 0 に状態を持つ (1 qubit 回路)
    // 出力回路は 3 qubit で qubit 0 の状態が 3 重化される
    inline QuantumCircuit encodeBitFlip(const QuantumCircuit& input) {
        if (input.state().size() != 2)
            throw std::invalid_argument("encodeBitFlip: input must be a 1-qubit circuit");

        Complex alpha = input.state()[0]; // |0> の振幅
        Complex beta  = input.state()[1]; // |1> の振幅

        // 3 qubit 回路を構築: α|000> + β|111>
        QuantumCircuit encoded(3);

        // 直接状態ベクトルを設定するため、
        // qubit 0 に入力状態を設定してから CNOT で拡散
        // |ψ> ⊗ |00> → CNOT(0,1) → CNOT(0,2)
        // これは QuantumCircuit の操作で実現する

        // まず qubit 0 に α|0> + β|1> を設定
        // QuantumCircuit は |000> で初期化されるので、
        // 状態ベクトルを直接操作する代わりにゲートで構築

        // α|0> + β|1> を作るには Ry ゲートを使う
        // ただし位相も必要なので、一般的には U3 ゲートが必要
        // ここでは状態ベクトルに直接書き込むアプローチを取る

        // 内部の状態ベクトルにアクセスできないため、
        // 入力状態から 3-qubit 符号語を構築
        QuantumCircuit result(3);
        // 状態ベクトル: |000>=0, |001>=1, ..., |111>=7
        // 符号語: α|000> + β|111>
        // |000> = index 0,  |111> = index 7
        auto& st = const_cast<StateVector&>(result.state());
        st[0] = alpha;
        st[7] = beta;

        return result;
    }

    // ビットフリップシンドローム検出
    // シンドロームを返す: 0=エラーなし, 1=qubit0, 2=qubit1, 3=qubit2
    inline int detectBitFlip(QuantumCircuit& encoded) {
        if (encoded.state().size() != 8)
            throw std::invalid_argument("detectBitFlip: encoded must be a 3-qubit circuit");

        const auto& st = encoded.state();

        // 各基底状態の確率を計算し、最も確率の高い基底を特定
        // 正しい符号語: α|000> + β|111>
        // bit-flip エラー:
        //   qubit 0 反転: α|001> + β|110>  → syndrome = 1
        //   qubit 1 反転: α|010> + β|101>  → syndrome = 2
        //   qubit 2 反転: α|100> + β|011>  → syndrome = 3

        // パリティチェック: qubit 0⊕1, qubit 1⊕2
        // syndrome (s1, s2): (0,0)=OK, (1,0)=q0, (1,1)=q1, (0,1)=q2

        // 各パターンの確率を集計
        double probNoError = 0.0;   // |000>, |111>
        double probFlip0 = 0.0;     // |001>, |110>
        double probFlip1 = 0.0;     // |010>, |101>
        double probFlip2 = 0.0;     // |100>, |011>

        probNoError = std::norm(st[0]) + std::norm(st[7]);  // 000, 111
        probFlip0   = std::norm(st[1]) + std::norm(st[6]);  // 001, 110
        probFlip1   = std::norm(st[2]) + std::norm(st[5]);  // 010, 101
        probFlip2   = std::norm(st[4]) + std::norm(st[3]);  // 100, 011

        // 最大確率のパターンがエラー位置
        double maxProb = probNoError;
        int syndrome = 0;

        if (probFlip0 > maxProb) { maxProb = probFlip0; syndrome = 1; }
        if (probFlip1 > maxProb) { maxProb = probFlip1; syndrome = 2; }
        if (probFlip2 > maxProb) { syndrome = 3; }

        return syndrome;
    }

    // ビットフリップ訂正: 検出されたシンドロームに基づいて X ゲートを適用
    inline void correctBitFlip(QuantumCircuit& encoded, int syndrome) {
        if (syndrome < 0 || syndrome > 3)
            throw std::invalid_argument("correctBitFlip: syndrome must be 0-3");

        if (syndrome == 0) return; // エラーなし

        // syndrome 1: qubit 0 を反転, 2: qubit 1, 3: qubit 2
        encoded.x(static_cast<size_t>(syndrome - 1));
    }

    // ---------------------------------------------------------------
    // Phase-flip code (3 qubit)
    // Hadamard 基底でのビットフリップ = 位相フリップ
    // 符号語: α|+++> + β|--->
    // ---------------------------------------------------------------

    // エンコード: 1 qubit → 3 qubit (Phase-flip code)
    inline QuantumCircuit encodePhaseFlip(const QuantumCircuit& input) {
        if (input.state().size() != 2)
            throw std::invalid_argument("encodePhaseFlip: input must be a 1-qubit circuit");

        // まずビットフリップ符号を作成
        QuantumCircuit result = encodeBitFlip(input);

        // 各 qubit に Hadamard を適用して位相フリップ符号に変換
        result.h(0);
        result.h(1);
        result.h(2);

        return result;
    }

    // ---------------------------------------------------------------
    // Shor code (9 qubit)
    // 外側: Phase-flip code (3 ブロック)
    // 内側: 各ブロックが Bit-flip code (3 qubit)
    // → ビットフリップと位相フリップの両方を訂正可能
    // ---------------------------------------------------------------

    // エンコード: 1 qubit → 9 qubit (Shor code)
    inline QuantumCircuit encodeShor(const QuantumCircuit& input) {
        if (input.state().size() != 2)
            throw std::invalid_argument("encodeShor: input must be a 1-qubit circuit");

        Complex alpha = input.state()[0];
        Complex beta  = input.state()[1];

        // Shor 符号語:
        // α|ψ_0> + β|ψ_1>
        // |ψ_0> = (|000> + |111>)(|000> + |111>)(|000> + |111>) / 2√2
        // |ψ_1> = (|000> - |111>)(|000> - |111>)(|000> - |111>) / 2√2

        QuantumCircuit result(9);
        auto& st = const_cast<StateVector&>(result.state());

        // 全状態を 0 で初期化
        std::fill(st.begin(), st.end(), Complex{0});

        // 9 qubit のインデックス: qubit 8..0 (MSB..LSB)
        // ブロック構成: [q0,q1,q2], [q3,q4,q5], [q6,q7,q8]
        // |000> = 各ブロックの 3 qubit がすべて 0
        // |111> = 各ブロックの 3 qubit がすべて 1

        // 各ブロックの |000> と |111> のビットパターン
        constexpr int blk0_000 = 0b000;        // qubits 0,1,2 = 000
        constexpr int blk0_111 = 0b111;        // qubits 0,1,2 = 111
        constexpr int blk1_000 = 0b000 << 3;   // qubits 3,4,5 = 000
        constexpr int blk1_111 = 0b111 << 3;   // qubits 3,4,5 = 111
        constexpr int blk2_000 = 0b000 << 6;   // qubits 6,7,8 = 000
        constexpr int blk2_111 = 0b111 << 6;   // qubits 6,7,8 = 111

        // |ψ_0> の 8 項: (|000>+|111>)^3 を展開
        // |ψ_1> の 8 項: (|000>-|111>)^3 を展開
        const double norm = 1.0 / (2.0 * std::sqrt(2.0));
        const int blk_patterns[2][2] = {
            {blk0_000, blk0_111},  // ブロック 0
        };

        // 全 2^3 = 8 組み合わせを生成
        for (int b2 = 0; b2 < 2; ++b2) {
            int bits2 = (b2 == 0) ? blk2_000 : blk2_111;
            int sign2_0 = 1;                      // |ψ_0> は常に +
            int sign2_1 = (b2 == 0) ? 1 : -1;    // |ψ_1> は |111> で -

            for (int b1 = 0; b1 < 2; ++b1) {
                int bits1 = (b1 == 0) ? blk1_000 : blk1_111;
                int sign1_0 = sign2_0;
                int sign1_1 = sign2_1 * ((b1 == 0) ? 1 : -1);

                for (int b0 = 0; b0 < 2; ++b0) {
                    int bits0 = (b0 == 0) ? blk0_000 : blk0_111;
                    int sign0_0 = sign1_0;
                    int sign0_1 = sign1_1 * ((b0 == 0) ? 1 : -1);

                    int idx = bits2 | bits1 | bits0;
                    // α * sign_0 * norm + β * sign_1 * norm
                    st[idx] = alpha * Complex{sign0_0 * norm}
                            + beta  * Complex{sign0_1 * norm};
                }
            }
        }

        return result;
    }

} // namespace calx::quantum

#endif // CALX_QUANTUM_QUANTUMECC_HPP
