#include "stral_score.hh"
#include "rna_data.hh"
#include "trace_controller.hh"
#include "free_endgaps.hh"

namespace LocARNA {
    template <class pf_score_t>
    PFGotoh<pf_score_t>::PFGotoh(const RnaData &rnaA,
                                 const RnaData &rnaB,
                                 const TraceController &trace_controller,
                                 const Matrix<double> &sim_mat,
                                 const Alphabet<char, 4> &alphabet,
                                 double gap_opening,
                                 double gap_extension,
                                 double pf_struct_weight,
                                 double temp,
                                 const FreeEndgaps &free_endgaps,
                                 bool flag_local)
        : lenA_(rnaA.length()),
          lenB_(rnaB.length()),
          temp_(temp),
          flag_local_(flag_local) {
        StralScore score(rnaA, rnaB, sim_mat, alphabet, pf_struct_weight,
                         gap_opening, gap_extension);

        // forward
        pf_gotoh(zM_, zA_, zB_, trace_controller, score, free_endgaps);

        // backward
        score.reverse();
        pf_gotoh(zMr_, zAr_, zBr_, trace_controller.reverse(),
                 score, // reversed !
                 free_endgaps.reverse());

        if (flag_local) {
            // for the local pf we need to sum over all Mij entries
            z_ = 1; // weight of the empty alignment
            for (size_type i = 0; i <= lenA_; i++) {
                for (size_type j = 0; j <= lenB_; j++) {
                    z_ += zM_(i, j);
                }
            }
        } else { // global
            z_ = zM_(lenA_, lenB_) + zA_(lenA_, lenB_) + zB_(lenA_, lenB_);

            // handle free ends (free ends & local does not make sense!)

            // free end gaps
            if (free_endgaps.allow_left_2()) {
                for (size_type i = 0; i <= lenA_; i++) {
                    z_ += zA_(i, lenB_);
                }
            }
            if (free_endgaps.allow_left_1()) {
                for (size_type j = 0; j <= lenB_; j++) {
                    z_ += zB_(lenA_, j);
                }
            }
        }
    }

    /* @brief Gotoh partition function
     *
     * Perform the partition variant of Gotoh on seqA and seqB
     * and fill the three matrices
     *
     * @note Local alignment (flag_local_):
     * -- additional case: one can always start a new local alignment in M
     *    --> add math_ij (to each entry M_ij, only in M matrix)
     *
     * @note Free endgaps:
     * -- if allow_left_1: add 1 to B_0j
     * -- if allow_left_2: add 1 to A_i0
     *
     */
    template <class pf_score_t>
    void
    PFGotoh<pf_score_t>::pf_gotoh(Matrix<pf_score_t> &zMl,
                                  Matrix<pf_score_t> &zAl,
                                  Matrix<pf_score_t> &zBl,
                                  const TraceController &tc,
                                  const StralScore &score,
                                  const FreeEndgaps &free_endgaps) {
        // Boltzman-weights for gap opening and extension
        double g_open = exp(score.indel_opening() / temp_);
        double g_ext = exp(score.indel() / temp_);

        // std::cout << "g_open: "<<g_open<<std::endl;
        // std::cout << "g_ext: "<<g_ext<<std::endl;

        zMl.resize(lenA_ + 1, lenB_ + 1);
        zAl.resize(lenA_ + 1, lenB_ + 1);
        zBl.resize(lenA_ + 1, lenB_ + 1);

        // initialization
        //

        // start by filling all matrices with 0s
        zMl.fill(0);
        zAl.fill(0);
        zBl.fill(0);

        // For the init of zM(0,0) one must distringuish global/local:
        // In global alignment, there is an empty alignment,
        // which does not exist as local alignment!
        if (tc.is_valid(0, 0)) {
            zMl(0, 0) = flag_local_ ? 0 : 1;
        }

        if (lenA_>0 && tc.is_valid(1, 0)) {
            zAl(1, 0) = g_open * g_ext;
        }
        if (lenB_>0 && tc.is_valid(0, 1)) {
            zBl(0, 1) = g_open * g_ext;
        }

        for (size_type i = 2; i <= lenA_; i++) {
            if (tc.min_col(i) > 0)
                break;
            zAl(i, 0) = zAl(i - 1, 0) * g_ext;
        }

        for (size_type j = std::max(tc.min_col(0), (size_t)2);
             j <= std::min(tc.max_col(0), lenB_); j++) {
            zBl(0, j) = zBl(0, j - 1) * g_ext;
        }

        // free end gaps
        if (free_endgaps.allow_left_2()) {
            for (size_type i = 1; i <= lenA_; i++) {
                zAl(i, 0) += 1;
            }
        }
        if (free_endgaps.allow_left_1()) {
            for (size_type j = 1; j <= lenB_; j++) {
                zBl(0, j) += 1;
            }
        }

        // recursion
        for (size_type i = 1; i <= lenA_; i++) {
            for (size_type j = std::max(tc.min_col(i), (size_t)1);
                 j <= std::min(tc.max_col(i), lenB_); j++) {
                double match_score_ij = score.sigma(i, j);

                // Boltzman-weight for match of i and j
                double match_ij = exp(match_score_ij / temp_);

                zMl(i, j) = zMl(i - 1, j - 1) * match_ij +
                    zAl(i - 1, j - 1) * match_ij +
                    zBl(i - 1, j - 1) * match_ij + (flag_local_ ? match_ij : 0);

                zAl(i, j) = zAl(i - 1, j) * g_ext +
                    zMl(i - 1, j) * g_open * g_ext +
                    zBl(i - 1, j) * g_open * g_ext;

                zBl(i, j) = zBl(i, j - 1) * g_ext +
                    zMl(i, j - 1) * g_open * g_ext +
                    zAl(i, j - 1) * g_open * g_ext;
            }
        }
    }

    template <class pf_score_t>
    PFMatchProbs<pf_score_t>::PFMatchProbs(
        const RnaData &rnaA,
        const RnaData &rnaB,
        const TraceController &trace_controller,
        const Matrix<double> &sim_mat,
        const Alphabet<char, 4> &alphabet,
        double gap_opening,
        double gap_extension,
        double pf_struct_weight,
        double temp,
        const FreeEndgaps &free_endgaps,
        bool flag_local)
        : PFGotoh<pf_score_t>(rnaA,
                              rnaB,
                              trace_controller,
                              sim_mat,
                              alphabet,
                              gap_opening,
                              gap_extension,
                              pf_struct_weight,
                              temp,
                              free_endgaps,
                              flag_local) {
        probs_.resize(this->lenA_ + 1, this->lenB_ + 1);
        probs_.fill(0);

        // in local alignment, we add 1 for the empty alignment;
        // for avoiding redundancy the weight of the empty alignment is
        // not included in any matrix entry
        double locality_add = (this->flag_local_ ? 1 : 0);

        for (size_type i = 1; i <= this->lenA_; i++) {
            for (size_type j = std::max(trace_controller.min_col(i), (size_t)1);
                 j <= std::min(trace_controller.max_col(i), this->lenB_); j++) {
                probs_(i, j) = (this->zM_(i, j) *
                                (this->zMr_(this->lenA_ - i, this->lenB_ - j) +
                                 this->zAr_(this->lenA_ - i, this->lenB_ - j) +
                                 this->zBr_(this->lenA_ - i, this->lenB_ - j) +
                                 locality_add)) /
                    this->z_;
                // std::cout <<i<<" "<<j<<": "<<probs_(i,j)<<std::endl;
            }
        }

        // std::cout << "Probs:" << std::endl << probs_ << std::endl;
    }

    template <class pf_score_t>
    PFTraceProbs<pf_score_t>::PFTraceProbs(
        const RnaData &rnaA,
        const RnaData &rnaB,
        const TraceController &trace_controller,
        const Matrix<double> &sim_mat,
        const Alphabet<char, 4> &alphabet,
        double gap_opening,
        double gap_extension,
        double pf_struct_weight,
        double temp,
        const FreeEndgaps &free_endgaps,
        bool flag_local)
        : PFGotoh<pf_score_t>(rnaA,
                              rnaB,
                              trace_controller,
                              sim_mat,
                              alphabet,
                              gap_opening,
                              gap_extension,
                              pf_struct_weight,
                              temp,
                              free_endgaps,
                              flag_local) {
        probs_.resize(this->lenA_ + 1, this->lenB_ + 1);
        probs_.fill(0);

        double locality_add = (this->flag_local_ ? 1 : 0);

        double g_open = exp(gap_opening / temp);

        for (size_type i = 0; i <= this->lenA_; i++) {
            for (size_type j = std::max(trace_controller.min_col(i), (size_t)0);
                 j <= std::min(trace_controller.max_col(i), this->lenB_); j++) {

		pf_score_t z_ij =
		    this->zM_(i, j) *
                        (this->zMr_(this->lenA_ - i, this->lenB_ - j) +
                         this->zAr_(this->lenA_ - i, this->lenB_ - j) +
                         this->zBr_(this->lenA_ - i, this->lenB_ - j) +
                         locality_add) +
                    this->zA_(i, j) *
                        (this->zMr_(this->lenA_ - i, this->lenB_ - j) +
                         this->zAr_(this->lenA_ - i, this->lenB_ - j) / g_open +
                         this->zBr_(this->lenA_ - i, this->lenB_ - j)) +
                    this->zB_(i, j) *
                        (this->zMr_(this->lenA_ - i, this->lenB_ - j) +
                         this->zAr_(this->lenA_ - i, this->lenB_ - j) +
                         this->zBr_(this->lenA_ - i, this->lenB_ - j) / g_open);

                this->probs_(i, j) = z_ij / this->z_;
                // std::cout <<i<<" "<<j<<": "<<probs_(i,j)<<std::endl;
            }
        }

        // std::cout << "Probs:" << std::endl << probs_ << std::endl;
    }
}
