#pragma once
#include <string.h>
#include <unordered_map>
#ifndef _WIN32
#include <sys/mman.h>
#endif

// It appears that on OSX MAP_ANONYMOUS is mapped to MAP_ANON
// https://github.com/leftmike/foment/issues/4
#ifdef __APPLE__
#define MAP_ANONYMOUS MAP_ANON
#endif

#include "array_parameters_dense.h"

class sparse_parameters;
typedef std::unordered_map<uint64_t, weight*> weight_map;

template <typename T>
class sparse_iterator
{
private:
	weight_map::iterator _iter;
	uint32_t _stride;

public:
	typedef std::forward_iterator_tag iterator_category;
	typedef T value_type;
	typedef ptrdiff_t difference_type;
	typedef  T* pointer;
	typedef  T& reference;

	sparse_iterator(weight_map::iterator& iter, uint32_t stride)
		: _iter(iter), _stride(stride)
	{ }

	sparse_iterator& operator=(const sparse_iterator& other)
	{
		_iter = other._iter;
		_stride = other._stride;
		return *this;

	}
	uint64_t index() { return _iter->first; }

	T& operator*() { return *(_iter->second); }

	sparse_iterator& operator++()
	{
		_iter++;
		return *this;
	}

	bool operator==(const sparse_iterator& rhs) const { return _iter == rhs._iter; }
	bool operator!=(const sparse_iterator& rhs) const { return _iter != rhs._iter; }
};


class sparse_parameters
{
private:
	weight_map _map;
	uint64_t _weight_mask;  // (stride*(1 << num_bits) -1)
	uint32_t _stride_shift;
	bool _seeded; // whether the instance is sharing model state with others
	bool _delete;
	void* default_data;
  float* default_value;
public:
	typedef sparse_iterator<weight> iterator;
	typedef sparse_iterator<const weight> const_iterator;
 private:
	void(*fun)(const weight*, void*);
 public:

	sparse_parameters(size_t length, uint32_t stride_shift = 0)
		: _map(),
		_weight_mask((length << stride_shift) - 1),
		_stride_shift(stride_shift),
		_seeded(false), _delete(false), default_data(nullptr),
    fun(nullptr)
	{ default_value = calloc_mergable_or_throw<weight>(stride());}

	sparse_parameters()
		: _map(), _weight_mask(0), _stride_shift(0), _seeded(false), _delete(false), default_data(nullptr), fun(nullptr)
	{ default_value = calloc_mergable_or_throw<weight>(stride());}

	bool not_null() { return (_weight_mask > 0 && !_map.empty()); }

	sparse_parameters(const sparse_parameters &other) { shallow_copy(other); }
	sparse_parameters(sparse_parameters &&) = delete;

	weight* first() { throw 1; } //TODO: Throw better exceptions. Allreduce currently not supported in sparse.

	//iterator with stride
	iterator begin() { weight_map::iterator i = _map.begin(); return iterator(i, stride()); }
	iterator end() { weight_map::iterator i = _map.end(); return iterator(i, stride()); }

	//const iterator
	const_iterator cbegin() { weight_map::iterator i = _map.begin(); return const_iterator(i,  stride()); }
	const_iterator cend() { weight_map::iterator i = _map.begin(); return const_iterator(i, stride()); }

	inline weight& operator[](size_t i)
	{   uint64_t index = i & _weight_mask;
    weight_map::iterator iter = _map.find(index);
		if (iter == _map.end())
		  {     _map.insert(std::make_pair(index, calloc_mergable_or_throw<weight>(stride())));
			iter = _map.find(index);
			if (fun != nullptr)
        fun(iter->second, default_data);
		}
		return *(iter->second);
	}

  inline const weight& operator[](size_t i) const
	{   uint64_t index = i & _weight_mask;
		weight_map::const_iterator iter = _map.find(index);
		if (iter == _map.end())
      return *default_value;
		return *(iter->second);
  }

	inline weight& strided_index(size_t index) { return operator[](index << _stride_shift); }

	void shallow_copy(const sparse_parameters& input)
	{
		// TODO: this is level-1 copy (weight* are stilled shared)
		if (!_seeded)
		{
		  for (auto iter = _map.begin(); iter != _map.end(); ++iter)
				free(iter->second);
		}
		_map = input._map;
		_weight_mask = input._weight_mask;
		_stride_shift = input._stride_shift;
    free(default_value);
    default_value = calloc_mergable_or_throw<weight>(stride());
    memcpy(default_value, input.default_value, stride());
    default_data = input.default_data;
		_seeded = true;
	}

	template<class R, class T> void set_default(R& info)
	{
	  R& new_R = calloc_or_throw<R>();
	  new_R = info;
	  default_data = &new_R;
	  fun = (void(*)(const weight*, void*))T::func;
    fun(default_value, default_data);
	}

	template<class T> void set_default() { fun = (void(*)(const weight*, void*))T::func; }

	void set_zero(size_t offset)
	{
		for (weight_map::iterator iter = _map.begin(); iter != _map.end(); ++iter)
			(&(*(iter->second)))[offset] = 0;
	}

	uint64_t mask()	const { return _weight_mask; }

	uint64_t seeded() const { return _seeded; }

	uint32_t stride() const { return 1 << _stride_shift; }

	uint32_t stride_shift()	const { return _stride_shift; }

	void stride_shift(uint32_t stride_shift) {
    _stride_shift = stride_shift;
    free(default_value);
    default_value = calloc_mergable_or_throw<weight>(stride());
    if (fun != nullptr)
      fun(default_value, default_data);
  }

#ifndef _WIN32
	void share(size_t length)
	{throw 1; //TODO: add better exceptions
	}
#endif

	~sparse_parameters()
	{if (!_delete && !_seeded)  // don't free weight vector if it is shared with another instance
		{
    for (auto iter = _map.begin(); iter != _map.end(); ++iter)
      free(iter->second);
    _map.clear();
		 _delete = true;
		}
    if (default_data != nullptr)
      free(default_data);
    free(default_value);
	}
};

class parameters {
 public:
  bool sparse;
  dense_parameters dense_weights;
  sparse_parameters sparse_weights;

  inline weight& operator[](size_t i)
  {
    if (sparse)
      return sparse_weights[i];
    else
      return dense_weights[i];
  }

  inline uint32_t stride_shift()
  {
    if (sparse)
      return sparse_weights.stride_shift();
    else
      return dense_weights.stride_shift();
  }

  inline uint32_t stride()
  {
    if (sparse)
      return sparse_weights.stride();
    else
      return dense_weights.stride();
  }

  inline uint64_t mask()
  {
    if (sparse)
      return sparse_weights.mask();
    else
      return dense_weights.mask();
  }

  inline uint64_t seeded()
  {
    if (sparse)
      return sparse_weights.seeded();
    else
      return dense_weights.seeded();
  }

  inline void shallow_copy(const parameters& input)
  {
    if (sparse)
      sparse_weights.shallow_copy(input.sparse_weights);
    else
      dense_weights.shallow_copy(input.dense_weights);
  }

  inline void set_zero(size_t offset)
  {
    if (sparse)
      sparse_weights.set_zero(offset);
    else
      dense_weights.set_zero(offset);
  }
#ifndef _WIN32
  inline void share(size_t length)
  {
    if (sparse)
      sparse_weights.share(length);
    else
      dense_weights.share(length);
  }
  #endif

  inline void stride_shift(uint32_t stride_shift)
  { if (sparse)
      sparse_weights.stride_shift(stride_shift);
    else
      dense_weights.stride_shift(stride_shift);
  }

  inline weight& strided_index(size_t index)
  {
    if (sparse)
      return sparse_weights.strided_index(index);
    else
      return dense_weights.strided_index(index);
  }

  inline bool not_null()
  {
    if (sparse)
      return sparse_weights.not_null();
    else
      return dense_weights.not_null();
  }
};
