#pragma once


/*
WARNING: DO NOT USE FOR WEIGHTED STUFF YET! ALL WEIGHTS MUST BE 1!
*/

#include <vector>
#include <algorithm>
#include <cassert>
#include <fstream>
#include <string>

#include <gudhi/Simplex_tree.h>

#include "sparse_vector.hpp"
#include "misc.hpp"

namespace Gudhi_laplacian
{
  typedef Gudhi::Simplex_tree<> ST;
  typedef std::vector<std::vector<ST::Simplex_handle> > Dimensional_stratification;

  // This function is not const in the simplex tree because a bunch
  // of boundary/coboundary query functions are not const.
  // UB if x is not a face of y.
  int face_sign(ST & st, ST::Simplex_handle x, ST::Simplex_handle y)
  {
    auto x_vers = st.simplex_vertex_range(x);
    auto y_vers = st.simplex_vertex_range(y);
    
    auto it = x_vers.begin();
    auto jt = y_vers.begin();
    
    int sign = (st.dimension(y) % 2 == 0) ? 1 : -1; // The vertex ranges run backwards.
    while (it != x_vers.end())
    {
      if (*it != *jt)
        return sign;
      
      ++it;
      ++jt;
      sign *= -1;
    }
    return sign;
  }
  
  Dimensional_stratification make_contiguous_keys(ST & st)
  {
    Dimensional_stratification stratification;
    std::vector<ST::Simplex_key> keys;
       
    for (auto s : st.complex_simplex_range())
    {
      int dim = st.dimension(s);
      
      while (stratification.size() <= dim)
        stratification.push_back(std::vector<ST::Simplex_handle>());
      while (keys.size() <= dim)
        keys.push_back(0);

      stratification[dim].push_back(s);
      st.assign_key(s, keys[dim]++);
    }
    stratification.push_back(std::vector<ST::Simplex_handle>());

    return stratification;
  }

  Sparse_vector<int> assemble_cbd_row(ST & st, ST::Simplex_handle s)
  {
    Sparse_vector<int> ret;

    // This is a specialization of face_sign.
    int sign = (st.dimension(s) % 2 == 0) ? 1 : -1;
    for (auto bd : st.boundary_simplex_range(s))
    {
      ret.add(ST::key(bd), sign);
      sign *= -1;
    }

    // Note: We don't want to compress the row.
    return ret;
  }
  
  int write_laplacians_mm(ST & st, const Dimensional_stratification & stratification, const std::vector<std::vector<double> > & weights, int from_degree, int to_degree, const std::string & prefix)
  {
    if (preexisting_files(prefix, std::string("petsc")))
      return 1;

    if (stratification.size() <= to_degree)
      return 2;

    std::vector<Sparse_vector<int> > cbd_same;
    std::vector<Sparse_vector<int> > cbd_same_t;
    std::vector<Sparse_vector<int> > cbd_below;
    std::vector<Sparse_vector<int> > cbd_below_t;
    
    for (int d = from_degree; d < to_degree; ++d)
    {
      unsigned int m = d+1 < stratification.size() ? stratification[d+1].size() : 0;
      unsigned int n = stratification[d].size();
      unsigned int o = d > 0 ? stratification[d-1].size() : 0;

      if (d == 0)
      {
        cbd_below = std::vector<Sparse_vector<int> >(n); // An nx0 matrix to avoid special-casing further down.
        cbd_below_t = Sparse_vector<int>::transpose(cbd_below, 0);
      }
      else
      {
        cbd_below = cbd_same;
        cbd_below_t = cbd_same_t;
      }
      assert(cbd_below.size() == n);
      assert(cbd_below_t.size() == o);

      cbd_same.clear();
      for (unsigned int i = 0; i < m; ++i)
      {
        cbd_same.push_back(assemble_cbd_row(st, stratification[d+1][i]));
      }
      cbd_same_t = Sparse_vector<int>::transpose(cbd_same, n);
      assert(cbd_same.size() == m);
      assert(cbd_same_t.size() == n);

      std::vector<Sparse_vector<double> > lap;
      unsigned int nnz = 0;
      
      for (unsigned int i = 0; i < n; ++i)
      {
        lap.push_back(Sparse_vector<double>());

        // Up part.
        for (auto it = cbd_same_t[i].cbegin(); it != cbd_same_t[i].cend(); ++it)
        {
          for (auto jt = cbd_same[it->first].cbegin(); jt != cbd_same[it->first].cend(); ++jt)
          {
            lap[i].add(jt->first, (weights[d+1][it->first]/weights[d][jt->first]) * (it->second * jt->second));
          }
        }

        // Down part.
        for (auto it = cbd_below[i].cbegin(); it != cbd_below[i].cend(); ++it)
        {
          for (auto jt = cbd_below_t[it->first].cbegin(); jt != cbd_below_t[it->first].cend(); ++jt)
          {
            lap[i].add(jt->first, (weights[d][jt->first]/weights[d-1][it->first]) * (it->second * jt->second));
          }
        }

        lap[i].compress();
        nnz += lap[i].nnz();
      }

      std::ofstream file;
      file.open(prefix + std::string(".") + std::to_string(d) + std::string(".petsc"), std::ios::out | std::ios::binary);
      if (!file.is_open())
        return 3;
      
      write_be<int32_t>(file, PETSC_MAT_FILE_CLASSID);
      write_be<int32_t>(file, n);
      write_be<int32_t>(file, n);
      write_be<int32_t>(file, nnz);
      for (unsigned int i = 0; i < n; ++i)
        write_be<int32_t>(file, lap[i].nnz());
      for (unsigned int i = 0; i < n; ++i)
      {
        for (auto it = lap[i].cbegin(); it != lap[i].cend(); ++it)
        {
          write_be<int32_t>(file, it->first);
        }
      }
      for (unsigned int i = 0; i < n; ++i)
      {
        for (auto it = lap[i].cbegin(); it != lap[i].cend(); ++it)
        {
          write_be<double>(file, it->second);
        }
      }
      
      file.close();
    }
    return 0;
  }

  Sparse_vector<double> assemble_laplacian_row(ST & st, const std::vector<std::vector<double> > & weights, ST::Simplex_handle s)
  {
    Sparse_vector<double> row;

    int dim = st.dimension(s);
    ST::Simplex_key s_key = ST::key(s);
    double s_weight = weights[dim][s_key];

    // Up part.
    ST::Cofaces_simplex_range cbds = st.cofaces_simplex_range(s, 1);
    for (auto cbd : cbds)
    {
      int s_cbd_sign = face_sign(st, s, cbd);
      ST::Simplex_key cbd_key = ST::key(cbd);
      double cbd_weight = weights[dim+1][cbd_key];
      ST::Boundary_simplex_range bds = st.boundary_simplex_range(cbd);
      int bd_cbd_sign = ((dim+1) % 2 == 0) ? 1 : -1;
      for (auto bd : bds)
      {
        row.add(ST::key(bd), s_cbd_sign * bd_cbd_sign * (cbd_weight/s_weight));
        bd_cbd_sign *= -1;
      }
    }

    // Down part.
    ST::Boundary_simplex_range bds = st.boundary_simplex_range(s);
    int bd_s_sign = (dim % 2 == 0) ? 1 : -1;
    for (auto bd : bds)
    {
      ST::Simplex_key bd_key = ST::key(bd);
      double bd_weight = weights[dim-1][bd_key];
      ST::Cofaces_simplex_range cbds = st.cofaces_simplex_range(bd, 1);
      for (auto cbd : cbds)
      {
        int bd_cbd_sign = face_sign(st, bd, cbd);
        ST::Simplex_key cbd_key = ST::key(cbd);
        double cbd_weight = weights[dim][cbd_key];
        row.add(cbd_key, bd_s_sign * bd_cbd_sign * (cbd_weight/bd_weight));
      }
      bd_s_sign *= -1;
    }

    row.compress();
    return row;
  }

  int write_cbdrys(ST & st, const Dimensional_stratification & stratification, int from_degree, int to_degree, const std::string & prefix)
  {
    if (preexisting_files(prefix, std::string("petsc")))
      return 1;

    if (stratification.size() <= to_degree)
      return 2;

    for (int d = from_degree; d < to_degree; ++d)
    {
      // Coboundary matrix in degree d.
      // Shape: (dim C_{d+1}) x (dim C_{d})
      int m = stratification[d+1].size();
      int n = stratification[d].size();
      
      std::ofstream file;
      file.open(prefix + std::string(".") + std::to_string(d) + std::string(".petsc"), std::ios::out | std::ios::binary);
      if (!file.is_open())
        return 3;

      write_be<int32_t>(file, PETSC_MAT_FILE_CLASSID);
      write_be<int32_t>(file, m);
      write_be<int32_t>(file, n);
      write_be<int32_t>(file, (d+2)*m);

      for (int i = 0; i < m; ++i)
        write_be<int32_t>(file, d+2);

      // For indices.
      for (auto it = stratification[d+1].cbegin(); it != stratification[d+1].cend(); ++it)
      {
        Sparse_vector<int> row = assemble_cbd_row(st, *it);
        assert(row.size() == d+2);
        for (auto jt = row.cbegin(); jt != row.cend(); ++jt)
        {
          write_be<int32_t>(file, jt->first);
        }
      }

      // For values.
      for (auto it = stratification[d+1].cbegin(); it != stratification[d+1].cend(); ++it)
      {
        Sparse_vector<int> row = assemble_cbd_row(st, *it);
        assert(row.size() == d+2);
        for (auto jt = row.cbegin(); jt != row.cend(); ++jt)
        {
          write_be<double>(file, jt->second);
        }
      }
      
      file.close();
    } 
    
    return 0;
  }

  int write_laplacians(ST & st, const Dimensional_stratification & stratification, const std::vector<std::vector<double> > & weights, int from_degree, int to_degree, const std::string & prefix)
  {
    if (preexisting_files(prefix, std::string("petsc")))
      return 1;

    if (stratification.size() <= to_degree)
      return 2;

    for (int d = from_degree; d < to_degree; ++d)
    {
      // Laplacian matrix in degree d.
      // Shape: (dim C_{d}) x (dim C_{d})
      int m = stratification[d].size();
      
      std::ofstream file;
      file.open(prefix + std::string(".") + std::to_string(d) + std::string(".petsc"), std::ios::out | std::ios::binary);
      if (!file.is_open())
        return 3;

      write_be<int32_t>(file, PETSC_MAT_FILE_CLASSID);
      write_be<int32_t>(file, m);
      write_be<int32_t>(file, m);

      int num_nonzero = 0;
      std::vector<int> num_nonzero_row;
      num_nonzero_row.reserve(m);

      // First traversal (counting).
      for (auto it = stratification[d].cbegin(); it != stratification[d].cend(); ++it)
      {
        Sparse_vector<double> row = assemble_laplacian_row(st, weights, *it);
        num_nonzero += row.size();
        num_nonzero_row.push_back(row.size());
      }
      
      write_be<int32_t>(file, num_nonzero);
      for (auto it = num_nonzero_row.cbegin(); it != num_nonzero_row.cend(); ++it)
        write_be<int32_t>(file, *it);

      // Second traversal (indices).
      for (auto it = stratification[d].cbegin(); it != stratification[d].cend(); ++it)
      {
        Sparse_vector<double> row = assemble_laplacian_row(st, weights, *it);
        for (auto jt = row.cbegin(); jt != row.cend(); ++jt)
        {
          write_be<int32_t>(file, jt->first);
        }
      }

      // Third traversal (values).
      for (auto it = stratification[d].cbegin(); it != stratification[d].cend(); ++it)
      {
        Sparse_vector<double> row = assemble_laplacian_row(st, weights, *it);
        for (auto jt = row.cbegin(); jt != row.cend(); ++jt)
        {
          write_be<double>(file, jt->second);
        }
      }
      
      file.close();
    } 
    
    return 0;
  }

  int write_weights(const std::vector<std::vector<double> > & weights, int from_degree, int to_degree, const std::string & prefix)
  {
    if (preexisting_files(prefix, std::string("petsc")))
      return 1;

    if (weights.size() <= to_degree)
      return 2;

    for (int d = from_degree; d < to_degree; ++d)
    {
      // Weights in degree d.
      int m = weights[d].size();
      
      std::ofstream file;
      file.open(prefix + std::string(".") + std::to_string(d) + std::string(".petsc"), std::ios::out | std::ios::binary);
      if (!file.is_open())
        return 3;

      write_be<int32_t>(file, PETSC_VEC_FILE_CLASSID);
      write_be<int32_t>(file, m);

      for (auto it = weights[d].cbegin(); it != weights[d].cend(); ++it)
        write_be<double>(file, *it);
      
      file.close();
    } 
    
    return 0;
  }

  int write_complex(ST & st, const Dimensional_stratification & stratification, int from_degree, int to_degree, const std::string & prefix)
  {
    if (preexisting_files(prefix, std::string("txt")))
      return 1;

    if (stratification.size() <= to_degree)
      return 2;

    for (int d = from_degree; d < to_degree; ++d)
    {
      std::ofstream file;
      file.open(prefix + std::string(".") + std::to_string(d) + std::string(".txt"), std::ios::out);

      if (!file.is_open())
        return 3;

      for (auto it = stratification[d].begin(); it != stratification[d].end(); ++it)
      {
        std::vector<ST::Vertex_handle> tmp(std::begin(st.simplex_vertex_range(*it)), std::end(st.simplex_vertex_range(*it)));
        std::reverse(std::begin(tmp), std::end(tmp));
        for (auto jt = std::begin(tmp); jt != std::end(tmp) - 1; ++jt)
        {
          file << std::to_string(*jt) << " ";
        }
        file << tmp.back() << std::endl;
      }
      file.close();
    }

    return 0;
  }

}