#pragma once

#include "Interface.hxx"

#include <elisa/ConfidenceMatrix.hxx>
#include <elisa/Elisa.hxx>
#include <elisa/ElisaFDR.hxx>

#include <pybind11/chrono.h>
#include <pybind11/numpy.h>
#include <pybind11/pybind11.h>
#include <pybind11/pytypes.h>
#include <pybind11/stl.h>

#include <stdexcept>

using namespace pybind11::literals;
using namespace sd::disc;

namespace sd::disc
{

class ElisaInterface : public Interface
{
public:
    virtual real_matrix pvalues_indep() const     = 0;
    virtual real_matrix log_pvalues_indep() const = 0;
};

template <typename Trait, typename Statistic, class Interface>
class ElisaImpl : public MultiClass<Trait, ElisaInterface>
{
public:
    using pattern_type = typename Trait::pattern_type;
    using float_type   = typename Trait::float_type;
    using base_type    = MultiClass<Trait, ElisaInterface>;

    using base_type::base_type;

    Interface interface;

    void do_fit(Dataset<pattern_type> x, std::vector<size_t> const& y) override
    {
        const auto& cfg         = this->cfg;
        auto&       composition = this->composition;

        composition      = Composition<Trait>();
        composition.data = PartitionedData<pattern_type>(std::move(x), y);
        initialize_model(composition, cfg);
        auto fn       = interface;
        this->elapsed = timeit([&] { discover_patterns_generic(composition, cfg, fn); });
        this->extract_explanation(composition);
    }

    real_matrix pvalues_indep() const override
    {
        real_matrix m;
        disc::calc_pvalues_indep(this->composition,
                                 this->cfg,
                                 m,
                                 [](auto p, auto q, auto n)
                                 { return Statistic::pvalue(p, q, n); });
        return m;
    }
    real_matrix log_pvalues_indep() const override
    {
        real_matrix m;
        disc::calc_pvalues_indep(this->composition,
                                 this->cfg,
                                 m,
                                 [](auto p, auto q, auto n)
                                 { return Statistic::log_pvalue(p, q, n); });
        return m;
    }
};

template <typename Trait,
          typename Statistic,
          typename Interface = IElisa<typename Trait::float_type, Statistic>()>
class ElisaImpl1 : public SingleClass<Trait, ElisaInterface>
{
public:
    using pattern_type = typename Trait::pattern_type;
    using float_type   = typename Trait::float_type;
    using base_type    = SingleClass<Trait, ElisaInterface>;

    using base_type::base_type;

    Interface interface;

    void do_fit(Dataset<pattern_type> x) override
    {
        const auto& cfg       = this->cfg;
        auto&       elapsed   = this->elapsed;
        auto&       component = this->component;

        component.data = std::move(x);
        initialize_model(component, cfg);
        auto fn = interface;
        elapsed = timeit([&] { discover_patterns_generic(component, cfg, fn); });
        this->extract_explanation(component);
    }

    real_matrix pvalues_indep() const override
    {
        real_matrix m;
        disc::calc_pvalues_indep(this->component,
                                 this->cfg,
                                 m,
                                 [](auto p, auto q, auto n)
                                 { return Statistic::pvalue(p, q, n); });
        return m;
    }
    real_matrix log_pvalues_indep() const override
    {
        real_matrix m;
        disc::calc_pvalues_indep(this->component,
                                 this->cfg,
                                 m,
                                 [](auto p, auto q, auto n)
                                 { return Statistic::log_pvalue(p, q, n); });
        return m;
    }
};

enum class ElisaMode
{
    FDR_LORD,
    FDR_LOND,
    FDR_LORD_SEPARATE,
    FDR_LOND_SEPARATE,
    FWER //,FWER_TIGHT
};

class Elisa
{
public:
    bool      sparse      = false;
    bool      two_tailed  = false;
    bool      upper       = true;
    bool      gauss       = false;
    bool      relaxed     = false;
    double    alpha       = 0.05;
    size_t    min_support = 2;
    ElisaMode mode        = ElisaMode::FWER;
    // ElisaKind kind = ElisaKind::two_tailed;

#if HAS_HIGH_PRECISION_FLOAT_TYPE
    bool high_precision = false;
#else
    constexpr static bool high_precision = false;
#endif

    Elisa& fit(py::object const& x, py::list const& y = py::list())
    {
        setup(x, py::cast<std::vector<size_t>>(y));
        return *this;
    }
    std::vector<size_t> predict(py::object const& x) const { return get().predict(x); }
    auto                predict_probabilities(py::object const& x) const
    {
        return to_numpy(get().predict_probabilities(x));
    }
    auto predict_log_probabilities(py::object const& x) const
    {
        return to_numpy(get().predict_log_probabilities(x));
    }
    auto pvalues_indep() const { return to_numpy(get().pvalues_indep()); }
    auto log_pvalues_indep() const { return to_numpy(get().log_pvalues_indep()); }

    std::vector<size_t> contrasting(std::vector<size_t> const& s, std::vector<size_t> const& t)
    {
        return get().contrasting(s, t);
    }
    std::vector<size_t> emerging(std::vector<size_t> const& s, std::vector<size_t> const& t)
    {
        return get().emerging(s, t);
    }
    std::vector<size_t> common(std::vector<size_t> const& s) { return get().common(s); }
    std::vector<size_t> characteristic(int i) { return get().characteristic(i); }
    std::vector<size_t> unique(std::vector<size_t> const& s) { return get().unique(s); }
    const std::vector<std::vector<size_t>>& patterns() const { return get().patterns(); }
    std::chrono::milliseconds               elapsed() const { return get().elapsed; }
    auto confidence() const { return to_numpy(get().confidence_matrix); }
    auto frequencies() const { return to_numpy(get().frequency_matrix); }
    auto log_pvalues() const
    {
        auto m = get().confidence_matrix;
        for (auto& x : m) { x = -x; }
        return to_numpy(std::move(m));
        // return get()
        //     .confidence_matrix.unaryExpr([](auto x) { return -x; })
        //     .template cast<double>();
    }
    auto pvalues() const
    {
        using std::exp;
        auto m = get().confidence_matrix;
        for (auto& x : m) { x = exp(-x); }
        return to_numpy(std::move(m));
        // return get()
        //     .confidence_matrix.unaryExpr([](auto x) { return exp(-x); })
        //     .template cast<double>();
    }

private:
    std::unique_ptr<ElisaInterface> impl = nullptr;

    ElisaInterface& get()
    {
        if (impl) return *impl;
        throw std::runtime_error("Model is not yet fitted!");
    }
    const ElisaInterface& get() const
    {
        if (impl) return *impl;
        throw std::runtime_error("Model is not yet fitted!");
    }

    template <typename Trait, typename Stat>
    void emplace_implementation(bool use_single_class)
    {
        using F = typename Trait::float_type;

        auto emplace_into_impl = [&](auto fn)
        {
            if (use_single_class)
                impl = std::make_unique<ElisaImpl1<Trait, Stat, decltype(fn)>>();
            else
                impl = std::make_unique<ElisaImpl<Trait, Stat, decltype(fn)>>();
        };

        switch (mode)
        {
        case ElisaMode::FDR_LORD: emplace_into_impl(IElisaFDR<F, Stat, fdr::LORD<F>>()); break;
        case ElisaMode::FDR_LOND: emplace_into_impl(IElisaFDR<F, Stat, fdr::LORD3<F>>()); break;
        case ElisaMode::FDR_LORD_SEPARATE:
            emplace_into_impl(IElisaFDRI<F, Stat, fdr::LORD<F>>());
            break;
        case ElisaMode::FDR_LOND_SEPARATE:
            emplace_into_impl(IElisaFDRI<F, Stat, fdr::LORD3<F>>());
            break;
        case ElisaMode::FWER: emplace_into_impl(IElisa<F, Stat>()); break;
        }
    }

    template <typename X>
    void setup(X const& x, std::vector<size_t> const& y)
    {
        bool use_single_class = y.empty();

        build_trait(sparse,
                    high_precision,
                    relaxed,
                    [&](auto t)
                    {
                        using Trait       = std::decay_t<decltype(t)>;
                        auto assign_pimpl = [&](auto stat)
                        {
                            emplace_implementation<Trait, decltype(stat)>(use_single_class);

                            // if (use_single_class)
                            //     impl = std::make_unique<ElisaImpl1<Trait, decltype(stat)>>();
                            // else
                            // {
                            //     if (separate_fdr_controls)
                            //         impl = std::make_unique<ElisaImpl<Trait, decltype(stat)>>();
                            //     else
                            //         impl = std::make_unique<ElisaImpl<Trait, decltype(stat)>>();
                            // }
                        };

                        if (upper)
                        {
                            if (gauss)
                                assign_pimpl(ElisaGaussianUpper{});
                            else
                                assign_pimpl(ElisaBinomialUpper{});
                        }
                        else if (two_tailed)
                        {
                            if (gauss)
                                assign_pimpl(ElisaGaussianTT{});
                            else
                                assign_pimpl(ElisaBinomialTT{});
                        }
                        else
                        {
                            if (gauss)
                                assign_pimpl(ElisaGaussian{});
                            else
                                assign_pimpl(ElisaBinomial{});
                        }
                    });

        if (!impl) throw std::domain_error("cannot fit");

        get().cfg.alpha       = alpha;
        get().cfg.min_support = min_support;
        get().fit(x, y);
        // get().fit(x, std::vector<size_t>(y.begin(), y.end()));
    }
};
void init_pymodule_spass(py::module& m)
{
    py::enum_<ElisaMode>(m, "Control")
        .value("FWER", ElisaMode::FWER)
        .value("FDR_LORD", ElisaMode::FDR_LORD)
        .value("FDR_LORD_PER_DATASET", ElisaMode::FDR_LORD_SEPARATE)
        .export_values();

    py::class_<Elisa>(m, "Spass")
        .def(py::init<>())
        // .def(py::init<size_t>())
        // .def(py::init<size_t, size_t>())
        // .def(py::init<size_t, size_t, size_t>())
        .def("fit", &Elisa::fit, "x"_a, "y"_a = py::list()) // py::list()
        // .def("predict", &Elisa::predict, "x"_a)
        // .def("predict_probabilities", &Elisa::predict_probabilities, "x"_a)
        .def("elapsed_time", &Elisa::elapsed)
        // .def("log_pvalues", &Elisa::log_pvalues)
        // .def("pvalues", &Elisa::pvalues)
        // .def("pvalues_indep", &Elisa::pvalues_indep)
        // .def("log_pvalues_indep", &Elisa::log_pvalues_indep)
        .def("frequencies", &Elisa::frequencies)
        .def("confidence", &Elisa::confidence)
        .def("contrasting", &Elisa::contrasting, "i"_a, "j"_a)
        .def("common", &Elisa::common, "i"_a)
        .def("characteristic", &Elisa::characteristic, "i"_a)
        .def("emerging", &Elisa::emerging, "i"_a, "j"_a)
        .def("patterns", &Elisa::patterns)
        .def("unique", &Elisa::unique, "i"_a)
#if HAS_HIGH_PRECISION_FLOAT_TYPE
        .def_readwrite("high_precision", &Elisa::high_precision)
#endif
        .def_readwrite("sparse", &Elisa::sparse)
        .def_readwrite("gaussian", &Elisa::gauss)
        .def_readwrite("upper", &Elisa::upper)
        .def_readwrite("control", &Elisa::mode)
        .def_readwrite("alpha", &Elisa::alpha)
        .def_readwrite("relaxed", &Elisa::relaxed);
}

} // namespace sd::disc
