#pragma once

#include <cstdint>
#include <fstream>
#include <limits>
#include <map>
#include <string>
#include <vector>

#include "misc.hpp"

template <typename T> class Interval
{
public:
  inline T length() const { return death - birth; }
  T birth;
  T death;
};

template <typename T> inline bool operator==(Interval<T> x, Interval<T> y) { return x.birth == y.birth && x.death == y.death; }
template <typename T> inline bool operator!=(Interval<T> x, Interval<T> y) { return !(x == y); }
template <typename T> inline bool operator<(Interval<T> x, Interval<T> y)
{
  return (x.length() < y.length()) ||
    ((x.length() == y.length()) && (x.birth < y.birth)) ||
    (((x.length() == y.length()) && (x.birth == y.birth) && (x.death < y.death)));
}

template <typename T> class PD
{
public:
  PD() : intervals()
  {
  };

  void add(T b, T d, unsigned int m)
  {
    if (m > 0 && b < d)
    {
      Interval<T> interval;
      interval.birth = b;
      interval.death = d;
      intervals.push_back(std::make_pair(interval, m));
    }
  }

  void clear() { intervals.clear(); }

  void finitize(T scale)
  {
    for (auto it = intervals.begin(); it != intervals.end(); ++it)
    {
      if (std::isinf(it->first.death))
        it->first.death = scale;
    }
  }

  void discretize(T delta)
  {
    if (delta <= 0)
      return;
    for (auto it = begin(); it != end(); ++it)
    {
      it->first.birth = std::round(it->first.birth/delta)*delta;
      it->first.death = std::round(it->first.death/delta)*delta;
    }
  } 

  void compress_and_sort()
  {
    std::map<Interval<T>, unsigned int> tmp;
    for (auto it = cbegin(); it != cend(); ++it)
    {
      auto existing = tmp.find(it->first);
      if (existing == tmp.end())
      {
        tmp.insert(*it);
      }
      else
      {
        existing->second += it->second;
      }
    }
    intervals = std::vector<std::pair<Interval<T>, unsigned int> >(tmp.cbegin(), tmp.cend());
  }

  unsigned int size() const
  {
    unsigned int ret = 0;
    for (auto it = cbegin(); it != cend(); ++it)
      ret += it->second;
    return ret;
  }

  inline typename std::vector<std::pair<Interval<T>, unsigned int> >::size_type size_2() const { return intervals.size(); }

  using Iterator = typename std::vector<std::pair<Interval<T>, unsigned int> >::iterator;
  using Const_iterator = typename std::vector<std::pair<Interval<T>, unsigned int> >::const_iterator;

  inline Const_iterator cbegin() const { return intervals.cbegin(); }
  inline Const_iterator cend() const { return intervals.cend(); }
  inline Iterator begin() { return intervals.begin(); }
  inline Iterator end() { return intervals.end(); }
  
private:
  std::vector<std::pair<Interval<T>, unsigned int> > intervals;
};

template <typename T> inline T sqdist(Interval<T> x, Interval<T> y)
{
  return (x.birth - y.birth)*(x.birth - y.birth) + (x.death - y.death)*(x.death - y.death);
}

template <typename T> T heat_kernel(T sigma, const PD<T> & a, const PD<T> & b)
{
  T ret = 0;

  for (auto it = a.cbegin(); it != a.cend(); ++it)
  {
    auto x = *it;
    Interval<T> xbar;
    xbar.birth = x.first.death;
    xbar.death = x.first.birth;

    for (auto jt = b.cbegin(); jt != b.cend(); ++jt)
    {
      auto y = *jt;
      ret += x.second*y.second*(std::exp(-sqdist(x.first, y.first)/(8*sigma)) - std::exp(-sqdist(xbar, y.first)/(8*sigma)));
    }
  }

  ret /= 8*sigma*PI;
  
  return ret;
}
  
int read_dipha_degree(const std::string & fname, unsigned int degree, PD<double> & pd)
{
  std::ifstream f;
  f.open(fname, std::ios::in | std::ios::binary);
  if (!f.is_open())
  {
    std::cerr << "Failed to open " << fname << ". Bailing." << std::endl;
    return 1;
  }

  if (read_le<int64_t>(f) != DIPHA_MAGIC)
  {
    std::cerr << "File " << fname << " is not DIPHA file. Bailing." << std::endl;
    f.close();
    return 1;
  }

  if (read_le<int64_t>(f) != DIPHA_PERSISTENCE_DIAGRAM)
  {
    std::cerr << "File " << fname << " is not persistence diagram. Bailing." << std::endl;
    f.close();
    return 1;
  }
  
  int64_t n = read_le<int64_t>(f);

  for (int64_t i = 0; i < n; ++i)
  {
    int64_t d = read_le<int64_t>(f);
    double birth = read_le<double>(f);
    double death = read_le<double>(f);
    if (d == degree)
      pd.add(birth, death, 1);
    else if (-d - 1 == degree)
      pd.add(birth, std::numeric_limits<double>::infinity(), 1);
  }
  
  f.close();
  return 0;
}