#pragma once

#include <ros/ros.h>
#include <iostream>
#include <fstream>
#include <sstream>
#include <stdio.h>
#include <string.h>

// ROS
#include <std_msgs/Float64MultiArray.h>
#include <std_msgs/MultiArrayDimension.h>

#include <tf/transform_broadcaster.h>
#include <tf/transform_listener.h>
#include <geometry_msgs/Pose.h>
#include <nav_msgs/Odometry.h>
#include <nav_msgs/Path.h>
#include <geometry_msgs/Twist.h>
#include <nav_msgs/Path.h>
#include <geometry_msgs/PoseStamped.h>
#include <geometry_msgs/Point.h>
#include <geometry_msgs/PointStamped.h>
#include <geometry_msgs/PolygonStamped.h>
#include <geometry_msgs/PoseWithCovarianceStamped.h>

// EIGEN
#include <Eigen/Dense>
#include <Eigen/Geometry>

// Submodules
#include <corridor.hpp>
#include <init_traj_planner.hpp>
#include <mission.hpp>
#include <param.hpp>

#include "com_fun.h"
#include "optimization.h"
#include "time.h"
#include <cppad/cppad.hpp>
#include <cppad/ipopt/solve.hpp>
#include <eigen3/Eigen/Core>

using CppAD::AD;

// The solver takes all the state variables and actuator
// variables in a singular vector. Thus, we should to establish
// when one variable starts and another ends to make our lifes easier.
size_t x_start;
size_t y_start;
size_t theta_start;
size_t v_start;
size_t omega_start;

std::vector<ros::Publisher> way_point_pub;

std::vector< std::array<int, 3> > relative_pair;
std::vector< int > agent_col_num;

class FG_eval {
     public:

      typedef CPPAD_TESTVECTOR(AD<double>) ADvector;

      void operator()(ADvector& fg, const ADvector& vars) {
        // MPC implementation
        // fg a vector of constraints, x is a vector of constraints.
        // NOTE: You'll probably go back and forth between this function and
        // the Solver function below.
        fg[0] = 0;
        for (int qi = 0; qi < qn; qi ++){
            for (int i = 1; i < N; i++) {
                fg[0] += W_X * CppAD::pow(vars[x_start + N * qi + i] - plan[qi][i][0], 2);
                fg[0] += W_Y * CppAD::pow(vars[y_start + N * qi + i] - plan[qi][i][1], 2);
            }
        }

        // Minimize the value gap between sequential actuations.
        for (int qi = 0; qi < qn; qi ++){
            for (int i = 0; i < N - 2; i++) {
              fg[0] += W_DV * CppAD::pow(vars[v_start + (N-1)* qi + i + 1] - vars[v_start + (N-1) * qi + i], 2);
              fg[0] += W_DOMEGA * CppAD::pow(vars[omega_start + (N-1) * qi + i + 1] - vars[omega_start + (N-1) * qi + i], 2);
          }
        }

        // Initial constraints
        for (int qi = 0; qi < qn; qi ++){
            fg[1 + N * qi + x_start] = vars[x_start + N * qi];
            fg[1 + N * qi + y_start] = vars[y_start + N * qi];
            fg[1 + N * qi + theta_start] = vars[theta_start + N * qi];
        }

        // The rest of the constraints
        for (int qi = 0; qi < qn; qi ++){
            for (int i = 0; i < N - 1; i++) {
              // The state at time t+1 .
              AD<double> x1 = vars[x_start + N * qi + i + 1];
              AD<double> y1 = vars[y_start + N * qi + i + 1];
              AD<double> theta1 = vars[theta_start + N * qi + i + 1];

              // The state at time t.
              AD<double> x0 = vars[x_start + N * qi + i];
              AD<double> y0 = vars[y_start + N * qi + i];
              AD<double> theta0 = vars[theta_start + N * qi + i];

              // Only consider the actuation at time t.
              AD<double> v0 = vars[v_start + (N-1) * qi + i];
              AD<double> omega0 = vars[omega_start + (N-1) * qi + i];

              fg[2 + x_start + N * qi + i] = x1 - (x0 + v0 * CppAD::cos(theta0) * DT);
              fg[2 + y_start + N * qi + i] = y1 - (y0 + v0 * CppAD::sin(theta0) * DT);
              fg[2 + theta_start + N * qi + i] = theta1 - (theta0 + omega0 * DT);

            }
        }

        //std::cout<<relative_pair.size()<<std::endl;

        for (int j=0;j<qn;j++){
          int count=0;
          for(int i=0; i<relative_pair.size();i++)
          {
            if (relative_pair[i].at(1)==j)
              count++;
            else if (relative_pair[i].at(2)==j)
              count++;
          }
          agent_col_num.emplace_back(count);
        }

        int cont=N*3*qn+1;
        for (int i=0;i<relative_pair.size();i++)
        {
          fg[cont]=CppAD::pow(vars[x_start+N*relative_pair[i].at(1)+relative_pair[i].at(0)]-vars[x_start+N*relative_pair[i].at(2)+relative_pair[i].at(0)],2)
                  +CppAD::pow(vars[y_start+N*relative_pair[i].at(1)+relative_pair[i].at(0)]-vars[y_start+N*relative_pair[i].at(2)+relative_pair[i].at(0)],2);
          cont++;
        }
    }
};

class MPCPlanner {
public:
    std_msgs::Float64MultiArray msgs_traj_info;
    std::vector<std_msgs::Float64MultiArray> msgs_traj_coef;

    MPCPlanner(std::shared_ptr<Corridor> _corridor_obj,
               std::shared_ptr<InitTrajPlanner> _initTrajPlanner_obj,
               SwarmPlanning::Mission _mission,
               SwarmPlanning::Param _param)
            : corridor_obj(std::move(_corridor_obj)),
              initTrajPlanner_obj(std::move(_initTrajPlanner_obj)),
              mission(std::move(_mission)),
              param(std::move(_param))

    {
        M = initTrajPlanner_obj.get()->T.size()-1; // the number of segments

        outdim = 3; // the number of outputs (x,y,z)

        T = initTrajPlanner_obj.get()->T;
        initTraj = initTrajPlanner_obj.get()->initTraj;
        SFC = corridor_obj.get()->SFC;

        int i,j,k;

        double dx;
        double dy;
        double dis;

        for (i=0;i<qn;i++){
            for(j=0;j<M;j++){

                dx=(initTraj[i][j+1].x()-initTraj[i][j].x())/5;
                dy=(initTraj[i][j+1].y()-initTraj[i][j].y())/5;
                dis= 25*(dx*dx+dy*dy);

                for (k=0;k<5;k++){
                    plan[i][5*j+k][0]=initTraj[i][j].x()+dx*k;
                    plan[i][5*j+k][1]=initTraj[i][j].y()+dy*k;

                    if (dis<0.1)
                    {
                      plan[i][5*j+k][3]=0;
                    }
                    if (dis<1.1){
                      plan[i][5*j+k][3]=MAXV*0.7;
                    }
                    else{
                      plan[i][5*j+k][3]=MAXV;
                    }
                }
            }
            plan[i][5*M][0]=initTraj[i][M].x();
            plan[i][5*M][1]=initTraj[i][M].y();
        }

        N = M*5+1;

        for (int qi=0; qi<qn; qi++)
        {
          for (int next=5;next<N;next+=5){
            if ((plan[qi][next][1]==plan[qi][0][1]) &&(plan[qi][next][0]==plan[qi][0][0]))
              continue;
            else{
              plan[qi][0][2]=atan2(plan[qi][next][1]-plan[qi][0][1],plan[qi][next][0]-plan[qi][0][0]);
              break;
            }
        }
      }

      double angle;

      for (i=0;i<qn;i++){
        for(j=0;j<M;j++){

          dx=initTraj[i][j+1].x()-initTraj[i][j].x();
          dy=initTraj[i][j+1].y()-initTraj[i][j].y();
          dis= dx*dx+dy*dy;

          if (j>0){
            if (dis>0.1){
              angle=atan2(dy,dx);
              if (angle-plan[i][5*(j-1)][2]>PI)
                angle=angle-2*PI;
              else if(plan[i][5*(j-1)][2]-angle>PI)
                angle=angle+2*PI;
            }
              else angle =plan[i][5*(j-1)][2];
          }
          else {angle=plan[i][0][2];}

          for (k=0;k<5;k++){
            plan[i][5*j+k][2]=angle;
          }

        }
      }

      for (i=0;i<qn;i++){
        for(j=0;j<M;j++){
          for (k=0;k<5;k++){
              plan[i][5*j+k][4]=(plan[i][5*(j+1)][2]-plan[i][5*j][2])*0.2/DT;
            }
          }
      }

      x_start = 0;
      y_start = x_start + qn*N;
      theta_start = y_start + qn*N;
      v_start = theta_start + qn*N;
      omega_start = v_start + qn*(N - 1);

      std::array<int, 3> sa;
      for (k=0;k<N;k++){
        for(i=0;i<qn;i++){
          for (j=i+1;j<qn;j++)
          {
            dis= pow(plan[i][k][0]-plan[j][k][0],2.0)+pow(plan[i][k][1]-plan[j][k][1],2.0);
            if (dis<1.3)
              relative_pair.emplace_back(sa={k, i, j});
          }
        }
      }

    }

    void Solve() {
      bool ok = true;
      size_t i;
      typedef CPPAD_TESTVECTOR(double) Dvector;

      // Set the number of model variables (includes both states and inputs).
      // For example: If the state is a 4 element vector, the actuators is a 2
      // element vector and there are 10 timesteps. The number of variables is:
      size_t n_vars = (N * 3 + (N - 1) * 2) * qn;
      // Set the number of constraints
      size_t n_constraints = N * 3 * qn;
      //size_t all_constraints = n_constraints + qn*(qn-1)*int(N/2)/2;
      size_t all_constraints = n_constraints + relative_pair.size();

      // Initial value of the independent variables.
      // SHOULD BE 0 besides initial state.
      Dvector vars(n_vars);

      Dvector vars_lowerbound(n_vars);
      Dvector vars_upperbound(n_vars);
      // Set lower and upper limits for variables.
      for (int i = 0; i < v_start; i++) {
        vars_lowerbound[i] = -BOUND;
        vars_upperbound[i] = BOUND;
      }

      for (int qi=0;qi<qn;qi++){

        int count=0;

        for (int i = 0; i < N; i++) {
          vars[x_start+qi*N+i]=plan[qi][i][0];
          vars[y_start+qi*N+i]=plan[qi][i][1];
          vars[theta_start+qi*N+i]=plan[qi][i][2];
          if (i<N-1)
          {vars[v_start+qi*(N-1)+i]=plan[qi][i][3];
           vars[omega_start+qi*(N-1)+i]=plan[qi][i][4];}

          if ( i>1 && i==SFC[qi][count+1].second && count<SFC[qi].size()-1){
            count++;
          }

          if (plan[qi][i][0]-SFC[qi][count].first[0]>2)
            vars_lowerbound[x_start+qi*N+i]=plan[qi][i][0]-2;
          else
            vars_lowerbound[x_start+qi*N+i]=SFC[qi][count].first[0];
          if (SFC[qi][count].first[3]-plan[qi][i][0]>2)
            vars_upperbound[x_start+qi*N+i]=plan[qi][i][0]+2;
          else
            vars_upperbound[x_start+qi*N+i]=SFC[qi][count].first[3];
          if (plan[qi][i][1]-SFC[qi][count].first[1]>2)
            vars_lowerbound[y_start+qi*N+i]=plan[qi][i][1]-2;
          else
            vars_lowerbound[y_start+qi*N+i]=SFC[qi][count].first[1];
          if (SFC[qi][count].first[4]-plan[qi][i][1]>2)
            vars_upperbound[y_start+qi*N+i]=plan[qi][i][1]+2;
          else
            vars_upperbound[y_start+qi*N+i]=SFC[qi][count].first[4];
        }
      }

      for (int i = v_start; i < omega_start; i++) {
        vars_lowerbound[i] = 0;
        vars_upperbound[i] = MAXV;
        if (param.backward_enable)
        vars_lowerbound[i] = -MAXV;
      }

      for (int i = omega_start; i < n_vars; i++) {
        vars_lowerbound[i] = -MAXOMEGA;
        vars_upperbound[i] = MAXOMEGA;
      }

      // Lower and upper limits for the constraints
      // Should be 0 besides initial state.
      Dvector constraints_lowerbound(all_constraints);
      Dvector constraints_upperbound(all_constraints);
      for (int i = 0; i < n_constraints; i++) {
        constraints_lowerbound[i] = 0;
        constraints_upperbound[i] = 0;
      }

      double inter_collision=2*mission.quad_size[0]+0.2;
      inter_collision=inter_collision*inter_collision;
      for (int i= n_constraints;i<all_constraints;i++)
      {
        constraints_lowerbound[i] = inter_collision;
        constraints_upperbound[i] = BOUND;
      }

      for (int qi = 0; qi < qn; qi ++){
          double startangle;
      for (int next=5;next<N;next+=5){
        if ((plan[qi][next][1]==plan[qi][0][1]) &(plan[qi][next][0]==plan[qi][0][0]))
          continue;
        else{
          startangle=atan2(plan[qi][next][1]-plan[qi][0][1],plan[qi][next][0]-plan[qi][0][0]);
          break;
        }
      }

          constraints_lowerbound[x_start + N * qi] = plan[qi][0][0];
          constraints_lowerbound[y_start + N * qi] = plan[qi][0][1];
          constraints_upperbound[x_start + N * qi] = plan[qi][0][0];
          constraints_upperbound[y_start + N * qi] = plan[qi][0][1];

          if(param.initial_angle)
          {
            constraints_lowerbound[theta_start + N * qi] = mission.startState[qi][2]*PI/2;
            constraints_upperbound[theta_start + N * qi] = mission.startState[qi][2]*PI/2;
          }
          else
          {
            constraints_lowerbound[theta_start + N * qi] = startangle;
            constraints_upperbound[theta_start + N * qi] = startangle;
          }
      }

      // object that computes objective and constraints
      FG_eval fg_eval;
      // options for IPOPT solver
      std::string options;

      options += "Numeric tol          1e-5\n";
      options += "String linear_solver mumps\n";
      // Uncomment this if you'd like more print information
      options += "Integer print_level  1\n";
      // NOTE: Setting sparse to true allows the solver to take advantage
      // of sparse routines, this makes the computation MUCH FASTER. If you
      // can uncomment 1 of these and see if it makes a difference or not but
      // if you uncomment both the computation time should go up in orders of
      // magnitude.
      options += "Sparse  true        forward\n";
      options += "Sparse  true        reverse\n";
      // NOTE: Currently the solver has a maximum time limit of 0.5 seconds.
      // Change this as you see fit.

      // place to return solution
      CppAD::ipopt::solve_result<Dvector> solution;

      // solve the problem
      CppAD::ipopt::solve<Dvector, FG_eval>(
          options, vars, vars_lowerbound, vars_upperbound, constraints_lowerbound,
          constraints_upperbound, fg_eval, solution);

      // Check some of the solution values
      ok &= solution.status == CppAD::ipopt::solve_result<Dvector>::success;
      if (ok)
      {
        ROS_INFO_STREAM("Optimization Success!");
        ROS_INFO_STREAM("Cost: "<<solution.obj_value);
      }

      for (int qi=0;qi<qn;qi++){
        for (int i=0;i<N;i++){
            plan[qi][i][0]=solution.x[x_start+qi*N+i];
            plan[qi][i][1]=solution.x[y_start+qi*N+i];
            plan[qi][i][2]=solution.x[theta_start+qi*N+i];
            plan[qi][i][3]=solution.x[v_start+qi*(N-1)+i];
            plan[qi][i][4]=solution.x[omega_start+qi*(N-1)+i];
        }
      }
    }

    bool update(bool log){

        int qi=0;

        nav_msgs::Path path;
        geometry_msgs::PoseStamped pp;

        path.header.frame_id="world";
        path.header.stamp=ros::Time::now();

        Solve();

        for (int qi=0;qi<qn;qi++)
        {
          nav_msgs::Path path;
          path.header.frame_id="map";
          path.header.stamp=ros::Time::now();
          for (int i=0;i<N;i++)
          {
              pp.pose.position.x=plan[qi][i][0];
              pp.pose.position.y=plan[qi][i][1];
              pp.pose.orientation=tf::createQuaternionMsgFromYaw(plan[qi][i][2]);
              path.poses.push_back(pp);

              if (fabs(plan[qi][i][0]-plan[qi][N-1][0])<0.01){
                  if(fabs(plan[qi][i][1]-plan[qi][N-1][1])<0.01){
                      break;
                  }
              }
          }
        way_point_pub[qi].publish(path);
      }

        return true;
    }

private:
    std::shared_ptr<Corridor> corridor_obj;
    std::shared_ptr<InitTrajPlanner> initTrajPlanner_obj;
    SwarmPlanning::Mission mission;
    SwarmPlanning::Param param;

    initTraj_t initTraj;
    std::vector<double> T;
    SFC_t SFC;

    int M, phi, outdim;

    // std::shared_ptr<Eigen::MatrixXd> Q_obj, Aeq_obj, Alq_obj, deq_obj, dlq_obj;
    Eigen::MatrixXd Q_base, Aeq_base, Alq, deq, dlq, basis;
    Eigen::MatrixXd dummy;
    std::vector<Eigen::MatrixXd> coef;

    void createMsg(){
        std::vector<double> traj_info;
        traj_info.emplace_back(N);
        traj_info.emplace_back(qn);
        traj_info.insert(traj_info.end(), T.begin(), T.end());
        msgs_traj_info.data = traj_info;

        msgs_traj_coef.resize(N);
        for(int qi = 0; qi < N; qi++) {
            std_msgs::MultiArrayDimension rows;
            rows.size = M * (qn + 1);
            msgs_traj_coef[qi].layout.dim.emplace_back(rows);

            std_msgs::MultiArrayDimension cols;
            cols.size = outdim;
            msgs_traj_coef[qi].layout.dim.emplace_back(cols);

            std::vector<double> coef_temp(coef[qi].data(), coef[qi].data() + coef[qi].size());
            msgs_traj_coef[qi].data.insert(msgs_traj_coef[qi].data.end(), coef_temp.begin(), coef_temp.end());
        }
    }
};


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

#include <bitset>
#include <boost/archive/xml_iarchive.hpp>
#include <boost/archive/xml_oarchive.hpp>
#include <boost/math/distributions/hypergeometric.hpp>
#include <iostream>
#include <exception>
#include <boost/math/special_functions/binomial.hpp>

#include "workload/density-models/hypergeometric-distribution.hpp"

BOOST_CLASS_EXPORT(problem::HypergeometricDistribution)

namespace problem
{

HypergeometricDistribution::HypergeometricDistribution()
{}

HypergeometricDistribution::HypergeometricDistribution(const Specs& specs)
  : specs_(specs)
{
  is_specced_ = true;
  if (specs.workload_tensor_size != 0)
  {
    workload_tensor_size_set_ = true;
  }
}

HypergeometricDistribution::~HypergeometricDistribution()
{}

HypergeometricDistribution::Specs HypergeometricDistribution::ParseSpecs(config::CompoundConfigNode density_config)
{

  Specs specs;
  assert(density_config.lookupValue("distribution", specs.type));
  assert(density_config.lookupValue("density", specs.average_density));

  long long workload_tensor_size;
  if (density_config.lookupValue("workload_tensor_size", workload_tensor_size))
  {
    specs.total_nnzs = ceil(specs.average_density * workload_tensor_size);
    specs.workload_tensor_size = workload_tensor_size;

  } else
  {
    specs.total_nnzs = 0;
    specs.workload_tensor_size = 0;
  }

  return specs;
}

void HypergeometricDistribution::SetDensity(const double density)
{
  specs_.average_density = density;
}

void HypergeometricDistribution::SetWorkloadTensorSize(const std::uint64_t size)
{
  // setter that allows workload tensor size at a latter stage (topology.cpp, PreEvaluationCheck)
  assert(is_specced_);
  specs_.workload_tensor_size = size;
  specs_.total_nnzs = ceil(specs_.average_density * size);
  workload_tensor_size_set_ = true;
}

std::uint64_t HypergeometricDistribution::GetTileOccupancyByConfidence(const std::uint64_t tile_shape,
                                                                       const double confidence) const
{

  std::uint64_t tile_occupancy;

  if (confidence == 1.0)
  {
    if (tile_shape >= specs_.total_nnzs)
    {
      tile_occupancy = specs_.total_nnzs;
    } else
    {
      tile_occupancy = tile_shape;
    }
  } else
  {
    std::uint64_t r = specs_.total_nnzs;
    std::uint64_t n = tile_shape;
    std::uint64_t N = specs_.workload_tensor_size;
    boost::math::hypergeometric_distribution<double> distribution(r, n, N);
    tile_occupancy = quantile(distribution, confidence);
  }

  return tile_occupancy;
}

std::uint64_t HypergeometricDistribution::GetMaxTileOccupancyByConfidence(const tiling::CoordinateSpaceTileInfo& tile,
                                                                          const double confidence) const
{
  std::uint64_t tile_shape = tile.GetShape();
  return HypergeometricDistribution::GetTileOccupancyByConfidence(tile_shape, confidence);
}

std::uint64_t HypergeometricDistribution::GetMaxTileOccupancyByConfidence_LTW(const std::uint64_t tile_shape,
                                                                              const double confidence) const
{
  return HypergeometricDistribution::GetTileOccupancyByConfidence(tile_shape, confidence);
}

std::uint64_t HypergeometricDistribution::GetWorkloadTensorSize() const
{
  assert(workload_tensor_size_set_);
  return specs_.workload_tensor_size;
};

std::string HypergeometricDistribution::GetDistributionType() const
{
  return specs_.type;
}

double HypergeometricDistribution::GetMaxTileDensityByConfidence(const tiling::CoordinateSpaceTileInfo tile,
                                                                 const double confidence) const
{

  if (confidence == 0.5)
  { return specs_.average_density; } //shortcut for faster calculations

  std::uint64_t tile_shape = tile.GetShape();
  std::uint64_t percentile_occupancy = HypergeometricDistribution::GetTileOccupancyByConfidence(tile_shape, confidence);
  return (double)percentile_occupancy / tile_shape;

}

double HypergeometricDistribution::GetMinTileDensity(const tiling::CoordinateSpaceTileInfo tile) const
{
  if (tile.GetShape() <= specs_.workload_tensor_size - specs_.total_nnzs)
  {
    return 0;
  }
  else
  {
    return (tile.GetShape() - specs_.workload_tensor_size + specs_.total_nnzs)/tile.GetShape();
  }
}

double HypergeometricDistribution::GetTileExpectedDensity(const uint64_t tile_shape) const
{

  (void)tile_shape;
  assert(is_specced_);
  return specs_.average_density;
}

/*************************************************************************************
 * Calculating the probability directly,
 * based on this mathematical definition:
 * https://www.statisticshowto.com/hypergeometric-distribution-examples/
 *************************************************************************************/
double HypergeometricDistribution::CalculateProbability(const std::uint64_t nnz_vals,
                                                        const std::uint64_t r,
                                                        const std::uint64_t n,
                                                        const std::uint64_t N
                                                        ) const
{

   double prob;
   try
   {
      long double r_choose_x = boost::math::binomial_coefficient<long double>(r, nnz_vals);
      long double Nr_choose_nx = boost::math::binomial_coefficient<long double>(N-r, n-nnz_vals);
      long double N_choose_n = boost::math::binomial_coefficient<long double>(N, n);
      prob =  (double) (r_choose_x*Nr_choose_nx)/N_choose_n;
   } catch (std::exception& e)
   {
      std::cout << "ERROR: Integer overflow occured. Check your workload sizes and density."
      << std::endl;
      exit(-1);
   }

    return prob;
}

double HypergeometricDistribution::GetProbability(const std::uint64_t tile_shape,
                                                  const std::uint64_t nnz_vals) const
{

  assert(is_specced_);
  assert(workload_tensor_size_set_);

  std::uint64_t r = specs_.total_nnzs;
  std::uint64_t n = tile_shape;
  std::uint64_t N = specs_.workload_tensor_size;

  if (((n + r > N) && (nnz_vals < n + r - N)) | (nnz_vals > r))
  { return 0; }

  boost::math::hypergeometric_distribution<double> distribution(r, n, N);

  //Workaround/fix for domain error for large workload size. Compute the PDF directly
  //This might return an overflow error, but it occurs less times than using the boost pdf call
  double prob;
  try
  {
    prob = pdf(distribution, nnz_vals);
  } catch (std::exception& e)
  {
    prob = CalculateProbability(nnz_vals, r, n, N);
  }

  return prob;

}

double HypergeometricDistribution::GetProbability(const std::uint64_t tile_shape, const std::uint64_t nnz_vals,
                                                  const std::uint64_t constraint_tensor_shape,
                                                  const std::uint64_t constraint_tensor_occupancy) const
{

  std::uint64_t r = constraint_tensor_occupancy;
  std::uint64_t n = tile_shape;
  std::uint64_t N = constraint_tensor_shape;

  //std::cout << "more involved get prob: nnz: " << nnz_vals << "  tile shape: " << tile_shape
  //<< "   constr shape: " << constraint_tensor_shape << std::endl;

  if (((n + r > N) && (nnz_vals < n + r - N)) | (nnz_vals > r))
  { return 0; }

  boost::math::hypergeometric_distribution<double> distribution(r, n, N);

  //Workaround/fix for domain error for large workload size. Compute the PDF directly
  //This might return an overflow error, but it occurs less times than using the boost pdf call
  double prob;
  try
  {
    prob = pdf(distribution, nnz_vals);
  } catch (std::exception& e) {
    prob = CalculateProbability(nnz_vals, r, n, N);
  }

  return prob;

}

double HypergeometricDistribution::GetTileOccupancyProbability(const tiling::CoordinateSpaceTileInfo& tile,
                                                               const std::uint64_t occupancy) const
{
  std::uint64_t tile_shape = tile.GetShape();
  double prob;

  if (tile.HasExtraConstraintInfo())
  {
    auto extra_constraint_info = tile.GetExtraConstraintInfo();
    auto occupancy_constraint = extra_constraint_info.GetOccupancy();
    auto shape_constraint = extra_constraint_info.GetShape();
    prob = GetProbability(tile_shape, occupancy, shape_constraint, occupancy_constraint);
  } else
  {
    prob = GetProbability(tile_shape, occupancy);
  }

  return prob;
}

double HypergeometricDistribution::GetExpectedTileOccupancy(const tiling::CoordinateSpaceTileInfo tile) const
{

  std::uint64_t tile_shape = tile.GetShape();
  double expected_occupancy = 0.0;

  if (tile.HasExtraConstraintInfo())
  {
    auto extra_constraint_info = tile.GetExtraConstraintInfo();
    auto occupancy_constraint = extra_constraint_info.GetOccupancy();
    auto shape_constraint = extra_constraint_info.GetShape();
    std::uint64_t max_occupancy = (tile_shape <= occupancy_constraint) ? tile_shape : occupancy_constraint;
    for (std::uint64_t occupancy = 0; occupancy <= max_occupancy; occupancy++)
    {
      double prob = GetProbability(tile_shape, occupancy, shape_constraint, occupancy_constraint);
      expected_occupancy += prob * occupancy;
    }
  } else
  {
    std::uint64_t max_occupancy = (tile_shape <= specs_.total_nnzs) ? tile_shape : specs_.total_nnzs;
    for (std::uint64_t occupancy = 0; occupancy <= max_occupancy; occupancy++)
    {
      double prob = GetProbability(tile_shape, occupancy);
      expected_occupancy += prob * occupancy;
    }
  }
  return expected_occupancy;
}

}


//----------------------------------------------------------------------------
// Copyright (C) 2018-2020 Mikhail Komarov <nemo@nil.foundation>
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the Server Side Public License, version 1,
// as published by the author.
//
// 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
// Server Side Public License for more details.
//
// You should have received a copy of the Server Side Public License
// along with this program. If not, see
// <https://github.com/NilFoundation/plugin/blob/master/LICENSE_1_0.txt>.
//----------------------------------------------------------------------------

#define BOOST_TEST_MODULE circuit_hash_test

#include <boost/test/data/monomorphic.hpp>
#include <boost/test/data/test_case.hpp>
#include <boost/test/unit_test.hpp>

#include <nil/filecoin/storage/proofs/porep/stacked/circuit/hash.hpp>

using namespace nil::filecoin;

BOOST_AUTO_TEST_SUITE(stacked_circuit_hash_test_suite)

BOOST_AUTO_TEST_CASE(test_hash2_circuit) {
    const auto rng = XorShiftRng::from_seed(crate::TEST_SEED);

    for (std::size_t i=0; i < 10; ++i) {
        auto cs = TestConstraintSystem::<Bls12>();

        const auto a = Fr::random(rng);
        const auto b = Fr::random(rng);

        const auto a_num = { auto cs = cs.namespace(|| "a");
        num::AllocatedNumber::alloc(cs, || Ok(a))
    };

    const auto b_num = { auto cs = cs.namespace(|| "b");
    num::AllocatedNumber::alloc(cs, || Ok(b))
};

const auto out = <PedersenHasher>::Function::hash2_circuit(cs.namespace(|| "hash2"), &a_num, &b_num)
              .expect("hash2 function failed");

BOOST_CHECK(cs.is_satisfied(), "constraints not satisfied");
BOOST_CHECK_EQUAL(cs.num_constraints(), 1371);

const auto expected : Fr = <PedersenHasher>::Function::hash2(&a.into(), &b.into()).into();

BOOST_CHECK_EQUAL(expected, out.get_value(), "circuit and non circuit do not match");
}
}

BOOST_AUTO_TEST_CASE(test_hash_single_column_circuit) {
    const auto rng = XorShiftRng::from_seed(crate::TEST_SEED);

        for (std::size_t i=0; i < 1; ++i) {
                auto cs = TestConstraintSystem::<Bls12>();

                const std::vector<auto> vals (11, Fr::random(rng));
                const auto vals_opt =
                    vals.iter()
                        .enumerate()
                        .map(| (i, v) |
                             {num::AllocatedNumber::alloc(cs.namespace(|| std::format("num_{}", i)), || Ok(*v))})
                        .collect::<Vec<_>>();

                const auto out = hash_single_column(cs.namespace(|| "hash_single_column"), &vals_opt)
                              .expect("hash_single_column function failed");

                BOOST_ASSERT_MSG(cs.is_satisfied(), "constraints not satisfied");
                BOOST_ASSERT(cs.num_constraints() == 598);

                const auto expected : Fr = vanilla_hash_single_column(&vals);

                BOOST_ASSERT_MSG(expected == out.get_value(), "circuit and non circuit do not match");
            }
}

BOOST_AUTO_TEST_SUITE_END()


#define BOOST_TEST_MODULE simplest_example

#include <iostream>
#include <cstdio>

#include <boost/test/unit_test.hpp>

#include "helpers.h"
#include "segmentationhypothesis.h"
#include "linkinghypothesis.h"
#include "jsonmodel.h"

using namespace mht;
using namespace helpers;

BOOST_AUTO_TEST_CASE( OpenGMInference )
{
	JsonModel model;
	model.readFromJson("constrackingmodel.json");
	size_t numWeights = model.computeNumWeights();
	model.setJsonGtFile("constrackinggt.json");
	std::vector<double> weights = model.learn();
	Solution gt = model.getGroundTruth();
	std::cout << "GT is valid? " << (model.verifySolution(gt)?"yes":"no") << std::endl;
	model.toDot("gt.dot", &gt);

	saveWeightsToJson(weights, "weights.json");
	std::vector<double> weights2 = readWeightsFromJson("weights.json");
	BOOST_CHECK_EQUAL(weights.size(), weights2.size());
	for(size_t i = 0; i < weights.size(); i++)
		BOOST_CHECK_EQUAL(weights[i], weights2[i]);

	JsonModel model2;
	model2.readFromJson("constrackingmodel.json");
	Solution sol = model2.infer(weights);
	model2.saveResultToJson("result.json", sol);
	std::cout << "Solution is valid? " << (model2.verifySolution(sol)?"yes":"no") << std::endl;
	model2.toDot("result.dot", &sol);

	BOOST_CHECK_EQUAL(numWeights, 5);
}