/*
 * prefilter_view_pairs.cc
 *
 *  Created on: Jul 12, 2016
 *      Author: LeonMing
 */

#include "theia/sfm/view_graph/prefilter_view_pairs.h"

#include <ceres/rotation.h>
#include <Eigen/Core>
#include <Eigen/Geometry>
#include <Eigen/SVD>
#include <glog/logging.h>
#include <unordered_map>
#include <unordered_set>
#include <cereal/access.hpp>
#include <iostream>
#include <fstream>
#include <algorithm>
#include <cmath>

#include "theia/math/util.h"
#include "theia/math/graph/minimum_spanning_tree.h"
#include "theia/util/hash.h"
#include "theia/util/map_util.h"
#include "theia/sfm/twoview_info.h"
#include "theia/sfm/types.h"
#include "theia/sfm/view_graph/view_graph.h"
#include "theia/util/filesystem.h"

#include <lemon/lp.h>
#include <lemon/lp_base.h>

namespace theia {

std::ofstream fileWriter("cycle_error_report.txt", std::ios::out);

bool PrefilterViewPairs(ViewGraph* viewGraph) {
  CHECK_NOTNULL(viewGraph);

// Declare MST extractor and MST and edges
  MinimumSpanningTree<ViewId, int> mst_extractor;
  std::unordered_set<ViewIdPair> mst;
  std::unordered_map<ViewIdPair, TwoViewInfo> edges;

// Create MST
  edges = viewGraph->GetAllEdges();

  for (const auto& edge : edges) {
    mst_extractor.AddEdge(edge.first.first, edge.first.second, 1);
  }

// Add values to MST
  mst_extractor.Extract(&mst);

// Preparing variables for creating chains
  std::vector<std::vector<ViewIdPair> > chainList;
  std::unordered_set<ViewIdPair> visited;
  std::vector<ViewIdPair> oldChain;
// Create chains
  CreateChains(chainList, mst, visited, oldChain, mst.begin()->first);

// Preparing variables for creating cycles
// Store cycles as vectors of ViewIdPairs, with a big vector holding all the cycles
  std::vector<std::vector<ViewIdPair> > cycleList;
// Create cycles
  CreateCycles(edges, chainList, mst, cycleList);

// Preparing variables for calculating errors
  std::unordered_set<ViewIdPair> allUsedEdges;
  std::vector<double> errors;
// Calculate and print cycle errors
 CalculateErrors(cycleList, edges, allUsedEdges, errors);

// Preparing variables for calculating translation errors
  std::unordered_set<ViewIdPair> allUsedTranslationEdges;
  std::vector<std::vector<ViewIdPair> > translationCycles;
  std::vector<double> translationErrors;
// Calculate translation cycle errors
  CalculateTranslationErrors(cycleList, translationCycles, edges, allUsedTranslationEdges, translationErrors);

// Preparing variables for inferencing
  std::unordered_map<int, ViewIdPair> number2Edge;
  std::unordered_map<ViewIdPair, int> edge2Number;
  int it = 0;
  for (auto& edge : allUsedEdges) {
	  number2Edge.emplace(it, edge);
	  edge2Number.emplace(edge, it);
	  ++it;
  }

  std::unordered_map<int, ViewIdPair> number2TranslationEdge;
  std::unordered_map<ViewIdPair, int> translationEdge2Number;
  int it2 = 0;
  for (auto& edge : allUsedTranslationEdges) {
	  number2TranslationEdge.emplace(it2, edge);
	  translationEdge2Number.emplace(edge, it2);
	  ++it2;
  }

  double rotSigma = 2*M_PI/180;
  double transSigma = 0.011;

// Edges that we need to remove
  std::unordered_set<ViewIdPair> viewPairsToRemove;
// Do inferences to figure out what edges to remove
  DoInference(rotSigma, cycleList, errors, edge2Number, number2Edge, allUsedEdges, viewPairsToRemove);
// Do inferences based on relative translations
  DoInference(transSigma, translationCycles, translationErrors, translationEdge2Number, number2TranslationEdge, allUsedTranslationEdges, viewPairsToRemove);

// Remove the edge if it is incorrect
  for (const ViewIdPair viewIdPair : viewPairsToRemove) {
    viewGraph->RemoveEdge(viewIdPair.first, viewIdPair.second);
  }
  VLOG(1) << "Removed " << viewPairsToRemove.size()
          << " view pairs by rotation prefiltering.";
  return true;
}

void CreateChains(std::vector<std::vector<ViewIdPair> >& chainList,
			      std::unordered_set<ViewIdPair>& mst,
				  std::unordered_set<ViewIdPair>& visited,
				  std::vector<ViewIdPair> oldChain,
				  ViewId root) {
	for (ViewIdPair e : mst) {
		if (visited.find(e) != visited.end()) {
			continue;
		}
		if (e.first == root) {
			oldChain.emplace_back(e);
			std::vector<ViewIdPair> chain = oldChain;
			oldChain.pop_back();
			chainList.emplace_back(chain);
			visited.insert(e);

			CreateChains(chainList, mst, visited, chain, e.second);
		}
		else if (e.second == root) {
			oldChain.emplace_back(e);
			std::vector<ViewIdPair> chain = oldChain;
			oldChain.pop_back();
			chainList.emplace_back(chain);
			visited.insert(e);

			CreateChains(chainList, mst, visited, chain, e.first);
		}
	}
}

void CreateCycles(std::unordered_map<ViewIdPair, TwoViewInfo>& edges,
				  std::vector<std::vector<ViewIdPair> >& chainList,
				  std::unordered_set<ViewIdPair>& mst,
				  std::vector<std::vector<ViewIdPair> >& cycleList) {

	std::unordered_set<ViewIdPair> otherEdges; // not in MST

	for (auto& e : edges) {
		if (mst.find(e.first) == mst.end()) {
			otherEdges.insert(e.first);
		}
	}

	for (ViewIdPair e : otherEdges) {
		if (e.first == mst.begin()->first || e.second == mst.begin()->first) {
			continue;
		}
		std::vector<ViewIdPair> cycle;
		cycle.emplace_back(e);
		std::vector<ViewIdPair> chain1;
		std::vector<ViewIdPair> chain2;

		for (std::vector<ViewIdPair> chain : chainList) {
			ViewIdPair lastE = chain.end()[-1];
			ViewIdPair secondtolastE = chain.end()[-2];
			if (
				((e.first == lastE.first || e.first == lastE.second)
				&& !(e.first == secondtolastE.first || e.first == secondtolastE.second))
				||
				((e.first == lastE.first || e.first == lastE.second)
				&& (chain.size() == 1))
				) {
				chain1 = chain;
				continue;
			}
			if (
				((e.second == lastE.first || e.second == lastE.second)
				&& !(e.second == secondtolastE.first || e.second == secondtolastE.second))
				||
				((e.second == lastE.first || e.second == lastE.second)
				&& (chain.size() == 1))
				) {
				chain2 = chain;
				continue;
			}
		}

		std::set_symmetric_difference(chain1.begin(), chain1.end(),
				                      chain2.begin(), chain2.end(),
				                      std::back_inserter(cycle));
		cycleList.emplace_back(cycle);
	}
}

void CalculateErrors(std::vector<std::vector<ViewIdPair> >& cycleList,
					 std::unordered_map<ViewIdPair, TwoViewInfo>& edges,
					 std::unordered_set<ViewIdPair>& allUsedEdges,
					 std::vector<double>& errors) {
	for (std::vector<ViewIdPair> cycle : cycleList) {

		double error;
		Eigen::Matrix3d combinedRotation;
		ceres::AngleAxisToRotationMatrix(
			edges[cycle[0]].rotation_2.data(),
			ceres::RowMajorAdapter3x3(combinedRotation.data()));

		std::vector<ViewIdPair> usedEdges;
		usedEdges.push_back(cycle.front());
		allUsedEdges.insert(cycle.front());
		std::vector<ViewId> nextEdge;
		nextEdge.push_back(cycle.front().second);
		while (usedEdges.size() < cycle.size()) {
			for (ViewIdPair edge : cycle) {

				TwoViewInfo info = edges[edge];
				Eigen::Matrix3d rotation_matrix;

				if (std::find(usedEdges.begin(), usedEdges.end(), edge) != usedEdges.end()) {
					continue;
				} else if (edge.first == nextEdge.back()) {
					ceres::AngleAxisToRotationMatrix(
						info.rotation_2.data(),
						ceres::RowMajorAdapter3x3(rotation_matrix.data()));
					combinedRotation *= rotation_matrix;
					usedEdges.push_back(edge);
					if (allUsedEdges.find(edge) == allUsedEdges.end()) {
						allUsedEdges.insert(edge);
					}
					nextEdge.push_back(edge.second);
					fileWriter << rotation_matrix;
					fileWriter << "\n\n";
				} else if (edge.second == nextEdge.back()) {
					ceres::AngleAxisToRotationMatrix(
						info.rotation_2.data(),
						ceres::RowMajorAdapter3x3(rotation_matrix.data()));
					combinedRotation *= rotation_matrix.transpose();
					usedEdges.push_back(edge);
					if (allUsedEdges.find(edge) == allUsedEdges.end()) {
						allUsedEdges.insert(edge);
					}
					nextEdge.push_back(edge.first);
					fileWriter << rotation_matrix.transpose();
					fileWriter << "\n\n";
				} else {
					continue;
				}
			}
		}
		fileWriter << "Combined Rotation:\n";
		fileWriter << combinedRotation;

		fileWriter << "\n\nError:\n";
		error = acos((combinedRotation.trace()-1)/2);
		fileWriter << error << "\n\n\n\n";
		errors.push_back(error);
	}
}

void CalculateTranslationErrors(std::vector<std::vector<ViewIdPair> >& cycleList,
								std::vector<std::vector<ViewIdPair> >& translationCycles,
		 	 	 	 	 	 	std::unordered_map<ViewIdPair, TwoViewInfo>& edges,
								std::unordered_set<ViewIdPair>& allUsedTranslationEdges,
								std::vector<double>& translationErrors) {
	for (auto& cycle : cycleList) {
		if (cycle.size() == 3) {
			translationCycles.push_back(cycle);
		}
	}
	std::cout << "Singular Values: " << "\n\n";
	for (auto& cycle : translationCycles) {
		Eigen::Matrix3d translationsMatrix;
		for (int i = 0; i < 3; ++i) {
			translationsMatrix(i, 0) = edges[cycle.front()].position_2[i];
		}
		std::vector<ViewIdPair> usedEdges;
		usedEdges.push_back(cycle.front());
		allUsedTranslationEdges.insert(cycle.front());
		std::vector<ViewId> nextEdge;
		nextEdge.push_back(cycle.front().second);
		while (usedEdges.size() < 3) {
			for (int j = 1; j < 3; ++j) {
				if (std::find(usedEdges.begin(), usedEdges.end(), cycle[j]) != usedEdges.end()) {
					continue;
				} else if (cycle[j].first == nextEdge.back()) {
					for (int i = 0; i < 3; ++i) {
						translationsMatrix(i, j) = edges[cycle[j]].position_2[i];
					}
					usedEdges.push_back(cycle[j]);
					if (allUsedTranslationEdges.find(cycle[j]) == allUsedTranslationEdges.end()) {
						allUsedTranslationEdges.insert(cycle[j]);
					}
					nextEdge.push_back(cycle[j].second);
				} else if (cycle[j].second == nextEdge.back()) {
					for (int i = 0; i < 3; ++i) {
						translationsMatrix(i, j) = -(edges[cycle[j]].position_2[i]);
					}
					usedEdges.push_back(cycle[j]);
					if (allUsedTranslationEdges.find(cycle[j]) == allUsedTranslationEdges.end()) {
						allUsedTranslationEdges.insert(cycle[j]);
					}
					nextEdge.push_back(cycle[j].first);
				}
			}
		}
//		std::cout << translationsMatrix << "\n";
		Eigen::JacobiSVD<Eigen::Matrix3d> svd(translationsMatrix);
		double error = svd.singularValues()[2];
		translationErrors.push_back(error);
	}
	for (int i = 0; i < translationErrors.size(); ++i) {
		std::cout << "Cycle " << i << "\n" << translationErrors[i] << "\n\n";
	}
}

void DoInference(double sigma,
				 std::vector<std::vector<ViewIdPair> >& cycleList,
				 std::vector<double>& errors,
				 std::unordered_map<ViewIdPair, int>& edge2Number,
				 std::unordered_map<int, ViewIdPair>& number2Edge,
				 std::unordered_set<ViewIdPair>& allUsedEdges,
				 std::unordered_set<ViewIdPair>& viewPairsToRemove) {
	const double rhoe = (-log(0.1)) - (-log(0.9));
	double rhol;

	const double sigmaInliers = sigma;
	const double eps = pow(2,-52);

	using namespace lemon;

	// Create an instance of the default LP solver class
	// (it will represent an "empty" problem at first)
	Lp lp;

	// Set up variables
	std::vector<Lp::Col> xe;
	std::vector<Lp::Col> xl;
	Lp::Expr totalEdges;
	Lp::Expr totalCycles;
	std::map<std::vector<ViewIdPair>, std::vector<int> > index;
	for (auto& cycle : cycleList) {
		std::vector<int> tempVector;
		for (auto& edge : cycle) {
			tempVector.push_back(edge2Number[edge]);
		}
		index.emplace(cycle, tempVector);
	}

	// Add columns (variables) to the problem
	for (int i = 0; i < allUsedEdges.size(); ++i) { // For every edge used in any cycle, make a new xe
		xe.push_back(lp.addCol());
	}
	for (int i = 0; i < cycleList.size(); ++i) { // For every cycle in the cycleList, make a new xl
		xl.push_back(lp.addCol());
	}

	// Add rows (constraints) to the problem
	for (int cycle = 0; cycle < cycleList.size(); ++cycle) { // For every cycle in the cycleList

		Lp::Expr total;
		for (int edge = 0; edge < cycleList[cycle].size(); ++edge) {
			lp.addRow(xl[cycle] >= xe[(index[cycleList[cycle]])[edge]]); // xl for this cycle is greater than each xe
			total += xe[(index[cycleList[cycle]])[edge]];
			totalEdges += total;
		}
		lp.addRow(xl[cycle] <= total);

		// Calculate rhol for this cycle
		double pErrorInliers = exp(-errors[cycle]/sigmaInliers)/(sigmaInliers*(1-exp(-M_PI/sigmaInliers)));
		double pErrorOutliers = (1-exp(-errors[cycle]/sigmaInliers))/(M_PI-sigmaInliers*(1-exp(-M_PI/sigmaInliers)));
		rhol = (-log(pErrorOutliers + eps))-(-log(pErrorInliers + eps));
		totalCycles += rhol * xl[cycle];
	}

	// Set lower and upper bounds for the columns (variables)
	for (int k = 0; k < xe.size(); ++k) {
		lp.colBounds(xe[k], 0, 1);
	}
	for (int k = 0; k < xl.size(); ++k) {
		lp.colBounds(xl[k], 0, 1);
	}

	// Specify the objective function
	lp.min();
	lp.obj(rhoe * totalEdges + totalCycles);

	// Solve the problem using the underlying LP solver
	lp.solve();

	// Print the results
	if (lp.primalType() == Lp::OPTIMAL) {
		std::cout << "Objective function value: " << lp.primal() << "\n\n";
		for (int i = 0; i < xe.size(); ++i) {
			std::cout << "Edge " << i << ": " << number2Edge[i].first << " " << number2Edge[i].second << ":\t" << lp.primal(xe[i]) << "\n";
			if (lp.primal(xe[i]) == 1) {
				viewPairsToRemove.insert(number2Edge[i]);
			}
		}
		for (int i = 0; i < xl.size(); ++i) {
			std::cout << "Cycle " << i << ": " << lp.primal(xl[i]) << "\n";
		}
	} else {
		std::cout << "Optimal solution not found." << "\n";
	}
}

}  // namespace theia


#include "syanten.hpp"

#include <boost/dll.hpp>
#include <spdlog/spdlog.h>

#include <fstream>

#include "bitutils.hpp"

namespace mahjong
{

SyantenCalculator::SyantenCalculator() { initialize(); }

/**
 * @brief 向聴数を計算する。
 *
 * @param[in] hand 手牌
 * @param[in] type 計算対象の向聴数の種類
 * @return std::tuple<SyantenType, int> (向聴数の種類, 向聴数)
 */
std::tuple<int, int> SyantenCalculator::calc(const Hand &hand, int type)
{
#ifdef CHECK_ARGUMENT
    if (type < 0 || type > 7) {
        spdlog::warn("Invalid type {} passed.", type);
        return -2;
    }
#endif

    std::tuple<int, int> ret = {SyantenType::Null, std::numeric_limits<int>::max()};

    if (type & SyantenType::Normal) {
        int syanten = calc_normal(hand);
        if (syanten < std::get<1>(ret))
            ret = {SyantenType::Normal, syanten};
    }
    if (type & SyantenType::Tiitoi) {
        int syanten = calc_tiitoi(hand);
        if (syanten < std::get<1>(ret))
            ret = {SyantenType::Tiitoi, syanten};
    }
    if (type & SyantenType::Kokusi) {
        int syanten = calc_kokusi(hand);
        if (syanten < std::get<1>(ret))
            ret = {SyantenType::Kokusi, syanten};
    }

    return ret;
}

/**
 * @brief 初期化する。
 *
 * @return 初期化に成功した場合は true、そうでない場合は false を返す。
 */
bool SyantenCalculator::initialize()
{
    boost::filesystem::path exe_path = boost::dll::this_line_location().parent_path();

#ifdef USE_UNORDERED_MAP
    if (s_tbl_.empty()) {
        boost::filesystem::path path = exe_path / "syupai_table.bin";
        std::ifstream ifs(path.string(), std::ios::binary);
        for (size_t i = 0; i < ShuupaiPatternSize; ++i) {
            unsigned int key;
            Pattern pattern;
            ifs.read((char *)&key, sizeof(unsigned int));
            ifs.read((char *)&pattern, sizeof(Pattern));

            s_tbl_[key] = pattern;
        }
    }

    if (z_tbl_.empty()) {
        boost::filesystem::path path = exe_path / "zihai_table.bin";
        std::ifstream ifs(path.string(), std::ios::binary);

        for (size_t i = 0; i < ZihaiPatternSize; ++i) {
            unsigned int key;
            Pattern pattern;
            ifs.read((char *)&key, sizeof(unsigned int));
            ifs.read((char *)&pattern, sizeof(Pattern));

            z_tbl_[key] = pattern;
        }
    }
#else
    if (s_tbl_.empty()) {
        s_tbl_.resize(ShuupaiTableSize);

        boost::filesystem::path path = exe_path / "syupai_table.bin";
        std::ifstream ifs(path.string(), std::ios::binary);
        for (size_t i = 0; i < ShuupaiPatternSize; ++i) {
            unsigned int key;
            ifs.read((char *)&key, sizeof(unsigned int));
            ifs.read((char *)&s_tbl_[key], sizeof(Pattern));
        }
    }

    if (z_tbl_.empty()) {
        z_tbl_.resize(ZihaiTableSize);

        boost::filesystem::path path = exe_path / "zihai_table.bin";
        std::ifstream ifs(path.string(), std::ios::binary);

        for (size_t i = 0; i < ZihaiPatternSize; ++i) {
            unsigned int key;
            ifs.read((char *)&key, sizeof(unsigned int));
            ifs.read((char *)&z_tbl_[key], sizeof(Pattern));
        }
    }
#endif

    return true;
}

/**
 * @brief 通常手の向聴数を計算する。
 *
 * @param[in] hand 手牌
 * @return int 向聴数
 */
int SyantenCalculator::calc_normal(const Hand &hand)
{
    // 制約条件「面子 + 候補 <= 4」で「面子数 * 2 + 候補」の最大値を計算する。
    int n_melds = int(hand.melds.size());
    int n_mentu_base = n_melds + s_tbl_[hand.manzu].n_mentu +
                       s_tbl_[hand.pinzu].n_mentu + s_tbl_[hand.sozu].n_mentu +
                       z_tbl_[hand.zihai].n_mentu;
    int n_kouho_base = s_tbl_[hand.manzu].n_kouho + s_tbl_[hand.pinzu].n_kouho +
                       s_tbl_[hand.sozu].n_kouho + z_tbl_[hand.zihai].n_kouho;

    // 雀頭なし
    int max = n_mentu_base * 2 + std::min(4 - n_mentu_base, n_kouho_base);

    if (s_tbl_[hand.manzu].head) {
        // 萬子の雀頭有り
        int n_mentu = n_mentu_base + s_tbl_[hand.manzu].n_mentu_diff;
        int n_kouho = n_kouho_base + s_tbl_[hand.manzu].n_kouho_diff;
        max = std::max(max, n_mentu * 2 + std::min(4 - n_mentu, n_kouho) + 1);
    }

    if (s_tbl_[hand.pinzu].head) {
        // 筒子の雀頭有り
        int n_mentu = n_mentu_base + s_tbl_[hand.pinzu].n_mentu_diff;
        int n_kouho = n_kouho_base + s_tbl_[hand.pinzu].n_kouho_diff;
        max = std::max(max, n_mentu * 2 + std::min(4 - n_mentu, n_kouho) + 1);
    }

    if (s_tbl_[hand.sozu].head) {
        // 索子の雀頭有り
        int n_mentu = n_mentu_base + s_tbl_[hand.sozu].n_mentu_diff;
        int n_kouho = n_kouho_base + s_tbl_[hand.sozu].n_kouho_diff;
        max = std::max(max, n_mentu * 2 + std::min(4 - n_mentu, n_kouho) + 1);
    }

    if (z_tbl_[hand.zihai].head) {
        // 字牌の雀頭有り
        int n_mentu = n_mentu_base + z_tbl_[hand.zihai].n_mentu_diff;
        int n_kouho = n_kouho_base + z_tbl_[hand.zihai].n_kouho_diff;
        max = std::max(max, n_mentu * 2 + std::min(4 - n_mentu, n_kouho) + 1);
    }

    return 8 - max;
}

/**
 * @brief 七対子手の向聴数を計算する。
 *
 * @param[in] hand 手牌
 * @return int 向聴数
 */
int SyantenCalculator::calc_tiitoi(const Hand &hand)
{
    // 牌の種類 (1枚以上の牌) を数える。
    int n_types = s_tbl_[hand.manzu].n_ge1 + s_tbl_[hand.pinzu].n_ge1 +
                  s_tbl_[hand.sozu].n_ge1 + z_tbl_[hand.zihai].n_ge1;
    // 対子の数 (2枚以上の牌) を数える。
    int n_toitu = s_tbl_[hand.manzu].n_ge2 + s_tbl_[hand.pinzu].n_ge2 +
                  s_tbl_[hand.sozu].n_ge2 + z_tbl_[hand.zihai].n_ge2;

    int syanten = 6 - n_toitu;
    if (n_types < 7)
        syanten += 7 - n_types; // 4枚持ちを考慮

    return syanten;
}

/**
 * @brief 国士無双の向聴数を計算する。
 *
 * @param[in] hand 手牌
 * @return int 向聴数
 */
int SyantenCalculator::calc_kokusi(const Hand &hand)
{
    // 老頭牌を抽出する。
    int manzu19 = hand.manzu & Bit::RotohaiMask;
    int pinzu19 = hand.pinzu & Bit::RotohaiMask;
    int sozu19 = hand.sozu & Bit::RotohaiMask;

    // 幺九牌の種類 (1枚以上の牌) を数える。
    int n_yaotyuhai = s_tbl_[manzu19].n_ge1 + s_tbl_[pinzu19].n_ge1 +
                      s_tbl_[sozu19].n_ge1 + z_tbl_[hand.zihai].n_ge1;

    // 幺九牌の対子があるかどうか
    int toitu = ((manzu19 & 0b110'000'000'000'000'000'000'000'110) |
                 (pinzu19 & 0b110'000'000'000'000'000'000'000'110) |
                 (sozu19 & 0b110'000'000'000'000'000'000'000'110) |
                 (hand.zihai & 0b110'110'110'110'110'110'110)) > 0;

    return 13 - toitu - n_yaotyuhai;
}

#ifdef USE_UNORDERED_MAP
std::unordered_map<unsigned int, SyantenCalculator::Pattern> SyantenCalculator::s_tbl_;
std::unordered_map<unsigned int, SyantenCalculator::Pattern> SyantenCalculator::z_tbl_;
#else
std::vector<SyantenCalculator::Pattern> SyantenCalculator::s_tbl_;
std::vector<SyantenCalculator::Pattern> SyantenCalculator::z_tbl_;
#endif

static SyantenCalculator inst;

} // namespace mahjong


#ifndef TVTML_FUNCTIONAL3D_HPP
#define TVTML_FUNCTIONAL3D_HPP

// Eigen includes
#include <Eigen/Core>
#include <Eigen/Sparse>

// video++ includes
#include <vpp/vpp.hh>

// system includes
#include <cmath>
#include <iostream>
#include <fstream>
#include <vector>

// own includes
#include "enumerators.hpp"
#include "manifold.hpp"
#include "data.hpp"

namespace tvmtl{

template <enum FUNCTIONAL_DISC disc, typename MANIFOLD, class DATA>
class Functional<FIRSTORDER, disc, MANIFOLD, DATA, 3>{

    public:
	// Manifold typedefs and constants
	static const MANIFOLD_TYPE mf_type = MANIFOLD::MyType;
	static const int value_dim = MANIFOLD::value_dim;
	static const int manifold_dim = MANIFOLD::manifold_dim;
	typedef typename MANIFOLD::scalar_type scalar_type;
	typedef typename MANIFOLD::value_type value_type;
	typedef typename MANIFOLD::ref_type ref_type;
	typedef typename MANIFOLD::cref_type cref_type;
	typedef typename MANIFOLD::deriv1_type deriv1_type;
	typedef typename MANIFOLD::deriv2_type deriv2_type;
	typedef typename MANIFOLD::restricted_deriv2_type restricted_deriv2_type;
	typedef typename MANIFOLD::tm_base_type tm_base_type;

	// Data typedef and constants
	static const int img_dim = DATA::img_dim;
	typedef typename DATA::storage_type img_type;
	typedef typename DATA::weights_type weights_type;
	typedef typename DATA::weights_mat weights_mat;
	typedef typename DATA::inp_mat inp_mat;

	// Functional parameters and return types
	static const FUNCTIONAL_DISC disc_type;
	typedef double param_type;
	typedef double result_type;

	// Tangent space transformation matrix types
	typedef vpp::imageNd<tm_base_type, img_dim> tm_base_mat_type;

	// Gradient and Hessian types
	typedef Eigen::Matrix<scalar_type, Eigen::Dynamic, 1> gradient_type;
	typedef vpp::imageNd<deriv2_type, img_dim> hessian_type;
	typedef Eigen::SparseMatrix<scalar_type> sparse_hessian_type;

	//Constructor
	Functional(param_type lambda, DATA& dat):
	    lambda_(lambda),
	    data_(dat)
	{
	    eps2_=1e-10;
	   static_assert(img_dim == 3, "Dimension of data and functional must match!");
	}

	void updateWeights();

	void updateTMBase();

	// Evaluation functions
	result_type evaluateJ();
	void  evaluateDJ();
	void  evaluateHJ();

	template <class IMG>
	void output_img(const IMG& img, const char* filename) const;
	template <class IMG>
	void output_matval_img(const IMG& img, const char* filename) const;

	// Getter and Setter
	inline param_type getlambda() const { return lambda_; }
	inline void setlambda(param_type lam) { lambda_=lam; }
	inline param_type geteps2() const { return eps2_; }
	inline void seteps2(param_type eps) { eps2_=eps; }

	inline const weights_mat& getweightsX() const { return weightsX_; }
	inline const weights_mat& getweightsY() const { return weightsY_; }
	inline const weights_mat& getweightsZ() const { return weightsZ_; }

	inline const gradient_type& getDJ() const { return DJ_; }
	inline const sparse_hessian_type& getHJ() const { return HJ_; }
	inline const tm_base_mat_type& getT() const { return T_; }

    private:
	DATA& data_;

	param_type lambda_, eps2_;
	weights_mat weightsX_, weightsY_, weightsZ_;

	tm_base_mat_type T_;
	gradient_type DJ_;
	sparse_hessian_type HJ_;
};

//--------Implementation FIRSTORDER-----/

template <enum FUNCTIONAL_DISC disc, typename MANIFOLD, class DATA >
const FUNCTIONAL_DISC Functional<FIRSTORDER, disc, MANIFOLD, DATA, 3>::disc_type = disc;

template <enum FUNCTIONAL_DISC disc, typename MANIFOLD, class DATA>
void Functional<FIRSTORDER, disc, MANIFOLD, DATA, 3>::updateWeights(){

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t Update Weights..." << std::endl;
    #endif

    // Neighbourhood box
    int ns = data_.img_.nslices();  // z
    int nr = data_.img_.nrows();    // y
    int nc = data_.img_.ncols();    // x

    // Subimage boxes
    vpp::box3d without_last_x(vpp::vint3(0,0,0), vpp::vint3(ns - 1, nr - 1, nc - 2)); // subdomain without last xslice
    vpp::box3d without_last_y(vpp::vint3(0,0,0), vpp::vint3(ns - 1, nr - 2, nc - 1)); // subdomain without last yslice
    vpp::box3d without_last_z(vpp::vint3(0,0,0), vpp::vint3(ns - 2, nr - 1, nc - 1)); // subdomain without last zslice
    vpp::box3d without_first_x(vpp::vint3(0,0,1), vpp::vint3(ns - 1, nr - 1, nc - 1)); // subdomain without first xlice
    vpp::box3d without_first_y(vpp::vint3(0,1,0), vpp::vint3(ns - 1, nr - 1, nc - 1)); // subdomain without first yslice
    vpp::box3d without_first_z(vpp::vint3(1,0,0), vpp::vint3(ns - 1, nr - 1, nc - 1)); // subdomain without first zslice

    weightsX_ = weights_mat(data_.img_.domain());
    weightsY_ = weights_mat(data_.img_.domain());
    weightsZ_ = weights_mat(data_.img_.domain());

    auto calc_dist = [&] (weights_type& w, const value_type i, const value_type n) {
	w = MANIFOLD::dist_squared(i, n);
    };

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...X neighbours " << std::endl;
    #endif

    // X Neighbours
    fill3d(weightsX_, 0.0);
    pixel_wise3d(calc_dist, weightsX_ | without_last_x, data_.img_ | without_last_x, data_.img_ | without_first_x );

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...Reweighting" << std::endl;
	std::cout << "\t\t...Y neighbours" << std::endl;
    #endif

    // Y Neighbours
    fill3d(weightsY_, 0.0);
    pixel_wise3d(calc_dist, weightsY_ | without_last_y, data_.img_ | without_last_y, data_.img_ | without_first_y );

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...Reweighting" << std::endl;
	std::cout << "\t\t...Z neighbours" << std::endl;
    #endif

    // Z Neighbours
    fill3d(weightsZ_, 0.0);
    pixel_wise3d(calc_dist, weightsZ_ | without_last_z, data_.img_ | without_last_z, data_.img_ | without_first_z );

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...Reweighting" << std::endl;
    #endif

    if(disc==ISO){
	auto g =  [&] (const weights_type& ew, weights_type& x, weights_type& y, weights_type& z) { x = ew / std::sqrt(x + y + z + eps2_); y = x; z = x; };
	pixel_wise3d(g, data_.edge_weights_, weightsX_, weightsY_, weightsZ_);
    }
    else{
	auto g =  [&] (const weights_type& ew, weights_type& w) { w = ew / std::sqrt(w+eps2_); };
	pixel_wise3d(g, data_.edge_weights_, weightsX_);
	pixel_wise3d(g, data_.edge_weights_, weightsY_);
	pixel_wise3d(g, data_.edge_weights_, weightsZ_);
    }
}

// Update the Tangent space ONB
template <enum FUNCTIONAL_DISC disc, typename MANIFOLD, class DATA >
void Functional<FIRSTORDER, disc, MANIFOLD, DATA, 3 >::updateTMBase(){

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\tUpdate tangent space basis..." << std::endl;
    #endif

   tm_base_mat_type T(data_.img_.domain());
   pixel_wise3d([&] (tm_base_type& t, const value_type& i) { MANIFOLD::tangent_plane_base(i,t); }, T, data_.img_);
   T_=T;

}

// Evaluation of J
template <enum FUNCTIONAL_DISC disc, typename MANIFOLD, class DATA >
typename Functional<FIRSTORDER, disc, MANIFOLD, DATA, 3 >::result_type Functional<FIRSTORDER, disc, MANIFOLD, DATA, 3 >::evaluateJ(){

    // sum d^2(img, img_noise)
    result_type J1, J2;
    J1 = J2 = 0.0;

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\tFunctional evaluation..." << std::endl;
	std::cout << "\t\t...Fidelity part" << std::endl;
    #endif

    if(data_.doInpaint()){
	auto f = [&] (const value_type& i, const value_type& n, const bool inp ) { J1 += MANIFOLD::dist_squared(i,n)*(1-inp); };
	pixel_wise3d_nothreads(f, data_.img_, data_.noise_img_, data_.inp_);
    }
    else{
	auto f = [&] (const value_type& i, const value_type& n) { J1 += MANIFOLD::dist_squared(i,n); };
	pixel_wise3d_nothreads(f, data_.img_, data_.noise_img_);
    }

	updateWeights();

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...TV part." << std::endl;
    #endif

    if(disc==ISO)
	pixel_wise3d_nothreads([&] (const weights_type& w) { J2 += 1.0/w;} ,weightsX_);
    else
	pixel_wise3d_nothreads([&] (const weights_type& wx, const weights_type& wy, const weights_type& wz) { J2 += 1.0 / wx + 1.0 / wy+ 1.0 / wz;}, weightsX_, weightsY_, weightsZ_);
    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "J1: " << J1 << std::endl;
	std::cout << "J2: " << J2 << std::endl;
    #endif

    return 0.5 * J1 + lambda_* J2;
}

// Evaluation of J'
template <enum FUNCTIONAL_DISC disc, typename MANIFOLD, class DATA >
void Functional<FIRSTORDER, disc, MANIFOLD, DATA, 3 >::evaluateDJ(){

    img_type grad = img_type(data_.img_.domain());
    int ns = data_.img_.nslices();
    int nr = data_.img_.nrows();
    int nc = data_.img_.ncols();

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\tGradient evaluation..." << std::endl;
	std::cout << "\t\t...Fidelity part" << std::endl;
    #endif
    //GRADIENT OF FIDELITY TERM
    if(data_.doInpaint()){
	auto f = [] (value_type& g, const value_type& i, const value_type& n, const bool inp ) { MANIFOLD::deriv1x_dist_squared(i,n,g); g*=(1-inp); };
	pixel_wise3d(f, grad, data_.img_, data_.noise_img_, data_.inp_);
    }
    else{
	auto f = [] (value_type& g, const value_type& i, const value_type& n) { MANIFOLD::deriv1x_dist_squared(i,n,g); };
	pixel_wise3d(f, grad, data_.img_, data_.noise_img_);
    }

    //GRADIENT OF TV TERM

    // Subimage boxes
    vpp::box3d without_last_x(vpp::vint3(0,0,0), vpp::vint3(ns - 1, nr - 1, nc - 2)); // subdomain without last xslice
    vpp::box3d without_last_y(vpp::vint3(0,0,0), vpp::vint3(ns - 1, nr - 2, nc - 1)); // subdomain without last yslice
    vpp::box3d without_last_z(vpp::vint3(0,0,0), vpp::vint3(ns - 2, nr - 1, nc - 1)); // subdomain without last zslice
    vpp::box3d without_first_x(vpp::vint3(0,0,1), vpp::vint3(ns - 1, nr - 1, nc - 1)); // subdomain without first xlice
    vpp::box3d without_first_y(vpp::vint3(0,1,0), vpp::vint3(ns - 1, nr - 1, nc - 1)); // subdomain without first yslice
    vpp::box3d without_first_z(vpp::vint3(1,0,0), vpp::vint3(ns - 1, nr - 1, nc - 1)); // subdomain without first zslice

    auto calc_first_arg_deriv = [&] (value_type& x, const weights_type& w, const value_type& i, const value_type& n) { MANIFOLD::deriv1x_dist_squared(i, n, x); x *= w; };
    auto calc_second_arg_deriv = [&] (value_type& y, const weights_type& w, const value_type& i, const value_type& n) { MANIFOLD::deriv1y_dist_squared(i, n, y); y *= w; };
    auto add_to_gradient = [&] (value_type& g, const value_type& d) { g+=d*lambda_; };

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\tGradient evaluation..." << std::endl;
	std::cout << "\t\t...TV part" << std::endl;
	std::cout << "\t\t...-> XD1" << std::endl;
    #endif
    // X neighbors and reweighting
    // ... w.r.t. to first argument
    { // Temporary image XD1 is deallocated after this scope
	img_type XD1 = img_type(without_last_x);
	pixel_wise3d(calc_first_arg_deriv, XD1, weightsX_ | without_last_x, data_.img_ | without_last_x, data_.img_ | without_first_x);
	pixel_wise3d(add_to_gradient, grad | without_last_x, XD1);
    }
    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...-> XD2" << std::endl;
    #endif
    // ... w.r.t. second argument
    {
	img_type XD2 = img_type(without_last_x);
	pixel_wise3d(calc_second_arg_deriv, XD2, weightsX_ | without_last_x, data_.img_ | without_last_x, data_.img_ | without_first_x);
	pixel_wise3d(add_to_gradient, grad | without_first_x,  XD2);
    }

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...-> YD1" << std::endl;
    #endif
    // Vertical derivatives and weighting
    // ... w.r.t. first argument
    {
	img_type YD1 = img_type(without_last_y);
	pixel_wise3d(calc_first_arg_deriv, YD1, weightsY_ | without_last_y, data_.img_ | without_last_y, data_.img_ | without_first_y);
	pixel_wise3d(add_to_gradient, grad | without_first_y, YD1);
    }

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...-> YD2" << std::endl;
    #endif
    // ... w.r.t second argument
    {
	img_type YD2 = img_type(without_last_y);
	pixel_wise3d(calc_second_arg_deriv, YD2, weightsY_ | without_last_y, data_.img_ | without_last_y, data_.img_ | without_first_y);
        pixel_wise3d(add_to_gradient, grad | without_first_y, YD2);
    }

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...-> ZD1" << std::endl;
    #endif
    // Vertical derivatives and weighting
    // ... w.r.t. first argument
    {
	img_type ZD1 = img_type(without_last_z);
	pixel_wise3d(calc_first_arg_deriv, ZD1, weightsZ_ | without_last_z, data_.img_ | without_last_z, data_.img_ | without_first_z);
	pixel_wise3d(add_to_gradient, grad | without_first_z, ZD1);
    }

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...-> ZD2" << std::endl;
    #endif
    // ... w.r.t second argument
    {
	img_type ZD2 = img_type(without_last_z);
	pixel_wise3d(calc_second_arg_deriv, ZD2, weightsZ_ | without_last_z, data_.img_ | without_last_z, data_.img_ | without_first_z);
        pixel_wise3d(add_to_gradient, grad | without_first_z, ZD2);
    }

    #ifdef TV_FUNC_DEBUG
	    data_.output_matval_img(grad,"3dgrad.csv");
    #endif

    DJ_ = gradient_type::Zero(ns*nr*nc*manifold_dim);

    updateTMBase();
    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...Local to global insert" << std::endl;
    #endif

   for(int s = 0; s < ns; ++s){
	//#pragma omp parallel for
	for(int r = 0; r < nr; ++r){
	    // Start of row pointers
	    value_type* p = &grad(s, r, 0);
	    tm_base_type* t = &T_(s, r, 0);
	    for(int c = 0; c < nc; ++c)
		DJ_.segment(manifold_dim * (nr * nc * s + nc * r + c), manifold_dim) = t[c].transpose() * Eigen::Map<const Eigen::VectorXd>(p[c].data(), p[c].size());
	//	DJ_.segment(manifold_dim * (s + ns * r + ns * nr * c), manifold_dim) = t[c].transpose() * Eigen::Map<const Eigen::VectorXd>(p[c].data(), p[c].size());
	}
    }

    #ifdef TV_FUNC_DEBUG
	std::fstream f;
	f.open("3dDJ.csv",std::fstream::out);
	f << DJ_;
	f.close();
    #endif
}

// Evaluation of Hessian J
template <enum FUNCTIONAL_DISC disc, typename MANIFOLD, class DATA >
void Functional<FIRSTORDER, disc, MANIFOLD, DATA, 3 >::evaluateHJ(){
    hessian_type hessian(data_.img_.domain());

    int ns = data_.img_.nslices();
    int nr = data_.img_.nrows();
    int nc = data_.img_.ncols();
    int sparsedim = ns*nr*nc*manifold_dim;

    //HESSIAN OF FIDELITY TERM
     #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\tHessian evaluation..." << std::endl;
	std::cout << "\t\t...Fidelity part" << std::endl;
    #endif

    sparse_hessian_type HF(sparsedim,sparsedim);

    //HF.reserve(Eigen::VectorXi::Constant(nc,manifold_dim));
    typedef Eigen::Triplet<double> Trip;
    std::vector<Trip> triplist;
    triplist.reserve(sparsedim*manifold_dim);

    if(data_.doInpaint()){
	auto f = [] (deriv2_type& h, const value_type& i, const value_type& n, const bool inp ) { MANIFOLD::deriv2xx_dist_squared(i,n,h); h*=(1-inp); };
	pixel_wise3d(f, hessian, data_.img_, data_.noise_img_, data_.inp_);
    }
    else{
	auto f = [] (deriv2_type& h, const value_type& i, const value_type& n) { MANIFOLD::deriv2xx_dist_squared(i,n,h); };
	pixel_wise3d(f, hessian, data_.img_, data_.noise_img_);
    }

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...Local to global insert" << std::endl;
    #endif
   for(int s = 0; s < ns; ++s){
    for(int r = 0; r < nr; ++r){
	    // Start of row pointers
	    deriv2_type* h = &hessian(s, r, 0);
	    tm_base_type* t = &T_(s, r, 0);
	    for(int c = 0; c < nc; ++c){
		int pos = manifold_dim * (nr * nc * s + nc * r + c); // rowwise flattening
		//int pos = manifold_dim * (s + ns * r + ns * nr * c); // columnwise flattening
		restricted_deriv2_type ht=t[c].transpose()*h[c]*t[c];

		for(int local_row = 0; local_row<ht.rows(); local_row++)
		    for(int local_col = local_row; local_col < ht.cols(); local_col++){
			scalar_type e = ht(local_row, local_col);
			    if(e!=0){
				int global_row = pos + local_row;
				int global_col = pos + local_col;
				triplist.push_back(Trip(global_row,global_col,e));
				if(global_row != global_col)
				    triplist.push_back(Trip(global_col,global_row,e));
			    }
		    }

	    }
	}
    }

     #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...Triplet list created" << std::endl;
    #endif
    HF.setFromTriplets(triplist.begin(),triplist.end());
    HF.makeCompressed();

    //HESSIAN OF TV TERM
    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\tHessian evaluation..." << std::endl;
	std::cout << "\t\t...TV part" << std::endl;
    #endif

    sparse_hessian_type HTV(sparsedim,sparsedim);

    //HTV.reserve(Eigen::VectorXi::Constant(nc,5*manifold_dim));
    triplist.clear();
    triplist.reserve(7 * sparsedim*manifold_dim);

    // Subimage boxes
    vpp::box3d without_last_x(vpp::vint3(0,0,0), vpp::vint3(ns - 1, nr - 1, nc - 2)); // subdomain without last xslice
    vpp::box3d without_last_y(vpp::vint3(0,0,0), vpp::vint3(ns - 1, nr - 2, nc - 1)); // subdomain without last yslice
    vpp::box3d without_last_z(vpp::vint3(0,0,0), vpp::vint3(ns - 2, nr - 1, nc - 1)); // subdomain without last zslice
    vpp::box3d without_first_x(vpp::vint3(0,0,1), vpp::vint3(ns - 1, nr - 1, nc - 1)); // subdomain without first xlice
    vpp::box3d without_first_y(vpp::vint3(0,1,0), vpp::vint3(ns - 1, nr - 1, nc - 1)); // subdomain without first yslice
    vpp::box3d without_first_z(vpp::vint3(1,0,0), vpp::vint3(ns - 1, nr - 1, nc - 1)); // subdomain without first zslice

    auto calc_xx_der = [&] (deriv2_type& x, const weights_type& w, const value_type& i, const value_type& n) { MANIFOLD::deriv2xx_dist_squared(i, n, x); x*=w; };
    auto calc_xy_der = [&] (deriv2_type& xy, const weights_type& w, const value_type& i, const value_type& n) { MANIFOLD::deriv2xy_dist_squared(i, n, xy); xy*=w; };
    auto calc_yy_der = [&] (deriv2_type& y, const weights_type& w, const value_type& i, const value_type& n) { MANIFOLD::deriv2yy_dist_squared(i, n, y); y*=w; };
    auto add_to_hessian =  [&] (deriv2_type& h, const deriv2_type& d) { h += d; };

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...->XD11" << std::endl;
    #endif

    // Horizontal Second Derivatives and weighting
    // ... w.r.t. first arguments
    { // Temporary image XD11 is deallocated after this scope
	hessian_type XD11(without_last_x);
        pixel_wise3d(calc_xx_der, XD11, weightsX_ | without_last_x, data_.img_ | without_last_x, data_.img_ | without_first_x);
	pixel_wise3d([&] (deriv2_type& h, const deriv2_type& d) { h=d; }, hessian | without_last_x, XD11);
	#ifdef TV_FUNC_DEBUG
	    data_.output_matval_img(XD11,"3dXD11.csv");
	#endif
    }

    for(int s=0; s < ns; ++s)
	for(int r=0; r<nr; ++r)
	    hessian(s,r,nc-1)=deriv2_type::Zero();

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...->XD22" << std::endl;
    #endif
    //... w.r.t. second arguments
    {
	hessian_type XD22(without_last_x);
        pixel_wise3d(calc_yy_der, XD22, weightsX_ | without_last_x, data_.img_ | without_last_x, data_.img_ | without_first_x);
	pixel_wise3d(add_to_hessian, hessian | without_first_x, XD22);
	#ifdef TV_FUNC_DEBUG
	    data_.output_matval_img(XD22,"3dXD22.csv");
	#endif
    }
    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...->YD11" << std::endl;
    #endif
    // Vertical Second Derivatives weighting
    //... w.r.t. first arguments
    {
	hessian_type YD11(without_last_y);
        pixel_wise3d(calc_xx_der, YD11, weightsY_ | without_last_y, data_.img_ | without_last_y, data_.img_ | without_first_y);
	pixel_wise3d(add_to_hessian, hessian | without_last_y, YD11);
	#ifdef TV_FUNC_DEBUG
	    data_.output_matval_img(YD11,"3dYD11.csv");
	#endif
    }

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...->YD22" << std::endl;
    #endif
    //... w.r.t. second arguments
    {
	hessian_type YD22(without_last_y);
        pixel_wise3d(calc_yy_der, YD22, weightsY_ | without_last_y, data_.img_ | without_last_y, data_.img_ | without_first_y);
	pixel_wise3d(add_to_hessian, hessian | without_first_y, YD22);
	#ifdef TV_FUNC_DEBUG
	    data_.output_matval_img(YD22,"3dYD22.csv");
	#endif
    }
    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...->ZD11" << std::endl;
    #endif
    // Vertical Second Derivatives weighting
    //... w.r.t. first arguments
    {
	hessian_type ZD11(without_last_z);
        pixel_wise3d(calc_xx_der, ZD11, weightsZ_ | without_last_z, data_.img_ | without_last_z, data_.img_ | without_first_z);
	pixel_wise3d(add_to_hessian, hessian | without_last_z, ZD11);
	#ifdef TV_FUNC_DEBUG
	    data_.output_matval_img(ZD11,"3dZD11.csv");
	#endif
    }

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...->ZD22" << std::endl;
    #endif
    //... w.r.t. second arguments
    {
	hessian_type ZD22(without_last_z);
        pixel_wise3d(calc_yy_der, ZD22, weightsZ_ | without_last_z, data_.img_ | without_last_z, data_.img_ | without_first_z);
	pixel_wise3d(add_to_hessian, hessian | without_first_z, ZD22);
	#ifdef TV_FUNC_DEBUG
	    data_.output_matval_img(ZD22,"3dZD22.csv");
	#endif
    }
    	#ifdef TV_FUNC_DEBUG
	    data_.output_matval_img(hessian,"3dhessian.csv");
	#endif
    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...Local to global insert" << std::endl;
    #endif
    // Insert elementwise into sparse Hessian
    // NOTE: Eventually make single version for both cases, including an offset
    // --> additional parameters sparse_mat, offset
    auto HTV_insert = [&]( hessian_type H, tm_base_mat_type T1, tm_base_mat_type T2, int row_offset, int col_offset){
	int Hns = H.nslices();
	int Hnr = H.nrows();
	int Hnc = H.ncols();

	for(int s = 0; s < Hns; ++s){
	    for(int r = 0; r < Hnr; ++r){
		// Start of row pointers
		deriv2_type* h = &H(s, r, 0);
		tm_base_type* t1 = &T1(s, r, 0);
		tm_base_type* t2 = &T2(s, r, 0);
		for(int c = 0; c < Hnc; ++c){
		    int pos = manifold_dim * (nr * nc * s + nc * r +c); // rowwise flattening
		   // int pos = manifold_dim * (s + ns * r + ns * nr * c); // columnwise flattening
		    restricted_deriv2_type ht = t1[c].transpose()*h[c]*t2[c];

		    for(int local_row = 0; local_row<ht.rows(); local_row++)
			for(int local_col = 0; local_col < ht.cols(); local_col++){
			    scalar_type e = ht(local_row, local_col);
				if(e!=0){
				    int global_row = pos + row_offset + local_row;
				    int global_col = pos + col_offset + local_col;
				    triplist.push_back(Trip(global_row,global_col,e));
				    if(row_offset > 0 || col_offset > 0)
					triplist.push_back(Trip(global_col,global_row,e));
				}
			}

		}
	    }
	}
    };

    auto HTV_single_insert = [&] (deriv2_type& h, tm_base_type& t1, tm_base_type& t2, int pos, int row_offset, int col_offset){
   		    restricted_deriv2_type ht = t1.transpose()*h*t2;
		    for(int local_row = 0; local_row<ht.rows(); local_row++)
			for(int local_col = 0; local_col < ht.cols(); local_col++){
			    scalar_type e = ht(local_row, local_col);
				if(e!=0){
				    int global_row = pos + row_offset + local_row;
				    int global_col = pos + col_offset + local_col;
				    triplist.push_back(Trip(global_row,global_col,e));
				    if(row_offset > 0 || col_offset > 0)
					triplist.push_back(Trip(global_col,global_row,e));
				}
			}
    };

    HTV_insert(hessian, T_, T_, 0, 0);

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...->ZD12" << std::endl;
    #endif
    // Z Neighbors Second Derivatives and weighting
    // ... w.r.t. first and second arguments
    {
	hessian_type ZD12(without_last_z);
	pixel_wise3d(calc_xy_der, ZD12, weightsZ_ | without_last_z, data_.img_ | without_last_z, data_.img_ | without_first_z );
	#ifdef TV_FUNC_DEBUG
	    data_.output_matval_img(ZD12,"3dZD12.csv");
	#endif
	#ifdef TV_FUNC_DEBUG_VERBOSE
		std::cout << "\t\t...Local to global insert:" << std::endl;
	#endif
	// Offsets for upper nx*ny-th subdiagonal
	int offset =  manifold_dim * nr * nc;
	HTV_insert(ZD12, T_ | without_last_z, T_ | without_first_z, 0, offset);
    }

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...->YD12" << std::endl;
    #endif
    // Y Neighbors Second Derivatives and weighting
    //... w.r.t. second arguments
    {
	hessian_type YD12(data_.img_.domain());
	pixel_wise3d(calc_xy_der, YD12 | without_last_y, weightsY_ | without_last_y, data_.img_ | without_last_y, data_.img_ | without_first_y );

	//Set last y-slice to zero
	#pragma omp parallel for
	for(int s = 0; s < ns ; ++s){
	    deriv2_type* row = &YD12(s, nr-1 ,0);
	    for(int c = 0; c < nc; ++c)
		row[c] = deriv2_type::Zero();
	}
	#ifdef TV_FUNC_DEBUG
	    data_.output_matval_img(YD12,"3dYD12.csv");
	#endif
	#ifdef TV_FUNC_DEBUG_VERBOSE
		std::cout << "\t\t...Local to global insert:" << std::endl;
	#endif
	// Offsets for first nxth subdiagonal
	int offset =  manifold_dim * nc;

	// Insert YD12 except last ns entries
	auto t1_it = T_.begin();
	auto t2_it = typename tm_base_mat_type::iterator(vpp::vint3(0,1,0), T_); // Iterator at beginning of second row
	bool do_break = false;
	int k = 0;
	int max_entry = ns * nr * nc - nc;

	for(int s=0; s < ns; ++s){
	    for(int r=0; r < nr; ++r){
		 deriv2_type* row = &YD12(s, r, 0);
		 for(int c=0; c < nc; ++c){
		    int pos = manifold_dim * (nr * nc * s + nc * r +c); // rowwise flattening
	//	    int pos = manifold_dim * (s + ns * r + ns * nr * c); // columnwise flattening
		    HTV_single_insert(row[c], *t1_it, *t2_it, pos, 0, offset);
		    t1_it.next();
		    t2_it.next();
		    #ifdef TV_FUNC_DEBUG_VERBOSE2
			std::cout << "\n\t\t\tPosition: (" << s << ","<< c << "," << r << ")" << std::endl;
			std::cout << "\t\t\tValue(YD12): " << row[c] << std::endl;
			std::cout << "\t\t\tEntry number: " << k << std::endl;
			std::cout << "\t\t\tPos: " << pos << std::endl;
		    #endif
		    ++k;
		    if(k >= max_entry) {do_break = true; break;}
		 }
		 if(do_break) break;
	    }
	    if(do_break) break;
	}

    }

    #ifdef TV_FUNC_DEBUG_VERBOSE
	std::cout << "\t\t...->XD12" << std::endl;
    #endif
    // Z neightbors Second Derivatives and weighting
    //... w.r.t. second arguments
    {
	hessian_type XD12(data_.img_.domain());
	pixel_wise3d(calc_xy_der, XD12 | without_last_x, weightsX_ | without_last_x, data_.img_ | without_last_x, data_.img_ | without_first_x);

	//Set last slice to zero, - actually not necessary but safer than keeping uninitialized data
	#pragma omp parallel for
	for(int s = 0; s < ns; ++s){
	    for(int r = 0; r < nr; ++r)
		XD12(s, r, nc-1) = deriv2_type::Zero();
	}

	#ifdef TV_FUNC_DEBUG
	    data_.output_matval_img(XD12,"3dXD12.csv");
	#endif
	#ifdef TV_FUNC_DEBUG_VERBOSE
		std::cout << "\t\t...Local to global insert:" << std::endl;
	#endif

	// Insert into first upper subdiagonal
	auto t1_it = T_.begin();
	auto t2_it = T_.begin(); t2_it.next();
	for(int s=0; s < ns; ++s){
	    for(int r=0; r < nr; ++r){
		 deriv2_type* row = &XD12(s, r, 0);
		 for(int c=0; c < nc; ++c){
		    int pos = manifold_dim * (nr * nc * s + nc * r +c); // rowwise flattening
		    //int pos = manifold_dim * (s + ns * r + ns * nr * c); // columnwise flattening
		    HTV_single_insert(row[c], *t1_it, *t2_it, pos, 0, manifold_dim);
		    t1_it.next();
		    t2_it.next();
		    if(s==ns-1 && r == nc-1 && c >= nc-2 ) break;
		 }
	    }
	}

	HTV.setFromTriplets(triplist.begin(),triplist.end());
	HTV.makeCompressed();
	triplist.clear();
    }
	#ifdef TV_FUNC_DEBUG_VERBOSE
		std::cout << "\t\t...Combine Fidelity and TV parts:" << std::endl;
	#endif

	HJ_= HF + lambda_*HTV;

	#ifdef TV_FUNC_DEBUG_VERBOSE
		std::cout << "\t\t...Output Hessian (stats):" << std::endl;
	#endif
    #ifdef TV_FUNC_DEBUG
	if (sparsedim<200){
	    if(sparsedim<80){
		std::cout << "\nFidelity\n" << HF << std::endl;
		std::cout << "\nTV\n" << HTV << std::endl;
		std::cout << "\nHessian\n" << HJ_ << std::endl;
	    }

	    std::fstream f;
	    f.open("H.csv",std::fstream::out);
	    Eigen::IOFormat CommaInitFmt(Eigen::StreamPrecision, Eigen::DontAlignCols, ", ", ", ", "", "", "", "");
	    //Eigen::MatrixXd M(HJ_);
	    //f << M.format(CommaInitFmt);
	    f << HJ_;
	    f.close();

	}
	else{
	    std::cout << "\nFidelity Non-Zeros: " << HJ_.nonZeros() << std::endl;
	    std::cout << "\nTV Non-Zeros: " << HJ_.nonZeros() << std::endl;
	    std::cout << "\nHessian Non-Zeros: " << HJ_.nonZeros() << std::endl;
	}
    // Test Solver:
	#ifdef TV_FUNC_DEBUG_VERBOSE
		std::cout << "\t\t...Test Solve" << std::endl;
	#endif
	gradient_type x;

	Eigen::SparseLU<sparse_hessian_type> solver;
	solver.analyzePattern(HJ_);
	solver.factorize(HJ_);
	x = solver.solve(DJ_);

	std::fstream f;
	f.open("Sol.csv",std::fstream::out);
	f << x;
	f.close();
    #endif
}

template <enum FUNCTIONAL_DISC disc, typename MANIFOLD, class DATA >
template < class IMG >
void Functional<FIRSTORDER, disc, MANIFOLD, DATA, 3 >::output_img(const IMG& img, const char* filename) const{
    //TODO
}

template <enum FUNCTIONAL_DISC disc, typename MANIFOLD, class DATA >
template < class IMG >
void Functional<FIRSTORDER, disc, MANIFOLD, DATA, 3 >::output_matval_img(const IMG& img, const char* filename) const{
    //TODO
}

}// end namespace tvtml

#endif


/**********************************************************************************
 *  The MIT License (MIT)                                                        *
 *                                                                                *
 *  Copyright (c) 2014 Carnegie Mellon University                                 *
 *                                                                                *
 *  Permission is hereby granted, free of charge, to any person obtaining a copy  *
 *  of this software and associated documentation files (the "Software"), to deal *
 *  in the Software without restriction, including without limitation the rights  *
 *  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell     *
 *  copies of the Software, and to permit persons to whom the Software is         *
 *  furnished to do so, subject to the following conditions:                      *
 *                                                                                *
 *  The above copyright notice and this permission notice shall be included in    *
 *  all copies or substantial portions of the Software.                           *
 *                                                                                *
 *  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR    *
 *  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,      *
 *  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE   *
 *  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER        *
 *  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, *
 *  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN     *
 *  THE SOFTWARE.                                                                 *
 **********************************************************************************/
/**
 * @file   cpdr.cpp
 * @brief  A pedestrian dead-reckoning system.  Integrates gyros and step counting
 * @author M. George
 */

#include <cmath>
#include <iostream>

#include <Eigen/Dense>

#include "cpdr.h"

using namespace BASIL;

/**************************************************************************//**
* cPDR Constructor.  Constant vectors and matrices (gravity,
* earth rate) are formed here.  PDR state is set to ALIGN initially.  PDR
* cannot be used until an alignment method is called.
* \arg \c params.  Boost::property_tree object and id of attached IMU.
******************************************************************************/
cPDR::cPDR(boost::property_tree::ptree config, unsigned int imuid) : mStepDetector(config)
{
    mImuId = imuid;

    mMessageCounter = 0;

    // Clear all the member vectors and matrices.
    mEarthRate           = CartesianVector::Zero();
    mGravity             = CartesianVector::Zero();
    mAlpha               = CartesianVector::Zero();
    mDeltaP              = CartesianVector::Zero();
    mDeltaV              = CartesianVector::Zero();
    mGravityReaction     = CartesianVector::Zero();

    mIdentity = CartesianMatrix::Identity();
    mA        = CartesianMatrix::Zero();

    mAlphaSkew     = SkewMatrix::Zero();
    mAlphaMagnitude = 0.0;

    mf1 = mf2 = mf3 = mf4 = 0.0;

    mCbn = RotationMatrix::Identity();

    mdt = 0;

    // INS initially stopped, time = 0, no messages and ready to calibrate.
    mTime = 0.0;
    mPreviousTime = 0.0;
    mPreviousEarthRateTime = 0.0;
    mAlignmentTime = config.get<double>("parameters.ins.alignmentperiod");
    mIMUCount = 0;
    mIMURate = config.get<unsigned int>("parameters.hardware.imus.xsensmti.rate");
    mEarthUpdateRate = config.get<unsigned int>("parameters.ins.earthupdaterate");

    if (mEarthUpdateRate < 0)
    {
        std::cerr << "Earth update rate must be >= 0" << std::endl;
        mApplyEarthRateCorrections = false;
        mEarthSubsample = 0;
    }
    else if (mEarthUpdateRate == 0)
    {
        mApplyEarthRateCorrections = false;
        mEarthSubsample = 0;
    }
    else
    {
        mApplyEarthRateCorrections = true;
        mEarthSubsample = mIMURate / mEarthUpdateRate;
    }

    // Get initial conditions
    mInitialHeading = config.get<double>("parameters.ins.initialheading")*BASIL::CONSTANTS::DEG_TO_RAD;
    mInitialLatitude = config.get<double>("parameters.ins.approximatelatitude")*BASIL::CONSTANTS::DEG_TO_RAD;

    // Step data
    mStepLength = config.get<double>("parameters.pdr.steplength");
    mStepIncrement = CartesianVector::Zero();

    // Clear initial INS data structure
    mInsData.Velocity    = CartesianVector::Zero();
    mInsData.Position    = Array::Zero();
    mInsData.Orientation = AngleArray::Zero();
    mInsData.InsMdl = BASIL::INS::GENERIC_INS;
    mInsData.Id = imuid;

    // Clear IMU calibration data structure
    mImuCalibrationData.AccelerationBias = CartesianVector::Zero();
    mImuCalibrationData.AccelerationScaleFactor = CartesianVector::Zero();
    mImuCalibrationData.AngularRateBias << 0.00414044772524217,0.00189662656009109,0.00726575764268722;
    mImuCalibrationData.AngularRateGBias = CartesianVector::Zero();
    mImuCalibrationData.AngularRateScaleFactor = CartesianVector::Zero();
    mImuCalibrationData.ImuMdl = BASIL::IMUS::GENERIC_IMU;
    mImuCalibrationData.Id = imuid;

    mCurrentState = ALIGNING;
}

/**************************************************************************//**
* cPDR Destructor
******************************************************************************/
cPDR::~cPDR()
{
}

/**************************************************************************//**
* Integrates the attitude forward using Savage type exact
* integration. Ref: Savage, P.  Strapdown Analytics, 2nd Ed.  Chapters 7, 19
* Steps are detected and position is dead-reckoned using heading data and
* average step length.
* \arg \c omega.  Vector of x,y,z gyro rate readings in rad/s
* \arg \c accel.  Vector of x,y,z accelerometer readings in m/s/s
* \arg \c time.   Time stamp of IMU readings
* \return \c void
******************************************************************************/
void cPDR::integratePVA(const cImuData& imudata)
{
    if (mCurrentState < cPDR::RUNNING)
    {
        std::cerr << "Error (cPDR): INS must be aligned before integrating." << std::endl;
	    return;
    }

    // For the first message we receive predict using nominal IMU rate
    if (mMessageCounter == 0)
    {
        mdt = 1.0/double(mIMURate);
    }
    else
    {
        mdt = imudata.Time.getSeconds() - mPreviousTime;
    }

    // Catch an out of order packet here
    if (mdt < 0)
    {
        //Some kind of standard error here
        std::cerr << "Error (cPDR): dt is negative in integratePVA()" << std::endl;
        return;
    }

    if (mdt > 2.0/double(mIMURate))
    {
        // Some kind of standard error here
        std::cerr << "Warning (cPDR): dt (" << mdt << "s) is large in integratePVA()" << std::endl;
        // Reset the counters so this doesn't persist.
        mInsData.Time = imudata.Time;
        mPreviousTime = imudata.Time.getSeconds();
    	mMessageCounter++;
        //  Returning here is equivalent to assuming no motion during period of missing data.
        return;
    }

    // Get rotation vector
    mAlpha = imudata.AngularRate*mdt;
    mAlphaMagnitude = mAlpha.norm();

    // Catch zero rotation vector magnitudes which happen with stationary sensors and quantization
    if (mAlphaMagnitude > 0.0)
    {
        // Skew symmetric matrix from rotation vector
        BASIL::MATH::skew(mAlpha, mAlphaSkew);

        // Integration constants
        mf1 = std::sin(mAlphaMagnitude)/mAlphaMagnitude;
        mf2 = (1.0 - std::cos(mAlphaMagnitude))/(mAlphaMagnitude*mAlphaMagnitude);
        mf3 = (1.0/(mAlphaMagnitude*mAlphaMagnitude))*(1.0-mf1);
        mf4 = (1.0/(mAlphaMagnitude*mAlphaMagnitude))*(0.5-mf2);

	    // Position update
	    if (mStepDetector.detect(imudata.Acceleration(BASIL::IMUZ)))
	    {
	        mStepIncrement(BASIL::NORTH) = mStepLength * std::cos(mInsData.Orientation(BASIL::YAW));
	        mStepIncrement(BASIL::EAST) = mStepLength * std::sin(mInsData.Orientation(BASIL::YAW));
	        mStepIncrement(BASIL::DOWN) = 0.0;
            mInsData.Position.noalias() += mStepIncrement;
            //std::cout << mMessageCounter << std::endl;
        }

        // Orientation update
        mA.noalias() = mIdentity + mf1*mAlphaSkew + mf2*mAlphaSkew*mAlphaSkew;
        mCbn *= mA;
        cbn2euler(mCbn,mInsData.Orientation);
    }
    else
    {
        // Position update
	    if (mStepDetector.detect(imudata.Acceleration(BASIL::IMUZ)))
	    {
	        mStepIncrement(BASIL::NORTH) = mStepLength * std::cos(mInsData.Orientation(BASIL::YAW));
	        mStepIncrement(BASIL::EAST) = mStepLength * std::sin(mInsData.Orientation(BASIL::YAW));
	        mStepIncrement(BASIL::DOWN) = 0.0;
            mInsData.Position.noalias() += mStepIncrement;
            //std::cout << mMessageCounter << std::endl;
        }
    }

    // Apply an earth rate correction at a lower frequency
    if ( mApplyEarthRateCorrections && (mMessageCounter % mEarthSubsample == 0) )
    {
        cPDR::earthRateCorrection(imudata.Time.getSeconds());
    }

    // Times are set here so that a user can bypass position and velocity integration
    // if desired, effectively using this class as a AHRS system with gyros only.
    mInsData.Time = imudata.Time;
    mPreviousTime = imudata.Time.getSeconds();

    mMessageCounter++;
}

/**************************************************************************//**
* Called at sub-intervals of the integreatePVA() function.  Adds correction
* to current attitude to account for Earth rotation effects.  Position and velocity
* corrections are taken into account in integratePVA() and as a natural consequence
* of this update.  Not used if no latitude input is used in alignment call.
* \arg \c void
* \return \c void
******************************************************************************/
void cPDR::earthRateCorrection(const double time)
{
    // First time in we use nominal earth correction rate for the dt.
    if (mPreviousEarthRateTime == 0.0)
    {
        mdt = 1.0 / (double)mEarthSubsample;
    }
    else
    {
        mdt = time - mPreviousEarthRateTime;
    }

    // Catch negative time increments
    if (mdt < 0)
    {
        //Some kind of standard error here
        std::cerr << "Error (cPDR): dt is negative in earthRateCorrection()" << std::endl;
        return;
    }

    mAlpha = -1.0*mdt*mEarthRate;
    mAlphaMagnitude = mAlpha.norm();
    mf1 = std::sin(mAlphaMagnitude)/mAlphaMagnitude;
    mf2 = (1.0 - std::cos(mAlphaMagnitude))/(mAlphaMagnitude*mAlphaMagnitude);
    BASIL::MATH::skew(mAlpha, mAlphaSkew);

    mA.noalias() = mIdentity + mf1*mAlphaSkew + mf2*mAlphaSkew*mAlphaSkew;
    mCbn = mA*mCbn;
    cbn2euler(mCbn,mInsData.Orientation);

    mPreviousEarthRateTime = time;
}

/**************************************************************************//**
* Local initial alignment with velocity and position in local navigation coordinates.
* Roll and pitch angles are coarsely aligned from an input average acceleration vector
* (2-3s recommended) with gravity.  Initial heading (North), velocity and position (Local)
* are user inputs.  User must also input an approximate latitude.  The global heading and
* approximate latitude allow reasonable gravity models and earth rate corrections to be fixed.
*\arg \c Pointer to 3x1 array of time averaged x,y,z accelerometer signals.
* 		 Suggested averaging time 2-3s.  IMU must be stationary during this time!
* \return \c void
******************************************************************************/
void cPDR::localAlign(const CartesianVector& gravityReactionForce,  const double heading, const CartesianVector& initialUTMPosition,
                      const CartesianVector& initialUTMVelocity, const double approximateLatitude)
{
    gravityAlignRollPitch(gravityReactionForce, mInsData.Orientation);
    mInsData.Orientation(BASIL::YAW) = heading;

    cPDR::euler2cbn(mInsData.Orientation,mCbn);

    mInsData.Velocity = initialUTMVelocity;
    mInsData.Position = initialUTMPosition;

    // Earth rate:
    getEarthRate(approximateLatitude, mEarthRate);

    // Gravity
    mGravity(BASIL::DOWN) = getGravityMagnitude(approximateLatitude);

    mMode = cPDR::LOCAL;
    // Earth rate corrections will be applied unless already turned off in constructor
    mApplyEarthRateCorrections = mApplyEarthRateCorrections && true;

    // Zero counters in case we need to re-align at some point
    mMessageCounter = 0;
    mdt = 0;
    mInsData.Time.setTime(TIME::FREE_RUNNING,0,0.0);
    mPreviousTime = 0.0;
    mPreviousEarthRateTime = 0.0;

    // We are now running and ready to call integrate*() functions
    mCurrentState = cPDR::RUNNING;
}

/**************************************************************************//**
* Local initial alignment with of roll and pitch angles only.
* Roll and pitch angles are coarsely aligned from an input average acceleration vector
* (2-3s recommended) with gravity.  Initial heading, velocity and position are all set to 0.
* User must also input an approximate latitude to allow a coarse estimate of gravity.
*\arg \c Pointer to 3x1 array of time averaged x,y,z accelerometer signals.
* 		 Suggested averaging time 2-3s.  IMU must be stationary during this time!
* \return \c void
******************************************************************************/
void cPDR::localAlign(const CartesianVector& gravityReactionForce, const double approximateLatitude)
{
    gravityAlignRollPitch(gravityReactionForce, mInsData.Orientation);

    euler2cbn(mInsData.Orientation,mCbn);

    // Earth rate:
    getEarthRate(approximateLatitude, mEarthRate);

    // Gravity
    mGravity(BASIL::DOWN) = getGravityMagnitude(approximateLatitude);

    mMode = cPDR::LOCAL;
    // Earth rate corrections will be applied unless already turned off in constructor
    mApplyEarthRateCorrections = mApplyEarthRateCorrections && true;

    // Zero counters in case we need to re-align at some point
    mMessageCounter = 0;
    mdt = 0;
    mInsData.Time.setTime(TIME::FREE_RUNNING,0,0.0);
    mPreviousTime = 0.0;
    mPreviousEarthRateTime = 0.0;

    // We are now running and ready to call integrate*() functions
    mCurrentState = cPDR::RUNNING;
}

/**************************************************************************//**
* Get the magnitude of gravity at our current latitude using a Defense Mapping
* Agency model.  Reference:  Jay A. Farrell.  Aided Inertial Navigation, p.33.
* \arg \c latitude.  Reference to latitude data (radians)
* \arg \c gravity.  Reference to gravity output.
* \return \c void
******************************************************************************/
double cPDR::getGravityMagnitude(const double latitude)
{
    double gravity = BASIL::WGS84::EQUATORIAL_GRAVITY *
                    ( (          1.0 + 0.0019318513530*std::sin(latitude)*std::sin(latitude)) /
                      (std::sqrt(1.0 - 0.0066943800229*std::sin(latitude)*std::sin(latitude))) );
    return gravity;
}

/**************************************************************************//**
* Correct the current position, velocity and orientation with errors.
* \arg \c position.  Position.
* \arg \c velocity.  Velocity.
* \arg \c eulers.  Euler angles.
* \return \c void
******************************************************************************/
void cPDR::correctPVA(const CartesianVector& positionerror,
                      const CartesianVector& velocityerror,
                      const AngleArray& skewangles)
{
    mInsData.Velocity -= velocityerror;
    mInsData.Position -= positionerror;

    BASIL::MATH::skew(skewangles, mAlphaSkew);
    mA = mIdentity - mAlphaSkew;
    mCbn = mA.inverse() * mCbn;

    cbn2euler(mCbn, mInsData.Orientation);
}

/**************************************************************************//**
* Get the components of earth rate at our current latitude.
* \arg \c latitude.  Reference to latitude data (radians)
* \arg \c gravity.  Pointer to earthRate vector for output.
* \return \c void
******************************************************************************/
void cPDR::getEarthRate(const double latitude, CartesianVector& earthRate)
{
    earthRate(NORTH) = BASIL::WGS84::EARTH_RATE * std::cos(latitude);
    earthRate(EAST) = 0;
    earthRate(DOWN) = -1.0 * BASIL::WGS84::EARTH_RATE*std::sin(latitude);
}

/**************************************************************************//**
* Get roll and pitch angles from a gravity vector.  Input vector should be averaged
* over a period of 2-3 seconds with a stationary IMU.  Aligns IMU x,y coordinates
* to Down gravity vector.  1 mg of accelerometer bias -> 1 mrad of roll, pitch error.
* Calibrate your accelerometer before averaging if possible.
* \arg \c latitude.  Reference to latitude data (radians)
* \arg \c gravity.  Pointer to earthRate vector for output.
* \return \c void
******************************************************************************/
void cPDR::gravityAlignRollPitch(const CartesianVector& gravityReactionForce, AngleArray& angles)
{
    angles(ROLL) = std::atan2(-gravityReactionForce(BODYY),
                              -gravityReactionForce(BODYZ));
    angles(PITCH) = std::atan2(gravityReactionForce(BASIL::BODYX),
                                           std::sqrt((gravityReactionForce(BODYY) * gravityReactionForce(BODYY))
                                                    +(gravityReactionForce(BODYZ) * gravityReactionForce(BODYZ))));
}

/**************************************************************************//**
* Calculate a direction cosine matrix, Cbm, relating body frame to global North,
* East, Down frame from an Euler angle vector.  Cbm is propagated directly by
* integrate so this function is most useful at alignment time for initial Cbm.
* \arg \c Cbm.  Empty Cbm matrix to be filled.
* \arg \c type.  Euler angle vector.
* \return \c void
******************************************************************************/
void cPDR::euler2cbn(const AngleArray& eulers, CartesianMatrix& Cbn)
{
    double roll = eulers(0);
    double pitch = eulers(1);
    double yaw = eulers(2);

    Cbn << std::cos(yaw)*std::cos(pitch),
         std::cos(yaw)*std::sin(pitch)*std::sin(roll) - std::sin(yaw)*std::cos(roll),
         std::cos(yaw)*std::sin(pitch)*std::cos(roll) + std::sin(roll)*std::sin(yaw),
         std::cos(pitch)*std::sin(yaw),
         std::cos(roll)*std::cos(yaw) + std::sin(roll)*std::sin(pitch)*std::sin(yaw),
         std::cos(roll)*std::sin(pitch)*std::sin(yaw) - std::sin(roll)*std::cos(yaw),
         -std::sin(pitch),
         std::sin(roll)*std::cos(pitch),
         std::cos(roll)*std::cos(pitch);
}

/**************************************************************************//**
* Set the Euler angles from the current Cbm value
* \arg \c Cbm.  Current Cbm matrix.
* \return \c void
******************************************************************************/
void cPDR::cbn2euler(const RotationMatrix& Cbn, AngleArray& eulers)
{
    eulers(0) = std::atan2(Cbn(2,1),Cbn(2,2));
    eulers(1) = std::asin(-Cbn(2,0));
    eulers(2) = std::atan2(Cbn(1,0),Cbn(0,0));
}

/**************************************************************************//**
* Copy the current Cbm into the input reference.
* \arg \c out.  Empty Cbm matrix to be filled.
* \return \c void
******************************************************************************/
void cPDR::getCbn(RotationMatrix& out)
{
    out = mCbn;
}

/**************************************************************************//**
* Copy the current euler angles into the input reference.
* \arg \c out.  Empty position vector to be filled.
* \return \c void
******************************************************************************/
void cPDR::getEulers(AngleArray& out)
{
    cbn2euler(mCbn,out);
}

/**************************************************************************//**
* Copy the current velocity into the input reference.
* \arg \c out.  Empty velocity vector to be filled.
* \return \c void
******************************************************************************/
void cPDR::getVelocity(CartesianVector& out)
{
    out = mInsData.Velocity;
}

/**************************************************************************//**
* Copy the current position into the input reference.
* \arg \c out.  Empty position vector to be filled.
* \return \c void
******************************************************************************/
void cPDR::getPosition(CartesianVector& out)
{
    out = mInsData.Position;
}

/**************************************************************************//**
* Copy the current gravity vector into the input reference.
* \arg \c out.  Empty gravity vector to be filled.
* \return \c void
******************************************************************************/
void cPDR::getGravity(CartesianVector& out)
{
    out = mGravity;
}

/**************************************************************************//**
* Copy the current earth rate vector into the input reference.
* \arg \c out.  Empty earth rate vector to be filled.
* \return \c void
******************************************************************************/
void cPDR::getEarthRate(CartesianVector& out)
{
    out = mEarthRate;
}

/**************************************************************************//**
* Copy the current position, velocity and euler angles out.
* \arg \c position.  Empty position vector to be filled.
* \arg \c velocity.  Empty velocity vector to be filled.
* \arg \c eulers.  Empty euler angles vector to be filled.
* \return \c void
******************************************************************************/
void cPDR::getPVA(CartesianVector& position, CartesianVector& velocity, CartesianVector& eulers)
{
    position = mInsData.Position;
    velocity = mInsData.Velocity;
    cbn2euler(mCbn,eulers);
}

/**************************************************************************//**
* Set the current Cbm from the input field.
* \arg \c out.  Cbm matrix input.
* \return \c void
******************************************************************************/
void cPDR::setCbn(const RotationMatrix& in)
{
    mCbn = in;
    cbn2euler(mCbn, mInsData.Orientation);
}

/**************************************************************************//**
* Set the current euler angles from the input field.
* \arg \c out.  Euler angles (r,p,y) input.
* \return \c void
******************************************************************************/
void cPDR::setEulers(const AngleArray& in)
{
    mInsData.Orientation = in;
    euler2cbn(mInsData.Orientation, mCbn);
}

/**************************************************************************//**
* Set the current velocity from the input field.
* \arg \c out.  Velocity vector input.
* \return \c void
******************************************************************************/
void cPDR::setVelocity(const CartesianVector& in)
{
    mInsData.Velocity = in;
}

/**************************************************************************//**
* Set the current position from the input field.
* \arg \c out.  Position vector input.
* \return \c void
******************************************************************************/
void cPDR::setPosition(const CartesianVector& in)
{
    mInsData.Position = in;
}

/**************************************************************************//**
* Set the current position, velocity and euler angles.
* \arg \c position.  Position.
* \arg \c velocity.  Velocity.
* \arg \c eulers.  Euler angles.
* \return \c void
******************************************************************************/
void cPDR::setPVA(const CartesianVector& position,
                  const CartesianVector& velocity,
                  const AngleArray& eulers)
{
    mInsData.Position = position;
    mInsData.Velocity = velocity;
    mInsData.Orientation = eulers;
    euler2cbn(mInsData.Orientation, mCbn);
}

/**************************************************************************//**
* Set the current gravity vector
* \arg \c in.  Gravity in NED.
* \return \c void
******************************************************************************/
void cPDR::setGravity(const CartesianVector& in)
{
    mGravity = in;
}

/**************************************************************************//**
* Reset INS Counters so that a break in data does not cause a spike in the solution
* \arg \c void
* \return \c void
******************************************************************************/
void cPDR::resetCounters()
{
    // Zero counters in case we need to re-align at some point
    mMessageCounter = 0;
    mdt = 0;
    mTime = 0.0;
    mPreviousTime = 0.0;
    mPreviousEarthRateTime = 0.0;
}


#ifdef COMPILATION_INSTRUCTIONS
(echo '#include"'$0'"'>$0.cpp)&&$CXX -Wall -Wextra -Wpedantic -Wfatal-errors -D_TEST_MULTI_ADAPTORS_LAPACK_CUDA $0.cpp -o $0x `pkg-config --libs blas` -lcudart -lcublas -lcusolver&&$0x&&rm $0x $0.cpp; exit
#endif
// © Alfredo A. Correa 2020

#ifndef MULTI_ADAPTORS_LAPACK_CUDA_HPP
#define MULTI_ADAPTORS_LAPACK_CUDA_HPP

#include "../../memory/adaptors/cuda/ptr.hpp"
#include "../../memory/adaptors/cuda/managed/ptr.hpp"
#include "../../memory/adaptors/cuda/managed/allocator.hpp"

#include "../../adaptors/cuda.hpp"

//#include<cublas_v2.h>
#include <cusolverDn.h>

//#include<iostream> // debug

#include <boost/log/trivial.hpp>

#include<complex>
#include<memory>

#include "../blas/filling.hpp"

namespace boost{
namespace multi{

namespace cusolver{

enum class status : typename std::underlying_type<cusolverStatus_t>::type{
	success                   = CUSOLVER_STATUS_SUCCESS, // "The operation completed successfully."
	not_initialized           = CUSOLVER_STATUS_NOT_INITIALIZED, // "The cuSolver library was not initialized. This is usually caused by the lack of a prior call, an error in the CUDA Runtime API called by the cuSolver routine, or an error in the hardware setup. To correct: call cusolverCreate() prior to the function call; and check that the hardware, an appropriate version of the driver, and the cuSolver library are correctly installed."
	allocation_failed         = CUSOLVER_STATUS_ALLOC_FAILED, // "Resource allocation failed inside the cuSolver library. This is usually caused by a cudaMalloc() failure. To correct: prior to the function call, deallocate previously allocated memory as much as possible."
	invalid_value             = CUSOLVER_STATUS_INVALID_VALUE, // "An unsupported value or parameter was passed to the function (a negative vector size, for example). To correct: ensure that all the parameters being passed have valid values."
	architecture_mismatch     = CUSOLVER_STATUS_ARCH_MISMATCH, // "The function requires a feature absent from the device architecture; usually caused by the lack of support for atomic operations or double precision. To correct: compile and run the application on a device with compute capability 2.0 or above."
	execution_failed          = CUSOLVER_STATUS_EXECUTION_FAILED, // "The GPU program failed to execute. This is often caused by a launch failure of the kernel on the GPU, which can be caused by multiple reasons. To correct: check that the hardware, an appropriate version of the driver, and the cuSolver library are correctly installed."
	internal_error            = CUSOLVER_STATUS_INTERNAL_ERROR, // "An internal cuSolver operation failed. This error is usually caused by a cudaMemcpyAsync() failure. To correct: check that the hardware, an appropriate version of the driver, and the cuSolver library are correctly installed. Also, check that the memory passed as a parameter to the routine is not being deallocated prior to the routine’s completion."
	matrix_type_not_supported = CUSOLVER_STATUS_MATRIX_TYPE_NOT_SUPPORTED // "The matrix type is not supported by this function. This is usually caused by passing an invalid matrix descriptor to the function. To correct: check that the fields in descrA were set correctly."
};

std::string inline status_string(enum status s){ //https://stackoverflow.com/questions/13041399/equivalent-of-cudageterrorstring-for-cublas
	switch(s){
		case status::success                   : return "The operation completed successfully.";
		case status::not_initialized           : return "The cuSolver library was not initialized. This is usually caused by the lack of a prior call, an error in the CUDA Runtime API called by the cuSolver routine, or an error in the hardware setup. To correct: call cusolverCreate() prior to the function call; and check that the hardware, an appropriate version of the driver, and the cuSolver library are correctly installed.";
		case status::allocation_failed         : return "Resource allocation failed inside the cuSolver library. This is usually caused by a cudaMalloc() failure. To correct: prior to the function call, deallocate previously allocated memory as much as possible.";
		case status::invalid_value             : return "An unsupported value or parameter was passed to the function (a negative vector size, for example). To correct: ensure that all the parameters being passed have valid values.";
		case status::architecture_mismatch     : return "The function requires a feature absent from the device architecture; usually caused by the lack of support for atomic operations or double precision. To correct: compile and run the application on a device with compute capability 2.0 or above.";
		case status::execution_failed          : return "The GPU program failed to execute. This is often caused by a launch failure of the kernel on the GPU, which can be caused by multiple reasons. To correct: check that the hardware, an appropriate version of the driver, and the cuSolver library are correctly installed.";
		case status::internal_error            : return "An internal cuSolver operation failed. This error is usually caused by a cudaMemcpyAsync() failure. To correct: check that the hardware, an appropriate version of the driver, and the cuSolver library are correctly installed. Also, check that the memory passed as a parameter to the routine is not being deallocated prior to the routine’s completion.";
		case status::matrix_type_not_supported : return "The matrix type is not supported by this function. This is usually caused by passing an invalid matrix descriptor to the function. To correct: check that the fields in descrA were set correctly.";
	}
	return "cublas status <unknown>";
}
struct error_category : std::error_category{
	char const* name() const noexcept override{return "cusolver wrapper";}
	std::string message(int err) const override{return cusolver::status_string(static_cast<enum cusolver::status>(err));}
	static error_category& instance(){static cusolver::error_category instance; return instance;}
};
inline std::error_code make_error_code(cusolver::status s) noexcept{
	return std::error_code(int(s), cusolver::error_category::instance());
}

struct version_t{
	int major = -1, minor =-1, patch=-1;
	friend std::ostream& operator<<(std::ostream& os, version_t const& self){
		return os<< self.major <<'.'<< self.minor <<'.'<<self.patch <<'\n';
	}
};

auto version(){
	version_t ret;
	cusolverGetProperty(MAJOR_VERSION, &ret.major);
	cusolverGetProperty(MINOR_VERSION, &ret.minor);
	cusolverGetProperty(PATCH_LEVEL, &ret.patch);
	return ret;
}

namespace dense{

using blas::filling;

struct context : std::unique_ptr<std::decay_t<decltype(*cusolverDnHandle_t{})>, decltype(&cusolverDnDestroy)>{
	context()  : std::unique_ptr<std::decay_t<decltype(*cusolverDnHandle_t{})>, decltype(&cusolverDnDestroy)>{
		[]{
			cusolverDnHandle_t h;
			auto s=cusolverDnCreate(&h); assert(CUSOLVER_STATUS_SUCCESS==s and h);
			return h;
		}(), &cusolverDnDestroy
	}{}
	template<class A> int potrf_buffer_size(filling uplo, A const& a);
};

template<typename T>
struct cusolverDn;//{
//	template<class... ArgsA3, class... ArgsB2>
//	static auto potrf_bufferSize(ArgsA3... argsa3, T* ptr, ArgsB2... argsb2);
//};

template<>
struct cusolverDn<float>{
	template<class... A3, class... B2>
	static auto potrf_bufferSize(A3... a3, double* ptr, B2... b2)
	->decltype(cusolverDnSpotrf_bufferSize(a3..., ptr, b2...)){
		return cusolverDnSpotrf_bufferSize(a3..., ptr, b2...);}
	template<class... A> static auto syev(A... a)
	->decltype(cusolverDnSsyevd_bufferSize(a...)){
		return cusolverDnSsyevd_bufferSize(a...);}
	template<class... A> static auto syev(A... a)
	->decltype(cusolverDnSsyevd(a...)){
		return cusolverDnSsyevd(a...);}
};

template<>
struct cusolverDn<double>{
	template<class... A3, class... B2>
	static auto potrf_bufferSize(A3... a3, double* ptr, B2... b2)
	->decltype(cusolverDnDpotrf_bufferSize(a3..., ptr, b2...)){
		return cusolverDnDpotrf_bufferSize(a3..., ptr, b2...);}
	template<class... A> static auto syevd_bufferSize(A... a)
//	->decltype(cusolverDnDsyevd_bufferSize(a...)){
	{	return cusolverDnDsyevd_bufferSize(a...);}
	template<class... A> static auto syevd(A... a)
	->decltype(cusolverDnDsyevd(a...))
	{	return cusolverDnDsyevd(a...);}
};

template<>
struct cusolverDn<std::complex<float>>{
	template<class... A3, class... B2>
	static auto potrf_bufferSize(A3... a3, std::complex<double>* ptr, B2... b2)
	->decltype(cusolverDnCpotrf_bufferSize(a3..., reinterpret_cast<cuComplex*>(ptr), b2...)){
		return cusolverDnCpotrf_bufferSize(a3..., reinterpret_cast<cuComplex*>(ptr), b2...);}
};

template<>
struct cusolverDn<std::complex<double>>{
	static auto translate(std::complex<double>* p){return reinterpret_cast<cuDoubleComplex*>(p);}
	template<class T> static auto translate(T t){return t;}
	template<class... A> static auto potrf_bufferSize(A... a)
	->decltype(cusolverDnZpotrf_bufferSize(translate(a)...)){
		return cusolverDnZpotrf_bufferSize(translate(a)...);}
	template<class... A> static auto potrf(A... a)
	->decltype(cusolverDnZpotrf(translate(a)...)){
		return cusolverDnZpotrf(translate(a)...);}
};

}
}

namespace memory{
namespace cuda{

template<class UL, class S, class PtrT, typename T = typename std::pointer_traits<PtrT>::element_type>
void potrf(UL ul, S n, PtrT A, S incx, int& info){
	boost::multi::cusolver::dense::context ctx; //BOOST_LOG_TRIVIAL(trace)<<"cuda::potrf called on size/stride "<< n <<' '<< incx <<'\n';
	int lwork = -1;
	{
		auto s = cusolver::dense::cusolverDn<T>::potrf_bufferSize(ctx.get(), ul=='U'?CUBLAS_FILL_MODE_UPPER:CUBLAS_FILL_MODE_LOWER, n, raw_pointer_cast(A), incx, &lwork);
		assert(s == CUSOLVER_STATUS_SUCCESS); assert(lwork >= 0);
	}
	multi::cuda::array<T, 1> work(lwork);
	multi::cuda::static_array<int, 0> devInfo;
	auto s = cusolver::dense::cusolverDn<T>::potrf(ctx.get(), ul=='U'?CUBLAS_FILL_MODE_UPPER:CUBLAS_FILL_MODE_LOWER, n, raw_pointer_cast(A), incx, raw_pointer_cast(base(work)), lwork, raw_pointer_cast(base(devInfo)) );
	assert(s == CUSOLVER_STATUS_SUCCESS);
	cudaDeviceSynchronize();
	info = devInfo();
}

// https://docs.nvidia.com/cuda/cusolver/index.html#cuds-lt-t-gt-syevd
template<class S, class PtrT, typename T = typename std::pointer_traits<PtrT>::element_type>
void syev(char jobz, char uplo, S n, PtrT a, S lda, PtrT w, PtrT /*work*/, S /*lwork*/, int& info){
	boost::multi::cusolver::dense::context ctx;
	int lwork_needed = -1;
	{
		auto s = cusolver::dense::cusolverDn<T>::syevd_bufferSize(
			ctx.get(), jobz=='V'?CUSOLVER_EIG_MODE_VECTOR:CUSOLVER_EIG_MODE_NOVECTOR,
			uplo=='U'?CUBLAS_FILL_MODE_UPPER:CUBLAS_FILL_MODE_LOWER,
			n,
			raw_pointer_cast(a),
			lda,
			raw_pointer_cast(w),
			&lwork_needed
		);
		assert(s == CUSOLVER_STATUS_SUCCESS); assert(lwork_needed >= 0);
	}
	multi::cuda::array<T, 1> tmp_work(lwork_needed); // buffers needs no-managed memory!
	multi::cuda::static_array<int, 0> devInfo;
	auto s = cusolver::dense::cusolverDn<T>::syevd(
		ctx.get(), jobz=='V'?CUSOLVER_EIG_MODE_VECTOR:CUSOLVER_EIG_MODE_NOVECTOR,
		uplo=='U'?CUBLAS_FILL_MODE_UPPER:CUBLAS_FILL_MODE_LOWER,
		n,
		raw_pointer_cast(a),
		lda,
		raw_pointer_cast(w),
		raw_pointer_cast(tmp_work.data_elements()),
		tmp_work.size(),
		raw_pointer_cast(base(devInfo))
	);
	if( s != CUSOLVER_STATUS_SUCCESS ) throw std::system_error{cusolver::make_error_code(static_cast<cusolver::status>(s)), "cannot call cusolver function "};
//	cudaDeviceSynchronize();
	info = devInfo();
}

namespace managed{
	template<class UL, class S, class PtrT, typename T = typename std::pointer_traits<PtrT>::element_type>
	auto potrf(UL ul, S n, PtrT A, S incx, int& info)
	->decltype(cuda::potrf(ul, n, cuda::ptr<T>(A), incx, info)){
		return cuda::potrf(ul, n, cuda::ptr<T>(A), incx, info);}

	template<class S, class PtrT, class P2, typename T = typename std::pointer_traits<PtrT>::element_type>
	auto syev(char jobz, char uplo, S n, PtrT a, S lda, PtrT w, P2 work, S lwork, int& info)
	->decltype(cuda::syev(jobz, uplo, n, cuda::ptr<T>(a), lda, cuda::ptr<T>(w), cuda::ptr<T>(work), lwork, info)){
		return cuda::syev(jobz, uplo, n, cuda::ptr<T>(a), lda, cuda::ptr<T>(w), cuda::ptr<T>(work), lwork, info);}

}

}
}

}}

namespace std{template<> struct is_error_code_enum<::boost::multi::cusolver::status> : true_type{};}

///////////////////////////////////////////////////////////////////////////////

#if _TEST_MULTI_ADAPTORS_LAPACK_CUDA

#include "../../array.hpp"
#include "../../utility.hpp"
#include<cassert>

namespace multi = boost::multi;

int main(){
	std::cout << "cusolver version " << multi::cusolver::version() << std::endl;
	multi::cusolver::dense::context c;
//	multi::cublas_context c;
//	assert( c.version() >= 10100 );
}

#endif
#endif