#pragma once

/**
\f$ \alpha^r=\displaystyle\sum_i n_i \delta^{d_i} \tau^{t_i}\f$
*/
class JustPowerEOSTerm {
public:
    Eigen::ArrayXd n, t, d;

    template<typename TauType, typename DeltaType>
    auto alphar(const TauType& tau, const DeltaType& delta) const {
        using result = std::common_type_t<TauType, DeltaType>;
        result r = 0.0, lntau = log(tau), lndelta = log(delta);
        for (auto i = 0; i < n.size(); ++i) {
            r = r + n[i] * exp(t[i] * lntau + d[i] * lndelta);
        }
        return forceeval(r);
    }
};

class PowerEOSTerm {
public:
    Eigen::ArrayXd n, t, d, c, l;
    Eigen::ArrayXi l_i;

    template<typename TauType, typename DeltaType>
    auto alphar(const TauType& tau, const DeltaType& delta) const {
        using result = std::common_type_t<TauType, DeltaType>;
        result r = 0.0, lntau = log(tau), lndelta = log(delta);
        for (auto i = 0; i < n.size(); ++i) {
            r = r + n[i] * exp(t[i]*lntau + d[i]*lndelta -c[i]*powi(delta, l_i[i]));
        }
        return forceeval(r);
    }
};

/**
\f$ \alpha^r=\displaystyle\sum_i n_i \delta^{d_i} \tau^{t_i} \exp(-\gamma_i\delta^{l_i})\f$
*/
class ExponentialEOSTerm {
public:
    Eigen::ArrayXd n, t, d, g, l;
    Eigen::ArrayXi l_i;

    template<typename TauType, typename DeltaType>
    auto alphar(const TauType& tau, const DeltaType& delta) const {
        using result = std::common_type_t<TauType, DeltaType>;
        result r = 0.0, lntau = log(tau), lndelta = log(delta);
        for (auto i = 0; i < n.size(); ++i) {
            r = r + n[i] * exp(t[i] * lntau + d[i] * lndelta - g[i] * powi(delta, l_i[i]));
        }
        return forceeval(r);
    }
};

/**
\f$ \alpha^r = \displaystyle\sum_i n_i \tau^{t_i}\delta^ {d_i} \exp(-\eta_i(\delta-\epsilon_i)^2 -\beta_i(\tau-\gamma_i)^2 }\f$
*/
class GaussianEOSTerm {
public:
    Eigen::ArrayXd n, t, d, eta, beta, gamma, epsilon;

    template<typename TauType, typename DeltaType>
    auto alphar(const TauType& tau, const DeltaType& delta) const {
        using result = std::common_type_t<TauType, DeltaType>;
        result r = 0.0, lntau = log(tau), lndelta = log(delta);
        auto square = [](auto x) { return x * x; };
        for (auto i = 0; i < n.size(); ++i) {
            r = r + n[i] * exp(t[i] * lntau + d[i] * lndelta - eta[i]*square(delta - epsilon[i]) - beta[i]*square(tau - gamma[i]));
        }
        return forceeval(r);
    }
};

/**
\f$ \alpha^r = \displaystyle\sum_i n_i \tau^{t_i}\delta^ {d_i} \exp(-\eta_i(\delta-\epsilon_i)^2 -\beta_i(\delta-\gamma_i) }\f$
*/
class GERG2004EOSTerm {
public:
    Eigen::ArrayXd n, t, d, eta, beta, gamma, epsilon;

    template<typename TauType, typename DeltaType>
    auto alphar(const TauType& tau, const DeltaType& delta) const {
        using result = std::common_type_t<TauType, DeltaType>;
        result r = 0.0, lntau = log(tau), lndelta = log(delta);
        auto square = [](auto x) { return x * x; };
        for (auto i = 0; i < n.size(); ++i) {
            r = r + n[i] * exp(t[i] * lntau + d[i] * lndelta - eta[i] * square(delta - epsilon[i]) - beta[i] * (delta - gamma[i]));
        }
        return forceeval(r);
    }
};


/**
\f$ \alpha^r = \displaystyle\sum_i n_i \delta^ { d_i } \tau^ { t_i } \exp(-\delta^ { l_i } - \tau^ { m_i })\f$
*/
class Lemmon2005EOSTerm {
public:
    Eigen::ArrayXd n, t, d, l, m;
    Eigen::ArrayXi l_i;

    template<typename TauType, typename DeltaType>
    auto alphar(const TauType& tau, const DeltaType& delta) const {
        using result = std::common_type_t<TauType, DeltaType>;
        result r = 0.0, lntau = log(tau), lndelta = log(delta);
        for (auto i = 0; i < n.size(); ++i) {
            r = r + n[i] * exp(t[i] * lntau + d[i] * lndelta - powi(delta, l[i]) - pow(tau, m[i]));
        }
        return forceeval(r);
    }
};

/**
\f$ \alpha^r = \displaystyle\sum_i n_i \tau^{t_i}\delta^ {d_i} \exp(-\eta_i(\delta-\epsilon_i)^2 + \frac{1}{\beta_i(\tau-\gamma_i)^2+b_i}\f$
*/
class GaoBEOSTerm {
public:
    Eigen::ArrayXd n, t, d, eta, beta, gamma, epsilon, b;

    template<typename TauType, typename DeltaType>
    auto alphar(const TauType& tau, const DeltaType& delta) const {

        using result = std::common_type_t<TauType, DeltaType>;
        result r = 0.0, lntau = log(tau), lndelta = log(delta);
        auto square = [](auto x) { return x * x; };
        for (auto i = 0; i < n.size(); ++i) {
            r = r + n[i] * exp(t[i] * lntau + d[i] * lndelta - eta[i] * square(delta - epsilon[i]) + 1.0 / (beta[i] * square(tau - gamma[i]) + b[i]));
        }
        return forceeval(r);
    }
};

/**
\f$ \alpha^r = 0\f$
*/
class NullEOSTerm {
public:
    template<typename TauType, typename DeltaType>
    auto alphar(const TauType& tau, const DeltaType& delta) const {
        return static_cast<std::common_type_t<TauType, DeltaType>>(0.0);
    }
};

class NonAnalyticEOSTerm {
public:
    Eigen::ArrayXd A, B, C, D, a, b, beta, n;

    template<typename TauType, typename DeltaType>
    auto alphar(const TauType& tau, const DeltaType& delta) const {
        // The non-analytic term
        auto square = [](auto x) { return x * x; };
        auto delta_min1_sq = square(delta - 1.0);

        using result = std::common_type_t<TauType, DeltaType>;
        result r = 0.0, lntau = log(tau), lndelta = log(delta);
        for (auto i = 0; i < n.size(); ++i) {
            auto Psi = exp(-C[i]*delta_min1_sq - D[i]*square(tau - 1.0));
            auto k = 1.0 / (2.0 * beta[i]);
            auto theta = (1.0 - tau) + A[i] * pow(delta_min1_sq, k);
            auto Delta = square(theta) + B[i]*pow(delta_min1_sq, a[i]);
            r = r + n[i]*pow(Delta, b[i])*delta*Psi;
        }
        result outval = forceeval(r);

        // If we are really, really close to the critical point (tau=delta=1), then the term will become undefined, so let's just return 0 in that case
        double dbl = static_cast<double>(getbaseval(outval));
        if (std::isfinite(dbl)) {
            return outval;
        }
        else {
            return static_cast<decltype(outval)>(0.0);
        }
    }
};


template<typename... Args>
class EOSTermContainer {  
private:
    using varEOSTerms = std::variant<Args...>;
    std::vector<varEOSTerms> coll;
public:

    auto size() const { return coll.size(); }

    template<typename Instance>
    auto add_term(Instance&& instance) {
        coll.emplace_back(std::move(instance));
    }

    template <class Tau, class Delta>
    auto alphar(const Tau& tau, const Delta& delta) const {
        std::common_type_t <Tau, Delta> ar = 0.0;
        for (const auto& term : coll) {
            auto contrib = std::visit([&](auto& t) { return t.alphar(tau, delta); }, term);
            ar = ar + contrib;
        }
        return ar;
    }
};

using EOSTerms = EOSTermContainer<JustPowerEOSTerm, PowerEOSTerm, GaussianEOSTerm, NonAnalyticEOSTerm, Lemmon2005EOSTerm, GaoBEOSTerm, ExponentialEOSTerm>;

using DepartureTerms = EOSTermContainer<JustPowerEOSTerm, PowerEOSTerm, GaussianEOSTerm, GERG2004EOSTerm, NullEOSTerm>;