#ifndef STAN_MATH_REV_MAT_FUN_TRACE_INV_QUAD_FORM_LDLT_HPP
#define STAN_MATH_REV_MAT_FUN_TRACE_INV_QUAD_FORM_LDLT_HPP

#include <stan/math/prim/mat/fun/Eigen.hpp>
#include <stan/math/rev/core.hpp>
#include <stan/math/rev/mat/fun/LDLT_alloc.hpp>
#include <stan/math/rev/mat/fun/LDLT_factor.hpp>
#include <boost/utility/enable_if.hpp>
#include <stan/math/prim/mat/err/check_multiplicable.hpp>
#include <stan/math/rev/scal/meta/is_var.hpp>

namespace stan {
  namespace math {

    namespace {
      template <typename T2, int R2, int C2, typename T3, int R3, int C3>
      class trace_inv_quad_form_ldlt_impl : public chainable_alloc {
      protected:
        inline void initializeB(const Eigen::Matrix<var, R3, C3> &B,
                                bool haveD) {
          Eigen::Matrix<double, R3, C3> Bd(B.rows(), B.cols());
          variB_.resize(B.rows(), B.cols());
          for (int j = 0; j < B.cols(); j++) {
            for (int i = 0; i < B.rows(); i++) {
              variB_(i, j) = B(i, j).vi_;
              Bd(i, j) = B(i, j).val();
            }
          }
          AinvB_ = ldlt_.solve(Bd);
          if (haveD)
            C_.noalias() = Bd.transpose()*AinvB_;
          else
            value_ = (Bd.transpose()*AinvB_).trace();
        }
        inline void initializeB(const Eigen::Matrix<double, R3, C3> &B,
                                bool haveD) {
          AinvB_ = ldlt_.solve(B);
          if (haveD)
            C_.noalias() = B.transpose()*AinvB_;
          else
            value_ = (B.transpose()*AinvB_).trace();
        }

        template<int R1, int C1>
        inline void initializeD(const Eigen::Matrix<var, R1, C1> &D) {
          D_.resize(D.rows(), D.cols());
          variD_.resize(D.rows(), D.cols());
          for (int j = 0; j < D.cols(); j++) {
            for (int i = 0; i < D.rows(); i++) {
              variD_(i, j) = D(i, j).vi_;
              D_(i, j) = D(i, j).val();
            }
          }
        }
        template<int R1, int C1>
        inline void initializeD(const Eigen::Matrix<double, R1, C1> &D) {
          D_ = D;
        }

      public:
        template<typename T1, int R1, int C1>
        trace_inv_quad_form_ldlt_impl(const Eigen::Matrix<T1, R1, C1> &D,
                                      const LDLT_factor<T2, R2, C2>
                                      &A,
                                      const Eigen::Matrix<T3, R3, C3> &B)
          : Dtype_(stan::is_var<T1>::value),
            ldlt_(A) {
          initializeB(B, true);
          initializeD(D);

          value_ = (D_*C_).trace();
        }

        trace_inv_quad_form_ldlt_impl(const LDLT_factor<T2, R2, C2>
                                      &A,
                                      const Eigen::Matrix<T3, R3, C3> &B)
          : Dtype_(2),
            ldlt_(A) {
          initializeB(B, false);
        }

        const int Dtype_;  // 0 = double, 1 = var, 2 = missing
        LDLT_factor<T2, R2, C2> ldlt_;
        Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> D_;
        Eigen::Matrix<vari*, Eigen::Dynamic, Eigen::Dynamic> variD_;
        Eigen::Matrix<vari*, R3, C3> variB_;
        Eigen::Matrix<double, R3, C3> AinvB_;
        Eigen::Matrix<double, C3, C3> C_;
        double value_;
      };

      template <typename T2, int R2, int C2, typename T3, int R3, int C3>
      class trace_inv_quad_form_ldlt_vari : public vari {
      protected:
        static inline
        void
        chainA(const double &adj,
               trace_inv_quad_form_ldlt_impl<double, R2, C2, T3, R3, C3>
               *impl) {
        }
        static inline
        void
        chainB(const double &adj,
               trace_inv_quad_form_ldlt_impl<T2, R2, C2, double, R3, C3>
               *impl) {
        }

        static inline
        void
        chainA(const double &adj,
               trace_inv_quad_form_ldlt_impl<var, R2, C2, T3, R3, C3> *impl) {
          Eigen::Matrix<double, R2, C2> aA;

          if (impl->Dtype_ != 2)
            aA.noalias() = -adj * (impl->AinvB_ * impl->D_.transpose()
                                   * impl->AinvB_.transpose());
          else
            aA.noalias() = -adj*(impl->AinvB_ * impl->AinvB_.transpose());

          for (int j = 0; j < aA.cols(); j++)
            for (int i = 0; i < aA.rows(); i++)
              impl->ldlt_.alloc_->variA_(i, j)->adj_ += aA(i, j);
        }
        static inline
        void
        chainB(const double &adj,
               trace_inv_quad_form_ldlt_impl<T2, R2, C2, var, R3, C3> *impl) {
          Eigen::Matrix<double, R3, C3> aB;

          if (impl->Dtype_ != 2)
            aB.noalias() = adj*impl->AinvB_*(impl->D_ + impl->D_.transpose());
          else
            aB.noalias() = 2*adj*impl->AinvB_;

          for (int j = 0; j < aB.cols(); j++)
            for (int i = 0; i < aB.rows(); i++)
              impl->variB_(i, j)->adj_ += aB(i, j);
        }

      public:
        explicit trace_inv_quad_form_ldlt_vari
        (trace_inv_quad_form_ldlt_impl<T2, R2, C2, T3, R3, C3> *impl)
          : vari(impl->value_), impl_(impl)
        { }

        virtual void chain() {
          // F = trace(D * B' * inv(A) * B)
          // aA = -aF * inv(A') * B * D' * B' * inv(A')
          // aB = aF*(inv(A) * B * D + inv(A') * B * D')
          // aD = aF*(B' * inv(A) * B)
          chainA(adj_, impl_);

          chainB(adj_, impl_);

          if (impl_->Dtype_ == 1) {
            for (int j = 0; j < impl_->variD_.cols(); j++)
              for (int i = 0; i < impl_->variD_.rows(); i++)
                impl_->variD_(i, j)->adj_ += adj_*impl_->C_(i, j);
          }
        }

        trace_inv_quad_form_ldlt_impl<T2, R2, C2, T3, R3, C3> *impl_;
      };

    }

    /**
     * Compute the trace of an inverse quadratic form.  I.E., this computes
     *       trace(B^T A^-1 B)
     * where the LDLT_factor of A is provided.
     **/
    template <typename T2, int R2, int C2, typename T3, int R3, int C3>
    inline typename
    boost::enable_if_c<stan::is_var<T2>::value ||
    stan::is_var<T3>::value,
                       var>::type
      trace_inv_quad_form_ldlt(const LDLT_factor<T2, R2, C2> &A,
                               const Eigen::Matrix<T3, R3, C3> &B) {
      check_multiplicable("trace_inv_quad_form_ldlt",
                          "A", A,
                          "B", B);

      trace_inv_quad_form_ldlt_impl<T2, R2, C2, T3, R3, C3> *impl_
        = new trace_inv_quad_form_ldlt_impl<T2, R2, C2, T3, R3, C3>(A, B);

      return var(new trace_inv_quad_form_ldlt_vari<T2, R2, C2, T3, R3, C3>
                 (impl_));
    }

  }
}
#endif
