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

// SequenceAlignment.hpp — Needleman-Wunsch / Smith-Waterman sequence alignment
#ifndef CALX_BIO_SEQUENCE_ALIGNMENT_HPP
#define CALX_BIO_SEQUENCE_ALIGNMENT_HPP

#include <string>
#include <vector>
#include <algorithm>

namespace calx::bio {

    // アラインメント結果
    struct AlignmentResult {
        std::string alignedSeq1, alignedSeq2;
        int score;
    };

    // Needleman-Wunsch (大域アラインメント)
    inline AlignmentResult needlemanWunsch(const std::string& seq1, const std::string& seq2,
                                            int match = 2, int mismatch = -1, int gap = -2) {
        const int m = static_cast<int>(seq1.size());
        const int n = static_cast<int>(seq2.size());

        // DP テーブル
        std::vector<std::vector<int>> dp(m + 1, std::vector<int>(n + 1, 0));
        for (int i = 0; i <= m; ++i) dp[i][0] = i * gap;
        for (int j = 0; j <= n; ++j) dp[0][j] = j * gap;

        for (int i = 1; i <= m; ++i) {
            for (int j = 1; j <= n; ++j) {
                int scoreDiag = dp[i - 1][j - 1] + ((seq1[i - 1] == seq2[j - 1]) ? match : mismatch);
                int scoreUp   = dp[i - 1][j] + gap;
                int scoreLeft = dp[i][j - 1] + gap;
                dp[i][j] = std::max({ scoreDiag, scoreUp, scoreLeft });
            }
        }

        // トレースバック
        std::string a1, a2;
        int i = m, j = n;
        while (i > 0 || j > 0) {
            if (i > 0 && j > 0 &&
                dp[i][j] == dp[i - 1][j - 1] + ((seq1[i - 1] == seq2[j - 1]) ? match : mismatch)) {
                a1 += seq1[i - 1];
                a2 += seq2[j - 1];
                --i; --j;
            } else if (i > 0 && dp[i][j] == dp[i - 1][j] + gap) {
                a1 += seq1[i - 1];
                a2 += '-';
                --i;
            } else {
                a1 += '-';
                a2 += seq2[j - 1];
                --j;
            }
        }

        std::reverse(a1.begin(), a1.end());
        std::reverse(a2.begin(), a2.end());

        return AlignmentResult{ std::move(a1), std::move(a2), dp[m][n] };
    }

    // Smith-Waterman (局所アラインメント)
    inline AlignmentResult smithWaterman(const std::string& seq1, const std::string& seq2,
                                          int match = 2, int mismatch = -1, int gap = -2) {
        const int m = static_cast<int>(seq1.size());
        const int n = static_cast<int>(seq2.size());

        // DP テーブル (0 以上を維持)
        std::vector<std::vector<int>> dp(m + 1, std::vector<int>(n + 1, 0));

        int bestScore = 0, bestI = 0, bestJ = 0;

        for (int i = 1; i <= m; ++i) {
            for (int j = 1; j <= n; ++j) {
                int scoreDiag = dp[i - 1][j - 1] + ((seq1[i - 1] == seq2[j - 1]) ? match : mismatch);
                int scoreUp   = dp[i - 1][j] + gap;
                int scoreLeft = dp[i][j - 1] + gap;
                dp[i][j] = std::max({ 0, scoreDiag, scoreUp, scoreLeft });

                if (dp[i][j] > bestScore) {
                    bestScore = dp[i][j];
                    bestI = i;
                    bestJ = j;
                }
            }
        }

        // トレースバック (スコアが 0 になるまで)
        std::string a1, a2;
        int i = bestI, j = bestJ;
        while (i > 0 && j > 0 && dp[i][j] > 0) {
            if (dp[i][j] == dp[i - 1][j - 1] + ((seq1[i - 1] == seq2[j - 1]) ? match : mismatch)) {
                a1 += seq1[i - 1];
                a2 += seq2[j - 1];
                --i; --j;
            } else if (dp[i][j] == dp[i - 1][j] + gap) {
                a1 += seq1[i - 1];
                a2 += '-';
                --i;
            } else {
                a1 += '-';
                a2 += seq2[j - 1];
                --j;
            }
        }

        std::reverse(a1.begin(), a1.end());
        std::reverse(a2.begin(), a2.end());

        return AlignmentResult{ std::move(a1), std::move(a2), bestScore };
    }

} // namespace calx::bio

#endif // CALX_BIO_SEQUENCE_ALIGNMENT_HPP
