// Copyright 2018 The Simons Foundation, Inc. - All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef NETKET_VARIATIONALMONTECARLO_HPP
#define NETKET_VARIATIONALMONTECARLO_HPP

#include <complex>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <string>
#include <vector>

#include <Eigen/Dense>
#include <Eigen/IterativeLinearSolvers>
#include <nonstd/optional.hpp>

#include "Machine/machine.hpp"
#include "Operator/abstract_operator.hpp"
#include "Optimizer/optimizer.hpp"
#include "Output/json_output_writer.hpp"
#include "Sampler/abstract_sampler.hpp"
#include "Stats/stats.hpp"
#include "Utils/parallel_utils.hpp"
#include "Utils/random_utils.hpp"
#include "common_types.hpp"
#include "matrix_replacement.hpp"

namespace netket {

// Variational Monte Carlo schemes to learn the ground state
// Available methods:
// 1) Stochastic reconfiguration optimizer
//   both direct and sparse version
// 2) Gradient Descent optimizer
class VariationalMonteCarlo {
  using GsType = Complex;
  using VectorT = Eigen::Matrix<Complex, Eigen::Dynamic, 1>;
  using MatrixT = Eigen::Matrix<Complex, Eigen::Dynamic, Eigen::Dynamic>;

  const AbstractOperator &ham_;
  AbstractSampler &sampler_;
  AbstractMachine &psi_;

  std::vector<std::vector<int>> connectors_;
  std::vector<std::vector<double>> newconfs_;
  std::vector<Complex> mel_;

  Eigen::VectorXcd elocs_;
  MatrixT Ok_;
  VectorT Okmean_;

  Eigen::MatrixXd vsamp_;

  Eigen::VectorXcd grad_;
  Eigen::VectorXcd gradprev_;

  double sr_diag_shift_;
  bool sr_rescale_shift_;
  bool use_iterative_;

  int totalnodes_;
  int mynode_;

  AbstractOptimizer &opt_;

  std::vector<AbstractOperator *> obs_;
  std::vector<std::string> obsnames_;
  ObsManager obsmanager_;

  bool dosr_;

  bool use_cholesky_;

  int nsamples_;
  int nsamples_node_;
  int ninitsamples_;
  int ndiscardedsamples_;

  Complex elocmean_;
  double elocvar_;
  int npar_;

 public:
  class Iterator {
   public:
    // typedefs required for iterators
    using iterator_category = std::input_iterator_tag;
    using difference_type = Index;
    using value_type = Index;
    using pointer_type = Index *;
    using reference_type = Index &;

   private:
    VariationalMonteCarlo &vmc_;
    Index step_size_;
    nonstd::optional<Index> n_iter_;

    Index cur_iter_;

   public:
    Iterator(VariationalMonteCarlo &vmc, Index step_size,
             nonstd::optional<Index> n_iter)
        : vmc_(vmc),
          step_size_(step_size),
          n_iter_(std::move(n_iter)),
          cur_iter_(0) {}

    Index operator*() const { return cur_iter_; }

    Iterator &operator++() {
      vmc_.Advance(step_size_);
      cur_iter_ += step_size_;
      return *this;
    }

    // TODO(C++17): Replace with comparison to special Sentinel type, since
    // C++17 allows end() to return a different type from begin().
    bool operator!=(const Iterator &) {
      return !n_iter_.has_value() || cur_iter_ < n_iter_.value();
    }
    // pybind11::make_iterator requires operator==
    bool operator==(const Iterator &other) { return !(*this != other); }

    Iterator begin() const { return *this; }
    Iterator end() const { return *this; }
  };

  VariationalMonteCarlo(const AbstractOperator &hamiltonian,
                        AbstractSampler &sampler, AbstractOptimizer &optimizer,
                        int nsamples, int discarded_samples = -1,
                        int discarded_samples_on_init = 0,
                        const std::string &method = "Sr",
                        double diag_shift = 0.01, bool rescale_shift = false,
                        bool use_iterative = false, bool use_cholesky = true)
      : ham_(hamiltonian),
        sampler_(sampler),
        psi_(sampler.GetMachine()),
        opt_(optimizer),
        elocvar_(0.) {
    Init(nsamples, discarded_samples, discarded_samples_on_init, method,
         diag_shift, rescale_shift, use_iterative, use_cholesky);
  }

  void Init(int nsamples, int discarded_samples, int discarded_samples_on_init,
            const std::string &method, double diagshift, bool rescale_shift,
            bool use_iterative, bool use_cholesky) {
    npar_ = psi_.Npar();

    opt_.Init(psi_.GetParameters());

    grad_.resize(npar_);
    Okmean_.resize(npar_);

    setSrParameters();

    MPI_Comm_size(MPI_COMM_WORLD, &totalnodes_);
    MPI_Comm_rank(MPI_COMM_WORLD, &mynode_);

    nsamples_ = nsamples;

    nsamples_node_ = int(std::ceil(double(nsamples_) / double(totalnodes_)));

    ninitsamples_ = discarded_samples_on_init;

    if (discarded_samples == -1) {
      ndiscardedsamples_ = 0.1 * nsamples_node_;
    } else {
      ndiscardedsamples_ = discarded_samples;
    }

    if (method == "Gd") {
      dosr_ = false;
    } else {
      setSrParameters(diagshift, rescale_shift, use_iterative, use_cholesky);
    }

    if (dosr_) {
      InfoMessage() << "Using the Stochastic reconfiguration method"
                    << std::endl;

      if (use_iterative_) {
        InfoMessage() << "With iterative solver" << std::endl;
      } else {
        if (use_cholesky_) {
          InfoMessage() << "Using Cholesky decomposition" << std::endl;
        }
      }
    } else {
      InfoMessage() << "Using a gradient-descent based method" << std::endl;
    }

    InfoMessage() << "Variational Monte Carlo running on " << totalnodes_
                  << " processes" << std::endl;

    MPI_Barrier(MPI_COMM_WORLD);
  }

  void AddObservable(AbstractOperator &ob, const std::string &obname) {
    obs_.push_back(&ob);
    obsnames_.push_back(obname);
  }

  void InitSweeps() {
    sampler_.Reset();

    for (int i = 0; i < ninitsamples_; i++) {
      sampler_.Sweep();
    }
  }

  void Sample() {
    sampler_.Reset();

    for (int i = 0; i < ndiscardedsamples_; i++) {
      sampler_.Sweep();
    }

    vsamp_.resize(nsamples_node_, psi_.Nvisible());

    for (int i = 0; i < nsamples_node_; i++) {
      sampler_.Sweep();
      vsamp_.row(i) = sampler_.Visible();
    }
  }

  /**
   * Computes the expectation values of observables from the currently stored
   * samples.
   */
  void ComputeObservables() {
    const Index nsamp = vsamp_.rows();
    for (const auto &obname : obsnames_) {
      obsmanager_.Reset(obname);
    }
    for (Index i_samp = 0; i_samp < nsamp; ++i_samp) {
      for (std::size_t i_obs = 0; i_obs < obs_.size(); ++i_obs) {
        const auto &op = obs_[i_obs];
        const auto &name = obsnames_[i_obs];
        obsmanager_.Push(name, ObsLocValue(*op, vsamp_.row(i_samp)).real());
      }
    }
  }

  void Gradient() {
    obsmanager_.Reset("Energy");
    obsmanager_.Reset("EnergyVariance");

    const int nsamp = vsamp_.rows();
    elocs_.resize(nsamp);
    Ok_.resize(nsamp, psi_.Npar());

    for (int i = 0; i < nsamp; i++) {
      elocs_(i) = ObsLocValue(ham_, vsamp_.row(i));
      Ok_.row(i) = psi_.DerLog(vsamp_.row(i));
      obsmanager_.Push("Energy", elocs_(i).real());
    }

    elocmean_ = elocs_.mean();
    SumOnNodes(elocmean_);
    elocmean_ /= double(totalnodes_);

    Okmean_ = Ok_.colwise().mean();
    SumOnNodes(Okmean_);
    Okmean_ /= double(totalnodes_);

    Ok_ = Ok_.rowwise() - Okmean_.transpose();

    elocs_ -= elocmean_ * Eigen::VectorXd::Ones(nsamp);

    for (int i = 0; i < nsamp; i++) {
      obsmanager_.Push("EnergyVariance", std::norm(elocs_(i)));
    }

    grad_ = 2. * (Ok_.adjoint() * elocs_);

    // Summing the gradient over the nodes
    SumOnNodes(grad_);
    grad_ /= double(totalnodes_ * nsamp);
  }

  /**
   * Computes the value of the local estimator of the operator `ob` in
   * configuration `v` which is defined by O_loc(v) = ⟨v|ob|Ψ⟩ / ⟨v|Ψ⟩.
   *
   * @param ob Operator representing the observable.
   * @param v Many-body configuration
   * @return The value of the local observable O_loc(v).
   */
  Complex ObsLocValue(const AbstractOperator &ob, const Eigen::VectorXd &v) {
    ob.FindConn(v, mel_, connectors_, newconfs_);

    assert(connectors_.size() == mel_.size());

    auto logvaldiffs = (psi_.LogValDiff(v, connectors_, newconfs_));

    assert(mel_.size() == std::size_t(logvaldiffs.size()));

    Complex obval = 0;

    for (int i = 0; i < logvaldiffs.size(); i++) {
      obval += mel_[i] * std::exp(logvaldiffs(i));
    }

    return obval;
  }

  double ElocMean() { return elocmean_.real(); }

  double Elocvar() { return elocvar_; }

  void Advance(Index steps = 1) {
    assert(steps > 0);
    for (Index i = 0; i < steps; ++i) {
      Sample();
      Gradient();
      UpdateParameters();
    }
  }

  Iterator Iterate(const nonstd::optional<Index> &n_iter = nonstd::nullopt,
                   Index step_size = 1) {
    assert(!n_iter.has_value() || n_iter.value() > 0);
    assert(step_size > 0);

    opt_.Reset();
    InitSweeps();

    Advance(step_size);
    return Iterator(*this, step_size, n_iter);
  }

  void Run(const std::string &output_prefix,
           nonstd::optional<Index> n_iter = nonstd::nullopt,
           Index step_size = 1, Index save_params_every = 50) {
    assert(n_iter > 0);
    assert(step_size > 0);
    assert(save_params_every > 0);

    nonstd::optional<JsonOutputWriter> writer;
    if (mynode_ == 0) {
      writer.emplace(output_prefix + ".log", output_prefix + ".wf",
                     save_params_every);
    }

    for (const auto step : Iterate(n_iter, step_size)) {
      ComputeObservables();

      // Note: This has to be called in all MPI processes, because converting
      // the ObsManager to JSON performs a MPI reduction.
      auto obs_data = json(obsmanager_);
      obs_data["Acceptance"] = sampler_.Acceptance();

      // writer.has_value() iff the MPI rank is 0, so the output is only
      // written once
      if (writer.has_value()) {
        writer->WriteLog(step, obs_data);
        writer->WriteState(step, psi_);
      }
      MPI_Barrier(MPI_COMM_WORLD);
    }
  }

  void UpdateParameters() {
    auto pars = psi_.GetParameters();

    if (dosr_) {
      const int nsamp = vsamp_.rows();

      Eigen::VectorXcd b = Ok_.adjoint() * elocs_;
      SumOnNodes(b);
      b /= double(nsamp * totalnodes_);

      if (!use_iterative_) {
        // Explicit construction of the S matrix
        Eigen::MatrixXcd S = Ok_.adjoint() * Ok_;
        SumOnNodes(S);
        S /= double(nsamp * totalnodes_);

        // Adding diagonal shift
        S += Eigen::MatrixXd::Identity(pars.size(), pars.size()) *
             sr_diag_shift_;

        Eigen::VectorXcd deltaP;
        if (use_cholesky_ == false) {
          Eigen::FullPivHouseholderQR<Eigen::MatrixXcd> qr(S.rows(), S.cols());
          qr.setThreshold(1.0e-6);
          qr.compute(S);
          deltaP = qr.solve(b);
        } else {
          Eigen::LLT<Eigen::MatrixXcd> llt(S.rows());
          llt.compute(S);
          deltaP = llt.solve(b);
        }
        // Eigen::VectorXcd deltaP=S.jacobiSvd(ComputeThinU |
        // ComputeThinV).solve(b);

        assert(deltaP.size() == grad_.size());
        grad_ = deltaP;

        if (sr_rescale_shift_) {
          Complex nor = (deltaP.dot(S * deltaP));
          grad_ /= std::sqrt(nor.real());
        }

      } else {
        Eigen::ConjugateGradient<MatrixReplacement, Eigen::Lower | Eigen::Upper,
                                 Eigen::IdentityPreconditioner>
            it_solver;
        // Eigen::GMRES<MatrixReplacement, Eigen::IdentityPreconditioner>
        // it_solver;
        it_solver.setTolerance(1.0e-3);
        MatrixReplacement S;
        S.attachMatrix(Ok_);
        S.setShift(sr_diag_shift_);
        S.setScale(1. / double(nsamp * totalnodes_));

        it_solver.compute(S);
        auto deltaP = it_solver.solve(b);

        grad_ = deltaP;
        if (sr_rescale_shift_) {
          auto nor = deltaP.dot(S * deltaP);
          grad_ /= std::sqrt(nor.real());
        }

        // if(mynode_==0){
        //   std::cerr<<it_solver.iterations()<<"
        //   "<<it_solver.error()<<std::endl;
        // }
        MPI_Barrier(MPI_COMM_WORLD);
      }
    }

    opt_.Update(grad_, pars);

    SendToAll(pars);

    psi_.SetParameters(pars);
    MPI_Barrier(MPI_COMM_WORLD);
  }

  void setSrParameters(double diagshift = 0.01, bool rescale_shift = false,
                       bool use_iterative = false, bool use_cholesky = true) {
    sr_diag_shift_ = diagshift;
    sr_rescale_shift_ = rescale_shift;
    use_iterative_ = use_iterative;
    dosr_ = true;
    use_cholesky_ = use_cholesky;
  }

  AbstractMachine &GetMachine() { return psi_; }
  const ObsManager &GetObsManager() const { return obsmanager_; }
};

}  // namespace netket

#endif


/*
 *                                                                                * Note: This license has also been called the "New BSD License"
 * or "Modified BSD License".
 *
 * Copyright (c) 2021 Electronics and Telecommunications Research
 * Institute All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * 1. Redistributions of source code must retain the above copyright notice,
 *    this list of conditions and the following disclaimer.
 *
 * 2. Redistributions in binary form must reproduce the above copyright notice,
 *    this list of conditions and the following disclaimer in the documentation
 *    and/or other materials provided with the distribution.
 *
 * 3. Neither the name of the copyright holder nor the names of its contributors
 *    may be used to endorse or promote products derived from this software
 *    without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
 * FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
 * COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
 * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
 * ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
*/

/**
 * CityGeodeticCalculator.cpp
 *
 * $Revision: 625 $
 * $LastChangedDate: 2019-07-19 15:18:25 +0900 (Fri, 19 Jul 2019) $
 */

#include <boost/math/constants/constants.hpp>
#include <string>
#include <vector>
#include <cmath>
#include "City.h"
#include "CityGeodeticCalculator.h"
#include "CityUtils.h"

namespace dtsim {

const double CityGeodeticCalculator::WGS84_a = 6378137;
const double CityGeodeticCalculator::WGS84_b = 6356752.314245;
const double CityGeodeticCalculator::WGS84_f = 1/298.257223563;
const double CityGeodeticCalculator::RADIUS_EARTH = 6371000.0;

CityGeodeticCalculator::CityGeodeticCalculator()
{
	if (City::instance()->containProperty("geodetic.distance.formula")) {
		std::string formulaStr = City::instance()->getStringProperty("geodetic.distance.formula");

		if ((formulaStr.compare("haversine") == 0) || (formulaStr.compare("Haversine") == 0))
			formula = GEODETIC_HAVERSINE;
		else if ((formulaStr.compare("vincenty") == 0) || (formulaStr.compare("Vincenty") == 0))
			formula = GEODETIC_VINCENTY;
		else
			formula = GEODETIC_HAVERSINE;
	} else {
		formula = GEODETIC_HAVERSINE;
	}
}

CityGeodeticCalculator::~CityGeodeticCalculator()
{
	_instance = 0;
}

CityGeodeticCalculator* CityGeodeticCalculator::_instance = 0;

CityGeodeticCalculator* CityGeodeticCalculator::instance()
{
	if (_instance == 0)
		_instance = new CityGeodeticCalculator();

	return _instance;
}

void CityGeodeticCalculator::haversine_getDestinationPoint(double startLon, double startLat,
														   double distance, double azimuth,
														   double& destLon, double& destLat)
{
	double delta = distance / RADIUS_EARTH;
	double theta = degreesToRadians(azimuth);

	double lambda1 = degreesToRadians(startLon);
	double phi1 = degreesToRadians(startLat);

	double sinPhi2 = (std::sin(phi1) * std::cos(delta)) +
					 (std::cos(phi1) * std::sin(delta) * std::cos(theta));
	double phi2 = std::asin(sinPhi2);

	double x = std::cos(delta) - (std::sin(phi1) * sinPhi2);
	double y = std::sin(theta) * std::sin(delta) * std::cos(phi1);
	double lambda2 = lambda1 + std::atan2(y, x);

	destLon = radiansToDegrees(lambda2);
	destLat = radiansToDegrees(phi2);
}

void CityGeodeticCalculator::vincenty_getDestinationPoint(double startLon, double startLat,
														  double distance, double azimuth,
														  double& destLon, double& destLat)
{
	double lambda1 = degreesToRadians(startLon);
	double phi1 = degreesToRadians(startLat);
	double alpha1 = degreesToRadians(azimuth);
	double s = distance;

	double sinAlpha1 = std::sin(alpha1);
	double cosAlpha1 = std::cos(alpha1);

	double tanU1 = (1 - WGS84_f) * std::tan(phi1);
	double cosU1 = 1 / std::sqrt((1 + tanU1*tanU1));
	double sinU1 = tanU1 * cosU1;

	double sigma1 = std::atan2(tanU1, cosAlpha1);
	double sinAlpha = cosU1 * sinAlpha1;
	double cosSqAlpha = 1 - sinAlpha*sinAlpha;
	double uSq = cosSqAlpha * (WGS84_a*WGS84_a - WGS84_b*WGS84_b) / (WGS84_b*WGS84_b);
	double A = 1 + uSq/16384*(4096+uSq*(-768+uSq*(320-175*uSq)));
	double B = uSq/1024 * (256+uSq*(-128+uSq*(74-47*uSq)));

	double cos2SigmaM, sinSigma, cosSigma, deltaSigma;

	double sigma = s / (WGS84_b*A);
	double sigmaP;
	int iterations = 0;
	do {
		cos2SigmaM = std::cos(2*sigma1 + sigma);
		sinSigma = std::sin(sigma);
		cosSigma = std::cos(sigma);
		deltaSigma = B * sinSigma * (cos2SigmaM + B/4*(cosSigma*(-1 + 2*cos2SigmaM*cos2SigmaM)-
					 B / 6*cos2SigmaM*(-3 + 4*sinSigma*sinSigma)*(-3 + 4*cos2SigmaM*cos2SigmaM)));
		sigmaP = sigma;
		sigma = s / (WGS84_b*A) + deltaSigma;
	} while ((std::abs(sigma - sigmaP) > 1e-12) && (++iterations < 1000));

	double chi = sinU1*sinSigma - cosU1*cosSigma*cosAlpha1;
	double phi2 = std::atan2(sinU1*cosSigma + cosU1*sinSigma*cosAlpha1, (1-WGS84_f)*std::sqrt(sinAlpha*sinAlpha + chi*chi));
	double lambda = std::atan2(sinSigma*sinAlpha1, cosU1*cosSigma - sinU1*sinSigma*cosAlpha1);
	double C = WGS84_f/16*cosSqAlpha*(4 + WGS84_f*(4 - 3*cosSqAlpha));
	double L = lambda - (1-C) * WGS84_f * sinAlpha * (sigma + C*sinSigma*(cos2SigmaM+C*cosSigma*(-1+2*cos2SigmaM*cos2SigmaM)));
	double lambda2 = lambda1 + L;

	double alpha2 = std::atan2(sinAlpha, -chi);

	destLon = radiansToDegrees(lambda2);
	destLat = radiansToDegrees(phi2);
}

void CityGeodeticCalculator::getDestinationPoint(double startLon, double startLat,
												 double distance, double azimuth,
												 double& destLon, double& destLat)
{
	if (formula == GEODETIC_HAVERSINE)
		haversine_getDestinationPoint(startLon, startLat, distance, azimuth, destLon, destLat);
	else
		vincenty_getDestinationPoint(startLon, startLat, distance, azimuth, destLon, destLat);
}

void CityGeodeticCalculator::getDestinationPoint(CityCoordinate& startPoint,
												 double distance, double azimuth,
												 CityCoordinate& destPoint)
{
	getDestinationPoint(startPoint.x, startPoint.y, distance, azimuth, destPoint.x, destPoint.y);
}

void CityGeodeticCalculator::getMidPoint(double startLon, double startLat,
										 double destLon, double destLat,
										 double& midLon, double& midLat)
{
	double lambda1 = degreesToRadians(startLon);
	double phi1 = degreesToRadians(startLat);
	double phi2 = degreesToRadians(destLat);

	double deltaLambda = degreesToRadians(destLon - startLon);

	std::vector<double> A;
	std::vector<double> B;
	std::vector<double> C;
	double x, y, z;

	x = std::cos(phi1);
	y = 0;
	z = std::sin(phi1);
	A.push_back(x);
	A.push_back(y);
	A.push_back(z);

	x = std::cos(phi2) * std::cos(deltaLambda);
	y = std::cos(phi2) * std::sin(deltaLambda);
	z = std::sin(phi2);
	B.push_back(x);
	B.push_back(y);
	B.push_back(z);

	x = A[0] + B[0];
	y = A[1] + B[1];
	z = A[2] + B[2];
	C.push_back(x);
	C.push_back(y);
	C.push_back(z);

	double phiM = std::atan2(C[2], std::sqrt((C[0]*C[0]) + (C[1]*C[1])));
	double lambdaM = lambda1 + std::atan2(C[1], C[0]);

	midLon = radiansToDegrees(lambdaM);
	midLat = radiansToDegrees(phiM);
}

void CityGeodeticCalculator::getMidPoint(CityCoordinate& startPoint, CityCoordinate& destPoint,
										 CityCoordinate& midPoint)
{
	getMidPoint(startPoint.x, startPoint.y, destPoint.x, destPoint.y, midPoint.x, midPoint.y);
}

void CityGeodeticCalculator::getIntermediatePoint(double startLon, double startLat, double fraction,
												  double destLon, double destLat,
												  double& intermediateLon, double& intermediateLat)
{
	if ((startLon == destLon) && (startLat == destLat)) {
		intermediateLon = startLon;
		intermediateLat = startLat;
		return;
	}

	double lambda1 = degreesToRadians(startLon);
	double phi1 = degreesToRadians(startLat);
	double lambda2 = degreesToRadians(destLon);
	double phi2 = degreesToRadians(destLat);

	double deltaPhi = phi2 - phi1;
	double deltaLambda = lambda2 - lambda1;

	double a = (std::sin(deltaPhi*0.5)*std::sin(deltaPhi*0.5)) +
			   (std::cos(phi1)*std::cos(phi2)*std::sin(deltaLambda*0.5)*std::sin(deltaLambda*0.5));
	double delta = 2 * std::atan2(std::sqrt(a), std::sqrt(1-a));

	double A = std::sin((1-fraction)*delta) / std::sin(delta);
	double B = std::sin(fraction*delta) / std::sin(delta);

	double x = (A * std::cos(phi1) * std::cos(lambda1)) + (B * std::cos(phi2) * std::cos(lambda2));
	double y = (A * std::cos(phi1) * std::sin(lambda1)) + (B * std::cos(phi2) * std::sin(lambda2));
	double z = (A * std::sin(phi1)) + (B * std::sin(phi2));

	double phi3 = std::atan2(z, std::sqrt((x*x) + (y*y)));
	double lambda3 = std::atan2(y, x);

	intermediateLon = radiansToDegrees(lambda3);
	intermediateLat = radiansToDegrees(phi3);
}

void CityGeodeticCalculator::getIntermediatePoint(CityCoordinate& startPoint, double fraction,
												  CityCoordinate& destPoint,
												  CityCoordinate& intermediatePoint)
{
	getIntermediatePoint(startPoint.x, startPoint.y, fraction, destPoint.x, destPoint.y,
						 intermediatePoint.x, intermediatePoint.y);
}

double CityGeodeticCalculator::haversine_getDistance(double startLon, double startLat,
													 double destLon, double destLat)
{
	double lambda1 = degreesToRadians(startLon);
	double phi1 = degreesToRadians(startLat);
	double lambda2 = degreesToRadians(destLon);
	double phi2 = degreesToRadians(destLat);

	double deltaPhi = phi2 - phi1;
	double deltaLambda = lambda2 - lambda1;

	double a = (std::sin(deltaPhi*0.5) * std::sin(deltaPhi*0.5)) +
			   (std::cos(phi1)*std::cos(phi2)*std::sin(deltaLambda*0.5)*std::sin(deltaLambda*0.5));
	double c = 2.0 * std::atan2(std::sqrt(a), std::sqrt(1-a));

	return RADIUS_EARTH * c;
}

double CityGeodeticCalculator::vincenty_getDistance(double startLon, double startLat,
													double destLon, double destLat, double& azimuth)
{
	double lambda1 = degreesToRadians(startLon);
	double phi1 = degreesToRadians(startLat);
	double lambda2 = degreesToRadians(destLon);
	double phi2 = degreesToRadians(destLat);

	double L = lambda2 - lambda1;
	double tanU1 = (1 - WGS84_f) * std::tan(phi1);
	double cosU1 = 1 / std::sqrt((1 + tanU1*tanU1));
	double sinU1 = tanU1 * cosU1;
	double tanU2 = (1 - WGS84_f) * std::tan(phi2);
	double cosU2 = 1 / std::sqrt((1 + tanU2*tanU2));
	double sinU2 = tanU2 * cosU2;

	double sinLambda, cosLambda, sinSigma = 0, cosSigma = 0, sinAlpha;
	double sinSqSigma, cosSqAlpha = 0, cos2SigmaM = 0, sigma = 0, C;

	double lambda = L;
	double lambdaP;
	int iterations = 0;

	bool antimeridian = std::abs(L) > boost::math::constants::pi<double>();

	do {
		sinLambda = std::sin(lambda);
		cosLambda = std::cos(lambda);

		sinSqSigma = (cosU2*sinLambda)*(cosU2*sinLambda) +
					 (cosU1*sinU2 - sinU1*cosU2*cosLambda) *
					 (cosU1*sinU2 - sinU1*cosU2*cosLambda);

		if (std::abs(sinSqSigma) < std::numeric_limits<double>::epsilon())
			break;

		sinSigma = std::sqrt(sinSqSigma);
		cosSigma = sinU1*sinU2 + cosU1*cosU2*cosLambda;

		sigma = std::atan2(sinSigma, cosSigma);
		sinAlpha = cosU1*cosU2*sinLambda / sinSigma;
		cosSqAlpha = 1 - sinAlpha*sinAlpha;
		cos2SigmaM = (cosSqAlpha != 0) ? (cosSigma - 2*sinU1*sinU2/cosSqAlpha) : 0;

		C = WGS84_f/16*cosSqAlpha*(4 + WGS84_f*(4 - 3*cosSqAlpha));
		lambdaP = lambda;
		lambda = L + (1-C)*WGS84_f*sinAlpha*(sigma + C*sinSigma*(cos2SigmaM + C*cosSigma*(-1 + 2*cos2SigmaM*cos2SigmaM)));
	} while ((std::abs(lambda - lambdaP) > 1e-12) && (++iterations < 1000));

	double uSq = cosSqAlpha * (WGS84_a*WGS84_a - WGS84_b*WGS84_b) / (WGS84_b*WGS84_b);
	double A = 1 + uSq/16384*(4096 + uSq*(-768 + uSq*(320  -175*uSq)));
	double B = uSq/1024 * (256 + uSq*(-128 + uSq*(74 - 47*uSq)));
	double deltaSigma = B*sinSigma*(cos2SigmaM + B/4*(cosSigma*(-1 + 2*cos2SigmaM*cos2SigmaM) -
						B/6*cos2SigmaM*(-3 + 4*sinSigma*sinSigma)*(-3 + 4*cos2SigmaM*cos2SigmaM)));

	double distance = WGS84_b*A*(sigma - deltaSigma);

	double alpha = std::atan2(cosU2*sinLambda, cosU1*sinU2 - sinU1*cosU2*cosLambda);
	double degrees = radiansToDegrees(alpha);
	if ((0 > degrees) || (degrees >= 360))
		degrees = std::fmod((std::fmod(degrees, 360.0) + 360.0), 360.0);

	azimuth = std::abs(distance) < std::numeric_limits<double>::epsilon() ? NAN : degrees;

	return distance;
}

double CityGeodeticCalculator::getDistance(double startLon, double startLat,
										   double destLon, double destLat)
{
	if (formula == GEODETIC_HAVERSINE) {
		return haversine_getDistance(startLon, startLat, destLon, destLat);
	} else {
		double azimuth;
		return vincenty_getDistance(startLon, startLat, destLon, destLat, azimuth);
	}
}

double CityGeodeticCalculator::getDistance(CityCoordinate& startPoint, CityCoordinate& destPoint)
{
	return getDistance(startPoint.x, startPoint.y, destPoint.x, destPoint.y);
}

double CityGeodeticCalculator::haversine_getAzimuth(double startLon, double startLat,
													double destLon, double destLat)
{
	double phi1 = degreesToRadians(startLat);
	double phi2 = degreesToRadians(destLat);

	double deltaLambda = degreesToRadians(destLon - startLon);

	double x = (std::cos(phi1) * std::sin(phi2)) -
			   (std::sin(phi1) * std::cos(phi2) * std::cos(deltaLambda));
	double y = std::sin(deltaLambda) * std::cos(phi2);
	double theta = std::atan2(y, x);

	double azimuth = radiansToDegrees(theta);

	if ((0 > azimuth) || (azimuth >= 360))
		azimuth = std::fmod((std::fmod(azimuth, 360.0) + 360.0), 360.0);

	return azimuth;
}

double CityGeodeticCalculator::vincenty_getAzimuth(double startLon, double startLat,
												   double destLon, double destLat)
{
	double azimuth;
	vincenty_getDistance(startLon, startLat, destLon, destLat, azimuth);

	return azimuth;
}

double CityGeodeticCalculator::getAzimuth(double startLon, double startLat,
										  double destLon, double destLat)
{
	if ((startLon == destLon) && (startLat == destLat))
		return NAN;

	if (formula == GEODETIC_HAVERSINE)
		return haversine_getAzimuth(startLon, startLat, destLon, destLat);
	else
		return vincenty_getAzimuth(startLon, startLat, destLon, destLat);
}

double CityGeodeticCalculator::getAzimuth(CityCoordinate& startPoint, CityCoordinate& destPoint)
{
	return getAzimuth(startPoint.x, startPoint.y, destPoint.x, destPoint.y);
}

} /* namespace dtsim */


/* Copyright (c) 2015, Julian Straub <jstraub@csail.mit.edu>
 * Licensed under the MIT license. See the license file LICENSE.
 */
#pragma once

#include <Eigen/Dense>
#include <iostream>

#include <boost/random/mersenne_twister.hpp>
#include <boost/shared_ptr.hpp>

#include "sphere.hpp"
#include "clusterer.hpp"
#include "dir.hpp"
#include "cat.hpp"

using namespace Eigen;
using std::cout;
using std::endl;

template<class T>
class KMeans : public Clusterer<T>
{
public:
  KMeans(const boost::shared_ptr<Matrix<T,Dynamic,Dynamic> >& spx, uint32_t K,
    boost::mt19937* pRndGen);
  virtual ~KMeans();

//  void initialize(const Matrix<T,Dynamic,Dynamic>& x);

  virtual void updateLabels();
  virtual void updateCenters();
  virtual MatrixXu mostLikelyInds(uint32_t n, Matrix<T,Dynamic,Dynamic>& deviates);
  virtual T avgIntraClusterDeviation();

  virtual T dist(const Matrix<T,Dynamic,1>& a, const Matrix<T,Dynamic,1>& b);
  virtual bool closer(T a, T b);
  virtual uint32_t indOfClosestCluster(int32_t i);
  virtual Matrix<T,Dynamic,1> computeCenter(uint32_t k);

protected:
  Sphere<T> S_;  // TODO this should not be here - needed for the empty cluster
};

// --------------------------- impl -------------------------------------------

template<class T>
KMeans<T>::KMeans(
    const boost::shared_ptr<Matrix<T,Dynamic,Dynamic> >& spx, uint32_t K,
    boost::mt19937* pRndGen)
  : Clusterer<T>(spx,K, pRndGen), S_(this->D_)
{
  Matrix<T,Dynamic,1> alpha(this->K_);
  alpha.setOnes(this->K_);
  Dir<Cat<T>,T> dir(alpha,this->pRndGen_);
  Cat<T> pi = dir.sample();
  cout<<"init pi="<<pi.pdf().transpose()<<endl;
  pi.sample(this->z_);

  this->ps_.setZero();
//  updateCenters();
//  for(uint32_t k=0; k<this->K_; ++k)
//    this->ps_.col(k) = S_.sampleUnif(this->pRndGen_);
//  cout<<"init centers"<<endl<<this->ps_<<endl;
}

template<class T>
KMeans<T>::~KMeans()
{}

template<class T>
T KMeans<T>::dist(const Matrix<T,Dynamic,1>& a, const Matrix<T,Dynamic,1>& b)
{
  return (a-b).norm();
};

template<class T>
bool KMeans<T>::closer(T a, T b)
{
  return a<b; // if dist a is smaller than dist b a is closer than b (Eucledian)
};

template<class T>
uint32_t KMeans<T>::indOfClosestCluster(int32_t i)
{
  T sim_closest = dist(this->ps_.col(0), this->spx_->col(i));
  uint32_t z_i = 0;
  for(uint32_t k=1; k<this->K_; ++k)
  {
    T sim_k = dist(this->ps_.col(k), this->spx_->col(i));
    if( closer(sim_k, sim_closest))
    {
      sim_closest = sim_k;
      z_i = k;
    }
  }
  return z_i;
};

template<class T>
void KMeans<T>::updateLabels()
{
#pragma omp parallel for
  for(uint32_t i=0; i<this->N_; ++i)
  {
    this->z_(i) = indOfClosestCluster(i);
  }
}

template<class T>
Matrix<T,Dynamic,1> KMeans<T>::computeCenter(uint32_t k)
{
  this->Ns_(k) = 0.0;
  Matrix<T,Dynamic,1> mean_k(this->D_);
  mean_k.setZero(this->D_);
  for(uint32_t i=0; i<this->N_; ++i)
    if(this->z_(i) == k)
    {
      mean_k += this->spx_->col(i);
      this->Ns_(k) ++;
    }
  return mean_k/this->Ns_(k);
}

template<class T>
void KMeans<T>::updateCenters()
{
#pragma omp parallel for
  for(uint32_t k=0; k<this->K_; ++k)
  {
    this->ps_.col(k) = computeCenter(k);
    if (this->Ns_(k) <= 0)
      this->ps_.col(k) = S_.sampleUnif(this->pRndGen_);
  }
}

template<class T>
MatrixXu KMeans<T>::mostLikelyInds(uint32_t n,
    Matrix<T,Dynamic,Dynamic>& deviates)
{
  MatrixXu inds = MatrixXu::Zero(n,this->K_);
  deviates = Matrix<T,Dynamic,Dynamic>::Ones(n,this->K_);

#pragma omp parallel for
  for (uint32_t k=0; k<this->K_; ++k)
  {
    for (uint32_t i=0; i<this->N_; ++i)
      if(this->z_(i) == k)
      {
        T deviate = dist(this->ps_.col(k), this->spx_->col(i));
//        T deviate = (this->ps_.col(k) - this->spx_->col(i)).norm();
        for (uint32_t j=0; j<n; ++j)
          if(closer(deviate, deviates(j,k)))
          {
            for(uint32_t l=n-1; l>j; --l)
            {
              deviates(l,k) = deviates(l-1,k);
              inds(l,k) = inds(l-1,k);
            }
            deviates(j,k) = deviate;
            inds(j,k) = i;
//            cout<<"after update "<<logLike<<endl;
//            Matrix<T,Dynamic,Dynamic> out(n,this->K_*2);
//            out<<logLikes.cast<T>(),inds.cast<T>();
//            cout<<out<<endl;
            break;
          }
      }
  }
  cout<<"::mostLikelyInds: deviates"<<endl;
  cout<<deviates<<endl;
  cout<<"::mostLikelyInds: inds"<<endl;
  cout<<inds<<endl;
  return inds;
};

template<class T>
T KMeans<T>::avgIntraClusterDeviation()
{
  Matrix<T,Dynamic,1> deviates(this->K_);
  deviates.setZero(this->K_);
#pragma omp parallel for
  for (uint32_t k=0; k<this->K_; ++k)
  {
    this->Ns_(k) = 0.0;
    for (uint32_t i=0; i<this->N_; ++i)
      if(this->z_(i) == k)
      {
        deviates(k) += dist(this->ps_.col(k), this->spx_->col(i));
//        deviates(k) += (this->ps_.col(k) - this->spx_->col(i)).norm();
        this->Ns_(k) ++;
      }
    if(this->Ns_(k) > 0.0) deviates(k) /= this->Ns_(k);
  }
  return deviates.sum()/static_cast<T>(this->K_);
}


/*
 * convlayer.cpp
 *
 * Feed-forward:
 *    z(l) = a(l-1) * w(l) + b(l) // where * is conv operation and b(l) is a scalar
 *    a(l) = activation(z(l))
 *
 * Back propagation:
 *    gradient(C, a(l-1)) = full_conv(delta(l), rotate180(w(l)))
 *
 *    delta(l) = elem_prod(gradient(C, a(l)), activation_derivative(z(l)))
 *    dC/dw(l) = conv(a(l), delta(l + 1))
 *    dC/db(l) = sum_elem(delta(l + 1))
 */
#include <boost/assert.hpp>

#include "core/utils.h"
#include "core/random.h"
#include "core/functions.h"
#include "contlayer.h"
#include "convlayer.h"

using namespace std;
using namespace boost;
using namespace yann;

namespace yann {

////////////////////////////////////////////////////////////////////////////////////////////////
//
// ConvolutionalLayerContext implementation
//
class ConvolutionalLayer_Context :
    public Layer::Context
{
  typedef Layer::Context Base;

  friend class ConvolutionalLayer;

public:
  ConvolutionalLayer_Context(
      const MatrixSize & output_size,
      const MatrixSize & batch_size) :
    Base(output_size, batch_size)
  {
    _zz.resizeLike(get_output());
  }
  ConvolutionalLayer_Context(const RefVectorBatch & output) :
    Base(output)
  {
    _zz.resizeLike(get_output());
  }

protected:
  VectorBatch _zz;
}; // class ConvolutionalLayer_Context

////////////////////////////////////////////////////////////////////////////////////////////////
//
// ConvolutionalLayer_TrainingContext implementation
//
class ConvolutionalLayer_TrainingContext :
  public ConvolutionalLayer_Context
{
  typedef ConvolutionalLayer_Context Base;

  friend class ConvolutionalLayer;

public:
  ConvolutionalLayer_TrainingContext(
      const MatrixSize & output_size,
      const MatrixSize & batch_size,
      const MatrixSize & filter_size,
      const unique_ptr<Layer::Updater> & updater) :
    Base(output_size, batch_size),
    _ww_rotated(filter_size, filter_size),
    _delta_ww(filter_size, filter_size),
    _delta_bb(0),
    _ww_updater(updater->copy()),
    _bb_updater(updater->copy())
  {
    YANN_CHECK_GT(filter_size, 0);
    YANN_CHECK_GT(output_size, 0);
    YANN_CHECK_GT(batch_size, 0);

    _delta.resizeLike(_zz);
    _sigma_derivative_zz.resizeLike(_zz);

    _ww_updater->init(filter_size, filter_size);
    _bb_updater->init(1, 1);
  }
  ConvolutionalLayer_TrainingContext(
      const RefVectorBatch & output,
      const MatrixSize & filter_size,
      const unique_ptr<Layer::Updater> & updater) :
    Base(output),
    _ww_rotated(filter_size, filter_size),
    _delta_ww(filter_size, filter_size),
    _delta_bb(0),
    _ww_updater(updater->copy()),
    _bb_updater(updater->copy())
  {
    YANN_CHECK_GT(filter_size, 0);
    YANN_CHECK_GT(yann::get_batch_size(output), 0);
    YANN_CHECK_GT(yann::get_batch_item_size(output), 0);

    _delta.resizeLike(_zz);
    _sigma_derivative_zz.resizeLike(_zz);

    _ww_updater->init(filter_size, filter_size);
    _bb_updater->init(1, 1);
  }

  // Layer::Context  overwrites
  virtual void start_epoch()
  {
    YANN_SLOW_CHECK(_ww_updater);
    YANN_SLOW_CHECK(_bb_updater);

    Base::start_epoch();

    _ww_updater->start_epoch();
    _bb_updater->start_epoch();
  }

  virtual void reset_state()
  {
    YANN_SLOW_CHECK(_ww_updater);
    YANN_SLOW_CHECK(_bb_updater);

    Base::reset_state();

    _delta_ww.setZero();
    _delta_bb = 0;

    _ww_updater->reset();
    _bb_updater->reset();
  }

private:
  Matrix _ww_rotated;
  Matrix _delta_ww;
  Value  _delta_bb;

  VectorBatch _delta;
  VectorBatch _sigma_derivative_zz;

  unique_ptr<Layer::Updater> _ww_updater;
  unique_ptr<Layer::Updater> _bb_updater;
}; // class ConvolutionalLayer_TrainingContext

}; // namespace yann

////////////////////////////////////////////////////////////////////////////////////////////////
//
// yann::ConvolutionalLayer implementation
//
MatrixSize yann::ConvolutionalLayer::get_input_size(
    const MatrixSize & input_rows,
    const MatrixSize & input_cols)
{
  return input_rows * input_cols;
}

MatrixSize yann::ConvolutionalLayer::get_conv_output_rows(
    const MatrixSize & input_rows,
    const MatrixSize & filter_rows)
{
  return (input_rows - filter_rows + 1);
}
MatrixSize yann::ConvolutionalLayer::get_conv_output_cols(
    const MatrixSize & input_cols,
    const MatrixSize & filter_cols)
{
  return (input_cols - filter_cols + 1);
}
MatrixSize yann::ConvolutionalLayer::get_conv_output_size(
    const MatrixSize & input_rows,
    const MatrixSize & input_cols,
    const MatrixSize & filter_rows,
    const MatrixSize & filter_cols)
{
  return get_conv_output_rows(input_rows, filter_rows)
      * get_conv_output_cols(input_cols, filter_cols);
}
MatrixSize yann::ConvolutionalLayer::get_conv_output_size(
    const MatrixSize & input_rows,
    const MatrixSize & input_cols,
    const MatrixSize & filter_size)
{
  return get_conv_output_size(input_rows, input_cols, filter_size, filter_size);
}

MatrixSize yann::ConvolutionalLayer::get_full_conv_output_rows(
    const MatrixSize & input_rows,
    const MatrixSize & filter_rows)
{
  return (input_rows + filter_rows - 1);
}
MatrixSize yann::ConvolutionalLayer::get_full_conv_output_cols(
    const MatrixSize & input_cols,
    const MatrixSize & filter_cols)
{
  return (input_cols + filter_cols - 1);
}
MatrixSize yann::ConvolutionalLayer::get_full_conv_output_size(
    const MatrixSize & input_rows,
    const MatrixSize & input_cols,
    const MatrixSize & filter_rows,
    const MatrixSize & filter_cols)
{
  return get_full_conv_output_rows(input_rows, filter_rows)
      * get_full_conv_output_cols(input_cols, filter_cols);
}
MatrixSize yann::ConvolutionalLayer::get_full_conv_output_size(
    const MatrixSize & input_rows,
    const MatrixSize & input_cols,
    const MatrixSize & filter_size)
{
  return get_full_conv_output_size(
      input_rows, input_cols,
      filter_size, filter_size);
}

void yann::ConvolutionalLayer::plus_conv(
    const RefConstMatrix & input,
    const RefConstMatrix & filter,
    RefMatrix output,
    bool clear_output)
{
  YANN_SLOW_CHECK_GT(filter.rows(), 0);
  YANN_SLOW_CHECK_GT(filter.cols(), 0);
  YANN_SLOW_CHECK_LE(filter.rows(), input.rows());
  YANN_SLOW_CHECK_LE(filter.cols(), input.cols());

  const auto input_rows = input.rows();
  const auto filter_rows = filter.rows();
  const auto filter_cols = filter.cols();
  const auto output_rows = get_conv_output_rows(input.rows(), filter_rows);
  const auto output_cols = get_conv_output_cols(input.cols(), filter_cols);

  if(clear_output) {
    output.setZero();
  }

  YANN_SLOW_CHECK_EQ(output.rows(), output_rows);
  YANN_SLOW_CHECK_EQ(output.cols(), output_cols);
  YANN_SLOW_CHECK_LE(filter_rows, input_rows);
  YANN_SLOW_CHECK_LE(output_rows, input_rows);

  auto apply_to_row = [&](const auto & ii, const auto & filter_start_row, const auto & filter_max_rows) mutable {
    const auto in_row = input.row(ii);
    // then through the filter by rows, account for "cutoff" at bottom rows
    for(MatrixSize kk = filter_start_row, ll = ii - filter_start_row; kk < filter_max_rows; ++kk, --ll) {
      const auto filter_row = filter.row(kk);
      auto out_row = output.row(ll);
      for(MatrixSize jj = 0; jj < output_cols; ++jj) {
          const auto in_block = in_row.segment(jj, filter_cols);
          out_row(jj) += (in_block.array() * filter_row.array()).sum();
      }
    }
  };

  // iterate through the input by rows
  MatrixSize filter_start_row = 0, filter_max_rows = 0;
  for(MatrixSize ii = 0; ii < input_rows; ++ii) {
    if(ii < filter_rows) {
      // filter at the top of the input; bottom filter rows
      // don't contribute to the output
      ++filter_max_rows;
    }
    if(ii >= output_rows) {
      // filter at the bottom of the input; top filter rows
      // don't contribute to the output
      ++filter_start_row;
    }

    // and now through each col and then filter row
    apply_to_row(ii, filter_start_row, filter_max_rows);
  }
}

void yann::ConvolutionalLayer::plus_conv(
    const RefConstVectorBatch & input,
    const MatrixSize & input_rows,
    const MatrixSize & input_cols,
    const RefConstMatrix & filter,
    RefVectorBatch output,
    bool clear_output)
{
  YANN_SLOW_CHECK_GT(filter.rows(), 0);
  YANN_SLOW_CHECK_GT(filter.cols(), 0);
  YANN_SLOW_CHECK_LE(filter.rows(), input_rows);
  YANN_SLOW_CHECK_LE(filter.cols(), input_cols);
  YANN_SLOW_CHECK_GT(get_batch_size(input), 0);
  YANN_SLOW_CHECK_EQ(get_batch_item_size(input), input_rows * input_cols);
  YANN_SLOW_CHECK_EQ(get_batch_size(output), get_batch_size(input));

  const auto filter_rows = filter.rows();
  const auto filter_cols = filter.cols();
  const auto output_rows = get_conv_output_rows(input_rows, filter_rows);
  const auto output_cols = get_conv_output_cols(input_cols, filter_cols);
  const auto batch_size = get_batch_size(input);
  YANN_SLOW_CHECK_EQ(get_batch_item_size(output), output_rows * output_cols);

  if(clear_output) {
    output.setZero();
  }

  // ATTENTION: this code operates on raw Matrix.data() and might be broken
  // if Matrix.data() layout changes
  for(MatrixSize ii = 0; ii < batch_size; ++ii) {
    const auto in_batch = get_batch(input, ii);
    auto out_batch = get_batch(output, ii);
    MapConstMatrix in(in_batch.data(), input_rows, input_cols);
    MapMatrix out(out_batch.data(), output_rows, output_cols);
    plus_conv(in, filter, out, false); // we already cleared output if needed
  }
}

void yann::ConvolutionalLayer::full_conv(
    const RefConstMatrix & input,
    const RefConstMatrix & filter,
    RefMatrix output)
{
  YANN_SLOW_CHECK_GT(filter.rows(), 0);
  YANN_SLOW_CHECK_GT(filter.cols(), 0);
  YANN_SLOW_CHECK_LE(filter.rows(), input.rows());
  YANN_SLOW_CHECK_LE(filter.cols(), input.cols());

  const auto input_rows = input.rows();
  const auto input_cols = input.cols();
  const auto filter_rows = filter.rows();
  const auto filter_cols = filter.cols();
  const auto output_cols = get_full_conv_output_cols(input_cols, filter_cols);

  YANN_SLOW_CHECK_EQ(output.rows(), get_full_conv_output_rows(input_rows, filter_rows));
  YANN_SLOW_CHECK_EQ(output.cols(), output_cols);

  // clear output
  output.setZero();

  MatrixSize input_row, output_col, output_last_row;
  auto apply_filter = [&](const auto & in_block, const auto & filter_block) mutable {
    // iterate through the filter by rows, account for "cutoff" at bottom rows
    for(MatrixSize kk = 0; kk < filter_rows; ++kk) {
      output(output_last_row - kk, output_col) += (in_block.array() * filter_block.row(kk).array()).sum();
    }
  };

  // iterate through the input by rows
  for(input_row = 0, output_last_row = filter_rows - 1;
      input_row < input_rows;
      ++input_row, ++output_last_row)
  {
    auto in_row = input.row(input_row);

    // iterate through output columns: we have 3 stages:
    // - filter on the left side of input with partial overlap
    // - filter over input with full overlap
    // - filter on the right side of input with partial overlap

    // on the left side from input
    MatrixSize in_start = 0, filter_start = filter_cols - 1, col_size = 1;
    for(output_col = 0; col_size < filter_cols; ++output_col, ++col_size, --filter_start) {
      // in_start == 0
      // col_size == output_col + 1
      // filter_start == filter_cols - col_size;
      auto in_block = in_row.segment(0, col_size);
      apply_filter(in_block, filter.rightCols(col_size));
    }

    // full overlap by columns
    for(; output_col < input_cols; ++output_col, ++in_start) {
      // in_start == output_col - filter_cols + 1;
      // col_size == filter_cols
      // filter_start == 0
      auto in_block = in_row.segment(in_start, col_size);
      // using here just filter instead of filter.rightCols() affects perf in a bad way
      apply_filter(in_block, filter.leftCols(col_size));
    }

    // on the right side from input
    for(--col_size; output_col < output_cols; ++output_col, ++in_start, --col_size) {
      // in_start == output_col - filter_cols + 1;
      // col_size == input_cols - in_start;
      // filter_start == 0
      auto in_block = in_row.segment(in_start, col_size);
      apply_filter(in_block, filter.leftCols(col_size));
    }
  }
}

void yann::ConvolutionalLayer::full_conv(
    const RefConstVectorBatch & input,
    const MatrixSize & input_rows,
    const MatrixSize & input_cols,
    const RefConstMatrix & filter,
    RefVectorBatch output)
{
  YANN_SLOW_CHECK_GT(filter.rows(), 0);
  YANN_SLOW_CHECK_GT(filter.cols(), 0);
  YANN_SLOW_CHECK_LE(filter.rows(), input_rows);
  YANN_SLOW_CHECK_LE(filter.cols(), input_cols);
  YANN_SLOW_CHECK_GT(get_batch_size(input), 0);
  YANN_SLOW_CHECK_EQ(get_batch_item_size(input), input_rows * input_cols);
  YANN_SLOW_CHECK_EQ(get_batch_size(output), get_batch_size(input));

  const auto filter_rows = filter.rows();
  const auto filter_cols = filter.cols();
  const auto output_rows = get_full_conv_output_rows(input_rows, filter_rows);
  const auto output_cols = get_full_conv_output_cols(input_cols, filter_cols);
  const auto batch_size = get_batch_size(input);
  YANN_SLOW_CHECK_EQ(get_batch_item_size(output), output_rows * output_cols);

  // ATTENTION: this code operates on raw Matrix.data() and might be broken
  // if Matrix.data() layout changes
  for(MatrixSize ii = 0; ii < batch_size; ++ii) {
    auto in_batch = get_batch(input, ii);
    auto out_batch = get_batch(output, ii);
    MapConstMatrix in(in_batch.data(), input_rows, input_cols);
    MapMatrix out(out_batch.data(), output_rows, output_cols);
    full_conv(in, filter, out);
  }
}

void yann::ConvolutionalLayer::rotate180(const Matrix & input, Matrix & output)
{
  output.resize(input.rows(), input.cols());
  for (MatrixSize ii = 0, size1 = input.rows(), size2 = input.cols(); ii < size1;
      ++ii) {
    for (MatrixSize jj = 0; jj < size2; ++jj) {
      output(ii, jj) = input(size1 - ii - 1, size2 - jj - 1);
    }
  }
}

// Convolutional layer: broadcast from input to all the conv layers
unique_ptr<BroadcastLayer> yann::ConvolutionalLayer::create_conv_bcast_layer(
    const size_t & output_frames_num,
    const MatrixSize & input_rows,
    const MatrixSize & input_cols,
    const MatrixSize & filter_size,
    const unique_ptr<ActivationFunction> & activation_function)
{
  YANN_CHECK_GT(output_frames_num, 0);

  auto bcast_layer = make_unique<BroadcastLayer>();
  YANN_CHECK(bcast_layer);
  for(auto ii = output_frames_num; ii > 0; --ii) {
    auto conv_layer = make_unique<ConvolutionalLayer>(
        input_rows,
        input_cols,
        filter_size);
    YANN_CHECK(conv_layer);

    if(activation_function) {
      conv_layer->set_activation_function(activation_function);
    }
    bcast_layer->append_layer(std::move(conv_layer));
  }

  return bcast_layer;
}

yann::ConvolutionalLayer::ConvolutionalLayer(
    const MatrixSize & input_rows,
    const MatrixSize & input_cols,
    const MatrixSize & filter_size) :
  _input_rows(input_rows),
  _input_cols(input_cols),
  _filter_size(filter_size),
  _ww(filter_size, filter_size),
  _bb(0),
  _activation_function(new SigmoidFunction())
{
  YANN_CHECK_GE(input_rows, filter_size);
  YANN_CHECK_GE(input_cols, filter_size);
  YANN_CHECK_GT(filter_size, 0);
}

yann::ConvolutionalLayer::~ConvolutionalLayer()
{
}

void yann::ConvolutionalLayer::set_activation_function(
    const unique_ptr<ActivationFunction> & activation_function)
{
  YANN_CHECK(activation_function);
  _activation_function = activation_function->copy();
}

void yann::ConvolutionalLayer::set_values(const Matrix & ww, const Value & bb)
{
  YANN_CHECK(is_same_size(ww, _ww));
  _ww = ww;
  _bb = bb;
}

// Layer overwrites
bool yann::ConvolutionalLayer::is_valid() const
{
  if(!Base::is_valid()) {
    return false;
  }
  if(!_activation_function) {
    return false;
  }
  return true;
}

std::string yann::ConvolutionalLayer::get_name() const
{
  return "ConvolutionalLayer";
}

string yann::ConvolutionalLayer::get_info() const
{
  YANN_CHECK(is_valid());

  ostringstream oss;
  oss << Base::get_info()
      << " activation: " << _activation_function->get_info()
      << ", input rows: " << _input_rows
      << ", input cols: " << _input_cols
      << ", filter: " << _filter_size
      << ", output rows: " << get_output_rows()
      << ", output cols: " << get_output_cols()
  ;
  return oss.str();
}

bool yann::ConvolutionalLayer::is_equal(const Layer& other, double tolerance) const
{
  if(!Base::is_equal(other, tolerance)) {
    return false;
  }
  auto * the_other = dynamic_cast<const ConvolutionalLayer*>(&other);
  if(the_other == nullptr) {
    return false;
  }
  // TOOD: add deep compare
  if(_activation_function->get_info() != the_other->_activation_function->get_info()) {
    return false;
  }
  if(!_ww.isApprox(the_other->_ww, tolerance)) {
    return false;
  }
  if(fabs(_bb - the_other->_bb) >= tolerance) {
    return false;
  }
  return true;
}

MatrixSize yann::ConvolutionalLayer::get_input_size() const
{
  return get_input_size(_input_rows, _input_cols);
}

MatrixSize yann::ConvolutionalLayer::get_output_size() const
{
  return get_conv_output_size(_input_rows, _input_cols, _filter_size);
}

MatrixSize yann::ConvolutionalLayer::get_output_rows() const
{
  return get_conv_output_rows(_input_rows, _filter_size);
}

MatrixSize yann::ConvolutionalLayer::get_output_cols() const
{
  return get_conv_output_cols(_input_cols, _filter_size);
}

unique_ptr<Layer::Context> yann::ConvolutionalLayer::create_context(
    const MatrixSize & batch_size) const
{
  return make_unique<ConvolutionalLayer_Context>(get_output_size(), batch_size);
}
unique_ptr<Layer::Context> yann::ConvolutionalLayer::create_context(
    const RefVectorBatch & output) const
{
  return make_unique<ConvolutionalLayer_Context>(output);
}
unique_ptr<Layer::Context> yann::ConvolutionalLayer::create_training_context(
    const MatrixSize & batch_size,
    const std::unique_ptr<Layer::Updater> & updater) const
{
  YANN_CHECK(updater);
  YANN_CHECK(is_valid());
  return make_unique<ConvolutionalLayer_TrainingContext>(
      get_output_size(), batch_size, _filter_size, updater);
}
unique_ptr<Layer::Context> yann::ConvolutionalLayer::create_training_context(
    const RefVectorBatch & output,
    const std::unique_ptr<Layer::Updater> & updater) const
{
  YANN_CHECK(updater);
  YANN_CHECK(is_valid());
  return make_unique<ConvolutionalLayer_TrainingContext>(
      output, _filter_size, updater);
}

void yann::ConvolutionalLayer::feedforward(
    const RefConstVectorBatch & input,
    Context * context,
    enum OperationMode mode) const
{
  auto ctx = dynamic_cast<ConvolutionalLayer_Context *>(context);
  YANN_CHECK(ctx);
  YANN_CHECK(is_valid());

  // z(l) = conv(a(l-1))*w(l) + b(l)
  plus_conv(input, _input_rows, _input_cols, _ww, ctx->_zz);
  ctx->_zz.array() += _bb;

  // a(l) = activation(z(l))
  YANN_CHECK(is_same_size(ctx->_zz, ctx->get_output()));
  _activation_function->f(ctx->_zz, ctx->get_output(), mode);
}

void yann::ConvolutionalLayer::feedforward(
    const RefConstSparseVectorBatch & input,
    Context * context,
    enum OperationMode mode) const
{
  throw runtime_error("ConvolutionalLayer::feedforward() is not implemented for sparse vectors");
}

void yann::ConvolutionalLayer::backprop(
    const RefConstVectorBatch & gradient_output,
    const RefConstVectorBatch & input,
    optional<RefVectorBatch> gradient_input,
    Context * context) const
{
  auto ctx = dynamic_cast<ConvolutionalLayer_TrainingContext*>(context);
  YANN_CHECK(ctx);
  YANN_CHECK(is_valid());
  YANN_CHECK_GT(get_batch_size(gradient_output), 0);
  YANN_CHECK_EQ(get_batch_item_size(gradient_output), get_output_size());
  YANN_CHECK_EQ(get_batch_size(input), get_batch_size(gradient_output));
  YANN_CHECK_EQ(get_batch_item_size(input), get_input_size());
  YANN_CHECK(!gradient_input || is_same_size(input, *gradient_input));

  // just to make it easier to read
  const auto & zz = ctx->_zz;
  auto & sigma_derivative_zz = ctx->_sigma_derivative_zz;
  auto & delta = ctx->_delta;
  auto & delta_ww = ctx->_delta_ww;
  auto & delta_bb = ctx->_delta_bb;
  const auto batch_size = get_batch_size(input);

  // delta(l) = elem_prod(gradient(C, a(l)), activation_derivative(z(l)))
  YANN_CHECK(is_same_size(zz, sigma_derivative_zz));
  YANN_CHECK(is_same_size(zz, gradient_output));
  _activation_function->derivative(zz, sigma_derivative_zz);
  delta.array() = gradient_output.array() * sigma_derivative_zz.array();

  // update deltas
  // dC/dw(l) = conv(a(l), delta(l + 1))
  // dC/db(l) = sum_elem(delta(l + 1))
  YANN_CHECK_EQ(batch_size, get_batch_size(delta));
  for (MatrixSize ii = 0; ii < batch_size; ++ii) {
    auto input_batch = get_batch(input, ii);
    auto delta_batch = get_batch(delta, ii);
    MapConstMatrix input_row(input_batch.data(), _input_rows, _input_cols);
    MapConstMatrix delta_row(delta_batch.data(), get_output_rows(), get_output_cols());
    plus_conv(input_row, delta_row, delta_ww, false); // delta_ww += conv()
  }
  delta_bb += delta.array().sum();

  // we don't need to calculate the first gradient(C, a(l)) for the actual inputs
  if(gradient_input) {
    // gradient(C, a(l-1)) = full_conv(delta(l), rotate180(w(l)))
    rotate180(_ww, ctx->_ww_rotated); // TODO: we can cache the rotation if _ww doesn't change
    full_conv(delta, get_output_rows(), get_output_cols(), ctx->_ww_rotated, *gradient_input);
  }
}

void yann::ConvolutionalLayer::backprop(
    const RefConstVectorBatch & gradient_output,
    const RefConstSparseVectorBatch & input,
    optional<RefVectorBatch> gradient_input,
    Context * context) const
{
  throw runtime_error("ConvolutionalLayer::backprop() is not implemented for sparse vectors");
}

void yann::ConvolutionalLayer::init(enum InitMode mode, optional<InitContext> init_context)
{
  switch (mode) {
  case InitMode_Zeros:
    _ww.setZero();
    _bb = 0;;
    break;
  case InitMode_Random:
    {
      unique_ptr<RandomGenerator> gen01 = RandomGenerator::normal_distribution(0, 1,
          init_context ? optional<Value>(init_context->seed()) : boost::none);
      gen01->generate(_ww);
      gen01->generate(_bb);
    }
    break;
  }
}

void yann::ConvolutionalLayer::update(Context * context, const size_t & tests_num)
{
  auto ctx = dynamic_cast<ConvolutionalLayer_TrainingContext *>(context);
  YANN_CHECK(ctx);
  YANN_CHECK(ctx->_ww_updater);
  YANN_CHECK(ctx->_bb_updater);
  YANN_CHECK(is_same_size(_ww, ctx->_delta_ww));

  ctx->_ww_updater->update(ctx->_delta_ww, tests_num, _ww);
  ctx->_bb_updater->update(ctx->_delta_bb, tests_num, _bb);
}

// the format is (w:<weights>,b:<bias>)
void yann::ConvolutionalLayer::read(std::istream & is)
{
  Base::read(is);

  read_char(is, '(');
  read_object(is, "w", _ww);
  read_char(is, ',');
  read_object(is, "b", _bb);
  read_char(is, ')');
}

// the format is (w:<weights>,b:<bias>)
void yann::ConvolutionalLayer::write(std::ostream & os) const
{
  Base::write(os);

  os << "(";
  write_object(os, "w", _ww);
  os << ",";
  write_object(os, "b", _bb);
  os << ")";
}


//
// Created by wang yu on 2018/6/20.
//
#include "mview.h"
#include <iostream>

#include <Eigen/Eigen>
#include <Eigen/Dense>
#include <opencv2/core.hpp>
#include <opencv/cv.hpp>
#include <opencv2/core/eigen.hpp>

static int debug_count = 0;

cv::Mat convertRgbToOpenCV(const RgbImage& rgb) {
	GrayImage r = rgb.unaryExpr([](Eigen::Vector3f rgb) { return rgb[0]; });
	GrayImage g = rgb.unaryExpr([](Eigen::Vector3f rgb) { return rgb[1]; });
	GrayImage b = rgb.unaryExpr([](Eigen::Vector3f rgb) { return rgb[2]; });

	cv::Mat r_mat;
    cv::eigen2cv(r,r_mat);
	cv::Mat g_mat;
    cv::eigen2cv(g,g_mat);
	cv::Mat b_mat;
    cv::eigen2cv(b,b_mat);

	cv::Mat arr[3] { r_mat, g_mat, b_mat };
	cv::Mat rgb_mat;
	cv::merge(arr, 3, rgb_mat);
	return rgb_mat;
}

RgbImage convertOpenCVToRgb(const cv::Mat imageMat){
    cv::Mat rgb[3];
    cv::split(imageMat, rgb);

	int width = imageMat.cols;
	int height = imageMat.rows;
    GrayImage r(height, width), g(height, width), b(height, width);

    cv::cv2eigen(rgb[0], r);
    cv::cv2eigen(rgb[1], g);
    cv::cv2eigen(rgb[2], b);

    RgbImage rgbImage (g.rows(), g.cols());

    for(int row=0;row<g.rows();row++){
        for(int col=0;col<g.cols();col++){
            rgbImage(row,col)<<r(row,col),g(row,col),b(row,col);
        }
    }

    return rgbImage;
}

void remap_rgb(cv::Mat& rgb_image, cv::Mat map1, cv::Mat map2) {
	cv::Mat rgbs[3];
	cv::split(rgb_image, rgbs);

	cv::Mat rgbs_out[3];
	for(int i = 0; i < 3; i++)
		cv::remap(rgbs[i], rgbs_out[i], map1, map2, cv::INTER_NEAREST, cv::BORDER_CONSTANT);

	cv::merge(rgbs_out, 3, rgb_image);
}

auto rectify(const Image& left, const Image& right) -> Rectified{
    //load data
    GrayImage left_gray_pixels=left.gray_pixels;
    RgbImage left_rgb_pixels=left.rgb_pixels;

    GrayImage right_gray_pixels=right.gray_pixels;
    RgbImage right_rgb_pixels=right.rgb_pixels;

	Eigen::Matrix4f left_to_right = right.extrinsics * left.extrinsics.inverse();
	for(int i = 0; i < 3; i++)
		assert(left_to_right(3, i) == 0);
	assert(left_to_right(3, 3) == 1.0);

    //rotation and transformation from left camera to right camera
    Eigen::Matrix3f R = left_to_right.block<3, 3>(0, 0);
    Eigen::Vector3f T = R.transpose() * left_to_right.block<3, 1>(0, 3);

    //convert eigen matrix to mat
    cv::Mat left_gray_mat, left_ground_mat;
    cv::eigen2cv(left_gray_pixels, left_gray_mat);
    cv::eigen2cv(left.ground_truth, left_ground_mat);
    cv::Mat left_rgb_mat = convertRgbToOpenCV(left_rgb_pixels);

    cv::Mat right_gray_mat, right_ground_mat;
    cv::eigen2cv(right_gray_pixels, right_gray_mat);
    cv::eigen2cv(right.ground_truth, right_ground_mat);
    cv::Mat right_rgb_mat = convertRgbToOpenCV(right_rgb_pixels);
    cv::Mat left_intrinsics_mat, right_intrinsics_mat, R_mat, T_mat;

    cv::eigen2cv(left.intrinsics,left_intrinsics_mat);
    cv::eigen2cv(right.intrinsics,right_intrinsics_mat);

    cv::eigen2cv(R,R_mat);
    cv::eigen2cv(T,T_mat);

	R_mat.assignTo(R_mat, CV_64F);
	T_mat.assignTo(T_mat, CV_64F);

    cv::Size imageSize = left_gray_mat.size();
	cv::Size targetSize = imageSize; // (left_gray_mat.cols/4., left_gray_mat.rows/4.);
    cv::Mat R1,R2,P1,P2,Q;
    cv::stereoRectify(left_intrinsics_mat,{},right_intrinsics_mat,{},imageSize,R_mat,T_mat,R1,R2,P1,P2,Q,
			cv::CALIB_ZERO_DISPARITY,
		   	1,targetSize,0,0);

    cv::Mat map1, map2, left_gray_out, left_ground_out;
    cv::initUndistortRectifyMap(left_intrinsics_mat,{},R1,P1,targetSize,CV_32FC1,map1,map2);
    cv::remap(left_gray_mat, left_gray_out, map1, map2, cv::INTER_NEAREST, cv::BORDER_CONSTANT);
    cv::remap(left_ground_mat, left_ground_out, map1, map2, cv::INTER_NEAREST, cv::BORDER_CONSTANT);
	left_gray_mat = left_gray_out;
	remap_rgb(left_rgb_mat, map1, map2);

    cv::Mat map3,map4, right_gray_out, right_ground_out;
    cv::initUndistortRectifyMap(right_intrinsics_mat,{},R2,P2,targetSize,CV_32FC1,map3,map4);
    cv::remap(right_gray_mat, right_gray_out, map3, map4, cv::INTER_NEAREST, cv::BORDER_CONSTANT);
    cv::remap(right_ground_mat, right_ground_out, map3, map3, cv::INTER_NEAREST, cv::BORDER_CONSTANT);
	right_gray_mat = right_gray_out;
	remap_rgb(right_rgb_mat, map3, map4);

	std::cout << "Debugging rectification results\n";
	std::cout << "R1 " << R1 << "\nR2 " << R2 << "\nP1 " << P1 << "\nP2" << P2 << "\nQ " << Q << std::endl;

#if !MVIEW_NDEBUG
	std::stringstream debug_name;
	std::string debug_left = ((debug_name << "debug.rectified" << debug_count << ".left.png"), debug_name.str()); debug_name.str("");
	std::string debug_right = ((debug_name << "debug.rectified" << debug_count << ".right.png"), debug_name.str()); debug_name.str("");
	std::string debug_depth = ((debug_name << "debug.rectified" << debug_count << ".depth.png"), debug_name.str()); debug_name.str("");
	debug_count++;

	cv::imwrite(debug_left, left_gray_mat*255.);
	cv::imwrite(debug_right, right_gray_mat*255.);
	cv::imwrite(debug_depth, left_ground_out*25.);
#endif

	Rectified rectified;

    //convert mat to eigen matrix
	rectified.pixel_left_gray = GrayImage(left_gray_mat.rows, left_gray_mat.cols);
	rectified.left_ground_truth = GrayImage(left_gray_mat.rows, left_gray_mat.cols);
    cv::cv2eigen(left_gray_mat, rectified.pixel_left_gray);
    cv::cv2eigen(left_ground_out, rectified.left_ground_truth);
	rectified.pixel_left_rgb = convertOpenCVToRgb(left_rgb_mat);

	rectified.pixel_right_gray = GrayImage (right_gray_mat.rows, right_gray_mat.cols);
	rectified.right_ground_truth = GrayImage(right_gray_mat.rows, right_gray_mat.cols);
    cv::cv2eigen(right_gray_mat, rectified.pixel_right_gray);
    cv::cv2eigen(right_ground_out, rectified.right_ground_truth);
	rectified.pixel_right_rgb = convertOpenCVToRgb(right_rgb_mat);

	rectified.extrinsics_left = left.extrinsics;
	rectified.extrinsics_right = right.extrinsics;

	rectified.R1 = R1;
	rectified.R2 = R2;
	rectified.P1 = P1;
	rectified.P2 = P2;
	rectified.Q = Q;

    return rectified;
}


// implements the solver class - Felipe Figueredo Rocha
#ifndef _solver_hpp
#define _solver_hpp

#include <iostream>
#include <iomanip>
#include <string>
#include <fstream>
#include <vector>
#include <stdlib.h>
#include <sstream>
#include <cmath>
#include "utils.hpp"
#include "linAlg.hpp"
#include "data.hpp"
#include "simulSettings.hpp"
#include "solvergp.h"
#include <armadillo>

//~ #define DEBUG if(false)
#define DEBUG if(true)

using namespace arma;
using namespace std;

class Solver{
	public:
	// see constructor for further explanation about the variabls
	Data *d;
	SimulSettings *s;
	mat *A; // stiffness matrix
	vec *b; // force vector
	vec *Sol; // force vector
	vector<intVec*> *Coupling; // Coupling vector, works similar to flagDirich
	Solver() {}
	void init();
	void prepareMatrices(); // alloc linear system matrices and vector
	void createCoupling();
	void linsolve(); // call the routine of solving the lin System
	void enforceDirichlet(SGPrealVec *vDirich, intVec *flagDirich, double pen); // modify matrices to enforce Dirichlet
	void assembly(SGPrealMat *A_L, SGPrealVec *b_L, int e); // assembly local in global matrices
	void assemblyCoupling(intVec *Cvec, intMat *Cmat_L, int e);
	void globalToLocal(SGPrealVec *XLL,SGPrealVec *Sol0E, SGPrealVec *Sol1E,SGPrealVec *ParamE, int e); // fills the e^{th} of the mesh in a auxiliary object
	void localToGlobal(SGPrealVec *ParamE,int e); // fills the e^{th} of the mesh in a auxiliary object
	void run();
};

void Solver:: init(){
	s = new SimulSettings;
	s->loadSettings();
	d = new Data;
	d->loadSGPmesh();
	d->loadSGPdirich(s->NSubsteps);
	d->loadSGPinifile();
	d->loadSGPparam();
	prepareMatrices();
	createCoupling();
}

void Solver:: prepareMatrices(){
	int n = d->Ndof*d->Nnodes;

 	A = new mat(n,n);
	b = new vec(n);
	Sol = new vec(n);

}

void Solver:: createCoupling(){
	int n = d->Ndof*d->Nnodes;
	int NodElt = d->eNNod->max();
	int MaxLRows = NodElt*d->Ndof;
	intMat *CouplingE = new intMat(MaxLRows,MaxLRows);
	ElemGroup elG;

	Coupling = new vector<intVec*>(s->NSubsteps);
	for(int i=0; i<s->NSubsteps; i++) Coupling->at(i) = new intVec(n);

	for(int k = 0; k< s->NSubsteps; k++){

		for(int e=0; e< d->Nelem; e++){
			int iGroup = (*d->eType)(e);

			elG = s->substeps->at(k).elGroup->at(iGroup);

			(*CouplingE)=0;

			if(elG.FlagElemLib<0){
				getLocalSymbolic(&elG.ElemLib, CouplingE->v, elG.CommonPar->v,&d->Ndof, &d->Ndim, &MaxLRows);
			}

			assemblyCoupling(Coupling->at(k),CouplingE,e);
		}

		for(int i=0; i<n; i++) (*Coupling->at(k))(i) = (*Coupling->at(k))(i) - 1;
	}
}

void Solver:: linsolve(){
	solve(*Sol,*A,*b);
}

void Solver :: enforceDirichlet(SGPrealVec *vDirich, intVec *flagDirich, double pen = 99999.9){

	for(int i = 0 ; i<b->size(); i++){
		if((*flagDirich)(i)==-1){
			(*b)(i) = pen*(*vDirich)(i);
			(*A)(i,i) = pen*1.0;
			for(int j = 0 ; j<b->size(); j++){ if(j!=i) (*A)(i,j) = 0.0;}
		}
	}
}
//~
void Solver :: assembly(SGPrealMat *A_L, SGPrealVec *b_L, int e){

	int ip,jp,kip,kjp;

	int iShiftElem = (*d->eNNod_acc)(e);

	for(int i=0; i<(*d->eNNod)(e); i++){ for(int ii=0; ii<d->Ndof; ii++){
		ip = i*d->Ndof + ii;
		kip = (*d->Elem)(iShiftElem + i)*d->Ndof + ii;
		(*b)(kip) += (*b_L)(ip);
		for(int j=0; j<(*d->eNNod)(e); j++){ for(int jj=0; jj<d->Ndof; jj++){
			jp = j*d->Ndof + jj;
			kjp = (*d->Elem)(iShiftElem + j)*d->Ndof + jj;
			(*A)(kip,kjp) += (*A_L)(ip,jp);
		}
		}
	}
	}
}

void Solver :: assemblyCoupling(intVec *Cvec, intMat *Cmat_L, int e){

	int ip,kip;
	int iShiftElem = (*d->eNNod_acc)(e);

	for(int i=0; i<(*d->eNNod)(e); i++){ for(int ii=0; ii<d->Ndof; ii++){
		ip = i*d->Ndof + ii;
		kip = (*d->Elem)(iShiftElem + i)*d->Ndof + ii;
		if((*Cmat_L)(ip,ip) == 1){
			(*Cvec)(kip) = 1;
		}
	}
	}
}

void Solver :: globalToLocal(SGPrealVec *XLL,SGPrealVec *Sol0E, SGPrealVec *Sol1E, SGPrealVec *ParamE, int e){
		int ipDim,kpDim,ipDof,kpDof;
		int iShiftElem = (*d->eNNod_acc)(e);
		int iShiftMat = (*d->eMat)(e);
		int iShiftParam = (*d->eParamSize_acc)(iShiftMat);

		(*XLL) = 0.0;
		(*Sol0E) = 0.0;
		(*Sol1E) = 0.0;
		(*ParamE) = 0.0;

		for(int i=0;i<(*d->eNNod)(e);i++){
			ipDim = i*d->Ndim;
			kpDim = (*d->Elem)(iShiftElem + i)*d->Ndim;
			ipDof = i*d->Ndof;
			kpDof = (*d->Elem)(iShiftElem + i)*d->Ndof;

			for(int j=0 ; j<d->Ndim; j++) (*XLL)(ipDim + j) = (*d->X)(kpDim +j);

			for(int j=0 ; j<d->Ndof; j++){
				(*Sol0E)(ipDof + j) = (*d->Sol0)(kpDof + j);
				(*Sol1E)(ipDof + j) = (*d->Sol1)(kpDof + j);
			}
		}

		for(int k=0; k<(*d->eParamSize)(iShiftMat) ; k++) (*ParamE)(k) = (*d->Param)(iShiftParam + k);

}

void Solver :: localToGlobal(SGPrealVec *ParamE, int e){
		int iShiftMat = (*d->eMat)(e);
		int iShiftParam = (*d->eParamSize_acc)(iShiftMat);

		for(int k=0; k<(*d->eParamSize)(iShiftMat) ; k++) (*d->Param)(iShiftParam + k) = (*ParamE)(k);
}

void Solver :: run(){
	int timeStep, iGroup;
	int id_Elem_Family, MaxLRows, iDofT, NodElt;
	intVec *JParam;
	double DTm, Time;
	SGPrealVec *BE, *XLL, *Sol0E, *Sol1E, *CommonParE, *ParamE;
	SGPrealMat *AE;
	ElemGroup elG;
	double error, tol, emax = 99999.0;
	int itNL, maxit;
	bool isNL;

	NodElt = d->eNNod->max();
	MaxLRows = NodElt*d->Ndof;

	AE = new SGPrealMat(MaxLRows,MaxLRows);
	BE = new SGPrealVec(MaxLRows);
	XLL = new SGPrealVec(NodElt*d->Ndim);
	Sol0E = new SGPrealVec(MaxLRows);
	Sol1E = new SGPrealVec(MaxLRows);
	ParamE = new SGPrealVec(d->eParamSize->max());
	JParam = new intVec(3); // temporary

	Time = s->Tini;
	DTm = s->DelT;
	timeStep = 0;

	d->writeSGPdataout(timeStep,Time,DTm);

	while(Time < s->Tmax){
		DEBUG{cout << "=> Beginning Timestep = " << timeStep << ", time = " << Time << "  =====\n " << endl;}

		for(int k = 0; k< s->NSubsteps; k++){

			error = emax;
			itNL = 0;
			isNL = s->substeps->at(k).isNL;
			if(isNL){
				maxit = s->substeps->at(k).maxit;
				tol = s->substeps->at(k).tol;
			}
			else{
				maxit = 1; // force to stop with just one iteration
				tol = 0.0;
			}

			DEBUG{cout << "===> Beginning SubStep = " << k << " ===========" << endl;}

			DEBUG{if(isNL) cout << "===> Starting NonLinear Convergence, tol = " << tol << endl;}
			while(error>tol && itNL<maxit){

				DEBUG{if(isNL) cout << "===> Starting NonLinear Internal Loop, it = " << itNL << " =======" << endl;}

				A->fill(0.0);
				b->fill(0.0);

				//~ cout << "=====> Assemblying Global Matrices ================ " << endl;
				for(int e=0; e< d->Nelem; e++){

					globalToLocal(XLL,Sol0E,Sol1E,ParamE,e);

					iGroup = (*d->eType)(e);

					elG = s->substeps->at(k).elGroup->at(iGroup);

					(*AE)=0.0;
					(*BE)=0.0;

					if(elG.FlagElemLib<0){
						getLocalMatrix(&elG.ElemLib, AE->v, BE->v, &MaxLRows, XLL->v, &d->Ndim, &d->Ndof, &NodElt, Sol0E->v, Sol1E->v,
										elG.CommonPar->v, ParamE->v, JParam->v, &s->DelT, &DTm, &Time);
					}

					assembly(AE,BE,e);

					localToGlobal(ParamE,e);
				}

				DEBUG{cout << "=====> Solving Linear System ================ " << endl;}
				enforceDirichlet( d->vDirich->at(k), d->flagDirich->at(k) , 1.0);
				enforceDirichlet( d->Sol1, Coupling->at(k) , 1.0);

				linsolve();
				if(isNL){
					error = computeError(*Sol,*d->Sol1);
					DEBUG{cout << "====> Ending NonLinear Internal Loop , it= " << itNL << " , error= " << error << " =========="<< endl;}
				}

				itNL ++;
				*d->Sol1 = *Sol; // uses the operator overloading
			}
			DEBUG{if(isNL) cout << "====> Convergence Achieved with it= " << itNL-1 << " , error= " << error << " =========="<< endl;}

			DEBUG{cout << "===> Ending SubStep = " << k << " =============== \n " << endl;}
		}

		DEBUG{cout << "=> Ending Timestep = " << timeStep << ", time = " << Time << " ===== \n" << endl;}

		Time += s->DelT;
		timeStep++;

		*d->Sol0 = *d->Sol1; // uses the operator overloading

		DEBUG{cout << "=> Writing Results in File ====== \n\n" << endl;}
		d->writeSGPdataout(timeStep,Time,DTm);
		d->writeSGPparam();
	}
}

#endif


/*

Copyright (c) 2005-2020, University of Oxford.
All rights reserved.

University of Oxford means the Chancellor, Masters and Scholars of the
University of Oxford, having an administrative office at Wellington
Square, Oxford OX1 2JD, UK.

This file is part of Chaste.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
 * Redistributions of source code must retain the above copyright notice,
   this list of conditions and the following disclaimer.
 * Redistributions in binary form must reproduce the above copyright notice,
   this list of conditions and the following disclaimer in the documentation
   and/or other materials provided with the distribution.
 * Neither the name of the University of Oxford nor the names of its
   contributors may be used to endorse or promote products derived from this
   software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE
GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT
OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

*/

#ifndef VOLUMEDEPENDENTAVERAGEDSOURCEELLIPTICPDE_HPP_
#define VOLUMEDEPENDENTAVERAGEDSOURCEELLIPTICPDE_HPP_

#include "ChasteSerialization.hpp"
#include <boost/serialization/base_object.hpp>

#include "NodeBasedCellPopulation.hpp"
#include "AveragedSourceEllipticPde.hpp"
#include "TetrahedralMesh.hpp"
#include "AbstractLinearEllipticPde.hpp"

/**
 * An elliptic PDE to be solved numerically using the finite element method, for
 * coupling to a cell-based simulation.
 *
 * This class inherits from AveragedSourceEllipticPde and may only be used with
 * a NodeBasedCellPopulation, since it assumes that each cell is associated with
 * a Node object.
 *
 * The PDE takes the form
 *
 * Grad.(D*Grad(u)) + k*u*rho(x) = 0,
 *
 * where the scalars D and k are specified by the members mDiffusionCoefficient and
 * mSourceCoefficient, respectively. Their values must be set in the constructor.
 *
 * The function rho(x) denotes the local density of non-apoptotic cells. This
 * quantity is computed for each element of a 'coarse' finite element mesh that is
 * passed to the method SetupSourceTerms() and stored in the member mCellDensityOnCoarseElements.
 *
 * For a point x, rho(x) is a weighted sum of the non-apoptotic cells whose centres
 * lie in each finite element containing that point, scaled by the area of that element.
 * The weighting assigned to each cell is given by the square of the radius of the
 * associated node, which is accessed using the GetRadius() method.
 *
 * \todo Consider creating a VolumeDependentAveragedSourceParabolicPde class
 */
template<unsigned DIM>
class VolumeDependentAveragedSourceEllipticPde : public AveragedSourceEllipticPde<DIM>
{
    friend class TestCellBasedEllipticPdes;

private:

    /** Needed for serialization.*/
    friend class boost::serialization::access;
    /**
     * Serialize the PDE and its member variables.
     *
     * @param archive the archive
     * @param version the current version of this class
     */
    template<class Archive>
    void serialize(Archive & archive, const unsigned int version)
    {
       archive & boost::serialization::base_object<AveragedSourceEllipticPde<DIM> >(*this);
    }

    /** Static cast of the NodeBasedCellPopulation. */
    NodeBasedCellPopulation<DIM>* mpStaticCastCellPopulation;

public:

    /**
     * Constructor.
     *
     * @param rCellPopulation reference to the cell population
     * @param coefficient the coefficient of consumption of nutrient by cells (defaults to 0.0)
     */
    VolumeDependentAveragedSourceEllipticPde(AbstractCellPopulation<DIM>& rCellPopulation, double coefficient=0.0);

    /**
     * Set up the source terms.
     *
     * @param rCoarseMesh reference to the coarse mesh
     * @param pCellPdeElementMap optional pointer to the map from cells to coarse elements
     */
    void SetupSourceTerms(TetrahedralMesh<DIM,DIM>& rCoarseMesh, std::map<CellPtr, unsigned>* pCellPdeElementMap=nullptr);
};

#include "SerializationExportWrapper.hpp"
EXPORT_TEMPLATE_CLASS_SAME_DIMS(VolumeDependentAveragedSourceEllipticPde)

namespace boost
{
namespace serialization
{
/**
 * Serialize information required to construct a VolumeDependentAveragedSourceEllipticPde.
 */
template<class Archive, unsigned DIM>
inline void save_construct_data(
    Archive & ar, const VolumeDependentAveragedSourceEllipticPde<DIM>* t, const unsigned int file_version)
{
    // Save data required to construct instance
    const AbstractCellPopulation<DIM>* p_cell_population = &(t->rGetCellPopulation());
    ar & p_cell_population;
}

/**
 * De-serialize constructor parameters and initialise a VolumeDependentAveragedSourceEllipticPde.
 */
template<class Archive, unsigned DIM>
inline void load_construct_data(
    Archive & ar, VolumeDependentAveragedSourceEllipticPde<DIM>* t, const unsigned int file_version)
{
    // Retrieve data from archive required to construct new instance
    AbstractCellPopulation<DIM>* p_cell_population;
    ar >> p_cell_population;

    // Invoke inplace constructor to initialise instance
    ::new(t)VolumeDependentAveragedSourceEllipticPde<DIM>(*p_cell_population);
}
}
} // namespace ...

#endif /*VOLUMEDEPENDENTAVERAGEDSOURCEELLIPTICPDE_HPP_*/


/*
 * s2slayer.cpp
 *
 * Feed-forward:
 *    z(l) = a(l-1) * w + b // where * (vector, matrix) multiplication
 *    a(l) = activation(z(l))
 *
 * Back propagation:
 *    delta(l) = elem_prod(gradient(C, a(l)), activation_derivative(z(l)))
 *    dC/db(l) = delta(l)
 *    dC/dw(l) = a(l-1) * delta(l)
 *    gradient(C, a(l - 1)) = transp(w) * delta(l)
 */
#include <boost/assert.hpp>

#include "core/utils.h"
#include "core/random.h"
#include "core/functions.h"
#include "lstmlayer.h"
#include "s2slayer.h"

using namespace std;
using namespace boost;
using namespace yann;

namespace yann {

////////////////////////////////////////////////////////////////////////////////////////////////
//
// Seq2SeqLayer_Context implementation
//
class Seq2SeqLayer_Context :
    public Layer::Context
{
  typedef Layer::Context Base;

  friend class Seq2SeqLayer;

public:
  Seq2SeqLayer_Context(
      unique_ptr<Layer::Context> encoder_ctx,
      unique_ptr<Layer::Context> decoder_ctx) :
    Base(decoder_ctx->get_output()) // decoder output is our output
  {
    YANN_CHECK(encoder_ctx);
    YANN_CHECK(decoder_ctx);
    _encoder_ctx = std::move(encoder_ctx);
    _decoder_ctx = std::move(decoder_ctx);
  }

  inline bool is_valid() const
  {
    return _encoder_ctx && _decoder_ctx;
  }

  // Layer::Context overwrites
  virtual void start_epoch()
  {
    YANN_CHECK(is_valid());

    Base::start_epoch();

    _encoder_ctx->start_epoch();
    _decoder_ctx->start_epoch();
  }

  virtual void reset_state()
  {
    YANN_CHECK(is_valid());

    Base::reset_state();

    _encoder_ctx->reset_state();
    _decoder_ctx->reset_state();
  }
protected:
  unique_ptr<Layer::Context> _encoder_ctx;
  unique_ptr<Layer::Context> _decoder_ctx;
}; // class Seq2SeqLayer_Context

////////////////////////////////////////////////////////////////////////////////////////////////
//
// Seq2SeqLayer_TrainingContext implementation
//
class Seq2SeqLayer_TrainingContext :
    public Seq2SeqLayer_Context
{
  typedef Seq2SeqLayer_Context Base;
  friend class Seq2SeqLayer;

public:
  Seq2SeqLayer_TrainingContext(
      unique_ptr<Layer::Context> encoder_ctx,
      unique_ptr<Layer::Context> decoder_ctx) :
    Base(std::move(encoder_ctx), std::move(decoder_ctx))
  {
    _decoder_input_gradient.resize(_decoder_ctx->get_output_size());
    _decoder_output_gradient.resize(_decoder_ctx->get_output_size());
  }

protected:
  Vector _decoder_input_gradient;
  Vector _decoder_output_gradient;
}; // class Seq2SeqLayer_TrainingContext

}; // namespace yann

////////////////////////////////////////////////////////////////////////////////////////////////
//
// yann::Seq2SeqLayer implementation
//
std::unique_ptr<Seq2SeqLayer> yann::Seq2SeqLayer::create_lstm(
      const MatrixSize & input_size,
      const MatrixSize & output_size,
      const std::unique_ptr<ActivationFunction> & gate_activation_function,
      const std::unique_ptr<ActivationFunction> & io_activation_function)
{
  YANN_CHECK_GT(input_size, 0);
  YANN_CHECK_GT(output_size, 0);
  YANN_CHECK(gate_activation_function);
  YANN_CHECK(io_activation_function);

  auto encoder = make_unique<LstmLayer>(input_size, output_size);
  YANN_CHECK(encoder);
  encoder->set_activation_functions(gate_activation_function, io_activation_function);

  auto decoder = make_unique<LstmLayer>(output_size, output_size);
  YANN_CHECK(decoder);
  decoder->set_activation_functions(gate_activation_function, io_activation_function);

  return make_unique<Seq2SeqLayer>(std::move(encoder), std::move(decoder));
}

std::unique_ptr<Seq2SeqLayer> yann::Seq2SeqLayer::create_lstm(
    const MatrixSize & input_size,
    const MatrixSize & output_size,
    const std::unique_ptr<ActivationFunction> & activation_function)
{
  YANN_CHECK(activation_function);
  return create_lstm(input_size, output_size, activation_function, activation_function);
}

yann::Seq2SeqLayer::Seq2SeqLayer(std::unique_ptr<Layer> encoder, std::unique_ptr<Layer> decoder) :
    _encoder(std::move(encoder)),
    _decoder(std::move(decoder))
{
  YANN_CHECK_EQ(_encoder->get_output_size(), _decoder->get_input_size());
}

yann::Seq2SeqLayer::~Seq2SeqLayer()
{
}

// Layer overwrites
bool yann::Seq2SeqLayer::is_valid() const
{
  if(!Base::is_valid()) {
    return false;
  }
  if(!_encoder || !_encoder->is_valid()) {
    return false;
  }
  if(!_decoder || !_decoder->is_valid()) {
    return false;
  }
  return true;
}

std::string yann::Seq2SeqLayer::get_name() const
{
  return "Seq2SeqLayer";
}

string yann::Seq2SeqLayer::get_info() const
{
  YANN_CHECK(is_valid());

  ostringstream oss;
  oss << Base::get_info()
      << " encoder: " << _encoder->get_info()
      << ", decoder: " << _decoder->get_info()
  ;
  return oss.str();
}

bool yann::Seq2SeqLayer::is_equal(const Layer & other, double tolerance) const
{
  if(!Base::is_equal(other, tolerance)) {
    return false;
  }
  auto the_other = dynamic_cast<const Seq2SeqLayer*>(&other);
  if(the_other == nullptr) {
    return false;
  }
  if(is_valid() != the_other->is_valid()) {
    return false;
  }
  if(_encoder && !_encoder->is_equal(*the_other->_encoder, tolerance)) {
    return false;
  }
  if(_decoder && !_decoder->is_equal(*the_other->_decoder, tolerance)) {
    return false;
  }
  return true;
}

MatrixSize yann::Seq2SeqLayer::get_input_size() const
{
  return _encoder ? _encoder->get_input_size() : 0;
}

MatrixSize yann::Seq2SeqLayer::get_output_size() const
{
  return _decoder ? _decoder->get_output_size() : 0;
}

unique_ptr<Layer::Context> yann::Seq2SeqLayer::create_context(const MatrixSize & batch_size) const
{
  YANN_CHECK(is_valid());

  auto encoder_ctx = _encoder->create_context(batch_size);
  YANN_CHECK(encoder_ctx);
  auto decoder_ctx = _decoder->create_context(batch_size);
  YANN_CHECK(decoder_ctx);

  return make_unique<Seq2SeqLayer_Context>(std::move(encoder_ctx), std::move(decoder_ctx));
}
unique_ptr<Layer::Context> yann::Seq2SeqLayer::create_context(const RefVectorBatch & output) const
{
  YANN_CHECK(is_valid());

  auto encoder_ctx = _encoder->create_context(get_batch_size(output));
  YANN_CHECK(encoder_ctx);
  auto decoder_ctx = _decoder->create_context(output); // decoder output is our output
  YANN_CHECK(decoder_ctx);

  return make_unique<Seq2SeqLayer_Context>(std::move(encoder_ctx), std::move(decoder_ctx));
}
unique_ptr<Layer::Context> yann::Seq2SeqLayer::create_training_context(
    const MatrixSize & batch_size, const std::unique_ptr<Layer::Updater> & updater) const
{
  YANN_CHECK(is_valid());
  YANN_CHECK(updater);

  auto encoder_ctx = _encoder->create_training_context(batch_size, updater);
  YANN_CHECK(encoder_ctx);
  auto decoder_ctx = _decoder->create_training_context(batch_size, updater);
  YANN_CHECK(decoder_ctx);

  return make_unique<Seq2SeqLayer_TrainingContext>(std::move(encoder_ctx), std::move(decoder_ctx));
}
unique_ptr<Layer::Context> yann::Seq2SeqLayer::create_training_context(
    const RefVectorBatch & output, const std::unique_ptr<Layer::Updater> & updater) const
{
  YANN_CHECK(is_valid());
  YANN_CHECK(updater);

  auto encoder_ctx = _encoder->create_training_context(get_batch_size(output), updater);
  YANN_CHECK(encoder_ctx);
  auto decoder_ctx = _decoder->create_training_context(output, updater); // decoder output is our output
  YANN_CHECK(decoder_ctx);

  return make_unique<Seq2SeqLayer_TrainingContext>(std::move(encoder_ctx), std::move(decoder_ctx));
}

template<typename InputType>
void yann::Seq2SeqLayer::feedforward_internal(
    const InputType & input,
    Context * context,
    enum OperationMode mode) const
{
  auto ctx = dynamic_cast<Seq2SeqLayer_Context *>(context);
  YANN_CHECK(ctx);
  YANN_CHECK(ctx->is_valid());
  YANN_CHECK(is_valid());

  YANN_CHECK_LT(0, get_batch_size(input));
  YANN_CHECK_LE(get_batch_size(input), get_batch_size(ctx->get_output()));
  YANN_CHECK_EQ(get_batch_item_size(input), get_input_size());

  //
  // we assume all input is given to us "at once" no incremental feedforward
  //
  auto & encoder_ctx = ctx->_encoder_ctx;
  auto & decoder_ctx = ctx->_decoder_ctx;

  // feed input into encoder, overwrite the output
  _encoder->feedforward(input, encoder_ctx.get(), Operation_Assign);
  auto encoder_output_size = get_batch_size(encoder_ctx->get_output());
  YANN_SLOW_CHECK_GT(encoder_output_size, 0);
  auto state = get_batch(encoder_ctx->get_output(), encoder_output_size - 1); // last output

  // feed state into decoder and then feed prev output
  YANN_CHECK_EQ(mode, Operation_Assign); // we don't support anything else because we re-use output buffer
  for(MatrixSize ii = 0; ii < ctx->get_batch_size(); ++ii) {
    if(ii > 0) {
      auto prev_output = get_batch(decoder_ctx->get_output(), ii - 1);
      _decoder->feedforward(prev_output, decoder_ctx.get(), Operation_Assign);
    } else {
      _decoder->feedforward(state, decoder_ctx.get(), Operation_Assign);
    }
  }
}

void yann::Seq2SeqLayer::feedforward(
    const RefConstVectorBatch & input,
    Context * context,
    enum OperationMode mode) const
{
  feedforward_internal(input, context, mode);
}

void yann::Seq2SeqLayer::feedforward(
    const RefConstSparseVectorBatch & input,
    Context * context,
    enum OperationMode mode) const
{
  feedforward_internal(input, context, mode);
}

template<typename InputType>
void yann::Seq2SeqLayer::backprop_internal(
    const RefConstVectorBatch & gradient_output,
    const InputType & input,
    optional<RefVectorBatch> gradient_input,
    Context * context) const
{
  auto ctx = dynamic_cast<Seq2SeqLayer_TrainingContext *>(context);
  YANN_CHECK(ctx);
  YANN_CHECK(ctx->is_valid());
  YANN_CHECK(is_valid());
  YANN_CHECK_GT(get_batch_size(gradient_output), 0);
  YANN_CHECK_EQ(get_batch_item_size(gradient_output), get_output_size());
  YANN_CHECK_EQ(get_batch_size(input), get_batch_size(gradient_output));
  YANN_CHECK_EQ(get_batch_item_size(input), get_input_size());
  YANN_CHECK_EQ(get_batch_item_size(input), get_input_size());
  YANN_CHECK(!gradient_input || is_same_size(input, *gradient_input));

  //
  // we assume all output is given to us "at once" no incremental backprop
  //
  auto & encoder_ctx = ctx->_encoder_ctx;
  YANN_SLOW_CHECK(encoder_ctx);
  auto & decoder_ctx = ctx->_decoder_ctx;
  YANN_SLOW_CHECK(decoder_ctx);

  auto encoder_output_size = get_batch_size(encoder_ctx->get_output());
  YANN_SLOW_CHECK_GT(encoder_output_size, 0);
  auto state = get_batch(encoder_ctx->get_output(), encoder_output_size - 1); // last output

  // first we propagate back the output for decoder
  auto & decoder_input_gradient = ctx->_decoder_input_gradient;
  auto & decoder_output_gradient = ctx->_decoder_output_gradient;
  YANN_SLOW_CHECK(is_same_size(decoder_input_gradient, decoder_output_gradient));

  decoder_input_gradient.setZero();
  decoder_output_gradient.setZero();

  const auto batch_size = get_batch_size(gradient_output);
  for(MatrixSize ii = batch_size - 1; ii >= 0; --ii) {
    decoder_output_gradient += get_batch(gradient_output, ii);
    if(ii > 0) {
      auto decoder_prev_output = get_batch(decoder_ctx->get_output(), ii - 1);
      _decoder->backprop(
          decoder_output_gradient, // gradient_output
          decoder_prev_output,     // input
          make_optional<RefVectorBatch>(decoder_input_gradient),  // gradient_input
          decoder_ctx.get());
      swap(decoder_input_gradient, decoder_output_gradient);
    } else {
      _decoder->backprop(
          decoder_output_gradient, // gradient_output
          state,                   // input
          make_optional<RefVectorBatch>(decoder_input_gradient),  // gradient_input
          decoder_ctx.get());
    }
    // decoder_input_gradient contains the state gradient
  }

  // now propagate the gradient for encoder
  decoder_output_gradient.setZero();
  for(MatrixSize ii = batch_size - 1; ii >= 0; --ii) {
    auto gradient_in = gradient_input ? make_optional<RefVectorBatch>(get_batch(*gradient_input, ii)) : boost::none;
    InputType in = input.block(ii, 0, 1, input.cols()); // TODO: RowMajor, switch to get_batch()

    _encoder->backprop(
        decoder_input_gradient, // gradient_output
        in,   // input
        gradient_in,            // gradient_input
        encoder_ctx.get());
    if(ii == batch_size - 1) {
      decoder_input_gradient.setZero(); // we don't care about encoder output except the last vector
    }
  }
}

void yann::Seq2SeqLayer::backprop(
    const RefConstVectorBatch & gradient_output,
    const RefConstVectorBatch & input,
    optional<RefVectorBatch> gradient_input,
    Context * context) const
{
 backprop_internal(gradient_output, input, gradient_input, context);
}

void yann::Seq2SeqLayer::backprop(
    const RefConstVectorBatch & gradient_output,
    const RefConstSparseVectorBatch & input,
    optional<RefVectorBatch> gradient_input,
    Context * context) const
{
  backprop_internal<RefConstSparseVectorBatch>(gradient_output, input, gradient_input, context);
}

void yann::Seq2SeqLayer::init(enum InitMode mode, boost::optional<InitContext> init_context)
{
  YANN_CHECK(is_valid());
  _encoder->init(mode, init_context);
  _decoder->init(mode, init_context);
}

void yann::Seq2SeqLayer::update(Context * context, const size_t & tests_num)
{
  auto ctx = dynamic_cast<Seq2SeqLayer_TrainingContext *>(context);
  YANN_CHECK(ctx);
  YANN_CHECK(ctx->is_valid());
  YANN_CHECK(is_valid());

  _encoder->update(ctx->_encoder_ctx.get(), tests_num);
  _decoder->update(ctx->_decoder_ctx.get(), tests_num);
}

// the format is (e:<encoder>,d:<decoder>)
void yann::Seq2SeqLayer::read(std::istream & is)
{
  YANN_CHECK(is_valid());

  Base::read(is);

  read_char(is, '(');
  read_object(is, "e", *_encoder);
  read_char(is, ',');
  read_object(is, "d", *_decoder);
  read_char(is, ')');
}

// the format is (e:<encoder>,d:<decoder>)
void yann::Seq2SeqLayer::write(std::ostream & os) const
{
  YANN_CHECK(is_valid());

  Base::write(os);

  os << "(";
  write_object(os, "e", *_encoder);
  os << ",";
  write_object(os, "d", *_decoder);
  os << ")";
}