#include <TdcpErrorEval.hpp>
#include <UnicycleErrorEval.hpp>

#include <cpo_utilities.hpp>
#include <CpoBackEnd.hpp>
#include <fstream>
#include <chrono>
#include <filesystem>
#include <Eigen/Core>

using Transformation = lgmath::se3::Transformation;
using TransformationWithCovariance = lgmath::se3::TransformationWithCovariance;
using TransformStateVar = steam::se3::TransformStateVar;
using TransformStateEvaluator = steam::se3::TransformStateEvaluator;
using TransformEvaluator = steam::se3::TransformEvaluator;
using SteamTrajVar = steam::se3::SteamTrajVar;
using SteamTrajInterface = steam::se3::SteamTrajInterface;
using VectorSpaceStateVar = steam::VectorSpaceStateVar;
using PositionEvaluator = steam::se3::PositionEvaluator;

CpoBackEnd::CpoBackEnd() : Node("cpo_back_end") {

  subscription_ = this->create_subscription<cpo_interfaces::msg::TDCP>("tdcp",
                                                                       10,
                                                                       std::bind(
                                                                           &CpoBackEnd::_tdcpCallback,
                                                                           this,
                                                                           std::placeholders::_1));
  enu_publisher_ =
      this->create_publisher<geometry_msgs::msg::PoseWithCovariance>("cpo_enu",
                                                                     10);

  this->declare_parameter("fixed_rate_publish", false);
  fixed_rate_publish_ = this->get_parameter("fixed_rate_publish").as_bool();

  if (fixed_rate_publish_) {
    this->declare_parameter("publish_frequency", 5.0);
    double freq = this->get_parameter("publish_frequency").as_double();
    double period = 1000 / freq;

    this->declare_parameter("publish_delay", 5.0);
    publish_delay_ = this->get_parameter("publish_delay").as_double();

    publish_timer_ = rclcpp::create_timer(this,
                                          this->get_clock(),
                                          std::chrono::milliseconds((long) period),
                                          std::bind(&CpoBackEnd::_timedCallback,
                                                    this));
  } else {
    publish_timer_ = nullptr;
  }

  query_traj_srv_ = this->create_service<QueryTrajectory>("query_trajectory",
                                                          std::bind(&CpoBackEnd::_queryCallback,
                                                                    this,
                                                                    std::placeholders::_1,
                                                                    std::placeholders::_2));

  // set up receiver-vehicle transform
  this->declare_parameter("r_veh_gps_inv",
                          std::vector<double>{-0.60, 0.00, -0.52});
  auto r = this->get_parameter("r_veh_gps_inv").as_double_array();

  Eigen::Matrix4d T_gps_vehicle_eigen = Eigen::Matrix4d::Identity();
  T_gps_vehicle_eigen(0, 3) = r[0];
  T_gps_vehicle_eigen(1, 3) = r[1];
  T_gps_vehicle_eigen(2, 3) = r[2];
  auto T_gps_vehicle = TransformationWithCovariance(T_gps_vehicle_eigen);
  T_gps_vehicle.setZeroCovariance();
  tf_gps_vehicle_ =
      steam::se3::FixedTransformEvaluator::MakeShared(T_gps_vehicle);

  getParams();
}

void CpoBackEnd::_tdcpCallback(const cpo_interfaces::msg::TDCP::SharedPtr msg_in) {

  uint n = msg_in->pairs.size();
  std::cout << "Received " << n << " satellite pairs." << std::endl;

  addMsgToWindow(msg_in);

  // currently, we only attempt optimization if we have at least 4 satellites to
  // ensure the problem's well-defined. However, some cases will work with fewer
  if (n >= 4) {
    initializeProblem();

    // set up steam problem

    // setup state variables using initial condition
    std::vector<SteamTrajVar> traj_states;
    std::vector<TransformStateVar::Ptr> statevars;

    // keep track of states composed with frame transform for use in TDCP terms
    std::vector<TransformEvaluator::ConstPtr> receiver_poses;
    std::vector<TransformEvaluator::ConstPtr> enu_poses;

    // set up T_0g state
    TransformStateVar::Ptr T_0g_statevar(new TransformStateVar(init_pose_));
    TransformEvaluator::ConstPtr
        T_0g = TransformStateEvaluator::MakeShared(T_0g_statevar);

    { // first pose in window gets locked (but not the first velocity)
      TransformStateVar::Ptr
          temp_statevar_a(new TransformStateVar(Transformation()));
      TransformStateEvaluator::Ptr
          temp_pose_a = TransformStateEvaluator::MakeShared(temp_statevar_a);
      VectorSpaceStateVar::Ptr temp_velocity_a =
          VectorSpaceStateVar::Ptr(new VectorSpaceStateVar(edges_.front().v_a));
      temp_statevar_a->setLock(true);
      SteamTrajVar temp_a(steam::Time((int64_t)
      edges_.front().msg.t_a), temp_pose_a, temp_velocity_a);
      statevars.push_back(temp_statevar_a);
      traj_states.push_back(temp_a);

      TransformEvaluator::Ptr
          rec_pose_a = steam::se3::compose(tf_gps_vehicle_, temp_pose_a);
      receiver_poses.emplace_back(rec_pose_a);
      enu_poses.emplace_back(steam::se3::compose(rec_pose_a, T_0g));
    }

    Transformation prev_T_k0_est = Transformation();
    // loop over window to add states for other poses
    for (auto &edge : edges_) {
      // k has been incremented so update T_k0
      Transformation T_k0_est(edge.T_ba * prev_T_k0_est);
      // prevents rotation matrix from drifting away from being orthogonal
      T_k0_est.reproject(true);

      TransformStateVar::Ptr temp_statevar(new TransformStateVar(T_k0_est));
      TransformStateEvaluator::Ptr
          temp_pose = TransformStateEvaluator::MakeShared(temp_statevar);
      VectorSpaceStateVar::Ptr temp_velocity =
          VectorSpaceStateVar::Ptr(new VectorSpaceStateVar(edge.v_b));
      SteamTrajVar temp(steam::Time((int64_t)
      edge.msg.t_b), temp_pose, temp_velocity);
      statevars.push_back(temp_statevar);
      traj_states.push_back(temp);

      // T_s0 = T_sv * T_v0
      TransformEvaluator::Ptr
          rec_pose = steam::se3::compose(tf_gps_vehicle_, temp_pose);
      receiver_poses.emplace_back(rec_pose); // T_k0
      // T_kg = T_k0 * T_0g
      enu_poses.emplace_back(steam::se3::compose(rec_pose, T_0g));

      prev_T_k0_est = T_k0_est;
    }

    // add TDCP terms
    for (uint k = 0; k < edges_.size(); ++k) {
      TransformEvaluator::ConstPtr T_k1k = steam::se3::compose(
          receiver_poses[k + 1],
          steam::se3::inverse(receiver_poses[k]));
      PositionEvaluator::ConstPtr r_ba_ina(new PositionEvaluator(T_k1k));

      // using constant covariance here for now
      steam::BaseNoiseModel<1>::Ptr
          tdcp_noise_model(new steam::StaticNoiseModel<1>(tdcp_cov_));

      // iterate through satellite pairs in msg and add TDCP costs
      for (const auto &pair : edges_[k].msg.pairs) {
        Eigen::Vector3d
            r_1a_ing_ata{pair.r_1a_a.x, pair.r_1a_a.y, pair.r_1a_a.z};
        Eigen::Vector3d
            r_1a_ing_atb{pair.r_1a_b.x, pair.r_1a_b.y, pair.r_1a_b.z};
        Eigen::Vector3d
            r_2a_ing_ata{pair.r_2a_a.x, pair.r_2a_a.y, pair.r_2a_a.z};
        Eigen::Vector3d
            r_2a_ing_atb{pair.r_2a_b.x, pair.r_2a_b.y, pair.r_2a_b.z};

        steam::TdcpErrorEval::Ptr tdcp_error(new steam::TdcpErrorEval(
            pair.phi_measured,
            r_ba_ina,
            enu_poses[k],   // T_kg
            r_1a_ing_ata,
            r_1a_ing_atb,
            r_2a_ing_ata,
            r_2a_ing_atb));
        auto tdcp_factor = steam::WeightedLeastSqCostTerm<1, 6>::Ptr(
            new steam::WeightedLeastSqCostTerm<1, 6>(
                tdcp_error,
                tdcp_noise_model,
                tdcp_loss_function_));
        tdcp_cost_terms_->add(tdcp_factor);
      }
    }

    // lock T_0g instead when we already have a good estimate of it
    if (init_pose_estimated_ && lock_first_pose_) {
      T_0g_statevar->setLock(true);
    } else {
      // add prior on initial pose to deal with roll uncertainty and constrain r^0g_g to zero
      steam::BaseNoiseModel<6>::Ptr
          sharedNoiseModel(new steam::StaticNoiseModel<6>(pose_prior_cov_));
      steam::TransformErrorEval::Ptr prior_error_func(
          new steam::TransformErrorEval(init_pose_, T_0g));
      auto pose_prior_factor_ = steam::WeightedLeastSqCostTerm<6, 6>::Ptr(
          new steam::WeightedLeastSqCostTerm<6, 6>(
              prior_error_func,
              sharedNoiseModel,
              pp_loss_function_));
      pose_prior_cost_->add(pose_prior_factor_);
    }

    steam::BaseNoiseModel<4>::Ptr nonholonomic_noise_model
        (new steam::StaticNoiseModel<4>(nonholonomic_cov_));
    trajectory_.reset(new SteamTrajInterface(smoothing_factor_information_,
                                             true));

    // loop through velocity state variables
    for (const auto &traj_state : traj_states) {
      if (!traj_state.getVelocity()->isLocked()) {
        // add nonholonomic costs
        steam::UnicycleErrorEval::Ptr non_holo_error_func(
            new steam::UnicycleErrorEval(traj_state.getVelocity()));
        auto non_holonomic_factor = steam::WeightedLeastSqCostTerm<4, 6>::Ptr(
            new steam::WeightedLeastSqCostTerm<4, 6>(
                non_holo_error_func,
                nonholonomic_noise_model,
                nonholonomic_loss_function_));
        nonholonomic_cost_terms_->add(non_holonomic_factor);

        // add smoothing costs
        steam::Time temp_time = traj_state.getTime();
        const TransformEvaluator::Ptr &temp_pose = traj_state.getPose();
        const steam::VectorSpaceStateVar::Ptr
            &temp_velocity = traj_state.getVelocity();
        trajectory_->add(temp_time, temp_pose, temp_velocity);

        // also add velocity state to problem
        problem_->addStateVariable(traj_state.getVelocity());
      }
    }
    trajectory_->appendPriorCostTerms(smoothing_cost_terms_);

    problem_->addCostTerm(tdcp_cost_terms_);
    problem_->addCostTerm(nonholonomic_cost_terms_);
    problem_->addCostTerm(smoothing_cost_terms_);
    problem_->addCostTerm(pose_prior_cost_);

    problem_->addStateVariable(T_0g_statevar);
    for (const auto &state : statevars) {
      problem_->addStateVariable(state);
    }

    if (costs_verbose_) {
      // print initial costs (for debugging/development)
      printCosts(false);
    }

    // setup solver and optimize
    steam::DoglegGaussNewtonSolver::Params params;
    params.verbose = steam_verbose_;
    params.maxIterations = steam_max_iterations_;
    params.absoluteCostChangeThreshold = 1e-2;
    solver_.reset(new steam::DoglegGaussNewtonSolver(problem_.get(), params));

    try {
      solver_->optimize();
    } catch (steam::unsuccessful_step &e) {
      // did any successful steps occur?
      if (solver_->getCurrIteration() <= 1) {
        // no: something is very wrong; we should start over. Should not occur frequently
        std::cout
            << "Steam has failed to optimize the problem! This is an ERROR."
            << std::endl;
        return;
      } else {
        // yes: just a marginal problem, let's use what we got
        std::cout
            << "Steam has failed due to an unsuccessful step. This should be OK if it happens infrequently."
            << std::endl;
      }
    } catch (steam::decomp_failure &e) {
      // Should not occur frequently
      std::cout
          << "Steam has encountered an LL^T decomposition error while optimizing! This is an ERROR."
          << std::endl;
      return;
    }

    if (costs_verbose_) {
      // print final costs (for debugging/development)
      printCosts(true);
    }

    // update with optimized transforms
    for (uint i = 0; i < edges_.size(); ++i) {
      edges_[i].T_ba = statevars[i + 1]->getValue()
          * statevars[i]->getValue().inverse();  // T_21 = T_20 * inv(T_10)
      edges_[i].T_ba.reproject(true);
      edges_[i].v_a = traj_states.at(i).getVelocity()->getValue();
      edges_[i].v_b = traj_states.at(i + 1).getVelocity()->getValue();
    }

    // update our orientation estimate
    if (!T_0g_statevar->isLocked()) {
      init_pose_ = T_0g_statevar->getValue();
    }

    init_pose_.reproject(true);

    double t_n = (double) edges_.back().msg.t_b * 1e-9;
    double t_0 = (double) edges_.front().msg.t_a * 1e-9;
    std::cout << "Last message time was: " << std::setprecision(12) << t_n;
    std::cout << "    Time at front of sliding window is: " << t_0
              << std::setprecision(6) << std::endl;

    // if not publishing based on timer, we'll publish edge leaving window
    if (!fixed_rate_publish_ && edges_.size() == window_size_) {
      // only publish if window is full
      publishPose(init_pose_);
      saveToFile(init_pose_, t_0);
    }

    init_pose_estimated_ = true;
    first_window_ = false;
  }
}

void CpoBackEnd::_timedCallback() {
  if (first_window_ || !init_pose_estimated_ || trajectory_ == nullptr) {
    std::cout
        << "Not publishing because don't currently have a valid pose estimate."
        << std::endl;
    return;
  }

  // grab times and extrapolate poses
  double t_last_msg = (double) edges_.back().msg.t_b * 1e-9;
  double t_n = get_clock()->now().seconds();      // current time
  double t_k = t_n - publish_delay_;              // time to save an estimate at

  if (t_n - t_last_msg > traj_timeout_limit_) {
    std::cout << "Latest carrier phase measurement is " << t_n - t_last_msg
              << " seconds old. "
              << "Will not publish because trajectory is no longer valid."
              << std::endl;
    return;
  }
  if (t_k > t_last_msg && publish_delay_ > 0) {
    // if here we have set non-zero value to delay publishing so we don't have
    // to extrapolate but there's a big enough observation gap we would have to
    std::cout
        << "Gap in TDCP observations. Choosing not to extrapolate trajectory. "
        << std::endl;
    return;
  }

  TransformationWithCovariance T_0g = init_pose_;
  T_0g.setZeroCovariance();
  TransformationWithCovariance
      T_k0 = trajectory_->getInterpPoseEval(t_k)->evaluate();
  T_k0.setZeroCovariance();
  TransformationWithCovariance T_kg = T_k0 * T_0g;

  // publish
  publishPose(T_kg);
  saveToFile(T_kg, t_k);
}

void CpoBackEnd::_queryCallback(const std::shared_ptr<QueryTrajectory::Request> request,
                                std::shared_ptr<QueryTrajectory::Response> response) {

  std::cout << "CPO's query trajectory service was called." << std::endl;

  if (first_window_ || !init_pose_estimated_ || trajectory_ == nullptr
      || solver_ == nullptr) {
    response->success = false;
    response->message = "CPO does not currently have a valid pose estimate.";
    return;
  }

  // get times of current window
#if 0
  uint64_t t_start = edges_.front().msg.t_a;
#else
  uint64_t t_start = edges_.front().msg.t_b;  // todo: currently have problems when one side of interpolation is locked
#endif
  uint64_t t_end = edges_.back().msg.t_b;

#if 0       // DEBUGGING
  std::cout << std::setprecision(19) << "QC win t_a " << t_start
            << std::endl;
  std::cout << "QC win t_b " << t_end << std::endl;
  std::cout << "QC req t_1 " << request->t_1 << std::endl;
  std::cout << "QC req t_2 " << request->t_2 << std::setprecision(6)
            << std::endl;
#endif

  // error checking
  if (request->t_1 > request->t_2) {
    response->success = false;
    response->message = "t_1 > t_2 but expected t_2 to be after t_1.";
    std::cout << "response->message: " << response->message << std::endl;
    return;
  } else if (request->t_1 < t_start) {
    response->success = false;
    response->message =
        "Requested time before CPO window. Trajectory history has not been implemented.";
    std::cout << "response->message: " << response->message << std::endl;
    return;
  } else if (request->t_2 > t_end + (long) (traj_timeout_limit_ * 1e9)) {
    response->success = false;
    response->message = "t_2 too long after last GPS msg to trust trajectory.";
    std::cout << "response->message: " << response->message << std::endl;
    return;
  }

  // requested times are reasonable so let's try to estimate a transform
  lgmath::se3::TransformationWithCovariance T_10 =
      trajectory_->getInterpPoseEval(steam::Time((int64_t) request->t_1))->evaluate();
  lgmath::se3::TransformationWithCovariance T_20 =
      trajectory_->getInterpPoseEval(steam::Time((int64_t) request->t_2))->evaluate();
  lgmath::se3::TransformationWithCovariance T_21 = T_20 / T_10;

  auto gn_solver =
      std::dynamic_pointer_cast<steam::GaussNewtonSolverBase>(solver_);
  trajectory_->setSolver(gn_solver);

  Eigen::Matrix<double, 6, 6> Cov_21 = trajectory_->getRelativeCovariance(steam::Time((int64_t) request->t_1), steam::Time((int64_t) request->t_2));
  T_21.setCovariance(Cov_21);
  response->tf_2_1 = toPoseMsg(T_21);

  response->success = true;
  response->message = "Trajectory query was successful.";
}

void CpoBackEnd::publishPose(const TransformationWithCovariance &T_0g) {
  geometry_msgs::msg::PoseWithCovariance enu_pose_msg = toPoseMsg(T_0g);
  enu_publisher_->publish(enu_pose_msg);
  std::cout << "Message published! " << std::endl;
}

void CpoBackEnd::saveToFile(const Transformation &T_kg, double t_k) const {
  // append latest estimate to file
  std::ofstream outstream;
  outstream.open(results_path_, std::ofstream::out | std::ofstream::app);

  const Eigen::Vector3d r_kg_g = T_kg.r_ba_ina();     // vehicle position
  Eigen::Vector3d r_sg_g = (tf_gps_vehicle_->evaluate() * T_kg).r_ba_ina();
  Eigen::Matrix4d T_sg = (tf_gps_vehicle_->evaluate() * T_kg).matrix();

  // save times and global position for easy plotting
  outstream << std::setprecision(12) << t_k << ", " << 0 << ", ";
  outstream << r_sg_g(0) << ", " << r_sg_g(1) << ", " << r_sg_g(2)
            << ", ";   // receiver position in ENU frame
  outstream << r_kg_g(0) << ", " << r_kg_g(1) << ", " << r_kg_g(2)
            << ", ";   // vehicle position in ENU frame

  // save full transformations of receiver frame as well (** column major)
  auto T_sg_flat = std::vector<double>(T_sg.data(), T_sg.data() + 16);
  for (auto entry : T_sg_flat) outstream << entry << ",";

  outstream << std::endl;
  outstream.close();
}

void CpoBackEnd::getParams() {
  this->declare_parameter("tdcp_cov", 0.1);
  this->declare_parameter("non_holo_y", 0.1);
  this->declare_parameter("non_holo_z", 0.1);
  this->declare_parameter("non_holo_roll", 0.1);
  this->declare_parameter("non_holo_pitch", 0.1);
  this->declare_parameter("lin_acc_std_dev_x", 1.0);
  this->declare_parameter("lin_acc_std_dev_y", 0.1);
  this->declare_parameter("lin_acc_std_dev_z", 0.1);
  this->declare_parameter("ang_acc_std_dev_x", 0.1);
  this->declare_parameter("ang_acc_std_dev_y", 0.1);
  this->declare_parameter("ang_acc_std_dev_z", 0.1);
  this->declare_parameter("prior_cov_x", 0.001);
  this->declare_parameter("prior_cov_y", 0.001);
  this->declare_parameter("prior_cov_z", 0.001);
  this->declare_parameter("prior_cov_ang1", 0.01);
  this->declare_parameter("prior_cov_ang2", 0.01);
  this->declare_parameter("prior_cov_ang3", 1.0);
  this->declare_parameter("window_size", 10);
  this->declare_parameter("lock_first_pose", true);
  this->declare_parameter("results_path", "~/cpo.csv");
  this->declare_parameter("solver_verbose", false);
  this->declare_parameter("costs_verbose", false);
  this->declare_parameter("solver_max_iterations", 5);
  this->declare_parameter("traj_timeout_limit", 5.0);

  tdcp_cov_ << this->get_parameter("tdcp_cov").as_double();

  Eigen::Array<double, 1, 4> non_holo_diag;
  non_holo_diag << this->get_parameter("non_holo_y").as_double(),
      this->get_parameter("non_holo_z").as_double(),
      this->get_parameter("non_holo_roll").as_double(),
      this->get_parameter("non_holo_pitch").as_double();
  nonholonomic_cov_.setZero();
  nonholonomic_cov_.diagonal() = non_holo_diag;

  Eigen::Array<double, 1, 6> Qc_diag;
  Qc_diag << this->get_parameter("lin_acc_std_dev_x").as_double(),
      this->get_parameter("lin_acc_std_dev_y").as_double(),
      this->get_parameter("lin_acc_std_dev_z").as_double(),
      this->get_parameter("ang_acc_std_dev_x").as_double(),
      this->get_parameter("ang_acc_std_dev_y").as_double(),
      this->get_parameter("ang_acc_std_dev_z").as_double();
  smoothing_factor_information_.setZero();
  smoothing_factor_information_.diagonal() = 1.0 / Qc_diag;

  pose_prior_cov_ = Eigen::Matrix<double, 6, 6>::Identity();
  pose_prior_cov_(0, 0) = this->get_parameter("prior_cov_x").as_double();
  pose_prior_cov_(1, 1) = this->get_parameter("prior_cov_y").as_double();
  pose_prior_cov_(2, 2) = this->get_parameter("prior_cov_z").as_double();
  pose_prior_cov_(3, 3) = this->get_parameter("prior_cov_ang1").as_double();
  pose_prior_cov_(4, 4) = this->get_parameter("prior_cov_ang2").as_double();
  pose_prior_cov_(5, 5) = this->get_parameter("prior_cov_ang3").as_double();

  window_size_ = this->get_parameter("window_size").as_int();
  lock_first_pose_ = this->get_parameter("lock_first_pose").as_bool();
  results_path_ = expandUser(this->get_parameter("results_path").as_string());
  steam_verbose_ = this->get_parameter("solver_verbose").as_bool();
  costs_verbose_ = this->get_parameter("costs_verbose").as_bool();
  steam_max_iterations_ = this->get_parameter("solver_max_iterations").as_int();
  traj_timeout_limit_ = this->get_parameter("traj_timeout_limit").as_double();
}

void CpoBackEnd::initializeProblem() {
  // setup loss functions
  tdcp_loss_function_.reset(new steam::CauchyLossFunc(2.0));
  nonholonomic_loss_function_.reset(new steam::L2LossFunc());
  pp_loss_function_.reset(new steam::L2LossFunc());

  // setup cost terms
  tdcp_cost_terms_.reset(new steam::ParallelizedCostTermCollection());
  nonholonomic_cost_terms_.reset(new steam::ParallelizedCostTermCollection());
  smoothing_cost_terms_.reset(new steam::ParallelizedCostTermCollection());
  pose_prior_cost_.reset(new steam::ParallelizedCostTermCollection());

  // set up the steam problem
  problem_.reset(new steam::OptimizationProblem());

  if (!init_pose_estimated_) {
    // estimate initial T_0g from code solutions
    Eigen::Vector3d r_k0_ing = toEigenVec3d(edges_.back().msg.enu_pos)
        - toEigenVec3d(edges_.front().msg.enu_pos);
    double theta = atan2(r_k0_ing.y(), r_k0_ing.x());

    Eigen::Matrix<double, 6, 1> init_pose_vec;
    init_pose_vec << toEigenVec3d(edges_.front().msg.enu_pos), 0, 0, -1 * theta;

    init_pose_ = Transformation(init_pose_vec);

    if (first_window_) {
      // save ENU origin to results file
      std::ofstream outstream;
      outstream.open(results_path_);
      Eigen::Vector3d enu_origin = toEigenVec3d(edges_.back().msg.enu_origin);
      outstream << std::setprecision(9) << enu_origin[0] << "," << enu_origin[1]
                << "," << enu_origin[2] << std::endl;
      outstream.close();
      if (!std::filesystem::exists(results_path_)) {
        std::cout << "Warning: Results path " << results_path_
                  << " does not exist." << std::endl;
      }
    } else {
      std::cout << "Warning: new initial pose was set midway through run."
                << "Estimates on either side of this time should not be compared."
                << std::endl;
    }
  }
}

void CpoBackEnd::resetEstimator() {
  std::cout << "Clearing window and resetting the estimator." << std::endl;
  edges_.clear();
  trajectory_ = nullptr;
  init_pose_estimated_ = false;
  std::cout << "WARNING: RESET ESTIMATOR !" << std::endl;
}

void CpoBackEnd::printCosts(bool final) {
  std::string stage_str = final ? "Final" : "Initial";
  std::cout << " === " << stage_str << " Costs === " << std::endl;
  std::cout << "Carrier Phase:       " << tdcp_cost_terms_->cost()
            << "        Terms:  " << tdcp_cost_terms_->numCostTerms()
            << std::endl;
  std::cout << "Nonholonomic:        " << nonholonomic_cost_terms_->cost()
            << "        Terms:  " << nonholonomic_cost_terms_->numCostTerms()
            << std::endl;
  std::cout << "Smoothing:           " << smoothing_cost_terms_->cost()
            << "        Terms:  " << smoothing_cost_terms_->numCostTerms()
            << std::endl;
  std::cout << "Pose Prior:          " << pose_prior_cost_->cost()
            << "        Terms:  " << pose_prior_cost_->numCostTerms()
            << std::endl;
}

void CpoBackEnd::addMsgToWindow(const cpo_interfaces::msg::TDCP::SharedPtr &msg) {

  if (!edges_.empty() && msg->t_a != edges_.back().msg.t_b) {
    // times don't align, so we've likely missed a msg. To be safe we'll clear
    std::cout << "Warning: mismatched times. Current t_a: " << msg->t_a
              << ". Previous t_b: " << edges_.back().msg.t_b
              << " Resetting estimator." << std::endl;
    resetEstimator();
  }

  // add latest message, setting initial values for states
  Transformation new_T_estimate;
  Eigen::Matrix<double, 6, 1> new_vel_a;
  Eigen::Matrix<double, 6, 1> new_vel_b;
  if (trajectory_ != nullptr) {
    Transformation T_b0 =
        trajectory_->getInterpPoseEval(steam::Time((int64_t) msg->t_b))->evaluate();
    Transformation T_a0 =
        trajectory_->getInterpPoseEval(steam::Time((int64_t) msg->t_a))->evaluate();
    new_T_estimate = T_b0 * T_a0.inverse();
    new_vel_a = trajectory_->getVelocity(steam::Time((int64_t) msg->t_a));
    new_vel_b = trajectory_->getVelocity(steam::Time((int64_t) msg->t_b));
  } else {
    // if we don't have a trajectory to extrapolate, try generic initial guesses
    double dist_since_last = edges_.empty() ? 0.0 : 1.0;
    Eigen::Vector3d r_ba_est{dist_since_last, 0.0, 0.0};
    new_T_estimate = Transformation(Eigen::Matrix3d::Identity(), r_ba_est);
    new_vel_a << -1.0, 0.0, 0.0, 0.0, 0.0, 0.0;
    new_vel_b << -1.0, 0.0, 0.0, 0.0, 0.0, 0.0;
  }

  edges_.emplace_back(CpoEdge{*msg, new_T_estimate, new_vel_a, new_vel_b});

  // if we have a full queue, discard the oldest msg
  while (edges_.size() > window_size_) {
    // incrementing indices so need to update our T_0g estimate
    init_pose_ = edges_.front().T_ba * init_pose_;
    init_pose_.reproject(true);

    edges_.pop_front();
  }
}

geometry_msgs::msg::PoseWithCovariance CpoBackEnd::toPoseMsg(
    TransformationWithCovariance T) {
  Eigen::Quaterniond q(T.C_ba());
  Eigen::Vector3d r_ab_inb = T.r_ab_inb();

  geometry_msgs::msg::PoseWithCovariance msg;
  msg.pose.position.set__x(r_ab_inb[0]);
  msg.pose.position.set__y(r_ab_inb[1]);
  msg.pose.position.set__z(r_ab_inb[2]);
  msg.pose.orientation.set__x(q.x());
  msg.pose.orientation.set__y(q.y());
  msg.pose.orientation.set__z(q.z());
  msg.pose.orientation.set__w(q.w());

  if (!T.covarianceSet())
    T.setZeroCovariance();

  Eigen::Matrix<double, 6, 6> T_cov = T.cov();
  std::array<double, 36> temp{};
  for (int i = 0; i < 6; ++i) {
    for (int j = 0; j < 6; ++j) {
      temp[6*i + j] = T_cov(i, j);
    }
  }
  msg.set__covariance(temp);

  return msg;
}

Eigen::Vector3d CpoBackEnd::toEigenVec3d(const geometry_msgs::msg::Vector3 &ros_vec) {
  return Eigen::Vector3d{ros_vec.x, ros_vec.y, ros_vec.z};
}


#include <ros/ros.h>
#include <pcl/io/pcd_io.h>
#include <pcl/point_types.h>
#include <pcl_conversions/pcl_conversions.h>
#include <pcl/PCLPointCloud2.h>
#include <pcl/common/common.h>
#include <pcl/common/centroid.h>
#include <pcl/common/pca.h>
#include <pcl/filters/statistical_outlier_removal.h>

#include <sensor_msgs/PointCloud2.h>
#include <cv_bridge/cv_bridge.h>
#include <sensor_msgs/image_encodings.h>
#include <sensor_msgs/CameraInfo.h>
#include <sensor_msgs/Image.h>
#include <tf/transform_listener.h>
#include <tf_conversions/tf_eigen.h>
#include <message_filters/subscriber.h>
#include <message_filters/synchronizer.h>
#include <message_filters/sync_policies/approximate_time.h>

#include <image_geometry/pinhole_camera_model.h>

#include <opencv2/highgui/highgui.hpp>

#include <apc_msgs/SegmentPointcloudFromLabels.h>
#include <apc_msgs/AlignedDimensions.h>

#include <apc_msgs/Camera_Info.h>

#include <boost/foreach.hpp>

#include <sstream>

#include <cmath>

/**
*Node for detecting 3D clusters and align them with corresponding 2D images
*/

using namespace message_filters;
using namespace std;
using namespace pcl;
using namespace cv;
using namespace ros;

typedef struct Proposal{
    cv::Rect bounding_box;
    std::vector<cv::Point2d> imagePoints;
    uint32_t label;
} proposal_t;

class ObjectSegment
{

public:

    ObjectSegment(ros::NodeHandle nh):
        nh_(nh),
        camera_initialized_(false),
        input_cloud_(new pcl::PointCloud<pcl::PointXYZRGB>),
        temp_cloud_(new pcl::PointCloud<pcl::PointXYZRGB>),
        output_cloud_(new pcl::PointCloud<pcl::PointXYZRGB>)
    {
        // Just do this once.
        std::string camera_info_topic_name = "/realsense_wrist/rgb/camera_info";
        sensor_msgs::CameraInfoConstPtr camera_info_ptr = ros::topic::waitForMessage<sensor_msgs::CameraInfo>(camera_info_topic_name);
        onCameraInfo(camera_info_ptr);

        segmentPCSrv_ = nh_.advertiseService("segment_pointcloud_from_labels", &ObjectSegment::segmentPointcloudFromLabels, this);

        sor_.setMeanK(50);
        sor_.setStddevMulThresh(1.0);
    }

    /**
    * @brief onCameraInfo gets the intrinsic camara info
    * @param camara_info the info
    */
    void onCameraInfo(const sensor_msgs::CameraInfoConstPtr &camara_info){

        if (camera_initialized_) return;
        else{
            ROS_INFO_STREAM("Camera Initialised!");
            model_.fromCameraInfo(camara_info);
            camera_initialized_=true;
        }
    }

    bool segmentPointcloudFromLabels(apc_msgs::SegmentPointcloudFromLabels::Request &req,
                                     apc_msgs::SegmentPointcloudFromLabels::Response &res)
    {
        pcl::fromROSMsg(req.input_cloud, *input_cloud_);

        cv_bridge::CvImagePtr cv_ptr;
        cv::Mat labels;
        decltype(ros::Time::now()) t1, t2;
        t1 = ros::Time::now();
        try{
            cv_ptr = cv_bridge::toCvCopy(req.classification, sensor_msgs::image_encodings::MONO8);
            labels = cv_ptr->image;
        }
        catch (cv_bridge::Exception& e){
            ROS_ERROR("cv_bridge exception: %s", e.what());
            return true;
        }
        t2 = ros::Time::now();
        ROS_INFO_STREAM("Getting image from message: " << (t2 - t1));

        t1 = ros::Time::now();
        double min_seg_d, max_seg_d;
        int max_seg;
        cv::minMaxIdx(labels, &min_seg_d, &max_seg_d);
        max_seg = static_cast<int>(max_seg_d);
        t2 = ros::Time::now();
        ROS_INFO_STREAM("Getting min/max from labels: " << (t2 - t1));

        ROS_INFO_STREAM("Number of segments is: " << max_seg << std::endl);

        // t1 = ros::Time::now();
        // sor_.setInputCloud(input_cloud_);
        // sor_.filter(*input_cloud_);
        // t2 = ros::Time::now();
        // ROS_INFO_STREAM("Setting SOR stuff: " << (t2 - t1));

        t1 = ros::Time::now();
        pc_segments_.resize(max_seg);
        for(int i = 0; i < max_seg; i++) {
            pc_segments_[i].resize(0);
            pc_segments_[i].reserve(input_cloud_->size());
        }
        t2 = ros::Time::now();
        ROS_INFO_STREAM("Adding PC segments: " << (t2 - t1));

        t1 = ros::Time::now();
        for(auto &point: *input_cloud_) {
            if(point.r == 0) {
                continue;
            }
            cv::Point3d xyz = cv::Point3d(point.x, point.y, point.z);
            cv::Point2d pt = model_.project3dToPixel(xyz);
            if (!std::isnan(pt.x) && !std::isnan(pt.y)){
                int label = (int)labels.at<uchar>((int)pt.y, (int)pt.x);
                if(label > 0) {
                    pc_segments_[label-1].push_back(point);
                }
            }
        }
        t2 = ros::Time::now();
        ROS_INFO_STREAM("Adding points to PC segments: " << (t2 - t1));

        t1 = ros::Time::now();
        res.segmented_pointclouds.reserve(max_seg);
        res.aligned_dimensions.reserve(max_seg);
        for(int i = 0; i < max_seg; i++) {
            // Push back the point clouds.
            sensor_msgs::PointCloud2 pc_msg;
            pcl::toROSMsg(pc_segments_[i], pc_msg);
            pc_msg.header = req.input_cloud.header;
            res.segmented_pointclouds.push_back(pc_msg);

            apc_msgs::AlignedDimensions dimensions_msg;

            // Compute PCA to create dimensions aligned along the 3 principal axes

            if(pc_segments_[i].size() < 10) {
                dimensions_msg.width.data = 0;
                dimensions_msg.height.data = 0;
                dimensions_msg.depth.data = 0;
                dimensions_msg.average_z.data = 0;
                dimensions_msg.min_z.data = 0;
                dimensions_msg.max_z.data = 0;
                res.aligned_dimensions.push_back(dimensions_msg);
                continue;
            }
            // Calculate the centroid
            Eigen::Vector4f centroid;
            pcl::compute3DCentroid(pc_segments_[i], centroid);
            pcl::demeanPointCloud<pcl::PointXYZRGB>(pc_segments_[i], centroid, *temp_cloud_);
            /// Compute PCA for the cloud
            pcl::PCA<pcl::PointXYZRGB> pca;
            pca.setInputCloud(temp_cloud_);
            Eigen::Matrix3d eigenvectors_temp = pca.getEigenVectors().cast<double>();
            Eigen::Matrix<double, 3, 3> eigenvectors = eigenvectors_temp;
            // Ensure z-axis direction satisfies right-hand rule
            eigenvectors.col(2) = eigenvectors.col(0).cross(eigenvectors.col(1));
            // Align the principal axes to the objects frame
            Eigen::Matrix<double, 4, 4> transformation(Eigen::Matrix<double, 4, 4>::Identity());
            transformation.template block<3, 3>(0, 0) = eigenvectors.transpose();
            pcl::transformPointCloud(*temp_cloud_, *output_cloud_, transformation);
            pcl::PointXYZRGB lowerBound;
            pcl::PointXYZRGB upperBound;
            pcl::getMinMax3D(*output_cloud_, lowerBound, upperBound);

            dimensions_msg.width.data = upperBound.x - lowerBound.x;
            dimensions_msg.height.data = upperBound.y - lowerBound.y;
            dimensions_msg.depth.data = upperBound.z - lowerBound.z;

            // Match our global reference frame.
            // Yaw is rotation around z.
            eigenvectors.col(0) = eigenvectors_temp.col(1);
            eigenvectors.col(1) = -1 * eigenvectors_temp.col(0);
            eigenvectors.col(2) = eigenvectors.col(0).cross(eigenvectors.col(1));

            eigenvectors = eigenvectors.normalized();
            tf::Matrix3x3 out;
            tf::matrixEigenToTF(eigenvectors, out);
            double roll, pitch, yaw;
            out.getRPY(roll, pitch, yaw);
            // align the yaw to a reachable direction.
            if(yaw > 1.57 || yaw < -1.57) {
                yaw = yaw + 3.14;
            }
            tf::Quaternion transformed_q = tf::createQuaternionFromRPY(0, 0, yaw);
            tf::quaternionTFToMsg(transformed_q, dimensions_msg.pca_rotation);

            // Calculate some other useful things
            double max_z = 0;
            double min_z = 10.0;
            double average_z = 0;
            for(auto &point: pc_segments_[i]) {
                if(point.z > max_z) {
                    max_z = point.z;
                }
                if(point.z < min_z) {
                    min_z = point.z;
                }
                average_z = average_z + point.z;
            }
            average_z = average_z/pc_segments_[i].size();

            dimensions_msg.average_z.data = average_z;
            dimensions_msg.min_z.data = min_z;
            dimensions_msg.max_z.data = max_z;

            res.aligned_dimensions.push_back(dimensions_msg);

        }
        t2 = ros::Time::now();
        ROS_INFO_STREAM("Everything else: " << (t2 - t1));
        return true;

    }

private:

    ros::NodeHandle nh_;
    image_geometry::PinholeCameraModel model_;
    bool camera_initialized_;

    pcl::PointCloud<pcl::PointXYZRGB>::Ptr input_cloud_;
    pcl::StatisticalOutlierRemoval<pcl::PointXYZRGB> sor_;

    std::vector<pcl::PointCloud<pcl::PointXYZRGB>> pc_segments_;

    pcl::PointCloud<pcl::PointXYZRGB>::Ptr temp_cloud_;
    pcl::PointCloud<pcl::PointXYZRGB>::Ptr output_cloud_;

    //pub sub
    // ros::Subscriber camera_info_sub_;
    // boost::shared_ptr<message_filters::Synchronizer<sync_policies::ApproximateTime<sensor_msgs::PointCloud2,sensor_msgs::Image> > > syncPtr_;
    // ros::Publisher proposal_pub_;
    // ros::Publisher overlay_image_pub_;
    // ros::Publisher object_pub_;

    //services
    // ros::ServiceServer enableProposalPubSrv_;
    // ros::ServiceServer enableObjectPubSrv_;
    // ros::ServiceServer enableOverlayPubSrv_;
    // ros::ServiceServer doObjectProposalSrv_;
    // ros::ServiceServer changeCameraInfoSrv_;
    ros::ServiceServer segmentPCSrv_;

    // ros::ServiceClient cnn_client_;

};

int main(int argc, char** argv){

    ros::init(argc,argv,"segment_pointcloud_node");
    ros::NodeHandle nh("~");
    ObjectSegment proposal(nh);

    // cv::namedWindow("Display window", cv::WINDOW_AUTOSIZE );

    ros::spin();

    ROS_INFO("Terminating segment_pointcloud_node ...");

    //cv::destroyAllWindows();
    return 0;
}


// graph-tool -- a general graph modification and manipulation thingy
//
// Copyright (C) 2006-2018 Tiago de Paula Peixoto <tiago@skewed.de>
//
// This program is free software; you can redistribute it and/or
// modify it under the terms of the GNU General Public License
// as published by the Free Software Foundation; either version 3
// of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.

#ifndef GRAPH_SBM_HH
#define GRAPH_SBM_HH

#include <tuple>
#include <iostream>
#include <boost/functional/hash.hpp>

#include "graph.hh"
#include "graph_filtering.hh"
#include "graph_util.hh"
#include "sampler.hh"
#include "urn_sampler.hh"

#include "random.hh"

#include "hash_map_wrap.hh"

namespace graph_tool
{
using namespace std;
using namespace boost;

template <bool micro_deg, class Graph, class VProp, class IVec, class FVec,
          class VDProp, class RNG>
void gen_sbm(Graph& g, VProp b, IVec& rs, IVec& ss, FVec probs, VDProp in_deg,
             VDProp out_deg, bool micro_ers, RNG& rng)
{
    typedef typename std::conditional_t<micro_deg,size_t,double> dtype;
    vector<vector<size_t>> rvs;
    vector<vector<dtype>> v_in_probs, v_out_probs;
    for (auto v : vertices_range(g))
    {
        size_t r = b[v];
        if (r >= v_out_probs.size())
        {
            if (graph_tool::is_directed(g))
                v_in_probs.resize(r+1);
            v_out_probs.resize(r+1);
            rvs.resize(r+1);
        }
        rvs[r].push_back(v);
        if (graph_tool::is_directed(g))
            v_in_probs[r].push_back(in_deg[v]);
        v_out_probs[r].push_back(out_deg[v]);
    }

    typedef std::conditional_t<micro_deg,
                               UrnSampler<size_t, false>,
                               Sampler<size_t>> vsampler_t;
    vector<vsampler_t> v_in_sampler_, v_out_sampler;
    for (size_t r = 0; r < rvs.size(); ++r)
    {
        if (graph_tool::is_directed(g))
            v_in_sampler_.emplace_back(rvs[r], v_in_probs[r]);
        v_out_sampler.emplace_back(rvs[r], v_out_probs[r]);
    }

    auto& v_in_sampler = (graph_tool::is_directed(g)) ? v_in_sampler_ : v_out_sampler;

    for (size_t i = 0; i < rs.shape()[0]; ++i)
    {
        size_t r = rs[i];
        size_t s = ss[i];
        auto p = probs[i];

        if (p == 0)
            continue;

        if (!graph_tool::is_directed(g) && r == s)
            p /= 2;

        auto& r_sampler = v_out_sampler[r];
        auto& s_sampler = v_in_sampler[s];

        size_t mrs;
        if (micro_ers)
        {
            mrs = p;
        }
        else
        {
            std::poisson_distribution<> poi(p);
            mrs = poi(rng);
        }

        size_t ers = (&r_sampler != &s_sampler) ? mrs : 2 * mrs;
        if (!r_sampler.has_n(ers) || !s_sampler.has_n(ers))
            throw GraphException("Inconsistent SBM parameters: node degrees do not agree with matrix of edge counts between groups");

        for (size_t j = 0; j < mrs; ++j)
        {
            size_t u = r_sampler.sample(rng);
            size_t v = s_sampler.sample(rng);
            add_edge(u, v, g);
        }
    }
}

} // graph_tool namespace

#endif // GRAPH_SBM_HH


// Copyright (c) 2015-2020 The Bitcoin Core developers
// Distributed under the MIT software license, see the accompanying
// file COPYING or https://opensource.org/licenses/mit-license.php.

#include <consensus/merkle.h>

#include <boost/test/unit_test.hpp>

BOOST_AUTO_TEST_SUITE(merkle_tests)

// Older version of the merkle root computation code, for comparison.
static uint256 BlockBuildMerkleTree(const CBlock& block, std::vector<uint256>& vMerkleTree)
{
    vMerkleTree.clear();
    for (auto const& tx : block.vtx)
        vMerkleTree.push_back(tx.GetHash());
    int j = 0;
    for (int nSize = block.vtx.size(); nSize > 1; nSize = (nSize + 1) / 2)
    {
        for (int i = 0; i < nSize; i += 2)
        {
            int i2 = std::min(i+1, nSize-1);
            vMerkleTree.push_back(Hash(BEGIN(vMerkleTree[j+i]),  END(vMerkleTree[j+i]),
                                        BEGIN(vMerkleTree[j+i2]), END(vMerkleTree[j+i2])));
        }
        j += nSize;
    }
    return (vMerkleTree.empty() ? uint256() : vMerkleTree.back());
}

static inline int ctz(uint32_t i) {
    if (i == 0) return 0;
    int j = 0;
    while (!(i & 1)) {
        j++;
        i >>= 1;
    }
    return j;
}

BOOST_AUTO_TEST_CASE(merkle_test)
{
    for (int i = 0; i < 32; i++) {
        // Try 32 block sizes: all sizes from 0 to 16 inclusive, and then 15 random sizes.
        int ntx = (i <= 16) ? i : 17 + (GetRand(4000));
        // Try up to 3 mutations.
        for (int mutate = 0; mutate <= 3; mutate++) {
            int duplicate1 = mutate >= 1 ? 1 << ctz(ntx) : 0; // The last how many transactions to duplicate first.
            if (duplicate1 >= ntx) break; // Duplication of the entire tree results in a different root (it adds a level).
            int ntx1 = ntx + duplicate1; // The resulting number of transactions after the first duplication.
            int duplicate2 = mutate >= 2 ? 1 << ctz(ntx1) : 0; // Likewise for the second mutation.
            if (duplicate2 >= ntx1) break;
            int ntx2 = ntx1 + duplicate2;
            int duplicate3 = mutate >= 3 ? 1 << ctz(ntx2) : 0; // And for the third mutation.
            if (duplicate3 >= ntx2) break;
            int ntx3 = ntx2 + duplicate3;
            // Build a block with ntx different transactions.
            CBlock block;
            block.vtx.resize(ntx);
            for (int j = 0; j < ntx; j++) {
                CTransaction tx;
                tx.nLockTime = j;
                block.vtx[j] = tx;
            }
            // Compute the root of the block before mutating it.
            bool unmutatedMutated = false;
            uint256 unmutatedRoot = BlockMerkleRoot(block, &unmutatedMutated);
            BOOST_CHECK(unmutatedMutated == false);
            // Optionally mutate by duplicating the last transactions, resulting in the same merkle root.
            block.vtx.resize(ntx3);
            for (int j = 0; j < duplicate1; j++) {
                block.vtx[ntx + j] = block.vtx[ntx + j - duplicate1];
            }
            for (int j = 0; j < duplicate2; j++) {
                block.vtx[ntx1 + j] = block.vtx[ntx1 + j - duplicate2];
            }
            for (int j = 0; j < duplicate3; j++) {
                block.vtx[ntx2 + j] = block.vtx[ntx2 + j - duplicate3];
            }
            // Compute the merkle root and merkle tree using the old mechanism.
            std::vector<uint256> merkleTree;
            uint256 oldRoot = BlockBuildMerkleTree(block, merkleTree);
            // Compute the merkle root using the new mechanism.
            bool newMutated = false;
            uint256 newRoot = BlockMerkleRoot(block, &newMutated);
            BOOST_CHECK(oldRoot == newRoot);
            BOOST_CHECK(newRoot == unmutatedRoot);
            BOOST_CHECK((newRoot == uint256()) == (ntx == 0));
            BOOST_CHECK(newMutated == !!mutate);
        }
    }
}

BOOST_AUTO_TEST_CASE(merkle_test_empty_block)
{
    bool mutated = false;
    CBlock block;
    uint256 root = BlockMerkleRoot(block, &mutated);

    BOOST_CHECK_EQUAL(root.IsNull(), true);
    BOOST_CHECK_EQUAL(mutated, false);
}

BOOST_AUTO_TEST_CASE(merkle_test_oneTx_block)
{
    bool mutated = false;
    CBlock block;

    block.vtx.resize(1);
    CTransaction tx;
    tx.nLockTime = 0;
    block.vtx[0] = tx;
    uint256 root = BlockMerkleRoot(block, &mutated);
    BOOST_CHECK(root == block.vtx[0].GetHash());
    BOOST_CHECK_EQUAL(mutated, false);
}

BOOST_AUTO_TEST_CASE(merkle_test_OddTxWithRepeatedLastTx_block)
{
    bool mutated;
    CBlock block, blockWithRepeatedLastTx;

    block.vtx.resize(3);

    for (std::size_t pos = 0; pos < block.vtx.size(); pos++) {
        CTransaction tx;
        tx.nLockTime = pos;
        block.vtx[pos] = tx;
    }

    blockWithRepeatedLastTx = block;
    blockWithRepeatedLastTx.vtx.push_back(blockWithRepeatedLastTx.vtx.back());

    uint256 rootofBlock = BlockMerkleRoot(block, &mutated);
    BOOST_CHECK_EQUAL(mutated, false);

    uint256 rootofBlockWithRepeatedLastTx = BlockMerkleRoot(blockWithRepeatedLastTx, &mutated);
    BOOST_CHECK(rootofBlock == rootofBlockWithRepeatedLastTx);
    BOOST_CHECK_EQUAL(mutated, true);
}

BOOST_AUTO_TEST_CASE(merkle_test_LeftSubtreeRightSubtree)
{
    CBlock block, leftSubtreeBlock, rightSubtreeBlock;

    block.vtx.resize(4);
    std::size_t pos;
    for (pos = 0; pos < block.vtx.size(); pos++) {
        CTransaction tx;
        tx.nLockTime = pos;
        block.vtx[pos] = tx;
    }

    for (pos = 0; pos < block.vtx.size() / 2; pos++)
        leftSubtreeBlock.vtx.push_back(block.vtx[pos]);

    for (pos = block.vtx.size() / 2; pos < block.vtx.size(); pos++)
        rightSubtreeBlock.vtx.push_back(block.vtx[pos]);

    uint256 root = BlockMerkleRoot(block);
    uint256 rootOfLeftSubtree = BlockMerkleRoot(leftSubtreeBlock);
    uint256 rootOfRightSubtree = BlockMerkleRoot(rightSubtreeBlock);
    std::vector<uint256> leftRight;
    leftRight.push_back(rootOfLeftSubtree);
    leftRight.push_back(rootOfRightSubtree);
    uint256 rootOfLR = ComputeMerkleRoot(leftRight);

    BOOST_CHECK(root == rootOfLR);
}

BOOST_AUTO_TEST_SUITE_END()


//==============================================================================
//         Copyright 2003 - 2012   LASMEA UMR 6602 CNRS/Univ. Clermont II
//         Copyright 2009 - 2012   LRI    UMR 8623 CNRS/Univ Paris Sud XI
//
//          Distributed under the Boost Software License, Version 1.0.
//                 See accompanying file LICENSE.txt or copy at
//                     http://www.boost.org/LICENSE_1_0.txt
//==============================================================================
#ifndef NT2_INTERPOL_FUNCTIONS_IDXY_BILINEAR_HPP_INCLUDED
#define NT2_INTERPOL_FUNCTIONS_IDXY_BILINEAR_HPP_INCLUDED
#include <nt2/include/functor.hpp>
#include <nt2/sdk/meta/value_as.hpp>
#include <nt2/core/container/dsl/size.hpp>
#include <nt2/core/container/dsl/value_type.hpp>
#include <boost/mpl/bool.hpp>
#include <nt2/sdk/meta/adapted_traits.hpp>

namespace nt2 { namespace tag
  {
   /*!
     @brief idxy_bilinear generic tag

     Represents the idxy_bilinear function in generic contexts.

     @par Models:
        Hierarchy
   **/
    struct idxy_bilinear_ : ext::unspecified_<idxy_bilinear_>
    {
      /// @brief Parent hierarchy
      typedef ext::unspecified_<idxy_bilinear_> parent;
      template<class... Args>
      static BOOST_FORCEINLINE BOOST_AUTO_DECLTYPE dispatch(Args&&... args)
      BOOST_AUTO_DECLTYPE_BODY( dispatching_idxy_bilinear_( ext::adl_helper(), static_cast<Args&&>(args)... ) )
    };
  }
  namespace ext
  {
   template<class Site>
   BOOST_FORCEINLINE generic_dispatcher<tag::idxy_bilinear_, Site> dispatching_idxy_bilinear_(adl_helper, boost::dispatch::meta::unknown_<Site>, ...)
   {
     return generic_dispatcher<tag::idxy_bilinear_, Site>();
   }
   template<class... Args>
   struct impl_idxy_bilinear_;
  }
  /*!
    Performs two dimensional idxy_bilinear interpolation given an
    array a and two vector idx and idy of "real" indices provides the
    interpolated values along the two chosen array dimensions by
    bilinear formula.

    @par Semantic:

    calls can be of the forms:

      - idxy_bilinear(a0, idx, idy)
      - idxy_bilinear(a0, idx, idy, true) allowing extrapolation
      - idxy_bilinear(a0, idx, idy, val1) putting val1 outside the bounds
      - idxy_bilinear(a0, idx, idy, val1, val2)}, putting val1 under the index bounds val2 over
      - idxy_bilinear(a0, idx, idy, val1x val2x, val1y, val2y)
      - idxy_bilinear(a0, idx, idy, _, dim1, dim2) are the dimensions of interpolation 2 (rows), 1 (columns) of a0 by default
      - idxy_bilinear(a0, idx, idy, val1, dim1, dim2)
      - idxy_bilinear(a0, idx, idy, val1, val2, dim1, dim2)
      - idxy_bilinear(a0, idx, idy, val1x val2x, val1y, val2y, dim1, dim2) can also be used

      typically if a is a matrix, idx is a value situated between i and i+1 and
      idy is a value situated between j and j+1 the result is:
      @code
      (a(i, j)*(idx-i)+a(i+1, j)*(i+1-idx))*(jdx-j)+
      (a(i, j+1)*(idx-i)+a(i+1, j+1)*(i+1-idx))*(j+1-jdx)
      @endcode

    @param a0
    @param a1
    @param a2
    @param a3
    @param a4
    @param a5
    @param a6
    @param a7
    @param a8
    @param a9

    @return a value of the same type as the parameter
  **/
  NT2_FUNCTION_IMPLEMENTATION(tag::idxy_bilinear_, idxy_bilinear, 10)
  /// @overload
  NT2_FUNCTION_IMPLEMENTATION(tag::idxy_bilinear_, idxy_bilinear, 3)
  /// @overload
  NT2_FUNCTION_IMPLEMENTATION(tag::idxy_bilinear_, idxy_bilinear, 4)
  /// @overload
  NT2_FUNCTION_IMPLEMENTATION(tag::idxy_bilinear_, idxy_bilinear, 5)
  /// @overload
  NT2_FUNCTION_IMPLEMENTATION(tag::idxy_bilinear_, idxy_bilinear, 6)
  /// @overload
  NT2_FUNCTION_IMPLEMENTATION(tag::idxy_bilinear_, idxy_bilinear, 7)
  /// @overload
  NT2_FUNCTION_IMPLEMENTATION(tag::idxy_bilinear_, idxy_bilinear, 9)
}

namespace nt2 { namespace ext
{
  template<class Domain, class Expr,  int N>
  struct size_of<tag::idxy_bilinear_, Domain, N, Expr>
  {
    typedef typename boost::proto::result_of::child_c<Expr&,N-2>::type   choice1_t;
    typedef typename boost::proto::result_of::child_c<Expr&,N-1>::type   choice2_t;

    typedef typename  boost::proto::result_of
                      ::child_c<Expr&, 1>::value_type::extent_type  result_type;

    result_type operator()(Expr& e) const
    {
      const choice1_t & v1   =  boost::proto::child_c<N-2>(e);
      const choice2_t & v2   =  boost::proto::child_c<N-1>(e);
      typedef typename nt2::meta::is_integral<choice1_t>::type c_t;
      std::size_t dim1 = 2, dim2 = 1;
      getdims(boost::proto::child_c<1>(e), v1, v2, dim1, dim2, c_t());
      result_type sizee = boost::proto::child_c<0>(e).extent();
      sizee[dim1-1] = numel(boost::proto::child_c<1>(e));
      sizee[dim2-1] = numel(boost::proto::child_c<2>(e));
      return sizee;
    }

    template<class X, class C1, class C2> static
    void getdims( const X &, const C1 &, const C2 &
                , std::size_t&, std::size_t&, const boost::mpl::false_&
                )
    {}

    template<class X, class C1, class C2> static
    void getdims( const X &, const C1& v1, const C2& v2
                , std::size_t& dim1, std::size_t& dim2, const boost::mpl::true_&
                )
    {
      dim1 = v1;
      dim2 = v2;
    }
  };

  template <class Domain, class Expr,  int N>
  struct  value_type < tag::idxy_bilinear_, Domain,N,Expr>
        : meta::value_as<Expr,0>
  {};
} }

#endif


#include "Stat.h"
#include <macgyver/Exception.h>

#include <boost/accumulators/accumulators.hpp>
#include <boost/accumulators/statistics/stats.hpp>
#include <boost/accumulators/statistics/weighted_mean.hpp>
#include <boost/accumulators/statistics/weighted_median.hpp>
#include <boost/accumulators/statistics/weighted_variance.hpp>

#include <cmath>
#include <iterator>
#include <numeric>
#include <stdexcept>

namespace SmartMet
{
namespace TimeSeries
{
namespace Stat
{

#define MODULO_VALUE_360 360.0
using DataIterator = DataVector::iterator;
using DataConstIterator = DataVector::const_iterator;

using namespace boost::accumulators;
using namespace std;

bool comp_value(const DataItem& data1, const DataItem& data2)
{
  try
  {
    return data1.value < data2.value;
  }
  catch (...)
  {
    throw Fmi::Exception::Trace(BCP, "Operation failed!");
  }
}

bool comp_time(const DataItem& data1, const DataItem& data2)
{
  try
  {
    return data1.time < data2.time;
  }
  catch (...)
  {
    throw Fmi::Exception::Trace(BCP, "Operation failed!");
  }
}

Stat::Stat(double theMissingValue /*= numeric_limits<double>::quiet_NaN()*/)
    : itsMissingValue(theMissingValue), itsWeights(true), itsDegrees(false)
{
}

Stat::Stat(const DataVector& theData,
           double theMissingValue /*= numeric_limits<double>::quiet_NaN()*/)
    : itsData(theData), itsMissingValue(theMissingValue), itsWeights(true), itsDegrees(false)
{
  try
  {
    calculate_weights();
  }
  catch (...)
  {
    throw Fmi::Exception::Trace(BCP, "Operation failed!");
  }
}

Stat::Stat(const std::vector<double>& theValues,
           double theMissingValue /*= std::numeric_limits<double>::quiet_NaN()*/)
    : itsMissingValue(theMissingValue), itsWeights(false), itsDegrees(false)
{
  try
  {
    addData(theValues);
  }
  catch (...)
  {
    throw Fmi::Exception::Trace(BCP, "Operation failed!");
  }
}

Stat::Stat(const LocalDateTimeValueVector& theValues,
           double theMissingValue /*= std::numeric_limits<double>::quiet_NaN()*/)
    : itsMissingValue(theMissingValue), itsWeights(true), itsDegrees(false)

{
  try
  {
    for (const LocalDateTimeValue& item : theValues)
    {
      itsData.push_back(DataItem(item.first.utc_time(), item.second));
    }
    calculate_weights();
  }
  catch (...)
  {
    throw Fmi::Exception::Trace(BCP, "Operation failed!");
  }
}

Stat::Stat(const PosixTimeValueVector& theValues,
           double theMissingValue /*= std::numeric_limits<double>::quiet_NaN(*/)
    : itsMissingValue(theMissingValue), itsWeights(true), itsDegrees(false)
{
  try
  {
    for (const PosixTimeValue& item : theValues)
    {
      itsData.push_back(DataItem(item.first, item.second));
    }
    calculate_weights();
  }
  catch (...)
  {
    throw Fmi::Exception::Trace(BCP, "Operation failed!");
  }
}

void Stat::setData(const DataVector& theData)
{
  try
  {
    itsData = theData;
    calculate_weights();
  }
  catch (...)
  {
    throw Fmi::Exception::Trace(BCP, "Operation failed!");
  }
}

void Stat::addData(double theValue)
{
  try
  {
    itsData.push_back(DataItem(not_a_date_time, theValue));
    calculate_weights();
  }
  catch (...)
  {
    throw Fmi::Exception::Trace(BCP, "Operation failed!");
  }
}

void Stat::addData(const boost::local_time::local_date_time& theTime, double theValue)
{
  try
  {
    itsData.push_back(DataItem(theTime.utc_time(), theValue));
    calculate_weights();
  }
  catch (...)
  {
    throw Fmi::Exception::Trace(BCP, "Operation failed!");
  }
}

void Stat::addData(const boost::posix_time::ptime& theTime, double theValue)
{
  try
  {
    itsData.push_back(DataItem(theTime, theValue));
    calculate_weights();
  }
  catch (...)
  {
    throw Fmi::Exception::Trace(BCP, "Operation failed!");
  }
}

void Stat::addData(const vector<double>& theData)
{
  try
  {
    for (double value : theData)
      itsData.push_back(DataItem(not_a_date_time, value));

    calculate_weights();
  }
  catch (...)
  {
    throw Fmi::Exception::Trace(BCP, "Operation failed!");
  }
}

void Stat::addData(const DataItem& theData)
{
  try
  {
    itsData.push_back(theData);
    calculate_weights();
  }
  catch (...)
  {
    throw Fmi::Exception::Trace(BCP, "Operation failed!");
  }
}

void Stat::clear()
{
  try
  {
    itsData.clear();
  }
  catch (...)
  {
    throw Fmi::Exception::Trace(BCP, "Operation failed!");
  }
}
// time in seconds is used as a weight
double Stat::integ(const boost::posix_time::ptime& startTime /*= not_a_date_time */,
                   const boost::posix_time::ptime& endTime /*= not_a_date_time */) const
{
  try
  {
#ifdef MYDEBUG
    std::cout << "integ(" << startTime << ", " << endTime << ")" << std::endl;
#endif
    DataVector subvector;

    if (!get_subvector(subvector, startTime, endTime))
      return itsMissingValue;

    accumulator_set<double, stats<tag::sum>, double> acc;

    for (const DataItem& item : subvector)
    {
      acc(item.value, weight = item.weight);
    }

    if (itsDegrees)
      return fmod(boost::accumulators::sum(acc), MODULO_VALUE_360) / 3600.0;

    return boost::accumulators::sum(acc) / 3600.0;
  }
  catch (...)
  {
    throw Fmi::Exception::Trace(BCP, "Operation failed!");
  }
}

// weight is alwaus 1.0
double Stat::sum(const boost::posix_time::ptime& startTime /*= not_a_date_time */,
                 const boost::posix_time::ptime& endTime /*= not_a_date_time */) const
{
  try
  {
#ifdef MYDEBUG
    std::cout << "sum(" << startTime << ", " << endTime << ")" << std::endl;
#endif
    DataVector subvector;

    if (!get_subvector(subvector, startTime, endTime))
      return itsMissingValue;

    accumulator_set<double, stats<tag::sum>, double> acc;

    for (const DataItem& item : subvector)
    {
      acc(item.value, weight = 1.0);
    }

    if (itsDegrees)
      return fmod(boost::accumulators::sum(acc), MODULO_VALUE_360);

    return boost::accumulators::sum(acc);
  }
  catch (...)
  {
    throw Fmi::Exception::Trace(BCP, "Operation failed!");
  }
}

double Stat::min(const boost::posix_time::ptime& startTime /*= not_a_date_time */,
                 const boost::posix_time::ptime& endTime /*= not_a_date_time */) const
{
  try
  {
#ifdef MYDEBUG
    std::cout << "min(" << startTime << ", " << endTime << ")" << std::endl;
#endif
    DataVector subvector;

    if (!get_subvector(subvector, startTime, endTime) || subvector.empty())
      return itsMissingValue;

    return std::min_element(subvector.begin(), subvector.end(), comp_value)->value;
  }
  catch (...)
  {
    throw Fmi::Exception::Trace(BCP, "Operation failed!");
  }
}

double Stat::mean(const boost::posix_time::ptime& startTime /*= not_a_date_time */,
                  const boost::posix_time::ptime& endTime /*= not_a_date_time */) const
{
  try
  {
#ifdef MYDEBUG
    std::cout << "mean(" << startTime << ", " << endTime << ")" << std::endl;
#endif
    DataVector subvector;

    if (!get_subvector(subvector, startTime, endTime))
      return itsMissingValue;

    if (itsDegrees)
    {
      double previousDirection = 0;
      double directionSum = 0;
      for (unsigned int i = 0; i < subvector.size(); i++)
      {
        if (i == 0)
        {
          directionSum = subvector[i].value;
          previousDirection = subvector[i].value;
        }
        else
        {
          double diff = subvector[i].value - previousDirection;
          double direction = previousDirection + diff;

          if (diff < -MODULO_VALUE_360 / 2.0)
            direction += MODULO_VALUE_360;
          else if (diff > MODULO_VALUE_360 / 2.0)
            direction -= MODULO_VALUE_360;

          directionSum += direction;
          previousDirection = direction;
        }
      }

      double mean = directionSum / subvector.size();
      mean -= (MODULO_VALUE_360 * floor(mean / MODULO_VALUE_360));
      return mean;
    }

    accumulator_set<double, stats<tag::mean>, double> acc;
    for (const DataItem& item : subvector)
    {
      acc(item.value, weight = (itsWeights ? item.weight : 1.0));
    }
    return boost::accumulators::mean(acc);
  }
  catch (...)
  {
    throw Fmi::Exception::Trace(BCP, "Operation failed!");
  }
}

double Stat::max(const boost::posix_time::ptime& startTime /*= not_a_date_time */,
                 const boost::posix_time::ptime& endTime /*= not_a_date_time */) const
{
  try
  {
#ifdef MYDEBUG
    std::cout << "max(" << startTime << ", " << endTime << ")" << std::endl;
#endif
    DataVector subvector;

    if (!get_subvector(subvector, startTime, endTime))
      return itsMissingValue;

    return std::max_element(subvector.begin(), subvector.end(), comp_value)->value;
  }
  catch (...)
  {
    throw Fmi::Exception::Trace(BCP, "Operation failed!");
  }
}

double Stat::change(const boost::posix_time::ptime& startTime /*= not_a_date_time */,
                    const boost::posix_time::ptime& endTime /*= not_a_date_time */) const
{
  try
  {
#ifdef MYDEBUG
    std::cout << "change(" << startTime << ", " << endTime << ")" << std::endl;
#endif
    DataVector subvector;

    if (!get_subvector(subvector, startTime, endTime) || subvector.empty())
      return itsMissingValue;

    if (itsDegrees)
    {
      double previousValue = 0;
      double cumulativeChange = 0;
      for (unsigned int i = 0; i < subvector.size(); i++)
      {
        if (i > 0)
        {
          double diff = subvector[i].value - previousValue;
          if (diff < -MODULO_VALUE_360 / 2.0)
            cumulativeChange += diff + MODULO_VALUE_360;
          else if (diff > MODULO_VALUE_360 / 2.0)
            cumulativeChange += diff - MODULO_VALUE_360;
          else
            cumulativeChange += diff;
        }
        previousValue = subvector[i].value;
      }
      return cumulativeChange;
    }

    double firstValue(subvector.begin()->value);
    double lastValue((subvector.end() - 1)->value);

    return lastValue - firstValue;
  }
  catch (...)
  {
    throw Fmi::Exception::Trace(BCP, "Operation failed!");
  }
}

double Stat::trend(const boost::posix_time::ptime& startTime /*= not_a_date_time */,
                   const boost::posix_time::ptime& endTime /*= not_a_date_time */) const
{
  try
  {
#ifdef MYDEBUG
    std::cout << "trend(" << startTime << ", " << endTime << ")" << std::endl;
#endif
    DataVector subvector;

    if (!get_subvector(subvector, startTime, endTime) || subvector.size() <= 1)
      return itsMissingValue;

    long positiveChanges(0);
    long negativeChanges(0);
    double lastValue = 0;

    for (auto iter = subvector.begin(); iter != subvector.end(); iter++)
    {
      if (iter != subvector.begin())
      {
        if (itsDegrees)
        {
          const double diff = iter->value - lastValue;
          if (diff < -MODULO_VALUE_360 / 2.0)
            ++positiveChanges;
          else if (diff > MODULO_VALUE_360 / 2.0)
            ++negativeChanges;
          else if (diff < 0)
            ++negativeChanges;
          else if (diff > 0)
            ++positiveChanges;
        }
        else
        {
          if (iter->value > lastValue)
            ++positiveChanges;
          else if (iter->value < lastValue)
            ++negativeChanges;
        }
      }
      lastValue = iter->value;
    }

    return static_cast<double>(positiveChanges - negativeChanges) /
           static_cast<double>(subvector.size() - 1) * 100.0;
  }
  catch (...)
  {
    throw Fmi::Exception::Trace(BCP, "Operation failed!");
  }
}

unsigned int Stat::count(double lowerLimit,
                         double upperLimit,
                         const boost::posix_time::ptime& startTime /*= not_a_date_time */,
                         const boost::posix_time::ptime& endTime /*= not_a_date_time */) const
{
  try
  {
#ifdef MYDEBUG
    std::cout << "count(" << lowerLimit << ", " << upperLimit << ", " << startTime << ", "
              << endTime << ")" << std::endl;
#endif
    DataVector subvector;

    if (!get_subvector(subvector, startTime, endTime, false))
      return static_cast<unsigned int>(itsMissingValue);

    unsigned int occurances(0);

    for (const DataItem& item : subvector)
      if (item.value >= lowerLimit && item.value <= upperLimit)
        occurances++;

    return occurances;
  }
  catch (...)
  {
    throw Fmi::Exception::Trace(BCP, "Operation failed!");
  }
}

double Stat::percentage(double lowerLimit,
                        double upperLimit,
                        const boost::posix_time::ptime& startTime /*= not_a_date_time */,
                        const boost::posix_time::ptime& endTime /*= not_a_date_time */) const
{
  try
  {
#ifdef MYDEBUG
    std::cout << "percentage(" << lowerLimit << ", " << upperLimit << ", " << startTime << ", "
              << endTime << ")" << std::endl;
#endif
    DataVector subvector;

    if (!get_subvector(subvector, startTime, endTime))
      return itsMissingValue;

    int occurances(0);
    int total_count(0);

    for (const DataItem& item : subvector)
    {
      if (item.value >= lowerLimit && item.value <= upperLimit)
        occurances += (itsWeights ? item.weight : 1);
      total_count += (itsWeights ? item.weight : 1);
    }

    if (occurances == 0)
      return 0.0;

    return static_cast<double>(
        (static_cast<double>(occurances) / static_cast<double>(total_count)) * 100.0);
  }
  catch (...)
  {
    throw Fmi::Exception::Trace(BCP, "Operation failed!");
  }
}

double Stat::median(const boost::posix_time::ptime& startTime /*= not_a_date_time */,
                    const boost::posix_time::ptime& endTime /*= not_a_date_time */) const
{
  try
  {
#ifdef MYDEBUG
    std::cout << "median(" << startTime << ", " << endTime << ")" << std::endl;
#endif
    DataVector subvector;

    if (!get_subvector(subvector, startTime, endTime) || subvector.empty())
      return itsMissingValue;

    if (subvector.size() == 1)
      return subvector[0].value;

    vector<double> double_vector;

    for (const DataItem& item : subvector)
    {
      double_vector.insert(double_vector.end(),
                           static_cast<std::size_t>(itsWeights ? item.weight : 1.0),
                           item.value);
    }

    std::size_t vector_size = double_vector.size();

    std::size_t middle_index = (vector_size / 2) + 1;

    auto middle_pos = double_vector.begin();
    std::advance(middle_pos, middle_index);
    partial_sort(double_vector.begin(), middle_pos, double_vector.end());

    if (double_vector.size() % 2 == 0)
    {
      return ((double_vector[static_cast<std::size_t>((vector_size / 2) - 1)] +
               double_vector[(vector_size / 2)]) /
              2.0);
    }

    return double_vector[static_cast<std::size_t>(((vector_size + 1) / 2) - 1)];
  }
  catch (...)
  {
    throw Fmi::Exception::Trace(BCP, "Operation failed!");
  }
}

double Stat::variance(const boost::posix_time::ptime& startTime /*= not_a_date_time */,
                      const boost::posix_time::ptime& endTime /*= not_a_date_time */) const
{
  try
  {
#ifdef MYDEBUG
    std::cout << "variance(" << startTime << ", " << endTime << ")" << std::endl;
#endif
    DataVector subvector;

    if (!get_subvector(subvector, startTime, endTime) || subvector.empty())
      return itsMissingValue;

    accumulator_set<double, stats<tag::variance>, double> acc;

    for (const DataItem& item : subvector)
    {
      acc(item.value, weight = (itsWeights ? item.weight : 1.0));
    }

    return boost::accumulators::variance(acc);
  }
  catch (...)
  {
    throw Fmi::Exception::Trace(BCP, "Operation failed!");
  }
}

double Stat::stddev(const boost::posix_time::ptime& startTime /*= not_a_date_time */,
                    const boost::posix_time::ptime& endTime /*= not_a_date_time */) const
{
  try
  {
#ifdef MYDEBUG
    std::cout << "stddev(" << startTime << ", " << endTime << ")" << std::endl;
#endif
    if (itsDegrees)
    {
      DataVector subvector;

      if (!get_subvector(subvector, startTime, endTime) || subvector.empty())
        return itsMissingValue;

      double sum = 0;
      double squaredSum = 0;
      double previousDirection = 0;
      for (unsigned int i = 0; i < subvector.size(); i++)
      {
        if (i == 0)
        {
          sum = subvector[i].value;
          squaredSum = subvector[i].value * subvector[i].value;
          previousDirection = subvector[i].value;
        }
        else
        {
          double diff = subvector[i].value - previousDirection;
          double dir = previousDirection + diff;
          if (diff < -MODULO_VALUE_360 / 2.0)
          {
            while (dir < MODULO_VALUE_360 / 2.0)
              dir += MODULO_VALUE_360;
          }
          else if (diff > MODULO_VALUE_360 / 2.0)
          {
            while (dir > MODULO_VALUE_360 / 2.0)
              dir -= MODULO_VALUE_360;
          }
          sum += dir;
          squaredSum += dir * dir;
          previousDirection = dir;
        }
      }
      double tmp = squaredSum - sum * sum / subvector.size();
      if (tmp < 0)
        return 0.0;

      return sqrt(tmp / subvector.size());
    }

    return sqrt(variance(startTime, endTime));
  }
  catch (...)
  {
    throw Fmi::Exception::Trace(BCP, "Operation failed!");
  }
}

double Stat::nearest(const boost::posix_time::ptime& timestep,
                     const boost::posix_time::ptime& startTime /*= not_a_date_time */,
                     const boost::posix_time::ptime& endTime /*= not_a_date_time */) const
{
  try
  {
#ifdef MYDEBUG
    std::cout << "nearest(" << timestep << ", " << startTime << ", " << endTime << ")" << std::endl;
#endif

    if (timestep == not_a_date_time)
      return itsMissingValue;

    DataVector subvector;

    if (!get_subvector(subvector, startTime, endTime, false) || subvector.empty())
      return itsMissingValue;

    if (subvector.size() == 1)
      return subvector[0].value;

    double nearest_value = itsMissingValue;
    long nearest_seconds = -1;

    for (unsigned int i = 0; i < subvector.size(); i++)
    {
      if (i == 0 || abs((subvector[i].time - timestep).total_seconds()) < nearest_seconds)
      {
        nearest_value = subvector[i].value;
        nearest_seconds = abs((subvector[i].time - timestep).total_seconds());
      }
    }

    return nearest_value;
  }
  catch (...)
  {
    throw Fmi::Exception::Trace(BCP, "Operation failed!");
  }
}

// Do linear inter-/extrapolation
double Stat::interpolate(const boost::posix_time::ptime& timestep,
                         const boost::posix_time::ptime& startTime /*= not_a_date_time */,
                         const boost::posix_time::ptime& endTime /*= not_a_date_time */) const
{
  try
  {
#ifdef MYDEBUG
    std::cout << "interpolate(" << timestep << ", " << startTime << ", " << endTime << ")"
              << std::endl;
#endif

    if (timestep == not_a_date_time)
      return itsMissingValue;

    DataVector subvector;

    if (!get_subvector(subvector, startTime, endTime, false))
      return itsMissingValue;

    // If there is only one timestep in the data and it is the same as requested we return the value
    if (subvector.size() == 1 && subvector.at(0).time == timestep)
      return subvector.at(0).value;

    // There must be at least two timesteps with valid values in order to interpolate
    if (subvector.size() == 2 &&
        (subvector.at(0).value == itsMissingValue || subvector.at(1).value == itsMissingValue))
      return itsMissingValue;

    std::vector<int> indicator_vector;  // -1 indicates earliter, + 1 later than requested timestep

    for (const auto& val : subvector)
    {
      if (val.value == itsMissingValue)
        continue;
      if (val.time < timestep)
        indicator_vector.push_back(-1);
      else if (val.time > timestep)
        indicator_vector.push_back(1);
      else if (val.time == timestep)
        return val.value;  // Exact timestep found
    }

    if (indicator_vector.size() < 2)
      return itsMissingValue;

    boost::posix_time::ptime first_time = boost::posix_time::not_a_date_time;
    boost::posix_time::ptime second_time = boost::posix_time::not_a_date_time;
    double first_value = itsMissingValue;
    double second_value = itsMissingValue;
    double time_diff_to_timestep_sec = 0.0;
    if (indicator_vector.back() == -1)
    {
      // All timesteps are earlier -> extrapolate to the future
      int last_index = indicator_vector.size() - 1;
      first_time = subvector.at(last_index - 1).time;
      second_time = subvector.at(last_index).time;
      first_value = subvector.at(last_index - 1).value;
      second_value = subvector.at(last_index).value;
      time_diff_to_timestep_sec = (timestep - first_time).total_seconds();
    }
    else if (indicator_vector.front() == 1)
    {
      // All timesteps are later -> extrapolate to the past
      first_time = subvector.at(0).time;
      second_time = subvector.at(1).time;
      first_value = subvector.at(0).value;
      second_value = subvector.at(1).value;
      time_diff_to_timestep_sec = (timestep - first_time).total_seconds();
    }
    else
    {
      // Interpolate between timesteps
      for (unsigned int i = 1; i < indicator_vector.size(); i++)
      {
        if (indicator_vector.at(i - 1) == -1 && indicator_vector.at(i) == 1)
        {
          first_time = subvector.at(i - 1).time;
          second_time = subvector.at(i).time;
          first_value = subvector.at(i - 1).value;
          second_value = subvector.at(i).value;
          time_diff_to_timestep_sec = (timestep - first_time).total_seconds();
          break;
        }
      }
    }

    double value_diff = second_value - first_value;
    double time_diff_sec = (second_time - first_time).total_seconds();
    double slope = value_diff / time_diff_sec;
    double interpolated_value = (first_value + (slope * time_diff_to_timestep_sec));

#ifdef MYDEBUG
    std::cout << "first_time: " << first_time << ",second_time: " << second_time
              << ", timestep: " << timestep << ", first_value: " << first_value
              << ", second_value: " << second_value << ",slope: " << slope
              << ", time_diff_to_timestep: " << time_diff_to_timestep_sec << " -> "
              << interpolated_value << std::endl;
#endif

    return interpolated_value;
  }
  catch (...)
  {
    throw Fmi::Exception::Trace(BCP, "Operation failed!");
  }
}

bool extract_subvector(const DataVector& itsData,
                       DataVector& subvector,
                       const boost::posix_time::ptime& startTime,
                       const boost::posix_time::ptime& endTime,
                       double itsMissingValue,
                       bool itsWeights)
{
  try
  {
    // if startTime is later than endTime return empty vector
    if ((startTime != not_a_date_time && endTime != not_a_date_time) && startTime > endTime)
      return false;

    boost::posix_time::ptime firstTimestamp(
        (startTime == not_a_date_time || startTime < itsData[0].time) ? itsData[0].time
                                                                      : startTime);
    boost::posix_time::ptime lastTimestamp(
        (endTime == not_a_date_time || endTime > itsData[itsData.size() - 1].time)
            ? itsData[itsData.size() - 1].time
            : endTime);

    time_period query_period(firstTimestamp, lastTimestamp + microseconds(1));

    // if there is only one element in the vector
    if (itsData.size() == 1)
    {
      if (query_period.contains(itsData[0].time))
      {
        if (itsData[0].value == itsMissingValue)
          return false;
        subvector.push_back(DataItem(itsData[0].time, itsData[0].value, 1.0));
      }
      return true;
    }

#ifdef MYDEBUG
    std::cout << "query_period: " << query_period << std::endl;
#endif

    for (auto iter = itsData.begin(); iter != itsData.end(); iter++)
    {
      if (iter->value == itsMissingValue)
      {
        subvector.clear();
        return false;
      }

      if (!itsWeights)
      {
        if (query_period.contains(iter->time))
          subvector.push_back(*iter);
        continue;
      }

      // iterate through the data vector and sort out periods and corresponding weights
      if ((iter + 1) != itsData.end())
      {
        // period between two timesteps
        time_period timestep_period(iter->time, (iter + 1)->time + microseconds(1));
        // if timestep period is inside the queried period handle it
        if (query_period.intersects(timestep_period))
        {
          // first find out intersecting period
          time_period intersection_period(query_period.intersection(timestep_period));
          if (intersection_period.length().total_seconds() == 0)
            continue;

          // value changes halfway of timestep, so we have to handle first half and second half of
          // separately
          boost::posix_time::ptime halfway_time(
              timestep_period.begin() + seconds(timestep_period.length().total_seconds() / 2));
          if (intersection_period.contains(halfway_time))
          {
            time_period first_part_period(intersection_period.begin(),
                                          halfway_time + microseconds(1));
            time_period second_part_period(halfway_time, intersection_period.end());
            subvector.push_back(
                DataItem(iter->time, iter->value, first_part_period.length().total_seconds()));
            subvector.push_back(DataItem(
                (iter + 1)->time, (iter + 1)->value, second_part_period.length().total_seconds()));
          }
          else
          {
            if (intersection_period.begin() >
                halfway_time)  // intersection_period is in the second half
              subvector.push_back(DataItem((iter + 1)->time,
                                           (iter + 1)->value,
                                           intersection_period.length().total_seconds()));
            else  // intersection_period must be in the first half
              subvector.push_back(
                  DataItem(iter->time, iter->value, intersection_period.length().total_seconds()));
          }
        }
      }
    }

    return !subvector.empty();
  }
  catch (...)
  {
    throw Fmi::Exception::Trace(BCP, "Operation failed!");
  }
}

bool Stat::get_subvector(DataVector& subvector,
                         const boost::posix_time::ptime& startTime /*= not_a_date_time*/,
                         const boost::posix_time::ptime& endTime /*= not_a_date_time*/,
                         bool useWeights /*= true*/) const
{
  try
  {
    if (invalid_timestamps())
    {
      if (startTime != not_a_date_time || endTime != not_a_date_time)
      {
        throw Fmi::Exception(
            BCP, "Error: startTime or endTime defined, but data contains invalid timestamps!");
      }

      for (const auto& val : itsData)
      {
        if (val.value == itsMissingValue)
          return false;
        subvector.push_back(val);
      }
      return true;
    }
    bool applyWeights = (useWeights == false ? false : itsWeights);
    return extract_subvector(itsData, subvector, startTime, endTime, itsMissingValue, applyWeights);
  }
  catch (...)
  {
    throw Fmi::Exception::Trace(BCP, "Operation failed!");
  }
}

void Stat::calculate_weights()
{
  try
  {
    if (itsData.size() == 1)
    {
      itsData[0].weight = 1.0;
      return;
    }

    bool invalidTimestampsFound = invalid_timestamps();

    if (!invalidTimestampsFound)
      sort(itsData.begin(), itsData.end(), comp_time);

    for (unsigned int i = 0; i < itsData.size(); i++)
    {
      if (invalidTimestampsFound)
      {
        itsData[i].weight = 1.0;
      }
      else
      {
        if (i == 0)
        {
          time_duration dur(itsData[i + 1].time - itsData[i].time);
          itsData[i].weight = dur.total_seconds() * 0.5;
        }
        else if (i == itsData.size() - 1)
        {
          time_duration dur(itsData[i].time - itsData[i - 1].time);
          itsData[i].weight = dur.total_seconds() * 0.5;
        }
        else
        {
          time_duration dur_prev(itsData[i].time - itsData[i - 1].time);
          time_duration dur_next(itsData[i + 1].time - itsData[i].time);
          itsData[i].weight = (dur_prev.total_seconds() * 0.5) + (dur_next.total_seconds() * 0.5);
        }
      }
    }
  }
  catch (...)
  {
    throw Fmi::Exception::Trace(BCP, "Operation failed!");
  }
}

bool Stat::invalid_timestamps() const
{
  try
  {
    for (const DataItem& item : itsData)
    {
      if (item.time == not_a_date_time)
      {
        return true;
      }
    }

    return false;
  }
  catch (...)
  {
    throw Fmi::Exception::Trace(BCP, "Operation failed!");
  }
}

std::ostream& operator<<(std::ostream& os, const DataItem& item)
{
  try
  {
    os << "timestamp = " << item.time << std::endl
       << "value = " << item.value << std::endl
       << "weight = " << item.weight << std::endl;

    return os;
  }
  catch (...)
  {
    throw Fmi::Exception::Trace(BCP, "Operation failed!");
  }
}

}  // namespace Stat
}  // namespace TimeSeries
}  // namespace SmartMet


/*Copyright (c) 2021 James Gayvert

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.
*/

/*
 * qchem_interface.cpp
 */

#include "qchem_interface.h"

#include <algorithm>
#include <chrono>
#include <Eigen/Dense>
#include <fstream>
#include <iostream>
#include <sstream>
#include <string>

#include "gto_ordering.h"
#include "opencap_exception.h"
#include "utils.h"

void qchem_parse_fchk_dms(std::string dmat_filename,std::vector<std::vector<Eigen::MatrixXd>> &alpha_opdms,
		std::vector<std::vector<Eigen::MatrixXd>> &beta_opdms, size_t nstates, size_t ntdm, BasisSet &bs,
		bool symmetric_rdm, bool do_spin)
{
	alpha_opdms = std::vector< std::vector<Eigen::MatrixXd>>(nstates, std::vector<Eigen::MatrixXd> (nstates));
	beta_opdms = std::vector< std::vector<Eigen::MatrixXd>>(nstates, std::vector<Eigen::MatrixXd> (nstates));
	int lt_number = bs.Nbasis*(bs.Nbasis+1)/2;
	int lt_tdm = nstates*(nstates+1)/2 - nstates;
	std::ifstream is(dmat_filename);
	std::vector<Eigen::MatrixXd> alpha_tdms;
	std::vector<Eigen::MatrixXd> beta_tdms;
    if (is.good())
    {
    	std::string line, rest;
    	std::getline(is, line);
    	auto start = std::chrono::high_resolution_clock::now();
    	for (size_t i=0;i<nstates;i++)
    	{
    			//alpha first, then beta
    			for(size_t spin=0;spin<=do_spin;spin++)
    			{
					while(line.find("State Density")== std::string::npos)
					{
						std::getline(is,line);
	    	    		if (is.peek()==EOF)
	    	    			opencap_throw("Error: Reached end of file before densities for "+
	    	    					std::to_string(nstates) + " states were found.");
					}
					//last part of line should be number of elements to read
					int num_elements = stoi(split(line,' ').back());
					if(sqrt(num_elements)!=bs.Nbasis && num_elements!=lt_number)
						opencap_throw("Error: dimensions of TDMs do not match specified basis set.");
					size_t lines_to_read = num_elements%5==0 ? (num_elements/5) : num_elements/5+1;
					std::vector<double> matrix_elements;
					for (size_t k=1;k<=lines_to_read;k++)
					{
						std::getline(is,line);
						std::vector<std::string> tokens = split(line,' ');
						for (auto token:tokens)
							matrix_elements.push_back(std::stod(token));
					}
					Eigen::MatrixXd st_opdm(bs.Nbasis,bs.Nbasis);
					if(symmetric_rdm)
						fill_LT<double>(matrix_elements,st_opdm);
					else
						fill_mat<double>(matrix_elements,st_opdm);
					to_opencap_ordering(st_opdm,bs,get_qchem_ids(bs));
					if(spin==0)
					{
						alpha_opdms[i][i]=st_opdm;
						if(!do_spin)
							beta_opdms[i][i]=st_opdm;
					}
					else
						beta_opdms[i][i]=st_opdm;
    			}
		}
        is.seekg (0, ios::beg);
    	for(size_t i=0;i<ntdm;i++)
    	{
			//alpha first, then beta
			for (size_t spin=0;spin<=do_spin;spin++)
			{
				while(line.find("Transition DM")== std::string::npos)
				{
					std::getline(is,line);
					if (is.peek()==EOF)
						opencap_throw("Error: Reached end of file before densities for "+
								std::to_string(nstates) + " states were found.");
				}
				//last part of line should be number of elements to read
				size_t num_elements = stoi(split(line,' ').back());
				if(sqrt(num_elements)!=bs.Nbasis && num_elements!=lt_number)
					opencap_throw("Error: dimensions of TDMs do not match specified basis set.");
				size_t lines_to_read = num_elements%5==0 ? (num_elements/5) : num_elements/5+1;
				std::vector<double> matrix_elements;
				for (size_t k=1;k<=lines_to_read;k++)
				{
					std::getline(is,line);
					std::vector<std::string> tokens = split(line,' ');
					for (auto token:tokens)
						matrix_elements.push_back(std::stod(token));
				}
				Eigen::MatrixXd st_opdm(bs.Nbasis,bs.Nbasis);
				fill_mat<double>(matrix_elements,st_opdm);
				to_opencap_ordering(st_opdm,bs,get_qchem_ids(bs));
				if(spin==0)
				{
					alpha_tdms.push_back(st_opdm);
					if(!do_spin)
						beta_tdms.push_back(st_opdm);
				}
				else
					beta_tdms.push_back(st_opdm);
			}
    	}
    }
    is.close();
    if(ntdm == lt_tdm)
    {
        std::cout << "Warning: TDM M-->N is assumed to be conjugate transpose of "
        << "TDM N-->M where M>N" << std::endl;
    	size_t dm_idx = 0;
    	for(size_t i=0;i<nstates;i++)
    	{
    		for(size_t j=i+1;j<nstates;j++)
    		{
    			alpha_opdms[i][j] = alpha_tdms[dm_idx];
    			alpha_opdms[j][i] = alpha_tdms[dm_idx].adjoint();
    			beta_opdms[i][j] = beta_tdms[dm_idx].adjoint();
    			beta_opdms[j][i] = beta_tdms[dm_idx];
				dm_idx++;
    		}
    	}
    }
    else
    {
    	size_t dm_idx = 0;
    	for(size_t i=0;i<nstates;i++)
    	{
    		for(size_t j=0;j<nstates;j++)
    		{
    			if(i!=j)
    			{
    				alpha_opdms[i][j] = alpha_tdms[dm_idx];
    				beta_opdms[i][j] = beta_tdms[dm_idx];
    				dm_idx++;
    			}
    		}
    	}
    }
}

void qchem_read_dms(std::vector<std::vector<Eigen::MatrixXd>> &alpha_dms,
		std::vector<std::vector<Eigen::MatrixXd>> &beta_dms,
		std::string fchk_filename, BasisSet &bs,size_t num_states)
{
	std::ifstream is(fchk_filename);
	size_t nstates = 0;
	size_t num_tdm = 0;
	bool sep_alpha_beta;
	bool symmetric_rdm = false;
	bool symmetric_tdm = false;
	std::string line, rest;
	if (is.good())
	{
		while(is.peek() != EOF )
		{
			std::getline(is, line);
			if(line.find("State Density")!= std::string::npos)
			{
				if(line.find("Alpha")!=std::string::npos)
				{
					sep_alpha_beta=true;
					int num_elements = stoi(split(line,' ').back());
					if(sqrt(num_elements)!=bs.Nbasis)
					{
						int lt_number = bs.Nbasis*(bs.Nbasis+1)/2;
						if(lt_number == num_elements)
							symmetric_rdm = true;
						else
							opencap_throw("Error: dimensions of DMs do not match specified basis set.");
					}
					break;
				}
				else
				{
					sep_alpha_beta=false;
					int num_elements = stoi(split(line,' ').back());
					if(sqrt(num_elements)!=bs.Nbasis)
					{
						int lt_number = bs.Nbasis*(bs.Nbasis+1)/2;
						if (lt_number == num_elements)
							symmetric_rdm = true;
						else
							opencap_throw("Error: dimensions of DMs do not match specified basis set.");
					}
					break;
				}
			}
		}
		is.seekg (0, ios::beg);
		while(is.peek()!=EOF)
		{
			std::getline(is, line);
			if((line.find("Alpha")!= std::string::npos && line.find("State Density")!= std::string::npos) && sep_alpha_beta)
				nstates++;
			else if(line.find("State Density")!= std::string::npos && !sep_alpha_beta)
				nstates++;
			if((line.find("Alpha")!= std::string::npos && line.find("Transition DM")!= std::string::npos) && sep_alpha_beta)
				num_tdm++;
			else if(line.find("Transition DM")!= std::string::npos && !sep_alpha_beta)
				num_tdm++;
		}
		is.close();
		if(nstates==0)
		{
			opencap_throw("Error: Unable to find any densities in:" +fchk_filename);
		}
		else if(nstates!=num_states)
		{
			opencap_throw("Error: number of states found: " + std::to_string(nstates) + ", does not match "
					"number of states specified in the input: " + std::to_string(num_states)+ ". Exiting...");
		}
		qchem_parse_fchk_dms(fchk_filename,alpha_dms,beta_dms,nstates,num_tdm,bs,symmetric_rdm,sep_alpha_beta);
	}
	else
	    opencap_throw("Error: I couldn't read:" + fchk_filename);
}

Eigen::MatrixXd qchem_read_overlap(std::string dmat_filename, BasisSet bs)
{
	size_t num_bf = bs.Nbasis;
    std::ifstream is(dmat_filename);
	Eigen::MatrixXd smat(num_bf,num_bf);
	smat=Eigen::MatrixXd::Zero(num_bf,num_bf);
    if (is.good())
    {
    	std::string line, rest;
    	while (line.find("Overlap Matrix")== std::string::npos)
    	{
        	std::getline(is, line);
    		if (is.peek()==EOF)
    			opencap_throw("Error: Reached end of file before Overlap Matrix was found.");
    	}
    	size_t num_elements = stoi(split(line,' ').back());
    	int lt_number = bs.Nbasis*(bs.Nbasis+1)/2;
    	if(num_elements!=lt_number)
    		opencap_throw("Error: Dimensions of overlap matrix do not match basis.");
		size_t lines_to_read = num_elements%5==0 ? (num_elements/5) : num_elements/5+1;
		std::vector<double> matrix_elements;
		for (size_t k=1;k<=lines_to_read;k++)
		{
			std::getline(is,line);
			std::vector<std::string> tokens = split(line,' ');
			for (auto token:tokens)
				matrix_elements.push_back(std::stod(token));
		}
		fill_LT<double>(matrix_elements,smat);
		to_opencap_ordering(smat,bs,get_qchem_ids(bs));
    }
    else
    	opencap_throw("Error: I couldn't read:" + dmat_filename);
    return smat;
}

Eigen::MatrixXd read_qchem_tddft_energies(size_t nstates,std::string output_file)
{
	Eigen::MatrixXd ZERO_ORDER_H(nstates,nstates);
	ZERO_ORDER_H=Eigen::MatrixXd::Zero(nstates,nstates);
	std::ifstream is(output_file);
    if (is.good())
    {
    	std::string line, rest;
    	std::getline(is, line);
    	size_t state_idx = 1;
    	while (state_idx<nstates)
    	{
    		if (is.peek()==EOF)
    			opencap_throw("Error: Reached end of file before "+ std::to_string(nstates) + " energies were found. "
    					"Only " + std::to_string(state_idx) + " states were found. Exiting...");
    		if (line.find("Total energy in the final basis set")!= std::string::npos)
    			ZERO_ORDER_H(0,0) = std::stod(split(line,' ')[8]);
			if (line.find("Total energy for state")!= std::string::npos)
			{
					ZERO_ORDER_H(state_idx,state_idx) = std::stod(split(line,' ')[5]);
					state_idx++;
					std::getline(is,line);
			}
    		else
    			std::getline(is,line);
    	}
    }
    else
    	opencap_throw("Error: I couldn't read:" + output_file);
    return ZERO_ORDER_H;
}

Eigen::MatrixXd read_qchem_eom_energies(size_t nstates,std::string output_file)
{
	Eigen::MatrixXd ZERO_ORDER_H(nstates,nstates);
	ZERO_ORDER_H=Eigen::MatrixXd::Zero(nstates,nstates);
	std::ifstream is(output_file);
    if (is.good())
    {
    	std::string line, rest;
    	std::getline(is, line);
    	size_t state_idx = 1;
    	while (state_idx<=nstates)
    	{
    		if (is.peek()==EOF)
    			opencap_throw("Error: Reached end of file before "+ std::to_string(nstates) + " energies were found. "
    					"Only " + std::to_string(state_idx) + " states were found. Exiting...");
			if (line.find("Total energy")!= std::string::npos && line.find("Excitation energy")!= std::string::npos)
			{
					ZERO_ORDER_H(state_idx-1,state_idx-1) = std::stod(split(line,' ')[3]);
					state_idx++;
					std::getline(is,line);
			}
    		else
    			std::getline(is,line);
    	}
    }
    else
    	opencap_throw("Error: I couldn't read:" + output_file);
    return ZERO_ORDER_H;
}

std::vector<Atom> read_geometry_from_fchk(std::string fchk_filename)
{
	std::vector<Atom> atoms;
	std::vector<size_t> atom_nums;
	std::vector<double> coords;
	std::ifstream is(fchk_filename);
    if (is.good())
    {
    	std::string line, rest;
    	std::getline(is, line);
		while(line.find("Atomic numbers")== std::string::npos)
		{
			std::getline(is,line);
    		if (is.peek()==EOF)
    			opencap_throw("Error: Reached end of file before atomic numbers were found.");
		}
		size_t num_elements = stoi(split(line,' ').back());
		size_t lines_to_read = num_elements%6==0 ? (num_elements/6) : num_elements/6+1;
		for (size_t k=1;k<=lines_to_read;k++)
		{
			std::getline(is,line);
			std::vector<std::string> tokens = split(line,' ');
			for (auto token:tokens)
				atom_nums.push_back(std::stoi(token));
		}
		while(line.find("Current cartesian coordinates")== std::string::npos)
		{
			std::getline(is,line);
    		if (is.peek()==EOF)
    			opencap_throw("Error: Reached end of file before Current Cartesian Coordinates were found.");
		}
		num_elements = stoi(split(line,' ').back());
		lines_to_read = num_elements%5==0 ? (num_elements/5) : num_elements/5+1;
		for (size_t k=1;k<=lines_to_read;k++)
		{
			std::getline(is,line);
			std::vector<std::string> tokens = split(line,' ');
			for (auto token:tokens)
				coords.push_back(std::stod(token));
		}
		//ok now lets populate our atoms
		for(size_t i=0;i<atom_nums.size();i++)
			atoms.push_back(Atom(atom_nums[i],coords[i*3],coords[i*3+1],coords[i*3+2]));
    }
    else
    	opencap_throw("Error: I couldn't read:" + fchk_filename);
    return atoms;
}

BasisSet read_basis_from_fchk(std::string fchk_filename, std::vector<Atom> atoms)
{
	auto start = std::chrono::high_resolution_clock::now();
	std::vector<int> shell_types;
	std::vector<int> prims_per_shell;
	std::vector<int> atom_ids;
	std::vector<double> exps;
	std::vector<double> coeffs;
	// needed for SP specification
	std::vector<double> p_coeffs;
	bool SP_basis_function = true;
	BasisSet bs;
	for(auto atm:atoms)
		bs.centers.push_back(atm.coords);
	std::vector<shell_id> ids;
	std::ifstream is(fchk_filename);
    if (is.good())
    {
    	std::string line, rest;
    	std::getline(is, line);
    	//First lets figure out if there are SP functions
    	//shell types
		while(line.find("Shell types")== std::string::npos)
		{
			std::getline(is,line);
    		if (is.peek()==EOF)
    			opencap_throw("Error: Reached end of file before shell types were found.");
		}
		size_t num_elements = stoi(split(line,' ').back());
		size_t lines_to_read = num_elements%6==0 ? (num_elements/6) : num_elements/6+1;
		for (size_t k=1;k<=lines_to_read;k++)
		{
			std::getline(is,line);
			std::vector<std::string> tokens = split(line,' ');
			for (auto token:tokens)
			{
				shell_types.push_back(std::stoi(token));
			}
		}
		//prims per shell
		while(line.find("Number of primitives per shell")== std::string::npos)
		{
			std::getline(is,line);
    		if (is.peek()==EOF)
    			opencap_throw("Error: Reached end of file before Number of primitives per shell were found.");
		}
		num_elements = stoi(split(line,' ').back());
		lines_to_read = num_elements%6==0 ? (num_elements/6) : num_elements/6+1;
		for (size_t k=1;k<=lines_to_read;k++)
		{
			std::getline(is,line);
			std::vector<std::string> tokens = split(line,' ');
			for (auto token:tokens)
				prims_per_shell.push_back(std::stoi(token));
		}
		//atom ids
		while(line.find("Shell to atom map")== std::string::npos)
		{
			std::getline(is,line);
    		if (is.peek()==EOF)
    			opencap_throw("Error: Reached end of file before Shell to atom map was found.");
		}
		num_elements = stoi(split(line,' ').back());
		lines_to_read = num_elements%6==0 ? (num_elements/6) : num_elements/6+1;
		for (size_t k=1;k<=lines_to_read;k++)
		{
			std::getline(is,line);
			std::vector<std::string> tokens = split(line,' ');
			for (auto token:tokens)
				atom_ids.push_back(std::stoi(token));
		}
		//prims
		while(line.find("Primitive exponents")== std::string::npos)
		{
			std::getline(is,line);
    		if (is.peek()==EOF)
    			opencap_throw("Error: Reached end of file before Primitive exponents were found.");
		}
		num_elements = stoi(split(line,' ').back());
		lines_to_read = num_elements%5==0 ? (num_elements/5) : num_elements/5+1;
		for (size_t k=1;k<=lines_to_read;k++)
		{
			std::getline(is,line);
			std::vector<std::string> tokens = split(line,' ');
			for (auto token:tokens)
				exps.push_back(std::stod(token));
		}
		//coeffs
		while(line.find("Contraction coefficients")== std::string::npos)
		{
			std::getline(is,line);
    		if (is.peek()==EOF)
    			opencap_throw("Error: Reached end of file before Contraction coefficients were found.");
		}
		num_elements = stoi(split(line,' ').back());
		lines_to_read = num_elements%5==0 ? (num_elements/5) : num_elements/5+1;
		for (size_t k=1;k<=lines_to_read;k++)
		{
			std::getline(is,line);
			std::vector<std::string> tokens = split(line,' ');
			for (auto token:tokens)
				coeffs.push_back(std::stod(token));
		}
		while(line.find("P(S=P) Contraction coefficients")== std::string::npos)
		{
			std::getline(is,line);
    		if (is.peek()==EOF)
    		{
    			SP_basis_function = false;
    			break;
    		}
		}
		if (SP_basis_function)
		{
			num_elements = stoi(split(line,' ').back());
			lines_to_read = num_elements%5==0 ? (num_elements/5) : num_elements/5+1;
			for (size_t k=1;k<=lines_to_read;k++)
			{
				std::getline(is,line);
				std::vector<std::string> tokens = split(line,' ');
				for (auto token:tokens)
					p_coeffs.push_back(std::stod(token));
			}
		}
    }
    else
    	opencap_throw("Error: I couldn't read:" + fchk_filename);
    size_t prim_idx=0;
    for(size_t i=0;i<shell_types.size();i++)
    {
    	//SP
    	if(shell_types[i]==-1)
    	{
    		if(p_coeffs.size()==0)
    			opencap_throw("Error: missing section P(S=P) Contraction coefficients.");
    		Shell s_shell(0,atoms[atom_ids[i]-1].coords);
    		Shell p_shell(1,atoms[atom_ids[i]-1].coords);
			int num_prims = prims_per_shell[i];
			for(int j=1;j<=num_prims;j++)
			{
				s_shell.add_primitive(exps[prim_idx],coeffs[prim_idx]);
				p_shell.add_primitive(exps[prim_idx],p_coeffs[prim_idx]);
				prim_idx++;
			}
			bs.add_shell(s_shell);
			bs.add_shell(p_shell);
    	}
    	else
    	{
			Shell new_shell(abs(shell_types[i]),atoms[atom_ids[i]-1].coords);
			if(shell_types[i]>0 && new_shell.l>1)
				new_shell.pure=false;
			int num_prims = prims_per_shell[i];
			for(int j=1;j<=num_prims;j++)
			{
				new_shell.add_primitive(exps[prim_idx],coeffs[prim_idx]);
				prim_idx++;
			}
			bs.add_shell(new_shell);
    	}
    }
    bs.normalize();
	auto stop = std::chrono::high_resolution_clock::now();
	auto total_time = std::chrono::duration<double>(stop-start).count();
    return bs;
}