// The contents of this file are licensed under the MIT license.
// See LICENSE.txt for more information.

/*
  This program takes an XML file with a list of photos brown bears, finds all the
  bear faces and the face landmarks then outputs an XML metadata file.

  The input XML file can be generated using "imglab -c". The output file can be
  viewed and edited using imglab.

	  The face detector uses a pretrained CNN from the dlib example:

  https://github.com/davisking/dlib/blob/master/examples/dnn_mmod_dog_hipsterizer.cpp

  While the model was trained for dogs, it works reasonably well for bears. The
  pretrained CNN data can be found here:

  http://dlib.net/files/mmod_dog_hipsterizer.dat.bz2
*/

#include <boost/filesystem.hpp>
#include <boost/algorithm/string.hpp>
#include <iostream>
#include <dlib/cmd_line_parser.h>
#include <dlib/dnn.h>
#include <dlib/data_io.h>
#include <dlib/image_processing.h>
// #include <dlib/gui_widgets.h>
#include <dlib/image_io.h>
// #include <ctype.h>
// #include <stdio.h>
// #include <stdlib.h>
#include <unistd.h>

using namespace std;
using namespace dlib;
using namespace boost::filesystem;

// -----------------------------------------------------------------------------

// A 3x3 conv layer that doesn't do any downsampling
template <long num_filters, typename SUBNET> using con3  = con<num_filters,3,3,1,1,SUBNET>;
// A 5x5 conv layer that does 2x downsampling
template <long num_filters, typename SUBNET> using con5d = con<num_filters,5,5,2,2,SUBNET>;
template <long num_filters, typename SUBNET> using con5  = con<num_filters,5,5,1,1,SUBNET>;

template <typename SUBNET> using downsampler  = relu<affine<con5d<32, relu<affine<con5d<32, relu<affine<con5d<16,SUBNET>>>>>>>>>;
template <typename SUBNET> using downsampler_bn  = relu<bn_con<con5d<32, relu<bn_con<con5d<32, relu<bn_con<con5d<16,SUBNET>>>>>>>>>;
template <typename SUBNET> using rcon5  = relu<affine<con5<45,SUBNET>>>;
template <typename SUBNET> using rcon5_bn  = relu<bn_con<con5<45,SUBNET>>>;

using net_type = loss_mmod<con<1,9,9,1,1,rcon5<rcon5<rcon5<downsampler<input_rgb_image_pyramid<pyramid_down<6>>>>>>>>;
using net_type_bn = loss_mmod<con<1,9,9,1,1,rcon5_bn<rcon5_bn<rcon5_bn<downsampler_bn<input_rgb_image_pyramid<pyramid_down<6>>>>>>>>;

// -----------------------------------------------------------------------------
// Define the 8x downsampling block with conv5d blocks.
// Use relu and batch normalization in the standard way.
template <typename SUBNET> using downsampler_t  = relu<bn_con<con5d<32, relu<bn_con<con5d<32, relu<bn_con<con5d<32,SUBNET>>>>>>>>>;

// The rest of the network will be 3x3 conv layers with batch normalization and
// relu.  So we define the 3x3 block we will use here.
template <typename SUBNET> using rcon3  = relu<bn_con<con3<32,SUBNET>>>;

// Finally, we define the entire network.   The special input_rgb_image_pyramid
// layer causes the network to operate over a spatial pyramid, making the detector
// scale invariant.
using net_type_t  = loss_mmod<con<1,6,6,1,1,rcon3<rcon3<rcon3<downsampler_t<input_rgb_image_pyramid<pyramid_down<6>>>>>>>>;

// -----------------------------------------------------------------------------

std::string g_mode;  // one of train_obj, train_sp, test, infer
unsigned MAX_LONG_SIDE = 2000;
unsigned MAX_SHORT_SIDE = 1500;
unsigned MAX_SIZE = MAX_LONG_SIDE*MAX_SHORT_SIDE;

std::vector<std::vector<double> > get_interocular_distances (
    const std::vector<std::vector<full_object_detection> >& objects
);

matrix<rgb_pixel> downscale_large_image_old (
  matrix<rgb_pixel>& img,
  float& pxRatio
)
{
  bool bDownscaled = false;
  // If the image is too big (this is a memory constraint), we need to downsample it
  pxRatio = 1;
  if (img.size() > (MAX_SIZE))
  {
    if (img.nc() > img.nr())
    {
      if (((float)MAX_LONG_SIDE / (float)img.nc()) < (float)MAX_SHORT_SIDE / (float)img.nr())
        pxRatio = (float)MAX_LONG_SIDE / (float)img.nc();
      else
        pxRatio = (float)MAX_SHORT_SIDE / (float)img.nr();
    }
    else
    {
      if (((float)MAX_LONG_SIDE / (float)img.nr()) < (float)MAX_SHORT_SIDE / (float)img.nc())
        pxRatio = (float)MAX_LONG_SIDE / (float)img.nr();
      else
        pxRatio = (float)MAX_SHORT_SIDE / (float)img.nc();
    }
    // cout << "File TOO BIG" << " Ratio: " << pxRatio << endl;
    matrix<rgb_pixel> smimg((int)(img.nr() * pxRatio), (int)(img.nc() * pxRatio));
    resize_image(img, smimg);
    bDownscaled = true;
    return smimg;
  }
  return img;
}

//------------------------------------------------------------------------------
//  scale image down to MAX_SIZE, fill in pxRatio
//------------------------------------------------------------------------------
matrix<rgb_pixel> downscale_image (
  matrix<rgb_pixel>& img,
  float& pxRatio
)
{
  pxRatio = 1;
  if (img.size() > (MAX_SIZE))
  {
	long orig_img_size = img.size ();
	pxRatio = sqrt ( (float) MAX_SIZE / (float) img.size());
    // cout << "File TOO BIG" << " Ratio: " << pxRatio << endl;
    matrix<rgb_pixel> smimg((int)(img.nr() * pxRatio), (int)(img.nc() * pxRatio));
    resize_image(img, smimg);
	long new_img_size = smimg.size ();
	float new_ratio = sqrt (orig_img_size / new_img_size);
	cout << "old size: " << orig_img_size << endl;
	cout << "new size: " << new_img_size << endl;
    return smimg;
  }
  return img;
}

//------------------------------------------------------------------------------
//  scale parts by ratio
//------------------------------------------------------------------------------
void scale_parts (
	std::map<std::string,point> &parts,
	float pxRatio
)
{
	std::map<std::string,point>::const_iterator itr;
	for (itr = parts.begin(); itr != parts.end(); ++itr)
	{
		parts[itr->first].x() = int((float)itr->second.x() * pxRatio);
		parts[itr->first].y() = int((float)itr->second.y() * pxRatio);
	}
}

//------------------------------------------------------------------------------
//  set rect
//------------------------------------------------------------------------------
void set_face_rect (
	rectangle &rect,
	rectangle detect_rect
)
{
	rect.set_left (detect_rect.left());
	rect.set_top (detect_rect.top());
	rect.set_right (detect_rect.right());
	rect.set_bottom (detect_rect.bottom());
}

//------------------------------------------------------------------------------
//  scale the rect by the ratio
//------------------------------------------------------------------------------
void scale_face_rect (
	rectangle &rect,
	float pxRatio
)
{
	// cout << "before scale: " << rect.left () << endl;
	rect.set_left((int)((float)rect.left() * pxRatio));
	rect.set_top((int)((float)rect.top() * pxRatio));
	rect.set_right((int)((float)rect.right() * pxRatio));
	rect.set_bottom((int)((float)rect.bottom() * pxRatio));
	// cout << "after scale: " << rect.left () << endl;
}

//------------------------------------------------------------------------------
//  set face
//------------------------------------------------------------------------------
void set_face_parts (
	image_dataset_metadata::box &face,
    full_object_detection shape
)
{
	std::vector<std::string> landmarks = {"head_top", "lear", "leye", "nose", "rear", "reye"};
	for (int i = 0; i < landmarks.size (); ++i)
	{
		face.parts[landmarks[i]].x() = shape.part(i).x();
		face.parts[landmarks[i]].y() = shape.part(i).y();
	}
}

//------------------------------------------------------------------------------
//  scale face parts by ratio
//------------------------------------------------------------------------------
void scale_face_parts (
	image_dataset_metadata::box &face,
	float pxRatio
)
{
	std::vector<std::string> landmarks = {"head_top", "lear", "leye", "nose", "rear", "reye"};
	for (int i = 0; i < landmarks.size(); ++i)
	{
		face.parts[landmarks[i]].x() *= pxRatio;
		face.parts[landmarks[i]].y() = face.parts[landmarks[i]].y() * pxRatio;
	}
}

//------------------------------------------------------------------------------
// clear labels from boxes so they won't be trained with images
//------------------------------------------------------------------------------
void clear_box_labels (std::vector<std::vector<mmod_rect>> &faces_list)
{
	for (unsigned long i=0; i < faces_list.size() ; ++i)
	{
		for (unsigned long j=0; j < faces_list[i].size() ; ++j)
			(faces_list[i][j]).label = "";
	}
}

//------------------------------------------------------------------------------
// downscale imgs and face boxes
//------------------------------------------------------------------------------
void downscale_imgs_and_faces (
	std::vector<matrix<rgb_pixel>> &imgs,
	std::vector<std::vector<mmod_rect>> &faces_list
	)
{
	matrix<rgb_pixel> img;
	float pxRatio = 1.0;
	std::vector<mmod_rect> faces;
	for (unsigned long i=0; i < imgs.size() ; ++i)
	{
		img = imgs[i];
		img = downscale_image (img, pxRatio);
		if (pxRatio != 1.0) { // downscaled.  apply to faces of img
			imgs[i] = img;
			faces = faces_list[i];
			for (unsigned long j=0; j < faces_list[i].size() ; ++j)
				scale_face_rect ((faces_list[i][j]).rect, pxRatio);
		}
	}
}

//------------------------------------------------------------------------------
// upscale image if it will remain under MAX_SIZE
//------------------------------------------------------------------------------
bool upscale_image (
  matrix<rgb_pixel>& img,
  float& pxRatio
)
{
  int upscaleRatio = 2;
	long origImgSize = img.size ();
    if (img.size() < (MAX_SIZE/(upscaleRatio*upscaleRatio)))
	{
	  cout << "upscaling..." << endl;
	  pyramid_up(img);
	  pxRatio = sqrt (img.size () / origImgSize);
	  cout << "Upscaled image size: " << img.size() << endl;
	  return true;
	}
	return false;
}

//------------------------------------------------------------------------------
// Find Faces
//------------------------------------------------------------------------------
std::vector<dlib::mmod_rect> find_faces (
  net_type& net,
  matrix<rgb_pixel>& img,
  float &sideRatio
)
{
  img = downscale_image (img, sideRatio);

  // Find faces
  cout << "Finding faces..." << endl;
  std::vector<dlib::mmod_rect> dets;
  // auto dets = net(img);
  dets = net(img);

  // if no faces, try with upscaling
  if (dets.size() == 0)
  {
	if (upscale_image (img, sideRatio))
	  dets = net(img);
  }
  return dets;
}
//------------------------------------------------------------------------------
// Find Faces and face landmarks
//------------------------------------------------------------------------------
void find_faces_and_landmarks (
  net_type& net,
  shape_predictor& sp,
  matrix<rgb_pixel>& img,
  std::vector<image_dataset_metadata::box>& faces,
  std::string bearID
)
{
  float pxRatio = 1.0;
  int numShapeParts = 6;
  // auto dets = net(img);

  auto dets = find_faces (net, img, pxRatio);

  // --- For each face, find the face landmarks ----
  // cout << "Finding landmarks..." << endl;
  for (auto&& d : dets)
  {
      // get the landmarks for this face
      auto shape = sp(img, d.rect);

      // fill in the data for faces
      image_dataset_metadata::box face;
	  float scaleRatio = 1/pxRatio;
	  // face: dlib::image_dataset_metadata::box

	  set_face_rect (face.rect, d);
	  scale_face_rect (face.rect, scaleRatio);
	  set_face_parts (face, shape);
	  scale_face_parts (face, scaleRatio);
	  face.label = bearID;
      faces.push_back(face);
  }
}

//------------------------------------------------------------------------------
// load a batch of images and their corresponding boxes
//------------------------------------------------------------------------------
void get_image_batch (
  size_t num_images,
  dlib::image_dataset_metadata::dataset data,
  std::vector<std::string>& files,
  std::vector<matrix<rgb_pixel>>& images_batch,
  std::vector<std::vector<mmod_rect>>& face_boxes_batch
)
{
  random_cropper cropper;
  std::mutex rnd_mutex;
  images_batch.clear();
  face_boxes_batch.clear();
  files.clear();
  cout << "num imagea per batch: " << num_images << endl;
  for (size_t i = 0; i < num_images; ++i)
  {
    matrix<rgb_pixel> img;
    size_t idx;
    idx = std::rand() % data.images.size();
    //{ std::lock_guard<std::mutex> lock(rnd_mutex);
    //    idx = rnd.get_random_64bit_number()%data.images.size();
    //}

    //cout << idx << " : " << data.images[idx].filename.c_str() << "..." << endl;
    load_image(img, data.images[idx].filename.c_str());
    std::vector<mmod_rect> boxes;
    for (unsigned long j = 0; j < data.images[idx].boxes.size(); ++j)
    {
      boxes.push_back(mmod_rect(data.images[idx].boxes[j].rect));
    }
    matrix<rgb_pixel> crop;
    std::vector<mmod_rect> crop_rects;
    cropper(img, boxes, crop, crop_rects);
    face_boxes_batch.push_back(crop_rects);
    images_batch.push_back(crop);
    files.push_back(data.images[idx].filename.c_str());
  }
}

// ----------------------------------------------------------------------------------------
//   copy of test_object_detection_function from dlib/dnn/validation.h
//		changes:
//		- load only one image at time due to memory limitation
//	  	- downscale image if beyond MAX_SIZE
//    	- if detect no faces, upscale if doesn't violate MAX_SIZE
// ----------------------------------------------------------------------------------------

template <
	typename SUBNET
	>
const matrix<double,1,3> my_test_object_detection_function (
	loss_mmod<SUBNET>& detector,
    dlib::image_dataset_metadata::dataset img_data,
	const std::vector<std::vector<mmod_rect>>& old_truth_dets,
	const test_box_overlap& overlap_tester = test_box_overlap(),
	const double adjust_threshold = 0,
	const test_box_overlap& overlaps_ignore_tester = test_box_overlap()
)
{
	double correct = 0;
	double correct_hits = 0;
	double total_true_targets = 0;

	std::vector<std::pair<double,bool> > all_dets;
	unsigned long missing_detections = 0;
	unsigned long prev_missed = 0;
	resizable_tensor temp;

	for (unsigned long i = 0; i < img_data.images.size(); ++i)
	{
		std::vector<mmod_rect> hits;
        matrix<rgb_pixel> img;
        // cout << img_data.images[i].filename.c_str() << "..." << endl;
        load_image(img, img_data.images[i].filename.c_str());

		float pxRatio = 1.0;
		//--------------------------------------------------
		//--------------------------------------------------
		img = downscale_image (img, pxRatio);
		detector.to_tensor(&img, &img+1, temp);
		detector.subnet().forward(temp);
		detector.loss_details().to_label(temp, detector.subnet(), &hits, adjust_threshold);
		if (hits.size () == 0)  // no faces found, try upscale image
		{
			if (upscale_image (img, pxRatio))
			{
				detector.to_tensor(&img, &img+1, temp);
				detector.subnet().forward(temp);
				detector.loss_details().to_label(temp, detector.subnet(), &hits, adjust_threshold);
			}
		}
		//--------------------------------------------------
		//--------------------------------------------------
        std::vector<mmod_rect> truth_dets;
		truth_dets.clear();
		// populate truth detects
		for (unsigned long j = 0; j < img_data.images[i].boxes.size(); ++j)
		{
			if (img_data.images[i].boxes[j].ignore)
			{
				cout << "Encountered unexpected ignore!" << endl;
			}
			else
			{
				if (pxRatio != 1) { // downscaled.  apply to truth recs
					scale_face_rect (img_data.images[i].boxes[j].rect, pxRatio);
				}
				truth_dets.push_back(mmod_rect(img_data.images[i].boxes[j].rect));
			}
		}
		// iterate through labels though we don't use labels
		for (auto& label : impl::get_labels(truth_dets, hits))
		{
			std::vector<full_object_detection> truth_boxes;
			std::vector<rectangle> ignore;
			std::vector<std::pair<double,rectangle>> boxes;
			// copy hits and truth_dets into the above three objects
			for (auto&& b : truth_dets)
			{
				if (b.ignore)
				{
					ignore.push_back(b);
				}
				else // if (b.label == label)
				{
					truth_boxes.push_back(full_object_detection(b.rect));
					++total_true_targets;
				}
			}
			for (auto&& b : hits)
			{
				// if (b.label == label)
					boxes.push_back(std::make_pair(b.detection_confidence, b.rect));
			}

			prev_missed = missing_detections;
			correct = impl::number_of_truth_hits(truth_boxes, ignore, boxes, overlap_tester, all_dets, missing_detections, overlaps_ignore_tester);
			correct_hits += correct;
			if (boxes.size () > truth_boxes.size ()) // # detects more than # truths
			{
				if (correct < truth_boxes.size()) // # detects == # truths but some didn't match
				{
					cout << img_data.images[i].filename.c_str() << " has extra detections and " << truth_boxes.size () - correct << " erroneous detection" << endl;
				}
				else
				{
					cout << img_data.images[i].filename.c_str() << " has extra detections" << endl;
				}
			}
			else if (boxes.size () < truth_boxes.size ()) // # detects less than # truths
			{
				if (correct < boxes.size ())
					cout << img_data.images[i].filename.c_str() << " has missing detections and " << boxes.size () - correct << " erroneous detection" << endl;
				else
					cout << img_data.images[i].filename.c_str() << " has missing detections" << endl;
			}
			else if (correct < boxes.size()) // # detects == # truths but some didn't match
			{
        		cout << img_data.images[i].filename.c_str() << " has " << boxes.size () - correct << " erroneous detections" << endl;
			}
		}
	}
	std::sort(all_dets.rbegin(), all_dets.rend());
	double precision, recall;
	double total_hits = all_dets.size();
	if (total_hits == 0)
		precision = 1;
	else
		precision = correct_hits / total_hits;
	if (total_true_targets == 0)
		recall = 1;
	else
		recall = correct_hits / total_true_targets;
	matrix<double, 1, 3> res;
	res = precision, recall, average_precision(all_dets, missing_detections);
	return res;
}

// ----------------------------------------------------------------
// initial copy from dlib/image_processing/shape_preditor.h
//   modified to take metadata instead of array of images
//	 due to memory limit.  load a single image and scale as needed.
// ----------------------------------------------------------------

    double my_test_shape_predictor (
        const shape_predictor& sp,
        dlib::image_dataset_metadata::dataset img_data,
        double & mean,
        double & stddev
    )
    {
	/*
		dlib::array<array2d<unsigned char> > images_test2;
		std::vector<std::vector<full_object_detection> > faces_test;
		load_image_dataset(images_test2, faces_test, imgs_file);
        cout << "shape predictor mean testing error:  " <<
            my_test_shape_predictor(sp, images_test2, faces_test, get_interocular_distances(faces_test)) << endl;
	*/

        running_stats<double> rs;
		for (unsigned long i = 0; i < img_data.images.size(); ++i)
		{
			matrix<rgb_pixel> img;
			// cout << i << ": " << img_data.images[i].filename.c_str() << "..." << endl;
			load_image(img, img_data.images[i].filename.c_str());

			float pxRatio = 1.0;
			// --- scale image -----
			img = downscale_image (img, pxRatio);
			// extract box rects from img_data
			auto objects = img_data.images[i].boxes;

            for (unsigned long j = 0; j < objects.size(); ++j)
            {
                // Just use a scale of 1 (i.e. no scale at all) if the caller didn't supply
                // any scales.
				if (pxRatio != 1.0)
				{
					scale_parts (objects[j].parts, pxRatio);
					scale_face_rect (objects[j].rect, pxRatio);
				}
                const double scale = length (objects[j].parts["leye"] - objects[j].parts["reye"]);

                full_object_detection det = sp(img, objects[j].rect);

				std::map<std::string,point>::const_iterator itr;
				unsigned long k = 0;
				for (itr = objects[j].parts.begin(); itr != objects[j].parts.end(); ++itr)
				{

					double score = length(det.part(k) - itr->second)/scale;
					// cout << "score for " << itr->first << ": " << score;
					// cout << " : " <<  det.part(k) << endl;
					rs.add(score);
					k++;
                }
            }
        }
        mean = rs.mean();
        stddev = rs.stddev();
        return mean;
    }

//----------------------------------------------------------------
//  train_obj
//----------------------------------------------------------------
net_type_bn run_train_obj (std::string train_file, int batch_size)
{
  cout << "\n\tTraining with file: " << train_file << endl;
  dlib::image_dataset_metadata::dataset data;
  load_image_dataset_metadata(data, train_file);

  std::vector<matrix<rgb_pixel>> images_train;
  std::vector<std::vector<mmod_rect>> face_boxes_train;
  //load_image_dataset(images_train, face_boxes_train, train_file);
  cout << "num training images: " << data.images.size() << endl;
  // --- scale images and face boxes ---
  //downscale_imgs_and_faces (images_train, face_boxes_train);
  for (unsigned long i = 0; i < data.images.size(); ++i)
  {
    std::vector<mmod_rect> boxes;
    for (unsigned long j = 0; j < data.images[i].boxes.size(); ++j)
    {
      boxes.push_back(mmod_rect(data.images[i].boxes[j].rect));
    }
    face_boxes_train.push_back(boxes);
  }

  clear_box_labels (face_boxes_train);

  // mmod_options options(face_boxes_train, 40,40); // doghip
  mmod_options options(face_boxes_train, 80,80); // bearface_network
  // mmod_options options(face_boxes_train, 40,40); // AnimalWeb
  // The detector will automatically decide to use multiple sliding windows if needed.
  // For the face data, only one is needed however.
  cout << "num detector windows: "<< options.detector_windows.size() << endl;
  for (auto& w : options.detector_windows)
    cout << "detector window width by height: " << w.width << " x " << w.height << endl;
  cout << "overlap NMS IOU thresh:             " << options.overlaps_nms.get_iou_thresh() << endl;
  cout << "overlap NMS percent covered thresh: " << options.overlaps_nms.get_percent_covered_thresh() << endl;

  // Now we are ready to create our network and trainer.
  net_type_bn net(options);

  // The MMOD loss requires that the number of filters in the final network layer equal
  // options.detector_windows.size().  So we set that here as well.
  net.subnet().layer_details().set_num_filters(options.detector_windows.size());
  dnn_trainer<net_type_bn> trainer(net);
  trainer.set_learning_rate(0.1);
  trainer.be_verbose();
  trainer.set_synchronization_file("mmod_sync", std::chrono::minutes(5));
  // trainer.set_iterations_without_progress_threshold(300); // doghip
  trainer.set_iterations_without_progress_threshold(8000);

  // Now let's train the network.  We are going to use mini-batches of 150
  // images.   The images are random crops from our training set (see
  // random_cropper_ex.cpp for a discussion of the random_cropper).
  std::vector<matrix<rgb_pixel>> mini_batch_samples;
  std::vector<std::vector<mmod_rect>> mini_batch_labels;
  std::vector<std::string> mini_batch_files;
  //random_cropper cropper;
  // cropper.set_chip_dims(200, 200);  // doghip
  // cropper.set_min_object_size(0.2);  // doghip

  // Run the trainer until the learning rate gets small.  This will probably take several
  // hours.

  // Set up data loaders
  dlib::pipe<std::vector<std::string>> qfiles(4);
  dlib::pipe<std::vector<matrix<rgb_pixel>>> qimages(4);
  dlib::pipe<std::vector<std::vector<mmod_rect>>> qlabels(4);
  auto data_loader = [&batch_size, &data, &qfiles, &qimages, &qlabels](time_t seed)
  {
    dlib::rand rnd(time(0)+seed);
    std::vector<matrix<rgb_pixel>> images;
    std::vector<std::vector<mmod_rect>> labels;
    std::vector<std::string> files;
    dlib::rand rnd_color(time(0)+seed);
    while(qimages.is_enabled())
    {
      try
      {
        get_image_batch(batch_size, data, files, images, labels);
        //get_image_batch(90, data, images, labels);
        //load_mini_batch(numid, numface, rnd, objs, images, labels);
        // We can also randomly jitter the colors and that often helps a detector
        // generalize better to new images.
        for (auto&& img : images)
          disturb_colors(img, rnd_color);

        //cout << "Queue batch" << endl;
        //for (size_t i = 0; i < files.size(); ++i)
        //{
        //  cout << i << " queue : " << files[i] << "..." << endl;
        //}

        qfiles.enqueue(files);
        qimages.enqueue(images);
        qlabels.enqueue(labels);
      }
      catch(std::exception& e)
      {
        cout << "EXCEPTION IN LOADING DATA" << endl;
        cout << e.what() << endl;
      }
    }
  };
  std::thread data_loader1([data_loader](){ data_loader(1); });
  std::thread data_loader2([data_loader](){ data_loader(2); });
  std::thread data_loader3([data_loader](){ data_loader(3); });
  std::thread data_loader4([data_loader](){ data_loader(4); });
  std::thread data_loader5([data_loader](){ data_loader(5); });

  cout << "\tRunning training ... \n" << endl;
  while(trainer.get_learning_rate() >= 1e-4)
  {
    // load a batch of images and their corresponding boxes from the queue
    qimages.dequeue(mini_batch_samples);
    qlabels.dequeue(mini_batch_labels);
    qfiles.dequeue(mini_batch_files);
    //cout << "Dequeue batch" << endl;
    //for (size_t i = 0; i < mini_batch_files.size(); ++i)
    //{
      //cout << "dequeue : " << mini_batch_files[i] << "..." << endl;
    //}
    //get_image_batch(75, data, mini_batch_samples, mini_batch_labels);
    // cropper(150, images_train, face_boxes_train, mini_batch_samples, mini_batch_labels);
    // crop size of 10
    //cropper(75, images_batch, face_boxes_batch, mini_batch_samples, mini_batch_labels);
    trainer.train_one_step(mini_batch_samples, mini_batch_labels);
  }
  // wait for training threads to stop
  trainer.get_net();
  cout << "done training" << endl;
  qimages.disable();
  qlabels.disable();
  qfiles.disable();
  data_loader1.join();
  data_loader2.join();
  data_loader3.join();
  data_loader4.join();
  data_loader5.join();

	return net;
}

//----------------------------------------------------------------
//  train_sp
//----------------------------------------------------------------
shape_predictor run_train_sp (std::string train_file)
{
	cout << "\n\tTraining with file: " << train_file << endl;
	// dlib::image_dataset_metadata::dataset data;
	// load_image_dataset_metadata(data, imgs_file);

	// std::vector<matrix<rgb_pixel>> images_train;
	// std::vector<std::vector<mmod_rect>> face_boxes_train;
	// load_image_dataset(images_train, face_boxes_train, train_file);
	// --- scale images and face boxes ---
	// downscale_imgs_and_faces (images_train, face_boxes_train);

//----------------------------------------------------------------------

	dlib::array<array2d<unsigned char> > images_train;
	std::vector<std::vector<full_object_detection> > faces_train;
	load_image_dataset(images_train, faces_train, train_file);
	cout << "num training images: " << images_train.size() << endl;
	shape_predictor_trainer trainer;
	trainer.set_oversampling_amount(300);
	trainer.set_nu(0.05);
	trainer.set_tree_depth(2);
	trainer.set_num_threads(2);
	// cascade depth=20, tree depth=5, padding=0.2
	trainer.set_cascade_depth (20);
	trainer.set_tree_depth (5);
	trainer.set_feature_pool_region_padding (0.2);

	trainer.be_verbose();
	cout << "\trunning shape predictor training ... \n";
	shape_predictor sp = trainer.train(images_train, faces_train);
	return sp;
	/*
	cout << "mean training error: "<<
		test_shape_predictor(sp, images_train, faces_train, get_interocular_distances(faces_train)) << endl;
	cout << "mean testing error:  "<<
		test_shape_predictor(sp, images_test, faces_test, get_interocular_distances(faces_test)) << endl;
	serialize("sp.dat") << sp;
	*/
}

//----------------------------------------------------------------
// get_output_network_name ()
//	 if specified in argument, ensure parent directory exist.
//	 else append current timestamp
//----------------------------------------------------------------
std::string get_output_network_name (command_line_parser& parser)
{
	std::string out_net;
	if (parser.option ("out_network"))
	{
		out_net = parser.option ("out_network").argument();
		boost::filesystem::path p_out_network (out_net);
		if (p_out_network.has_parent_path () &&
			!boost::filesystem::exists(p_out_network.parent_path()))
			boost::filesystem::create_directories(p_out_network.parent_path());
	}
	else
	{
		time_t rawtime;
		struct tm * timeinfo;
		char buffer[80];

		time (&rawtime);
		timeinfo = localtime(&rawtime);

		strftime(buffer,sizeof(buffer),"%Y%m%d%I%M",timeinfo);
		out_net = "network_";
		out_net.append (buffer);
		out_net.append (".dat");
	}
	return out_net;
}

// ----------------------------------------------------------------------------------------

// ----------------------------------------------------------------------------------------

int main(int argc, char** argv)
{try
{
		time_t timeStart = time(NULL);
		command_line_parser parser;

		parser.add_option("h","Display this help message.");
		parser.add_option("train_obj","Train object detector.", 1);
		parser.add_option("size_max", "scale images to max size.", 1);
		parser.add_option("xy_max", "scale images to max x and y.", 2);
		parser.add_option("out_network","Newly trained network.", 1);
		parser.add_option("train_sp","Train shape predictor.", 1);
		parser.add_option("test_obj","Test object detector.", 1);
		parser.add_option("test_sp","Test shape predictor.", 1);
		parser.add_option("infer","Detect faces using network file.", 1);
		parser.add_option("batch_size","Images to load into GPU for each train step. Defaults to 90.", 1);
		parser.add_option("load_all","Load all images into memory. Not yet implemented.");
		parser.parse(argc, argv);

		const char* one_time_opts[] = {"h", "train_obj", "train_sp", "test_obj", "test_sp", "infer"};
		parser.check_one_time_options(one_time_opts); // Can't give an option more than once

		if (parser.option("h") || parser.number_of_arguments () != 1)
		{
			cout << "\n\t bearface is used to detect a bear face in an image using a\n";
			cout << "\tnetwork file with the --infer option.";
			cout << "\tIt can be trained using the --train_obj or --train_sp \n";
			cout << "\tand tested with the --test flags.\n\n";
			cout << "\nUsage  : bearface --<infer|test|train_sp|train_obj> <network_file> <img_xml>\n";
			cout << "\nUsage  : bearface --infer bearface_network.dat imgs.xml\n";
			cout << "\nExample: bearface --test bearface_network.dat imgs.xml\n\n";
			cout << "\nExample: bearface --train_obj bearface_network.dat -out new_network.dat imgs.xml\n\n";
			parser.print_options();

			return EXIT_SUCCESS;
		}
		std::string network;
		if (parser.option ("size_max"))
		{
			MAX_SIZE = stoi (parser.option ("size_max").argument());
		}
		if (parser.option ("xy_max"))
		{
			MAX_LONG_SIDE = stoi (parser.option ("xy_max").argument(0));
			MAX_SHORT_SIDE = stoi (parser.option ("xy_max").argument(1));
			MAX_SIZE = MAX_LONG_SIDE*MAX_SHORT_SIDE;
		}
		if (parser.option("train_obj"))
		{
			g_mode = "train_obj";
			network = parser.option ("train_obj").argument();
		}
		else if (parser.option("train_sp"))
		{
			g_mode = "train_sp";
			network = parser.option ("train_sp").argument();
		}
		else if (parser.option("test_obj"))
		{
			g_mode = "test_obj";
			network = parser.option ("test_obj").argument();
		}
		else if (parser.option("test_sp"))
		{
			g_mode = "test_sp";
			network = parser.option ("test_sp").argument();
		}
		else if (parser.option("infer"))
		{
			g_mode = "infer";
			network = parser.option ("infer").argument();
		}

	//----------------------------------------------------------------------
    //image_window win_wireframe;

	char *lvalue = NULL;
	int index;
	int c;
	int batch_size = 90;
	std::string bearID;
	bool bLabelFixed = false;

	std::string imgs_file = parser[0];
	if (parser.option("batch_size"))
		batch_size = stoi (parser.option ("batch_size").argument());

    // load the models as well as glasses and mustache.
    net_type net;
    shape_predictor sp;
    matrix<rgb_alpha_pixel> glasses, mustache;
    deserialize(network) >> net >> sp >> glasses >> mustache;

	// doing testing
	if ((g_mode == "test_sp") || (g_mode == "test_obj")) {
		cout << "\tTesting with " << imgs_file << endl;
		std::vector<matrix<rgb_pixel>> images_test;
		std::vector<std::vector<mmod_rect>> face_boxes_test;
		dlib::image_dataset_metadata::dataset data;
		load_image_dataset_metadata(data, imgs_file);
		cout << "num testing images:  " << data.images.size() << endl;
		// ---  running test on object detection -----
		if (g_mode == "test_obj") {
			cout << "\tstarting object detector test..." << endl;
			matrix<double, 1, 3> res = my_test_object_detection_function(net, data, face_boxes_test);
			cout << "precision =  correct hits / total hits         " << endl;
			cout << "recall    =  correct hits / total true targets " << endl;
			cout << "av precision  " << endl;
			cout << "                   precision    average precision" << endl;
			cout << "                           recall " << endl;
			cout << "testing results:  " << res << endl;
		} else {
			// doing shape predictor testing
			cout << "\tstarting shape predictor test..." << endl;
			dlib::array<array2d<unsigned char> > images_test2;
			std::vector<std::vector<full_object_detection> > faces_test;
			double sp_mean, sp_stddev;

			my_test_shape_predictor(sp, data, sp_mean, sp_stddev);
			cout << "shape predictor mean testing error:  " <<
				sp_mean << " +/- " << sp_stddev << endl;

			cout << "on " << data.images.size() << " tests" << endl;
			cout << endl;
		}
		return 0;
	}
	//----------------------------------------------------------------------
	//
	//----------------------------------------------------------------------
	if (g_mode == "train_sp") {
		std::string out_net = get_output_network_name (parser);
		shape_predictor new_sp = run_train_sp (imgs_file);

		cout << "\n\tWriting to network: " << out_net+"_bn" << endl;
		net.clean();
		net_type anet = net;
		serialize(out_net) << anet << new_sp << glasses << mustache;
		cout << "\n\tWriting to network: " << out_net << endl;

		return 0;
	}
	else if (g_mode == "train_obj") {
		//std::string out_network = "";
    std::string out_net = get_output_network_name (parser);
    net_type_bn net = run_train_obj (imgs_file, batch_size);

		// Save the network to disk
		net.clean();

		serialize(out_net+"_bn") << net << sp << glasses << mustache;
		cout << "\n\tWriting to network: " << out_net+"_bn" << endl;
		net_type anet = net;
		serialize(out_net) << anet << sp << glasses << mustache;
		cout << "\n\tWriting to network: " << out_net << endl;

		return 0;
	}

	// doing inferencing

    int total_faces = 0;
	int files_with_faces = 0;
	int files_without_faces = 0;
	cout << "\nDoing inferencing ...\n" << endl;
    // Load XML metadata file
    dlib::image_dataset_metadata::dataset data;
    load_image_dataset_metadata(data, imgs_file);

    //Handle list of images
	std::vector <string> fields;
    for (int i = 0; i < data.images.size(); ++i)
    {
        matrix<rgb_pixel> img;

        cout << "processing: " << data.images[i].filename.c_str() << "...";
        load_image(img, data.images[i].filename.c_str());

		if (!bLabelFixed)
		{
			std::string fullpathfile = data.images[i].filename;
			boost::split( fields, fullpathfile, boost::is_any_of( "/" ));
			bearID = fields[fields.size() - 2];
		}
        std::vector<image_dataset_metadata::box> faces;
        find_faces_and_landmarks (net, sp, img, faces, bearID);
        data.images[i].boxes = faces;

        cout << "faces found: " << to_string(faces.size()) << endl;
        total_faces += faces.size();
		if (faces.size () == 0)
			files_without_faces++;
		else
			files_with_faces++;
    }
    cout << "Files with detected faces   : " << files_with_faces << endl;
    cout << "Files without detected faces: " << files_without_faces << endl;
    cout << "Total faces found           : " << total_faces << endl;
	path orig_path(imgs_file);
	std::string faces_file;
	if (orig_path.has_parent_path ())
	{
		faces_file = orig_path.parent_path().string () + "/";
	}
	faces_file += orig_path.stem().string() + "_faces.xml";
	cout << "faces_file: " << faces_file << endl;
    save_image_dataset_metadata(data, faces_file);
    if (!bLabelFixed)
    {
      cout << "\n\tGenerated " << total_faces << " faces in file: " << faces_file << "\n" << endl;
    }
    else
    {
      cout << "\n\tGenerated " << total_faces << " faces with label " << bearID << " in file: " << faces_file << "\n" << endl;
    }
}
catch(std::exception& e)
{
    cout << e.what() << endl;
}
}

// ----------------------------------------------------------------------------------------

double interocular_distance (
    const full_object_detection& det
)
{
    dlib::vector<double,2> l, r;
    // Find the center of the left eye by averaging the points around
    // the eye.
	l = det.part(2);

    // Find the center of the right eye by averaging the points around
    // the eye.
	r = det.part(5);

    // Now return the distance between the centers of the eyes
    return length(l-r);
}

std::vector<std::vector<double> > get_interocular_distances (
    const std::vector<std::vector<full_object_detection> >& objects
)
{
    std::vector<std::vector<double> > temp(objects.size());
    for (unsigned long i = 0; i < objects.size(); ++i)
    {
        for (unsigned long j = 0; j < objects[i].size(); ++j)
        {
            temp[i].push_back(interocular_distance(objects[i][j]));
        }
    }
    return temp;
}


#include <pcl/point_cloud.h>
#include <pcl/point_types.h>
#include <pcl/kdtree/kdtree_flann.h>
#include <pcl/search/kdtree.h>
#include <pcl_conversions/pcl_conversions.h>
#include <iostream>

#include <geometry_msgs/PoseStamped.h>
#include <geometry_msgs/Vector3.h>
#include <math.h>
#include <nav_msgs/Odometry.h>
#include <ros/console.h>
#include <ros/ros.h>
#include <sensor_msgs/PointCloud2.h>
#include <sensor_msgs/LaserScan.h>

#include <Eigen/Eigen>
#include <random>

using namespace std;
using namespace Eigen;

const int MAX_SWRAM_NUM = 40;

ros::Publisher _all_map_pub;
ros::Timer global_pcl_timer;
ros::Publisher pub_cloud[MAX_SWRAM_NUM+1];
ros::Subscriber _odom_sub[MAX_SWRAM_NUM+1];
nav_msgs::Odometry _odom[MAX_SWRAM_NUM+1];
ros::Timer local_sensing_timer[MAX_SWRAM_NUM+1];

ros::Publisher pub_cloud_ugv[MAX_SWRAM_NUM+1];
ros::Subscriber _odom_ugv_sub[MAX_SWRAM_NUM+1];
nav_msgs::Odometry _odom_ugv[MAX_SWRAM_NUM+1];
ros::Timer local_sensing_ugv_timer[MAX_SWRAM_NUM+1];

pcl::KdTreeFLANN<pcl::PointXYZ> kdtreeLocalMap;
vector<int> pointIdxRadiusSearch;
vector<float> pointRadiusSquaredDistance;

random_device rd;
default_random_engine eng(rd());
uniform_real_distribution<double> rand_x;
uniform_real_distribution<double> rand_y;

vector<double> _state;
int map_flag, swarm_num, swarm_num_ugv, uav_id[MAX_SWRAM_NUM+1], ugv_id[MAX_SWRAM_NUM+1];
std::string uav_name[MAX_SWRAM_NUM+1],ugv_name[MAX_SWRAM_NUM+1];
double _x_size, _y_size, _z_size;
double _x_l, _x_h, _y_l, _y_h;
double _sensing_range, sensing_horizon, _resolution, _sense_rate, _init_x, _init_y;
double _min_dist;

bool _map_ok = false;
bool _has_odom[MAX_SWRAM_NUM+1] = {false};
bool _has_odom_ugv[MAX_SWRAM_NUM+1] = {false};

// 障碍物参数 - 圆柱体、正方体
int cylinder_num;
double cylinder_radius, cylinder_height;
int sqaure_num;
double sqaure_size;

sensor_msgs::PointCloud2 globalMap_pcd;
sensor_msgs::PointCloud2 localMap_pcd;

pcl::PointCloud<pcl::PointXYZ> cloudMap,local_map;

void MapGenerateAMOV_1()
{
  // 待存入全局点云的点
  pcl::PointXYZ pt_random;

  double x, y,z;

  // 生成圆柱体障碍物
  for (int i = 1; i <= 21; i++)
  {
    for (int j = 1; j <= 21; j++)
    {

      x = -18 + 2*i;
      y = -18 + 2*j;

      if(j%2 == 1)
      {
        // y轴单数生成圆柱体
        x = floor(x / _resolution) * _resolution + _resolution / 2.0;
        y = floor(y / _resolution) * _resolution + _resolution / 2.0;

        int widNum = ceil((cylinder_radius) / _resolution);
        int heiNum = ceil(cylinder_height / _resolution);

        // 生成圆柱体的点云
        for (int r = -widNum; r < widNum; r++)
          for (int s = -widNum; s < widNum; s++) {
            for (int t = -1.0; t < heiNum; t++) {
              double temp_x = x + (r + 0.5) * _resolution + 1e-2;
              double temp_y = y + (s + 0.5) * _resolution + 1e-2;
              double temp_z = (t + 0.5) * _resolution + 1e-2;
              if ( (Eigen::Vector2d(temp_x,temp_y) - Eigen::Vector2d(x,y)).norm() <= cylinder_radius )
              {
                pt_random.x = temp_x;
                pt_random.y = temp_y;
                pt_random.z = temp_z;
                cloudMap.points.push_back(pt_random);
              }
            }
          }

      }else
      {
        // y轴单数生成正方体
        z = sqaure_size / 2;
        x = floor(x / _resolution) * _resolution + _resolution / 2.0;
        y = floor(y / _resolution) * _resolution + _resolution / 2.0;
        z = floor(z / _resolution) * _resolution + _resolution / 2.0;

        int widNum = ceil((sqaure_size) / _resolution);

        // 生成正方体的点云
        for (int r = -widNum / 2.0; r < widNum / 2.0; r++)
          for (int s = -widNum / 2.0; s < widNum / 2.0; s++) {
            for (int t = -widNum / 2.0; t < widNum / 2.0; t++) {
              double temp_x = x + (r + 0.5) * _resolution + 1e-2;
              double temp_y = y + (s + 0.5) * _resolution + 1e-2;
              double temp_z = z + (t + 0.5) * _resolution + 1e-2;
              pt_random.x = temp_x;
              pt_random.y = temp_y;
              pt_random.z = temp_z;
              cloudMap.points.push_back(pt_random);
            }
          }
      }

    }
  }

  cloudMap.width = cloudMap.points.size();
  cloudMap.height = 1;
  cloudMap.is_dense = true;

  ROS_INFO("\033[1;33;41m----> Finished generate AMOV map 1.\033[0m");

  kdtreeLocalMap.setInputCloud(cloudMap.makeShared());

  _map_ok = true;
}

void generate_cylinder(double x,double y)
{
  // 必须在地图范围内
  if(x<_x_l || x>_x_h || y<_y_l ||  y>_y_h)
  {
    return;
  }

  pcl::PointXYZ pt_random;
  x = floor(x/_resolution)*_resolution + _resolution/2.0;
  y = floor(y/_resolution)*_resolution + _resolution/2.0;
  int widNum = ceil((cylinder_radius)/_resolution);
  int heiNum = ceil(cylinder_height/_resolution);
  for(int r = -widNum; r < widNum; r++)
    for(int s = -widNum; s < widNum; s++)
      for(int t = -1.0; t < heiNum; t++) {
        double temp_x = x + (r + 0.5)*_resolution + 1e-2;
        double temp_y = y + (s + 0.5)*_resolution + 1e-2;
        double temp_z = (t + 0.5)*_resolution + 1e-2;
        if((Eigen::Vector2d(temp_x,temp_y) - Eigen::Vector2d(x,y)).norm() <= cylinder_radius)
        {
          pt_random.x = temp_x;
          pt_random.y = temp_y;
          pt_random.z = temp_z;
          cloudMap.points.push_back(pt_random);
        }
      }
}

void MapGenerateAMOV_2()
{
  // 待存入全局点云的点
  pcl::PointXYZ pt_random;

  double x, y,z;

  // 生成圆柱体障碍物
  generate_cylinder(2.0,0.0);
  generate_cylinder(4.0,0.0);
  generate_cylinder(6.0,0.0);

  cloudMap.width = cloudMap.points.size();
  cloudMap.height = 1;
  cloudMap.is_dense = true;

  ROS_INFO("\033[1;33;41m----> Finished generate AMOV map 2.\033[0m");

  kdtreeLocalMap.setInputCloud(cloudMap.makeShared());

  _map_ok = true;
}

void generate_cube(double x,double y)
{
  // 必须在地图范围内
  if(x<_x_l || x>_x_h || y<_y_l ||  y>_y_h)
  {
    return;
  }

  pcl::PointXYZ pt_random;
  double z = sqaure_size/2.0;
  x = floor(x/_resolution)*_resolution + _resolution/2.0;
  y = floor(y/_resolution)*_resolution + _resolution/2.0;
  z = floor(z/_resolution)*_resolution + _resolution/2.0;
  int widNum = ceil((sqaure_size)/_resolution);
  for(int r = -widNum / 2.0; r < widNum / 2.0; r++)
    for(int s = -widNum / 2.0; s < widNum / 2.0; s++)
      for(int t = -widNum / 2.0; t < widNum / 2.0; t++) {
        double temp_x = x + (r + 0.5)*_resolution + 1e-2;
        double temp_y = y + (s + 0.5)*_resolution + 1e-2;
        double temp_z = z + (t + 0.5)*_resolution + 1e-2;
        pt_random.x = temp_x;
        pt_random.y = temp_y;
        pt_random.z = temp_z;
        cloudMap.points.push_back(pt_random);
      }
}

int sign(int num, double mid,bool zero_flag)
{
  if(num>mid) return -1;
  if(zero_flag && num==mid) return 0;
  else return 1;
}

void MapGenerateAMOV_3()
{
  double x,y;
  for(int num = 0; num < 10; num++)
  {
    // format1
    for(int ci = 0; ci < 5; ci++)
    {
      x = sign(num,4,0) * (ci*2 + 3);
      y = (num*4*2)%40 - 13;
      generate_cylinder(x,y);
      // printf("<model name='cylinder_%d'>\n\t<include>\n\t\t<uri>model://cylinder</uri>\n\t\t<pose>",num*5+ci);
      // printf("%.1f %.1f 1.5 0 0 0</pose>\n\t</include>\n</model>\n",x,y);
    }
    for(int qi = 0; qi < 6; qi++)
    {
      x = sign(num,4,0) * (qi*2 + 2);
      y = (num*4*2)%40 - 13 + sign(qi%2,0,0)*0.5;
      generate_cube(x,y);
      // printf("<model name='square_%d'>\n\t<include>\n\t\t<uri>model://square</uri>\n\t\t<pose>",num*6+qi);
      // printf("%.1f %.1f 0.2 0 0 0</pose>\n\t</include>\n</model>\n",x,y);
    }
    // format2
    for(int ci = 0; ci < 5; ci++)
    {
      x = sign(num,4,0) * (ci*2 + 3);
      y = (num*4*2)%40 - 14.5 + sign(ci,2,1)*(ci-2)*0.5;
      generate_cylinder(x,y);
      // printf("<model name='cylinder_%d'>\n\t<include>\n\t\t<uri>model://cylinder</uri>\n\t\t<pose>",num*5+ci+50);
      // printf("%.1f %.1f 1.5 0 0 0</pose>\n\t</include>\n</model>\n",x,y);
    }
    for(int qi = 0; qi < 6; qi++)
    {
      x = sign(num,4,0) * (qi*2 + 2);
      y = (num*4*2)%40 - 15 - sign(qi%2,0,0)*0.5;
      generate_cube(x,y);
      // printf("<model name='square_%d'>\n\t<include>\n\t\t<uri>model://square</uri>\n\t\t<pose>",num*6+qi+60);
      // printf("%.1f %.1f 0.2 0 0 0</pose>\n\t</include>\n</model>\n",x,y);
    }
    // format3
    for(int ci = 0; ci < 5; ci++)
    {
      x = sign(num,4,0) * (ci*2 + 3);
      y = (num*4*2)%40 - 17 + sign(ci%2,0,0)*0.5;
      generate_cylinder(x,y);
      // printf("<model name='cylinder_%d'>\n\t<include>\n\t\t<uri>model://cylinder</uri>\n\t\t<pose>",num*5+ci+100);
      // printf("%.1f %.1f 1.5 0 0 0</pose>\n\t</include>\n</model>\n",x,y);
    }
    for(int qi = 0; qi < 6; qi++)
    {
      x = sign(num,4,0) * (qi*2 + 2);
      y = (num*4*2)%40 - 17;
      generate_cube(x,y);
      // printf("<model name='square_%d'>\n\t<include>\n\t\t<uri>model://square</uri>\n\t\t<pose>",num*6+qi+120);
      // printf("%.1f %.1f 0.2 0 0 0</pose>\n\t</include>\n</model>\n",x,y);
    }
    // format4
    for(int ci = 0; ci < 5; ci++)
    {
      x = sign(num,4,0) * (ci*2 + 3);
      y = (num*4*2)%40 - 19 + (ci-2)/2*0.5;
      generate_cylinder(x,y);
      // printf("<model name='cylinder_%d'>\n\t<include>\n\t\t<uri>model://cylinder</uri>\n\t\t<pose>",num*5+ci+150);
      // printf("%.1f %.1f 1.5 0 0 0</pose>\n\t</include>\n</model>\n",x,y);
    }
    for(int qi = 0; qi < 6; qi++)
    {
      x = sign(num,4,0) * (qi*2 + 2);
      y = (num*4*2)%40 - 19 + sign(qi%2,0,0)*0.5;
      generate_cube(x,y);
      // printf("<model name='square_%d'>\n\t<include>\n\t\t<uri>model://square</uri>\n\t\t<pose>",num*6+qi+180);
      // printf("%.1f %.1f 0.2 0 0 0</pose>\n\t</include>\n</model>\n",x,y);
    }
  }
  cloudMap.width = cloudMap.points.size();
  cloudMap.height = 1;
  cloudMap.is_dense = true;

  ROS_INFO("\033[1;33;41m----> Finished generate AMOV map 3.\033[0m");

  kdtreeLocalMap.setInputCloud(cloudMap.makeShared());

  _map_ok = true;
}

void MapGenerateAMOV_4()
{
  double x,y;
  for(int num = 0; num < 10; num++)
  {
    // format1
    for(int ci = 0; ci < 3; ci++)
    {
      x = sign(num,4,0) * (ci*4 + 4);
      y = (num*4*2)%40 - 13;
      generate_cylinder(x,y);
      // printf("<model name='cylinder_%d'>\n\t<include>\n\t\t<uri>model://cylinder</uri>\n\t\t<pose>",num*5+ci);
      // printf("%.1f %.1f 1.5 0 0 0</pose>\n\t</include>\n</model>\n",x,y);
    }
    for(int qi = 0; qi < 3; qi++)
    {
      x = sign(num,4,0) * (qi*4 + 2);
      y = (num*4*2)%40 - 13 + sign(qi%2,0,0)*0.5;
      generate_cube(x,y);
      // printf("<model name='square_%d'>\n\t<include>\n\t\t<uri>model://square</uri>\n\t\t<pose>",num*6+qi);
      // printf("%.1f %.1f 0.2 0 0 0</pose>\n\t</include>\n</model>\n",x,y);
    }
    // format2
    for(int ci = 0; ci < 3; ci++)
    {
      x = sign(num,4,0) * (ci*4 + 4);
      y = (num*4*2)%40 - 14.5 + sign(ci,2,1)*(ci-2)*0.5;
      generate_cylinder(x,y);
      // printf("<model name='cylinder_%d'>\n\t<include>\n\t\t<uri>model://cylinder</uri>\n\t\t<pose>",num*5+ci+50);
      // printf("%.1f %.1f 1.5 0 0 0</pose>\n\t</include>\n</model>\n",x,y);
    }
    for(int qi = 0; qi < 3; qi++)
    {
      x = sign(num,4,0) * (qi*4 + 2);
      y = (num*4*2)%40 - 15 - sign(qi%2,0,0)*0.5;
      generate_cube(x,y);
      // printf("<model name='square_%d'>\n\t<include>\n\t\t<uri>model://square</uri>\n\t\t<pose>",num*6+qi+60);
      // printf("%.1f %.1f 0.2 0 0 0</pose>\n\t</include>\n</model>\n",x,y);
    }
    // format3
    for(int ci = 0; ci < 3; ci++)
    {
      x = sign(num,4,0) * (ci*4 + 4);
      y = (num*4*2)%40 - 17 + sign(ci%2,0,0)*0.5;
      generate_cylinder(x,y);
      // printf("<model name='cylinder_%d'>\n\t<include>\n\t\t<uri>model://cylinder</uri>\n\t\t<pose>",num*5+ci+100);
      // printf("%.1f %.1f 1.5 0 0 0</pose>\n\t</include>\n</model>\n",x,y);
    }
    for(int qi = 0; qi < 3; qi++)
    {
      x = sign(num,4,0) * (qi*4 + 2);
      y = (num*4*2)%40 - 17;
      generate_cube(x,y);
      // printf("<model name='square_%d'>\n\t<include>\n\t\t<uri>model://square</uri>\n\t\t<pose>",num*6+qi+120);
      // printf("%.1f %.1f 0.2 0 0 0</pose>\n\t</include>\n</model>\n",x,y);
    }
    // format4
    for(int ci = 0; ci < 3; ci++)
    {
      x = sign(num,4,0) * (ci*4 + 4);
      y = (num*4*2)%40 - 19 + (ci-2)/2*0.5;
      generate_cylinder(x,y);
      // printf("<model name='cylinder_%d'>\n\t<include>\n\t\t<uri>model://cylinder</uri>\n\t\t<pose>",num*5+ci+150);
      // printf("%.1f %.1f 1.5 0 0 0</pose>\n\t</include>\n</model>\n",x,y);
    }
    for(int qi = 0; qi < 3; qi++)
    {
      x = sign(num,4,0) * (qi*4 + 2);
      y = (num*4*2)%40 - 19 + sign(qi%2,0,0)*0.5;
      generate_cube(x,y);
      // printf("<model name='square_%d'>\n\t<include>\n\t\t<uri>model://square</uri>\n\t\t<pose>",num*6+qi+180);
      // printf("%.1f %.1f 0.2 0 0 0</pose>\n\t</include>\n</model>\n",x,y);
    }
  }
  cloudMap.width = cloudMap.points.size();
  cloudMap.height = 1;
  cloudMap.is_dense = true;

  ROS_INFO("\033[1;33;41m----> Finished generate AMOV map 4.\033[0m");

  kdtreeLocalMap.setInputCloud(cloudMap.makeShared());

  _map_ok = true;
}

void RandomMapGenerateAMOV()
{
  // 待存入全局点云的点
  pcl::PointXYZ pt_random;

  // 障碍物位置容器
  vector<Eigen::Vector2d> obs_position;

  // 将无人机初始位置存入容器（防止在无人机初始位置周围生成障碍物）
  obs_position.push_back( Eigen::Vector2d(_init_x,_init_y) );

  // 障碍物位置随机生成器
  rand_x = uniform_real_distribution<double>(_x_l, _x_h);
  rand_y = uniform_real_distribution<double>(_y_l, _y_h);

  // 生成圆柱体障碍物
  for (int i = 0; i < cylinder_num && ros::ok(); i++)
  {
    double x, y;
    // 随机生成障碍物位置：[x,y]
    x = rand_x(eng);
    y = rand_y(eng);

    // 检查是否与已经生成的障碍物重叠：两个障碍物之间距离大于 _min_dist，否则生成失败
    bool flag_continue = false;
    for ( auto p : obs_position )
      if ( (Eigen::Vector2d(x,y) - p).norm() < _min_dist /*metres*/ )
      {
        i--; // i--代表此次生成失败，重新生成障碍物
        flag_continue = true;
        break;
      }
    if ( flag_continue ) continue;

    // 将本次生成的障碍物位置推入容器
    obs_position.push_back( Eigen::Vector2d(x,y) );

    // floor 是向下取整函数
    x = floor(x / _resolution) * _resolution + _resolution / 2.0;
    y = floor(y / _resolution) * _resolution + _resolution / 2.0;

    int widNum = ceil((cylinder_radius) / _resolution);
    int heiNum = ceil(cylinder_height / _resolution);

    // 生成圆柱体的点云
    for (int r = -widNum; r < widNum; r++)
      for (int s = -widNum; s < widNum; s++) {
        for (int t = -1.0; t < heiNum; t++) {
          double temp_x = x + (r + 0.5) * _resolution + 1e-2;
          double temp_y = y + (s + 0.5) * _resolution + 1e-2;
          double temp_z = (t + 0.5) * _resolution + 1e-2;
          if ( (Eigen::Vector2d(temp_x,temp_y) - Eigen::Vector2d(x,y)).norm() <= cylinder_radius )
          {
            pt_random.x = temp_x;
            pt_random.y = temp_y;
            pt_random.z = temp_z;
            cloudMap.points.push_back(pt_random);
          }
        }
      }
  }

  // 生成正方形障碍物
  for (int i = 0; i < sqaure_num; ++i)
  {
    double x, y, z;
    // 随机生成障碍物位置：[x,y]
    x = rand_x(eng);
    y = rand_y(eng);
    z = sqaure_size / 2;

    // 检查是否与已经生成的障碍物重叠：两个障碍物之间距离大于 _min_dist，否则生成失败
    bool flag_continue = false;
    for ( auto p : obs_position )
      if ( (Eigen::Vector2d(x,y) - p).norm() < _min_dist /*metres*/ )
      {
        i--; // i--代表此次生成失败，重新生成障碍物
        flag_continue = true;
        break;
      }
    if ( flag_continue ) continue;

    // 将本次生成的障碍物位置推入容器
    obs_position.push_back( Eigen::Vector2d(x,y) );

    x = floor(x / _resolution) * _resolution + _resolution / 2.0;
    y = floor(y / _resolution) * _resolution + _resolution / 2.0;
    z = floor(z / _resolution) * _resolution + _resolution / 2.0;

    int widNum = ceil((sqaure_size) / _resolution);

    // 生成正方体的点云
    for (int r = -widNum / 2.0; r < widNum / 2.0; r++)
      for (int s = -widNum / 2.0; s < widNum / 2.0; s++) {
        for (int t = -widNum / 2.0; t < widNum / 2.0; t++) {
          double temp_x = x + (r + 0.5) * _resolution + 1e-2;
          double temp_y = y + (s + 0.5) * _resolution + 1e-2;
          double temp_z = z + (t + 0.5) * _resolution + 1e-2;
          pt_random.x = temp_x;
          pt_random.y = temp_y;
          pt_random.z = temp_z;
          cloudMap.points.push_back(pt_random);
        }
      }
  }

  cloudMap.width = cloudMap.points.size();
  cloudMap.height = 1;
  cloudMap.is_dense = true;

  ROS_INFO("\033[1;33;41m----> Finished generate random map.\033[0m");

  kdtreeLocalMap.setInputCloud(cloudMap.makeShared());

  _map_ok = true;
}

void rcvOdometryCallbck(const nav_msgs::Odometry::ConstPtr& odom, int uav_id)
{
  _has_odom[uav_id] = true;

  _odom[uav_id] = *odom;

  _state = {odom->pose.pose.position.x,
            odom->pose.pose.position.y,
            odom->pose.pose.position.z,
            odom->twist.twist.linear.x,
            odom->twist.twist.linear.y,
            odom->twist.twist.linear.z,
            0.0,
            0.0,
            0.0};
}

void rcvUGVOdometryCallbck(const nav_msgs::Odometry::ConstPtr& odom, int ugv_id)
{
  _has_odom_ugv[ugv_id] = true;

  _odom_ugv[ugv_id] = *odom;

  _state = {odom->pose.pose.position.x,
            odom->pose.pose.position.y,
            odom->pose.pose.position.z,
            odom->twist.twist.linear.x,
            odom->twist.twist.linear.y,
            odom->twist.twist.linear.z,
            0.0,
            0.0,
            0.0};
}

void pubGlobalPoints(const ros::TimerEvent& event)
{
  if(_map_ok)
  {
    // 发布全局点云
    pcl::toROSMsg(cloudMap, globalMap_pcd);
    globalMap_pcd.header.frame_id = "world";
    _all_map_pub.publish(globalMap_pcd);
  }else
  {
    return;
  }
}

void renderSensedPoints(const ros::TimerEvent& event, int uav_id)
{
  if (!_map_ok || !_has_odom[uav_id]) return;

  Eigen::Quaterniond q;
  q.x() = _odom[uav_id].pose.pose.orientation.x;
  q.y() = _odom[uav_id].pose.pose.orientation.y;
  q.z() = _odom[uav_id].pose.pose.orientation.z;
  q.w() = _odom[uav_id].pose.pose.orientation.w;

  Eigen::Matrix3d rot;
  rot = q;
  Eigen::Vector3d yaw_vec = rot.col(0);

  local_map.points.clear();
  pcl::PointXYZ searchPoint(_odom[uav_id].pose.pose.position.x,
                            _odom[uav_id].pose.pose.position.y,
                            _odom[uav_id].pose.pose.position.z);
  pointIdxRadiusSearch.clear();
  pointRadiusSquaredDistance.clear();

  pcl::PointXYZ pt;
  if (kdtreeLocalMap.radiusSearch(searchPoint, sensing_horizon,
                                   pointIdxRadiusSearch,
                                   pointRadiusSquaredDistance) > 0) {
    for (size_t i = 0; i < pointIdxRadiusSearch.size(); ++i) {
      pt = cloudMap.points[pointIdxRadiusSearch[i]];

      // if ((fabs(pt.z - _odom.pose.pose.position.z) / (pt.x - _odom.pose.pose.position.x)) >
      //     tan(M_PI / 12.0))
      //   continue;
      if ((fabs(pt.z - _odom[uav_id].pose.pose.position.z) / sensing_horizon) >
          tan(M_PI / 6.0))
        continue;

      Vector3d pt_vec(pt.x - _odom[uav_id].pose.pose.position.x,
                      pt.y - _odom[uav_id].pose.pose.position.y,
                      pt.z - _odom[uav_id].pose.pose.position.z);

      if (pt_vec.normalized().dot(yaw_vec) < 0.5) continue;

      local_map.points.push_back(pt);
    }
  } else {
    return;
  }

  local_map.width = local_map.points.size();
  local_map.height = 1;
  local_map.is_dense = true;

  pcl::toROSMsg(local_map, localMap_pcd);
  localMap_pcd.header.frame_id = "/world";

  pub_cloud[uav_id].publish(localMap_pcd);
}

void renderSensedPointsUGV(const ros::TimerEvent& event, int ugv_id)
{
  if (!_map_ok || !_has_odom_ugv[ugv_id]) return;

  Eigen::Quaterniond q;
  q.x() = _odom_ugv[ugv_id].pose.pose.orientation.x;
  q.y() = _odom_ugv[ugv_id].pose.pose.orientation.y;
  q.z() = _odom_ugv[ugv_id].pose.pose.orientation.z;
  q.w() = _odom_ugv[ugv_id].pose.pose.orientation.w;

  Eigen::Matrix3d rot;
  rot = q;
  Eigen::Vector3d yaw_vec = rot.col(0);

  local_map.points.clear();
  pcl::PointXYZ searchPoint(_odom_ugv[ugv_id].pose.pose.position.x,
                            _odom_ugv[ugv_id].pose.pose.position.y,
                            _odom_ugv[ugv_id].pose.pose.position.z);
  pointIdxRadiusSearch.clear();
  pointRadiusSquaredDistance.clear();

  pcl::PointXYZ pt;
  if (kdtreeLocalMap.radiusSearch(searchPoint, sensing_horizon,
                                   pointIdxRadiusSearch,
                                   pointRadiusSquaredDistance) > 0) {
    for (size_t i = 0; i < pointIdxRadiusSearch.size(); ++i) {
      if(abs(pt.z - _odom_ugv[ugv_id].pose.pose.position.z)>5*_resolution+1e-2) continue; //qi comment
      pt = cloudMap.points[pointIdxRadiusSearch[i]];
      // pt.z = _odom_ugv[ugv_id].pose.pose.position.z; //qi comment

      // if ((fabs(pt.z - _odom_ugv.pose.pose.position.z) / (pt.x - _odom_ugv.pose.pose.position.x)) >
      //     tan(M_PI / 12.0))
      //   continue;
      if ((fabs(pt.z - _odom_ugv[ugv_id].pose.pose.position.z) / sensing_horizon) >
          tan(M_PI / 6.0))
        continue;

      Vector3d pt_vec(pt.x - _odom_ugv[ugv_id].pose.pose.position.x,
                      pt.y - _odom_ugv[ugv_id].pose.pose.position.y,
                      pt.z - _odom_ugv[ugv_id].pose.pose.position.z);

      if (pt_vec.normalized().dot(yaw_vec) < 0.5) continue;

      local_map.points.push_back(pt);
    }
  } else {
    return;
  }

  local_map.width = local_map.points.size();
  local_map.height = 1;
  local_map.is_dense = true;

  pcl::toROSMsg(local_map, localMap_pcd);
  localMap_pcd.header.frame_id =  "/world";

  pub_cloud_ugv[ugv_id].publish(localMap_pcd);
}

int main(int argc, char** argv)
{
  ros::init(argc, argv, "random_map_sensing_new");
  ros::NodeHandle n("~");

  n.param("swarm_num", swarm_num, 0);
  n.param("swarm_num_ugv", swarm_num_ugv, 0);
  // 0代表随机生成地图，1使用1号地图，2使用2号地图
  n.param("map_generator/map_flag", map_flag, 0);
  // 地图尺寸
  n.param("map_generator/x_size", _x_size, 30.0);
  n.param("map_generator/y_size", _y_size, 30.0);
  n.param("map_generator/z_size", _z_size, 2.0);
  // 点云分辨率
  n.param("map_generator/resolution", _resolution, 0.1);
  // 障碍物间最小距离
  n.param("map_generator/min_distance", _min_dist, 1.0);
  // 障碍物参数
  n.param("map_generator/cylinder_num", cylinder_num, 30);
  n.param("map_generator/sqaure_num", sqaure_num, 30);
  n.param("map_generator/cylinder_radius", cylinder_radius, 0.2);
  n.param("map_generator/cylinder_height", cylinder_height, 3.0);
  n.param("map_generator/sqaure_size", sqaure_size, 0.25);
  // 无人机初始位置
  n.param("init_state_x", _init_x, 0.0);
  n.param("init_state_y", _init_y, 0.0);
  // 感知范围与感知频率
  n.param("map_generator/sensing_range", _sensing_range, 10.0);
  n.param("map_generator/sensing_horizon", sensing_horizon, 5.0);
  n.param("map_generator/sense_rate", _sense_rate, 10.0);

  _x_l = -_x_size / 2.0;
  _x_h = +_x_size / 2.0;

  _y_l = -_y_size / 2.0;
  _y_h = +_y_size / 2.0;

  if( (cylinder_num + sqaure_num) > _x_size*8)
  {
    ROS_ERROR("The map can't put all the obstacles, remove some.");
    cylinder_num = _x_size*4;
    sqaure_num = _x_size*4;
  }

  // 发布 全局地图
  _all_map_pub = n.advertise<sensor_msgs::PointCloud2>("/map_generator/global_cloud", 1);
  global_pcl_timer = n.createTimer(ros::Duration(10.0/_sense_rate), pubGlobalPoints);

  // 无人机
  for(int i = 1 ; i <= swarm_num; i++)
  {
    uav_id[i] = i;
    uav_name[i] = "/uav"+std::to_string(i);
    pub_cloud[i] = n.advertise<sensor_msgs::PointCloud2>("/uav"+std::to_string(i)+"/map_generator/local_cloud", 1);
    _odom_sub[i] = n.subscribe<nav_msgs::Odometry>("/uav"+std::to_string(i)+"/prometheus/drone_odom", 10, boost::bind(&rcvOdometryCallbck,_1,uav_id[i]));
    local_sensing_timer[i] = n.createTimer(ros::Duration(1.0/_sense_rate), boost::bind(&renderSensedPoints,_1,uav_id[i]));
  }

  // 无人车
  for(int i = 1 ; i <= swarm_num_ugv; i++)
  {
    ugv_id[i] = i;
    ugv_name[i] = "/ugv"+std::to_string(i);
    pub_cloud_ugv[i] = n.advertise<sensor_msgs::PointCloud2>("/ugv"+std::to_string(i)+"/map_generator/local_cloud", 1);
    _odom_ugv_sub[i] = n.subscribe<nav_msgs::Odometry>("/ugv"+std::to_string(i)+"/odom", 10, boost::bind(&rcvUGVOdometryCallbck,_1,ugv_id[i]));
    local_sensing_ugv_timer[i] = n.createTimer(ros::Duration(1.0/_sense_rate), boost::bind(&renderSensedPointsUGV,_1,ugv_id[i]));
  }

  unsigned int seed = rd();
  // unsigned int seed = 2433201515;
  cout << "seed=" << seed << endl;
  eng.seed(seed);

  if(map_flag == 0) RandomMapGenerateAMOV();
  else if(map_flag == 1) MapGenerateAMOV_1();
  else if(map_flag == 2) MapGenerateAMOV_2();
  else if(map_flag == 3) MapGenerateAMOV_3();
  else if(map_flag == 4) MapGenerateAMOV_4();
  else ROS_ERROR("[Map server] wrong map_flag.");

  ros::Rate loop_rate(100.0);

  while (ros::ok()) {
    ros::spinOnce();
    loop_rate.sleep();
  }
}


#include <boost/geometry/strategies/geographic/distance_cross_track.hpp>


/*
* Software License Agreement (BSD License)
*
*  Copyright (c) 2015, Fabian Girrbach, Social Robotics Lab, University of Freiburg
*  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 the copyright holder nor the names of its contributors
*    may be used to endorse or promote products derived from this software
*    without specific prior written permission.
*
*  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
*  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
*  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
*  DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
*  FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
*  DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
*  SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
*  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
*  OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
*  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

#include <srl_nearest_neighbor_tracker/base/stl_helpers.h>
#include <srl_nearest_neighbor_tracker/ros/params.h>
#include <srl_nearest_neighbor_tracker/imm_filter.h>

#include <tf/tf.h>
#include <ros/ros.h>
#include <Eigen/LU>
#include <limits>
#include <map>

#include <std_msgs/ColorRGBA.h>
#include <std_msgs/Float32MultiArray.h>
#include <geometry_msgs/Point.h>
#include <visualization_msgs/Marker.h>

#define foreach BOOST_FOREACH

namespace srl_nnt {

IMMFilter::IMMFilter(const ros::NodeHandle& nodeHandle, const ros::NodeHandle& privateNodeHandle)
: m_nodeHandle(nodeHandle),
  m_privateNodeHandle(privateNodeHandle),
  m_immVisualization(false),
  m_sendDebugInformation(false)
{
    // TODO Initialization of filter with parameters and prefixes
    m_numberModels = Params::get<int>("number_of_models", 2);
    ROS_INFO_STREAM_NAMED("IMM", "Initializing IMM filter with " << m_numberModels << " models.");

    for (unsigned int i = 0; i < m_numberModels ; i++)
    {
        double checksum = 0.0;
        for (unsigned int j = 0; j < m_numberModels ; j++)
        {
            stringstream markovParamName;
            markovParamName << "markov_trans_prob_" << i << "_"<< j;
            m_markovTransitionProbabilities(i,j) = Params::get<double>(markovParamName.str(),0.0);
            checksum += m_markovTransitionProbabilities(i,j);
            ROS_DEBUG_STREAM("Markov transition probability from " << i << " to " << j << " set to " << m_markovTransitionProbabilities(i,j));

        }
        if(std::abs(checksum - 1.0) > 0.01) ROS_ERROR_STREAM_NAMED("IMM", "Markov transition probabilities have to sum up to one!");
    }

    for (unsigned int i = 0; i < m_numberModels ; i++)
    {
        stringstream currrentParameterPrefix;
        currrentParameterPrefix << "IMM" << i << Params::get<string>("prefix_model_parameter","_");
        EKF::Ptr ekf (new EKF(currrentParameterPrefix.str()));
        m_kalmanFilters.push_back(ekf);
    }

    privateNodeHandle.getParam("use_imm_visualization", m_immVisualization);
    if (m_immVisualization)
        m_immVisualizationPublisher = m_privateNodeHandle.advertise<visualization_msgs::Marker>("imm_visualization", 10);
    privateNodeHandle.getParam("send_imm_debug_information", m_sendDebugInformation);
    if(m_sendDebugInformation){
        m_immDebugPublisher = m_privateNodeHandle.advertise<spencer_tracking_msgs::ImmDebugInfos>("imm_debug_infos",10);
        m_debugMessages.header.seq = 0;
        m_debugMessages.header.frame_id = "odom";
    }

}

void IMMFilter::setTransitionMatrix(const StateVector& x,const double deltaT)
{
    ROS_DEBUG_NAMED("IMM", "Setting IMM transitions Matrix to all Kalman filters");
    foreach(EKF::Ptr ekf, m_kalmanFilters) {
        ekf->setTransitionMatrix(x,deltaT);
    }
}

FilterState::Ptr IMMFilter::initializeTrackState(Observation::ConstPtr observation, const VelocityVector& initialVelocity)
{
    IMMState::Ptr immState (new IMMState);
    ROS_DEBUG_NAMED("IMM", "New IMM state created");

    foreach(EKF::Ptr ekf, m_kalmanFilters) {
        //create new hypothesis for each model
        IMMHypothesis::Ptr hypothesis (new IMMHypothesis);
        ekf->initializeTrackState(hypothesis, observation, initialVelocity);
        //initialize probability of hypothesis equally depending on number of models
        hypothesis->m_probability = 1.0/(double) m_numberModels;
        immState->m_hypotheses.push_back(hypothesis);
        //ROS_INFO_STREAM("Initialize ekf with hypothesis C" << hypothesis->C() << " x " << hypothesis->x());
        m_debugMessages.header.frame_id = "odom";

    }

    // Setting current state to first hypothesis should be therefore the most probable model
    immState->m_currentHypothesis = immState->m_hypotheses.front();
    immState->m_currentHypothesisIdx =0;
    updateStateEstimate(immState);

    return immState;
}

FilterState::Ptr IMMFilter::initializeTrackStateFromLogicInitiator(InitiatorCandidate::Ptr candidate)
{
    IMMState::Ptr immState (new IMMState);
    ROS_DEBUG_NAMED("IMM", "New IMM state created");

    foreach(EKF::Ptr ekf, m_kalmanFilters) {
        // create new hypothesis for each model
        IMMHypothesis::Ptr hypothesis (new IMMHypothesis);
        hypothesis->m_x = candidate->state->x();
        hypothesis->m_C = candidate->state->C();

        // initialize probability of hypothesis equally depending on number of models
        hypothesis->m_probability = 1.0/(double) m_numberModels;
        immState->m_hypotheses.push_back(hypothesis);

        //ROS_INFO_STREAM("Initialize ekf with hypothesis C" << hypothesis->C() << " x " << hypothesis->x());
    }

    //Setting current state to first hypothesis should be therefore the most probable model
    immState->m_currentHypothesis = immState->m_hypotheses.front();
    immState->m_currentHypothesisIdx =0;
    updateStateEstimate(immState);
    return immState;
}

void IMMFilter::predictTrackState(FilterState::Ptr state, double deltatime)
{
    sendDebugInformation();
    ROS_DEBUG_NAMED("IMM", "Predict IMM track states!");

    IMMState::Ptr imm = boost::dynamic_pointer_cast<IMMState>(state);
    doMixing(imm);

    // Do the traditional kalman filter predictions
    for (int i = 0 ; i < m_kalmanFilters.size(); i++)
    {
        m_kalmanFilters.at(i)->predictTrackState(imm->m_hypotheses.at(i),deltatime);
        //ROS_WARN_STREAM("Prediction for hypothesis " << i << "with probability  " << imm->m_hypotheses.at(i)->m_probability << " Cov =  \n" << imm->m_hypotheses.at(i)->C()
        //      << "\n CovP =  \n" << imm->m_hypotheses.at(i)->Cp() << "\n Transition Matrix A \n" << m_kalmanFilters.at(i)->m_A);

    }
    mixPredictions(imm);
}

void IMMFilter::mixPredictions(IMMState::Ptr state)
{
    state->m_xp = StateVector::Zero();
    for (int i = 0; i < state->m_hypotheses.size(); i++)
    {
        IMMHypothesis::Ptr hyp = state->m_hypotheses.at(i);
        state->m_xp += hyp->m_xp * hyp->m_probability;
    }

    state->m_Cp = StateMatrix::Zero();
    for (int i = 0; i < state->m_hypotheses.size(); i++)
    {
        IMMHypothesis::Ptr hyp = state->m_hypotheses.at(i);
        StateVector diff = hyp->m_xp - state->m_xp;
        state->m_Cp += (hyp->m_Cp + (diff * diff.transpose())) * hyp->m_probability;
    }
}

void IMMFilter::updateMatchedTrack(FilterState::Ptr state, Pairing::ConstPtr pairingTracker)
{
    // cast to IMM state
    IMMState::Ptr imm = boost::dynamic_pointer_cast<IMMState>(state);

    // Do the traditional kalman filter updating
    for (int i = 0 ; i < m_kalmanFilters.size(); i++)
    {
        //get kalman filter state for current hypothesis
        IMMHypothesis::Ptr hyp = imm->m_hypotheses.at(i) ;

        // Create a new pairing for later update
        Pairing::Ptr pairing( new Pairing );

        pairing->observation = pairingTracker->observation;
        pairing->validated = pairingTracker->validated;

        // Calculate innovation v and inverse of innovation covariance S
        pairing->v = pairing->observation->z - hyp->zp();
        ObsMatrix S = hyp->H() * hyp->Cp() * hyp->H().transpose() + pairing->observation->R;
        Eigen::FullPivLU<ObsMatrix> lu(S);
        double ln_det_S = log(lu.determinant());

        // Calculate inverse of innovation covariance if possible
        const double MATRIX_LN_EPS = -1e8;

        if (ln_det_S > MATRIX_LN_EPS) {
            pairing->Sinv = lu.inverse();
            pairing->d = (pairing->v.transpose() * pairing->Sinv * pairing->v)(0,0);
            pairing->singular = pairing->d < 0.0;
        }
        else {
            pairing->Sinv = ObsMatrix::Constant(OBS_DIM, OBS_DIM, numeric_limits<double>::quiet_NaN());
            pairing->d = numeric_limits<double>::quiet_NaN();
            pairing->singular = true;

            ROS_WARN_STREAM("Singular pairing encountered!\nTrack  measurement prediction:\n" << state->zp() << "\nTrack covariance prediction:\n" << hyp->Cp()
                            << "\nObservation " << pairing->observation->id << " mean:\n" << pairing->observation->z << "\nObservation covariance:\n" << pairing->observation->R
                            << "\nH:\n" << hyp->H() << "\nS:\n" << S << "\nR:\n" << pairing->observation->R);
        }

        // update kalman filter with pairing for current filter
        ROS_DEBUG_NAMED("IMM", "Updating kalman filter of IMM");
        m_kalmanFilters.at(i)->updateMatchedTrack(hyp, pairing);

        ROS_DEBUG_NAMED("IMM", "Calculating likelihood");

        // save likelihood for hypothesis
        imm->m_hypotheses.at(i)->m_likelihood = calcLikelihood(pairing->d, lu.determinant());
    }

    modeProbabilityUpdate(imm);
    updateStateEstimate(imm);
    getDebugInformation(imm,pairingTracker);
    updateCurrentHypothesis(imm, pairingTracker->track);
}

double IMMFilter::calcLikelihood(double d, double detS)
{
    //REVIEW showed that this is wrong: NMODELS should not be here and also the square of mahlanobis distance d is shit
    //likelihood = exp(-(d * d) / 2.0) / (sqrt(pow((2 * M_PI), N_MODELS) * detS ));
    //ROS_INFO_STREAM("Exp likelihood " << likelihood);
    //Java implementation does it like this--> we can get rid of the pi factor because its just normalizing and constant for all hypotheses
    double likelihood = exp(-d/2.0) / sqrt(detS);
    return likelihood;
}

void IMMFilter::updateOccludedTrack(FilterState::Ptr state)
{
    // cast to IMM state
    IMMState::Ptr imm = boost::dynamic_pointer_cast<IMMState>(state);

    // Do the traditional kalman filter updating
    for (int i = 0 ; i < m_kalmanFilters.size(); i++)
    {
        IMMHypothesis::Ptr hyp = imm->getHypothesis(i);
        m_kalmanFilters.at(i)->updateOccludedTrack(hyp);
        //ROS_WARN_STREAM("Occluded update for hypothesis " << i << " Cov =  \n" << hyp->C());
    }

    updateStateEstimate(imm);
    //ROS_ERROR_STREAM("Occluded update mixed Cov= \n" << imm->C());
}

void IMMFilter::doMixing(IMMState::Ptr state)
{
    ROS_DEBUG_NAMED("IMM", "Model mixing");
    computeMixingProbabilities(state);
    computeMixedMean(state);
    computeMixedCovariance(state);
    state->useMixedValues();
}

void IMMFilter::computeMixingProbabilities(IMMState::Ptr state) {

    //ROS_DEBUG_NAMED("IMM", "Before Loop");
    for( int j = 0; j < state->getHypotheses().size(); j++ )
    {
        //FIXME Matrix representation should be possible
        double normalizer = 0.0;
        IMMHypothesis::Ptr hyp = state->getHypothesis(j);

        // Calculate normalizer which is constant for one target hypothesis
        for( int i = 0; i < state->m_hypotheses.size(); i++ )
        {
            normalizer += m_markovTransitionProbabilities(i,j) * state->getHypothesis(i)->getProbability();
        }
        hyp->m_cNormalizer = normalizer;

        // Calculate Mixed Probability for each hypothesis
        for( int i = 0; i < state->m_hypotheses.size(); i++ )
        {
            state->m_mixingProbabilities(i,j) = m_markovTransitionProbabilities(i,j) * state->getHypothesis(i)->getProbability() / hyp->m_cNormalizer;
            ROS_DEBUG_STREAM_NAMED("IMM", "Mixing probability (" << i<< ";" << j << ")=" << state->m_mixingProbabilities(i,j));
        }
    }
}

void IMMFilter::computeMixedMean(IMMState::Ptr state)
{
    // Compute mixed initial condition for all hypothesis
    for( int j = 0; j < state->getHypotheses().size(); j++ )
    {
        state->getHypothesis(j)->m_xMixed = StateVector::Zero();
        for( int i = 0; i < state->getHypotheses().size(); i++ )
        {
            state->getHypothesis(j)->m_xMixed += state->getHypothesis(i)->x() * state->m_mixingProbabilities(i,j);
        }
        ROS_DEBUG_STREAM_NAMED("IMM", "Unmixed Mean x = " << state->getHypothesis(j)->x() << " Mixed Mean x = " << state->getHypothesis(j)->xMixed() );
    }
}

void IMMFilter::computeMixedCovariance(IMMState::Ptr state)
{
    // Compute mixed covariance taken from book p.456
    // Matrix operations in java are quite strange
    for( int j = 0; j < state->getHypotheses().size(); j++ )
    {
        state->getHypothesis(j)->m_CMixed = StateMatrix::Zero();
        for( int i = 0; i < state->getHypotheses().size(); i++ )
        {
            StateVector diff = state->getHypothesis(i)->x() - state->getHypothesis(j)->xMixed();
            state->getHypothesis(j)->m_CMixed += state->m_mixingProbabilities(i,j) * (state->getHypothesis(i)->C()
                    + diff * diff.transpose());
        }
        ROS_DEBUG_STREAM_NAMED("IMM", "Unmixed Cov C = " << state->getHypothesis(j)->C() << " Mixed Cov Cmixed = " << state->getHypothesis(j)->CMixed() );
    }
}

void IMMFilter::updateStateEstimate(IMMState::Ptr state)
{
    state->m_x = StateVector::Zero();
    for (int i = 0; i < state->m_hypotheses.size(); i++)
    {
        IMMHypothesis::Ptr hyp = state->m_hypotheses.at(i);
        state->m_x += hyp->m_x * hyp->m_probability;
    }

    state->m_C = StateMatrix::Zero();
    for (int i = 0; i < state->m_hypotheses.size(); i++)
    {
        IMMHypothesis::Ptr hyp = state->m_hypotheses.at(i);
        StateVector diff = hyp->m_x - state->m_x;
        state->m_C += (hyp->m_C + (diff * diff.transpose())) * hyp->m_probability;
    }
}

    void IMMFilter::updateCurrentHypothesis(IMMState::Ptr state, Track::Ptr track)
    {
        ROS_DEBUG_NAMED("IMM", "Updating current hypothesis");

        double bestHypothesis = 0.0;
        unsigned int bestHypothesisIdx = 0;
        for (unsigned int i = 0; i < state->m_hypotheses.size(); i++)
        {
            if (state->getHypothesis(i)->m_probability > bestHypothesis)
            {
                bestHypothesis = state->getHypothesis(i)->m_probability;
                bestHypothesisIdx = i;
            }
        }

        if (state->m_currentHypothesis != state->getHypothesis(bestHypothesisIdx))
            ROS_ERROR_STREAM_NAMED("IMM", "Best Hypothesis switched to model "<< bestHypothesisIdx);

        state->m_currentHypothesis = state->getHypothesis(bestHypothesisIdx);
        state->m_currentHypothesisIdx =bestHypothesisIdx;
        track->model_idx = bestHypothesisIdx;
        //track->state = state->m_currentHypothesis;
    }

    void IMMFilter::modeProbabilityUpdate(IMMState::Ptr state)
    {
        ROS_DEBUG_NAMED("IMM", "Updating mode probability");

        double normalizer = 0.0;
        for (int i = 0; i < state->m_hypotheses.size(); i++)
        {
            IMMHypothesis::Ptr hyp = state->m_hypotheses.at(i);
            normalizer += hyp->m_likelihood * hyp->m_cNormalizer;
        }
        for (int i = 0; i < state->m_hypotheses.size(); i++)
        {
            IMMHypothesis::Ptr hyp = state->m_hypotheses.at(i);
            hyp->m_probability = hyp->m_likelihood*hyp->m_cNormalizer / normalizer;
            ROS_DEBUG_STREAM_NAMED("IMM", "Mode Probability of hypothesis " << i << "is " << hyp->m_probability);
        }
    }

    void IMMFilter::visualizeFilterProperties(const ros::Time& time, const string& trackerFrame, const Tracks tracks)
    {
        // check if visualization is enabled
        if (m_immVisualization)
        {
            // Create marker lists for vertical and horizontal probability bars
            visualization_msgs::Marker marker_top, marker_side;
            marker_top.header.stamp = marker_side.header.stamp =time;
            marker_top.header.frame_id = marker_side.header.frame_id =trackerFrame;
            marker_top.id = marker_side.id =0;
            marker_top.scale.x = marker_side.scale.x =0.1;
            marker_top.pose.orientation.w = marker_side.pose.orientation.w =1.0;
            marker_top.action = marker_side.action =visualization_msgs::Marker::ADD;
            marker_top.type = marker_side.type =visualization_msgs::Marker::LINE_LIST;
            marker_top.lifetime = marker_side.lifetime =ros::Duration();
            marker_top.ns = "imm_probabilites_vertical";
            marker_side.ns = "imm_probabilites_horizontal";

            // SRL People Tracker colors
            const size_t NUM_COLOR = 10, NUM_BW = 4;
            const unsigned int rainbow_colors[NUM_COLOR + NUM_BW] = {
                                                                     0xaf1f90, 0x000846, 0x00468a, 0x00953d, 0xb2c908, 0xfcd22a, 0xffa800, 0xff4500, 0xe0000b, 0xb22222,
                                                                     0xffffff, 0xb8b8b8, 0x555555, 0x000000
            };
            unsigned int id =0;
            foreach(Track::Ptr track, tracks){
                IMMState::Ptr imm = boost::dynamic_pointer_cast<IMMState>(track->state);

                const double spacer = 0.4;
                double offset = (imm->getHypotheses().size()-1)/2.0 * (-1*spacer);
                std_msgs::ColorRGBA color ;

                double x = track->state->x()(STATE_X_IDX);
                double y = track->state->x()(STATE_Y_IDX);

                double vx = track->state->x()(STATE_VX_IDX);
                double vy = track->state->x()(STATE_VY_IDX);
                double yaw =atan2(vy, vx);

                // get bar values for each hypothesis and add them to line list
                for (int i =0; i < imm->getHypotheses().size(); i++)
                {
                    //unsigned int rgb= spencer_colors[i*NUM_SRL_COLOR_SHADES];
                    unsigned int rgb= rainbow_colors[i];

                    // get color
                    color.r = ((rgb >> 16) & 0xff) / 255.0f;
                    color.g = ((rgb >> 8)  & 0xff) / 255.0f;
                    color.b = ((rgb >> 0)  & 0xff) / 255.0f;
                    color.a = 0.8;
                    marker_top.colors.push_back(color);
                    marker_top.colors.push_back(color);
                    marker_side.colors.push_back(color);
                    marker_side.colors.push_back(color);

                    // set position and orientation for horizontal and vertical bar
                    geometry_msgs::Point p;
                    p.x = x + cos(yaw+ M_PI_2)* offset;
                    p.y = y + sin(yaw+ M_PI_2)* offset;
                    p.z = 1.8;
                    marker_top.points.push_back(p);
                    marker_side.points.push_back(p);

                    p.z += imm->getHypothesis(i)->getProbability();
                    marker_top.points.push_back(p);

                    p.x += cos(yaw) * imm->getHypothesis(i)->getProbability();
                    p.y += sin(yaw) * imm->getHypothesis(i)->getProbability();
                    p.z = 1.8;
                    marker_side.points.push_back(p);

                    offset += spacer;

                }

                // for each track visualize most probable model as colored dot
                visualization_msgs::Marker activModelMarker;
                activModelMarker.header.stamp = time; ;
                activModelMarker.header.frame_id = trackerFrame;
                activModelMarker.id =   time.toNSec() + id++;
                activModelMarker.action =       visualization_msgs::Marker::ADD;
                activModelMarker.type =     visualization_msgs::Marker::SPHERE;
                activModelMarker.lifetime =     ros::Duration().fromSec(20);
                activModelMarker.ns =       "imm_active_model_path";
                activModelMarker.scale.x = activModelMarker.scale.y = activModelMarker.scale.z = 0.1;
                unsigned int rgb= rainbow_colors[imm->getCurrentHypothesisIndex()];

                color.r = ((rgb >> 16) & 0xff) / 255.0f;
                color.g = ((rgb >> 8)  & 0xff) / 255.0f;
                color.b = ((rgb >> 0)  & 0xff) / 255.0f;
                color.a = 0.8;

                activModelMarker.color = color;
                activModelMarker.pose.position.x = x;
                activModelMarker.pose.position.y = y;
                m_immVisualizationPublisher.publish(activModelMarker);
            }
            m_immVisualizationPublisher.publish(marker_top);
            m_immVisualizationPublisher.publish(marker_side);
        }
    }

    void IMMFilter::getDebugInformation(IMMState::Ptr state, Pairing::ConstPtr pairingTracker)
    {
        //calculate innovation from Mixed state estimates
        if (m_sendDebugInformation)
        {
            ROS_WARN_STREAM("packing new debug message");
            spencer_tracking_msgs::ImmDebugInfo debug;
            debug.track_id = pairingTracker->track->id;
            debug.CXX = state->m_C(STATE_X_IDX, STATE_X_IDX);
            debug.CYY = state->m_C(STATE_Y_IDX, STATE_Y_IDX);
            debug.CpXX = state->m_Cp(STATE_X_IDX, STATE_X_IDX);
            debug.CpYY = state->m_Cp(STATE_Y_IDX, STATE_Y_IDX);
            ObsVector innovation = state->m_xp.head(OBS_DIM) - pairingTracker->observation->z;
            debug.innovation = innovation.norm();
            for (int i = 0; i < state->m_hypotheses.size(); i++)
            {
                debug.modeProbabilities.push_back(state->m_hypotheses.at(i)->getProbability());
            }
            m_debugMessages.infos.push_back(debug);
        }
    }

    void IMMFilter::sendDebugInformation()
    {
        //calculate innovation from Mixed state estimates
        if (m_sendDebugInformation && m_debugMessages.infos.size() > 0){
            m_debugMessages.header.stamp = ros::Time::now();
            m_debugMessages.header.seq++;
            ROS_WARN_STREAM("Publishing " << m_debugMessages.infos.size() << " info messages");
            m_immDebugPublisher.publish(m_debugMessages);
            m_debugMessages.infos.clear();
        }

    }

}


/*
 * imageData.hh
 *
 *  Created on: 23.06.2014
 *      Author: Andreas Schuster
 */

#ifndef IMAGEDATA_HH_
#define IMAGEDATA_HH_

#define COSINE_WEIGHTED 1
// system
#include <boost/foreach.hpp>
#include <boost/algorithm/string.hpp>
#include <vector>
#include <cstdlib>
#include <cmath>
#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>
#include <limits>
#include <ctime>
#include <iostream>
#include <sstream>
#include <fstream>

#include <ImathBox.h>
#include <half.h>
#include <ImathVec.h>
#include <ImathMatrix.h>

#include <ImfInputFile.h>
#include <ImfArray.h>
#include <ImfPixelType.h>

namespace blech {

class ImageData
{

public:
    ImageData();
    virtual ~ImageData();

    bool loadImage( const std::string &i_path );
    bool saveImage( const std::string &i_path );
    void verifyZValues();
    bool processSSAO( double i_radius );
    bool processFog( double i_start, double i_end, double i_density );

    double getAverageDistance();
    void getDepthDistances( double &io_min, double &io_averg, double &io_max );
    void setQuiet( bool i_v ) { m_quiet = i_v; }
    bool getQuiet() { return m_quiet; }

private:
    void initDirections();
    void initSteps();
    bool fileExists( const std::string &i_path );
    bool readHeader(const std::string &i_path);
    half occlusion( const int i_x, const int i_y, float i_radius);
    // half occlusionRay( const int i_x, const int i_y, float i_radius );
    Imath::V3f getEyePos(float i_u, float i_v);

    Imath::V3f min_diff(Imath::V3f P, Imath::V3f P1, Imath::V3f P2);
    float tangent(Imath::V3f T);
    float tangent(Imath::V3f P, Imath::V3f S);
    Imath::V3f tangent_vector(Imath::V2f deltaUV, Imath::V3f dPdu, Imath::V3f dPdv);
    Imath::V2f snap_uv_offset(Imath::V2f uv);
    Imath::V2f snap_uv_coord(Imath::V2f uv);
    float rsqrt(float x);
    float invlength(Imath::V2f v);
    float tan_to_sin(float x);
    half sampleZChannel( int i_x, int i_y);
    Imath::V3f fetchEyePos( const Imath::V3f &i_eyePos );

public:

private:
    bool m_quiet;

    // openExr
    Imf::InputFile* m_imageHandleExr;
    Imf::Array2D<half> m_rPixels;
    Imf::Array2D<half> m_gPixels;
    Imf::Array2D<half> m_bPixels;
    Imf::Array2D<half> m_aPixels;
    Imf::Array2D<half> m_zPixels;
    Imf::Array2D<float> m_occlPixels;
    Imath::Box2i m_dw;

    bool m_hasCameraMatrix;
    bool m_hasScreenMatrix;

    int m_width;
    int m_height;

    float m_samples;

    Imath::M33f m_cameraMatrix;
    Imath::M33f m_cameraToWorldMatrix;
    Imath::M33f m_screenMatrix;
    Imath::M33f m_screenToWorldMatrix;

    Imath::M44f m_WorldToScreen;

    const short m_numDirs;
    const short m_numRays;
    const short m_numSteps;

    Imath::V2f *m_dirs;
    double *m_steps;

    // channelTypes
    Imf::PixelType m_rType;
    Imf::PixelType m_gType;
    Imf::PixelType m_bType;
    Imf::PixelType m_aType;
    Imf::PixelType m_zType;

};

}

#endif /* IMAGEDATA_HH_ */