#ifndef LOCARNA_MAIN_HELPER
#define LOCARNA_MAIN_HELPER

/**
 * define data structures and code common to
 * the driver programs of locarna, locarna_p, sparse, and exparna-p
 */

#include <string>
#include <map>
#include <memory>
#include "sequence.hh"
#include "edge_probs.hh"
#include "rna_data.hh"
#include "ribosum.hh"
#include "ribofit.hh"
#include "multiple_alignment.hh"
#include "alignment.hh"
#include "rna_ensemble.hh"
#include "free_endgaps.hh"

#include "LocARNA/ribosum85_60.icc"

namespace LocARNA {

    class ArcMatches;
    class ScoringParams;

    namespace MainHelper {

        // standard command line parameters common to locarna, locarna_p, sparse
        struct std_command_line_parameters {
            std::map<std::string, std::string> help_text = {
#include      "std_help_text.ihh"
            };

            // ----------------------------------------
            // command only parameters

            bool help; //!< whether to print help

            bool galaxy_xml; //!< whether to print a galaxy xml wrapper for the
                             //!parameters

            bool version; //!< whether to print version

            bool verbose; //!< whether to print verbose output

            bool quiet; //!< whether to stay quiet (overrides verbose)

            // ----------------------------------------
            // Scoring parameters

            int indel; //!< indel extension score

            int indel_opening; //!< indel opening score

            std::string ribosum_file; //!< ribosum_file

            bool use_ribosum; //!< use_ribosum

            bool ribofit;

            int match; //!< match score

            int mismatch; //!< mismatch score

            int unpaired_penalty; //!< penalty for unpaired bases

            int struct_weight; //!< structure weight

            double
                exp_prob; //!< expected probability of a base pair (null-model)
            bool exp_prob_given; //!< whether exp prob is given

            int tau;

            int exclusion;

            bool stacking; //!< whether to use special stacking arcmatch score

            bool new_stacking; //!< whether to use new stacking contributions

            // ----------------------------------------
            // partition function representation

            /** @brief Extended precision for partition function values
             *
             * If true, use an extended precision type for partition function values (pf_score_t).
             */
            bool extended_pf;

            /** @brief Quad precision for partition function values
             *
             * If true, use quad precision type for partition function values
             * (pf_score_t) if available; override
             * extended pf.
             */
            bool quad_pf;

            // ----------------------------------------
            // Locality

            bool struct_local;       //!< turns on/off structure locality
            bool struct_local_given; //!< is struct-local mode specified
                                     //!explicitely

            bool sequ_local; //!< sequence local alignment; maximize alignment
                             //!of subsequence
            bool sequ_local_given; //!< is sequ-local mode specified explicitely

            std::string free_endgaps; //!< free endgaps specification

            // ----------------------------------------
            // Output

            int width;

            std::string clustal; //!< name of clustal output file
            bool clustal_given;  //!< whether to write clustal output to file

            std::string stockholm; //!< name of stockholm output file
            bool stockholm_given; //!< whether to write stockholm output to file

            std::string pp; //!< name of pp output file
            bool pp_given;  //!< whether to write pp output to file

            bool alifold_consensus_dp; //!< whether to compute consensus dp by
                                       //!alifold

            std::string cons_struct_type; //! type of consensus structure

            double consensus_gamma; //! base pair weight for consensus mea

            bool local_output; //!< whether to write local output

            bool local_file_output; //!< whether to write local output to files
                                    //!(pp, aln)

            bool pos_output; //!< whether to output positions

            bool write_structure; //!< whether to write structure

            bool stopwatch; //!< whether to print verbose output

            // ----------------------------------------
            // Heuristics

            double min_prob;

            double max_bps_length_ratio;

            int max_diff_am;

            int max_diff;

            int max_diff_at_am;

            //! reference alignment for max-diff heuristic, name of clustalw
            //! format file
            std::string max_diff_alignment_file;

            //! pairwise reference alignment for max-diff heuristic,
            //! separator &
            std::string max_diff_pw_alignment;

            //! use relaxed variant of max diff with reference alignment
            bool max_diff_relax;

            //! minimum sequence alignment trace probability
            double min_trace_probability;

            // ----------------------------------------
            // Constraints

            bool no_lonely_pairs; //!< no lonely pairs option

            int max_bp_span; //!< maximum base pair span

            bool relaxed_anchors; //!< strict or relaxed anchor constraints

            // ------------------------------------------------------------
            // File arguments

            //! first input file
            std::string fileA;

            //! second input file
            std::string fileB;

        };

        struct mea_command_line_parameters {
            // ----------------------------------------
            // partition functions, probabilities and mea stuff

            bool mea_alignment; //!< whether to perform mea alignment

            int match_prob_method; //!< method for computing match probabilities

            std::string probcons_file; //!< probcons_file
            bool probcons_file_given;  //!< whether to probcons_file

            int temperature_alipf; //!< temperature for alignment partition
                                   //!functions

            int pf_struct_weight; //!< pf_struct_weight

            bool mea_gapcost; //!< whether to use mea gapcost

            int mea_alpha; //!< mea alpha

            int mea_beta; //!< mea beta

            int mea_gamma; //!< mea gamma

            int probability_scale; //!< probability scale

            bool write_matchprobs; //!< whether to write_matchprobs

            bool write_traceprobs; //!< whether to write_traceprobs

            bool read_matchprobs; //!< whether to read_matchprobs

            std::string matchprobs_infile;  //!< matchprobs input file
            std::string matchprobs_outfile; //!< matchprobs output file
            std::string traceprobs_outfile; //!< traceprobs output file

            bool write_arcmatch_scores; //!< whether to write arcmatch scores

            bool read_arcmatch_scores; //!< whether to read arcmatch scores

            bool
                read_arcmatch_probs; //!< whether to read arcmatch probabilities

            std::string arcmatch_scores_infile;  //!< arcmatch scores file
            std::string arcmatch_scores_outfile; //!< arcmatch scores file
        };

        //! @brief write input summary
        void
        report_input(const Sequence &seqA,
                     const Sequence &seqB,
                     const ArcMatches &arc_matches);

        //! @brief average sequence identity
        double
        average_pairwise_identity(const Sequence &seqA, const Sequence &seqB);

        //! @brief initialize ribosum/ribofit matrices
        template <class CLP>
        void
        init_ribo_matrix(const CLP &clp,
                         std::unique_ptr<RibosumFreq> &ribosum,
                         std::unique_ptr<Ribofit> &ribofit) {
            if (clp.ribofit) {
                if (clp.verbose) {
                    std::cout << "Use ribofit." << std::endl;
                }
                ribofit = std::make_unique<Ribofit_will2014>();
            } else if (clp.use_ribosum) {
                if (clp.ribosum_file == "RIBOSUM85_60") {
                    if (clp.verbose) {
                        std::cout << "Use built-in ribosum." << std::endl;
                    }
                    ribosum = std::make_unique<Ribosum85_60>();
                } else {
                    try {
                        ribosum = std::make_unique<RibosumFreq>(clp.ribosum_file);
                    } catch (std::ifstream::failure &e) {
                        std::cerr << "ERROR: " << e.what() << std::endl;
                        exit(-1);
                    } catch (failure &e) {
                        std::cerr << "ERROR: " << e.what() << std::endl;
                        exit(-1);
                    }
                }
            }
        }


        template <class CLP>
        void
        select_basematch_scores(CLP &clp,
                                const Sequence &seqA,
                                const Sequence &seqB,
                                const Ribosum *ribosum,
                                const Ribofit *ribofit,
                                Matrix<double> &basematch_scores,
                                Alphabet<char,4> &alphabet) {
            if (ribosum != nullptr) {
                alphabet = ribosum->alphabet();
                basematch_scores = ribosum->get_basematch_scores();
            } else if (ribofit != nullptr) {
                double avg_identity =
                    average_pairwise_identity(seqA, seqB);
                alphabet = ribofit->alphabet();
                ribofit->get_basematch_scores(avg_identity,
                                              basematch_scores);
            } else {
                alphabet = Alphabet<char,4>("ACGU");
                basematch_scores.resize(alphabet.size(), alphabet.size());
                basematch_scores.fill(clp.mismatch);
                for(size_t i=0; i<alphabet.size(); ++i) {
                    basematch_scores(i,i) = clp.match;
                }
            }
        }

        template <class CLP,class pf_score_t>
        std::unique_ptr<MatchProbs>
        init_match_probs(CLP &clp,
                         const RnaData *rna_dataA,
                         const RnaData *rna_dataB,
                         const TraceController *trace_controller,
                         const Ribosum *ribosum,
                         const Ribofit *ribofit,
                         const pf_score_t &dummy) {

            const Sequence &seqA = rna_dataA->sequence();
            const Sequence &seqB = rna_dataB->sequence();

            if (clp.read_matchprobs) {
		std::ifstream in(clp.matchprobs_infile);
                return std::make_unique<MatchProbs>(in,
						    seqA.length(),
						    seqB.length());
            } else {
                if (clp.match_prob_method == 1) {
                    if (!clp.probcons_file_given) {
                        std::cerr << "Probcons parameter file required for "
                                     "pairHMM-style computation"
                                  << " of basematch probabilities."
                                  << std::endl;
                        std::cerr << std::endl;
                        std::exit(-1);
                    }
                    if (clp.verbose) {
                        std::cout
                            << "Compute match probabilities using pairHMM."
                            << std::endl;
                    }

                    return std::make_unique<PairHMMMatchProbs>(
                        seqA, seqB,
                        PairHMMMatchProbs::PairHMMParams(clp.probcons_file));
                } else {
                    if (clp.verbose) {
                        std::cout << "Compute match probabilities using PF "
                                     "sequence alignment."
                                  << std::endl;
                    }

                    Alphabet<char,4> alphabet;
                    Matrix<double> basematch_scores;

                    select_basematch_scores(clp, seqA, seqB, ribosum, ribofit,
                                            basematch_scores, alphabet);

                    return std::make_unique<PFMatchProbs<pf_score_t>>(
                        *rna_dataA, *rna_dataB, *trace_controller,
                        basematch_scores, alphabet, clp.indel_opening / 100.0,
                        clp.indel / 100.0, clp.pf_struct_weight / 100.0,
                        clp.temperature_alipf / 100.0,
                        FreeEndgaps(clp.free_endgaps),
                        clp.sequ_local || (clp.match_prob_method == 3));
                }
            }
        }

        template <class CLP, class pf_score_t>
        PFTraceProbs<pf_score_t>
        make_trace_probs(CLP &clp,
                         const RnaData *rna_dataA,
                         const RnaData *rna_dataB,
                         const Ribosum *ribosum,
                         const Ribofit *ribofit,
                         TraceController *trace_controller,
                         const pf_score_t &dummy) {
            const auto &seqA = rna_dataA->sequence();
            const auto &seqB = rna_dataB->sequence();

            Alphabet<char,4> alphabet;
            Matrix<double> basematch_scores;

            select_basematch_scores(clp, seqA, seqB, ribosum, ribofit,
                                    basematch_scores, alphabet);

            PFTraceProbs<pf_score_t> trace_probs(*rna_dataA, *rna_dataB,
                                                 *trace_controller, basematch_scores,
                                                 alphabet, clp.indel_opening / 100.0,
                                                 clp.indel / 100.0,
                                                 clp.struct_weight / 100.0,
                                                 clp.temperature_alipf / 100.0,
                                                 FreeEndgaps(clp.free_endgaps),
                                                 clp.sequ_local);

            if (trace_probs.fail()) {
                std::cerr << "WARNING: no valid trace probabilities could be "
                             "computed with the requested floating point precision. "
                             "Consider flag --extended-pf (or even --quad-pf) to increase the precision."
                          << std::endl;
            }

            return trace_probs;
        }

        template <class pf_score_t,class CLP>
        void
        restrict_trace_by_probabilities(CLP &clp,
                                        const RnaData *rna_dataA,
                                        const RnaData *rna_dataB,
                                        const Ribosum *ribosum,
                                        const Ribofit *ribofit,
                                        TraceController *trace_controller,
                                        const pf_score_t &dummy
                                        ) {
            if (clp.min_trace_probability > 0.0) {
                auto trace_probs =
                    make_trace_probs(clp, rna_dataA, rna_dataB, ribosum,
                                     ribofit, trace_controller, pf_score_t());

                trace_controller->restrict_by_trace_probabilities(
                    trace_probs, clp.min_trace_probability);
            }
        }


        template <class CLP>
        void
        write_match_probs(const CLP &clp, const MatchProbs *match_probs) {
            if (clp.verbose) {
                std::cout << "Write match probabilities to file "
                          << clp.matchprobs_outfile << "." << std::endl;
            }

	    std::ofstream out(clp.matchprobs_outfile);
            match_probs->write_sparse(out,
                                      1.0 / clp.probability_scale);
        }

        template <class CLP, class pf_score_t>
        void
        write_trace_probs(const CLP &clp,
                          const RnaData *rna_dataA,
                          const RnaData *rna_dataB,
                          const Ribosum *ribosum,
                          const Ribofit *ribofit,
                          TraceController *trace_controller,
                          const pf_score_t &dummy
                          ) {
            if (clp.verbose) {
                std::cout << "Write trace probabilities to file "
                          << clp.traceprobs_outfile << "." << std::endl;
            }

            auto trace_probs =
                make_trace_probs(clp, rna_dataA, rna_dataB, ribosum, ribofit,
                                 trace_controller, pf_score_t());

            std::ofstream out(clp.traceprobs_outfile);
            trace_probs.write_sparse(out, 1.0 / clp.probability_scale);
        }

        using consensus_pair_t =
            std::pair<std::unique_ptr<RnaData>, std::string>;

        /**
         * @brief consensus ensemble
         */
        template <class CLP>
        consensus_pair_t
        consensus(const CLP &clp,
                  const PFoldParams &pfparams,
                  double my_exp_probA,
                  double my_exp_probB,
                  const RnaData *rna_dataA,
                  const RnaData *rna_dataB,
                  const Alignment &alignment,
                  bool local
                  ) {

            stopwatch.start("consensus");

            std::unique_ptr<RnaData> consensus;

            std::string consensus_structure="";

            if (clp.alifold_consensus_dp || clp.cons_struct_type == "alifold") {
                MultipleAlignment ma(alignment, local);

                // compute cutoff probability as geometric mean
                double min_prob = sqrt(rna_dataA->arc_cutoff_prob() *
                                       rna_dataB->arc_cutoff_prob());

                RnaEnsemble ens(ma, pfparams, false,
                                true); // alifold the alignment

                if (clp.cons_struct_type == "alifold") {
                    consensus_structure = ens.min_free_energy_structure();
                }

                if (clp.alifold_consensus_dp) {
                    consensus = std::make_unique<RnaData>(
                        ens, min_prob,
                        0,         // don't filter by max_bps_length_ratio
                        pfparams); // construct rna data from ensemble
                }
            }

            if (!clp.alifold_consensus_dp || clp.cons_struct_type == "mea") {
                // compute averaged consensus base pair probabilities

                std::unique_ptr<RnaData> mea_consensus =
                    std::make_unique<RnaData>(*rna_dataA, *rna_dataB, alignment,
                                              my_exp_probA, my_exp_probB,
                                              local);

                if (clp.cons_struct_type == "mea") {
                    consensus_structure = mea_consensus->mea_structure(clp.consensus_gamma);
                }

                if (!clp.alifold_consensus_dp) {
                    consensus = std::move(mea_consensus);
                }
            }

            stopwatch.stop("consensus");

            return std::make_pair(std::move(consensus), consensus_structure);
        }

        /** @brief write output to file/s (optionally)
         */
        template <class CLP, class DelayedCP>
        int
        write_alignment(const CLP &clp,
                        infty_score_t score,
                        DelayedCP &delayed_cp,
                        const Alignment &alignment,
                        const MultipleAlignment *multiple_ref_alignment) {

            const Sequence &seqA = alignment.seqA();
            const Sequence &seqB = alignment.seqB();

            int return_code = 0;

            // write MultipleAlignment deviation, if reference alignment given
            // and verbose
            if (clp.verbose && multiple_ref_alignment != nullptr) {
                MultipleAlignment resultMA(alignment);
                std::cout << "Deviation to reference: "
                          << multiple_ref_alignment->deviation(resultMA)
                          << std::endl;
            }

            if (clp.clustal_given) {
                std::ofstream out(clp.clustal.c_str());
                if (out.good()) {
                    MultipleAlignment ma(alignment, clp.local_file_output);

                    out << "CLUSTAL W --- " << PACKAGE_STRING;

                    // for legacy, clustal files of pairwise alignments contain
                    // the score
                    if (seqA.num_of_rows() == 1 && seqB.num_of_rows() == 1)
                        out << " --- Score: " << score;
                    out << std::endl << std::endl;

                    if (clp.write_structure) {
                        // annotate multiple alignment with structures
                        ma.prepend(MultipleAlignment::SeqEntry(
                            "", alignment.dot_bracket_structureA(
                                    clp.local_file_output)));
                        ma.append(MultipleAlignment::SeqEntry(
                            "", alignment.dot_bracket_structureB(
                                    clp.local_file_output)));
                    }

                    ma.write(out, clp.width,
                             MultipleAlignment::FormatType::CLUSTAL);

                } else {
                    std::cerr << "ERROR: Cannot write to " << clp.clustal << "."
                              << std::endl;
                    return_code = -1;
                }
            }

            if (clp.stockholm_given) {
                std::ofstream out(clp.stockholm.c_str());
                if (out.good()) {
                    MultipleAlignment ma(alignment, clp.local_file_output);

                    const auto &cp = delayed_cp.get();

                    if (cp.second != "") {
                        ma.set_annotation(
                            MultipleAlignment::AnnoType::consensus_structure,
                            SequenceAnnotation(cp.second));
                    }

                    out << "# STOCKHOLM 1.0" << std::endl;
                    out << "#=GF CC "
                        << "Generated by " << PACKAGE_STRING << std::endl;
                    out << "#=GF SQ " << ma.num_of_rows() << std::endl;
                    out << std::endl;

                    ma.write(out, clp.width,
                             MultipleAlignment::FormatType::STOCKHOLM);
                } else {
                    std::cerr << "ERROR: Cannot write to " << clp.stockholm << "."
                              << std::endl;
                    return_code = -1;
                }
            }

            if (clp.pp_given) {
                std::ofstream out(clp.pp.c_str());
                if (out.good()) {
                    const auto &cp = delayed_cp.get();
                    if (cp.first) {
                        cp.first->write_pp(out);
                    } else {
                        std::cerr << "ERROR: Cannot write pp, since "
                                     "consensus cannot be computed."
                                  << std::endl;
                        return_code = -1;
                    }
                } else {
                    std::cerr << "ERROR: Cannot write to " << clp.pp << std::endl;
                    return_code = -1;
                }
            }

            return return_code;
        }

        void
        report_input(const Sequence &seqA,
                     const Sequence &seqB,
                     const ArcMatches &arc_matches) {
            const BasePairs &bpsA = arc_matches.get_base_pairsA();
            const BasePairs &bpsB = arc_matches.get_base_pairsB();

            // ----------------------------------------
            // report on input in verbose mode
            std::cout << "Sequence A: " << std::endl;
            seqA.write(std::cout, MultipleAlignment::FormatType::CLUSTAL);
            std::cout << " (Length:" << seqA.length()
                      << ", Basepairs:" << bpsA.num_bps() << ")" << std::endl;

            std::cout << "Sequence B: " << std::endl;
            seqB.write(std::cout, MultipleAlignment::FormatType::CLUSTAL);
            std::cout << " (Length:" << seqB.length()
                      << ", Basepairs:" << bpsB.num_bps() << ")" << std::endl;

            std::cout << std::endl
                      << "Base Pair Matches: " << arc_matches.num_arc_matches()
                      << "." << std::endl;
            // std::cout << "Base Identity:
            // "<<(seq_identity(seqA,seqB)*100)<<std::endl;
        }

        double
        average_pairwise_identity(const Sequence &seqA, const Sequence &seqB) {
            double avg_identity = 0;
            for (size_t i = 0; i < seqA.num_of_rows(); i++) {
                for (size_t j = 0; j < seqB.num_of_rows(); j++) {
                    avg_identity += sequence_identity(seqA.seqentry(i).seq(),
                                                      seqB.seqentry(j).seq());
                }
            }
            return avg_identity / (seqA.num_of_rows() * seqB.num_of_rows());
        }
    }
}

#endif // LOCARNA_MAIN_HELPER
