#include <fstream>
#include <iostream>
#include <unordered_set>

#include "glass.hpp"

const unsigned Empty = std::numeric_limits<unsigned>::max();

class BiroliState {
public:
  unsigned maximumNeighbors;
  unsigned occupiedNeighbors;

  BiroliState() {
    maximumNeighbors = Empty;
    occupiedNeighbors = 0;
  }

  bool empty() const { return maximumNeighbors == Empty; }
  bool frustrated() const { return occupiedNeighbors > maximumNeighbors; }
  bool operator==(const BiroliState& s) const {
    return s.maximumNeighbors == maximumNeighbors;
  }
  void remove() {
    maximumNeighbors = Empty;
  }
};

template <int D> class BiroliSystem : public System<D, BiroliState> {
public:
  std::vector<unsigned> N;
  std::unordered_set<Vertex<D, BiroliState>*> occupants;

  unsigned types() const { return N.size() - 1; }

  unsigned occupancy() const {
    return BiroliSystem::size() - N[0];
  }

  double density() const { return (double)occupancy() / BiroliSystem::size(); }

  bool valid(const Vertex<D, BiroliState>& v) const {
    unsigned o = 0;

    for (const Vertex<D, BiroliState>& vn : v.neighbors) {
      if (!vn.empty()) {
        o++;
      }
    }

    return o == v.occupiedNeighbors;
  }

  bool compatible() const {
    for (const Vertex<D, BiroliState>& v : BiroliSystem::vertices) {
      if (v.state.frustrated()) {
        return false;
      }
    }

    return true;
  }

  bool valid() const {
    for (const Vertex<D, BiroliState>& v : BiroliSystem::vertices) {
      if (!valid(v)) {
        return false;
      }
    }

    return true;
  }

  BiroliSystem(unsigned L, unsigned T) : System<D, BiroliState>(L), N(T + 1, 0) {
    N[0] = BiroliSystem::size();
  }

  std::list<std::reference_wrapper<Vertex<D, BiroliState>>> overlaps(Vertex<D, BiroliState>& v) {
    std::list<std::reference_wrapper<Vertex<D, BiroliState>>> o;

    if (v.empty()) { // an empty site cannot be frustrating anyone
      return o;
    }

    bool selfFrustrated = v.frustrated();
    bool anyNeighborFrustrated = false;

    for (Vertex<D, BiroliState>& vn : v.neighbors) {
      if (!vn.empty()) {
        bool thisNeighborFrustrated = vn.frustrated();

        if (thisNeighborFrustrated) {
          anyNeighborFrustrated = true;
        }

        if (selfFrustrated || thisNeighborFrustrated) {
          o.push_back(vn);
        }
      }
    }

    if (selfFrustrated || anyNeighborFrustrated) {
      o.push_back(v);
    }

    return o;
  }

  std::list<std::reference_wrapper<Vertex<D, BiroliState>>> overlaps(const std::list<std::reference_wrapper<Vertex<D, BiroliState>>>& vs) {
    std::list<std::reference_wrapper<Vertex<D, BiroliState>>> o;

    for (Vertex<D, BiroliState>& v : vs) {
      o.splice(o.begin(), overlaps(v));
    }

    return o;
  }

  bool canReplaceWith(const Vertex<D, BiroliState>& v, unsigned t) const {
    if (t == Empty) {
      return true;
    }

    if (t < v.state.occupiedNeighbors) {
      return false;
    }

    if (v.empty()) {
      for (const HalfEdge<D, BiroliState>& e : v.adjacentEdges) {
        const Vertex<D, BiroliState>& vn = e.neighbor;
        if (vn.state.maximumNeighbors < vn.state.occupiedNeighbors + 1) {
          return false;
        }
      }
    }

    return true;
  }

  void setInitial() {
    for (Vertex<D, BiroliState>& v : BiroliSystem::vertices) {
      v.setInitial();
    }
  }

  bool insert(Vertex<D, BiroliState>& v, unsigned t, bool force = false) {
    if (force || (v.empty() && canReplaceWith(v, t))) {
      v.state.maximumNeighbors = t;
      for (HalfEdge<D, BiroliState>& e : v.adjacentEdges) {
        e.neighbor.state.occupiedNeighbors++;
      }
      N[t]++;
      N[0]--;
      if (t > 0) {
        occupants.insert(&v);
      }
      return true;
    } else {
      return false;
    }
  }

  bool remove(Vertex<D, BiroliState>& v) {
    if (v.empty()) {
      return false;
    } else {
      if (v.state.maximumNeighbors > 0) {
        occupants.erase(&v);
      }
      N[v.state.maximumNeighbors]--;
      N[0]++;
      v.state.maximumNeighbors = Empty;
      for (HalfEdge<D, BiroliState>& e : v.adjacentEdges) {
        e.neighbor.state.occupiedNeighbors--;
      }
      return true;
    }
  }

  bool randomLocalExchange(Rng& r) {
    Vertex<D, BiroliState>& v = r.pick(BiroliSystem::vertices);

    return exchange(v, r.pick(v.neighbors));
  }

  void swap(Vertex<D, BiroliState>& v1, Vertex<D, BiroliState>& v2) {
    if (v1.state.maximumNeighbors != v2.state.maximumNeighbors) {
      if (!v1.empty() && !v2.empty()) {
        std::swap(v1.state.maximumNeighbors, v2.state.maximumNeighbors);
      } else if (v1.empty()) {
        unsigned t = v2.state.maximumNeighbors;
        remove(v2);
        insert(v1, t, true);
      } else {
        unsigned t = v1.state.maximumNeighbors;
        remove(v1);
        insert(v2, t, true);
      }
      std::swap(v1.initialPosition, v2.initialPosition);
    }
  }

  bool exchange(Vertex<D, BiroliState>& v1, Vertex<D, BiroliState>& v2) {
    if (canReplaceWith(v1, v2.state.maximumNeighbors) && canReplaceWith(v2, v1.state.maximumNeighbors)) {
      swap(v1, v2);
      return true;
    } else {
      return false;
    }
  }

  std::complex<double> selfIntermediateScattering(unsigned k, unsigned t = 2) const {
    std::complex<double> F = 0;

    using namespace std::complex_literals;

    for (const Vertex<D, BiroliState>& v : BiroliSystem::vertices) {
      if (v.maximumNeighbors == t) {
        for (unsigned d = 0; d < D; d++) {
          F += exp(1i * 2.0 * M_PI * (double)k * (double)(v.position(d) - v.initialPosition(d)) / (double)BiroliSystem::L);
        }
      }
    }

    F /= D * N[t];

    return F;
  }

  double overlap(const BiroliSystem<D>& s, unsigned t = 2) const {
    int o = 0;

    for (const Vertex<D, BiroliState>& v1 : BiroliSystem::vertices) {
      const Vertex<D, BiroliState>& v2 = s.BiroliSystem::vertices[vectorToIndex(BiroliSystem::orientation.inverse().apply(v1.position))];
      if (v1.maximumNeighbors == v2.maximumNeighbors) {
        o++;
      }
    }

    unsigned nn = 0;

    for (unsigned n : N) {
      nn += iPow(n, 2);
    }

    return ((double)o / (double)BiroliSystem::size() - (double)nn / pow((double)BiroliSystem::size(), 2)) / (1 - (double)nn / pow((double)BiroliSystem::size(), 2));
  }

  bool randomExchange(Rng& r) {
    Vertex<D, BiroliState>& v1 = r.pick(BiroliSystem::vertices);
    Vertex<D, BiroliState>& v2 = r.pick(BiroliSystem::vertices);

    return exchange(v1, v2);
  }

  const std::array<double, 3> ratios = {0.1, 0.5, 0.4};

  void stepGrandCanonical(double z, Rng& r) {
    if (1.0 / 11.0 < r.uniform(0.0, 1.0)) {
      unsigned t = r.pick({1, 2, 2, 2, 2, 2, 3, 3, 3, 3});
      double pIns = BiroliSystem::size() * z / (occupancy() + 1);

      if (pIns > r.uniform(0.0, 1.0)) {
        Vertex<D, BiroliState>& v = r.pick(BiroliSystem::vertices);
        insert(v, t);
      }
    } else {
      double pDel = density() / z;

      if (pDel > r.uniform(0.0, 1.0)) {
        Vertex<D, BiroliState>* v = r.pick(occupants);
        remove(*v);
      }
    }
  }

  void sweepGrandCanonical(double z, Rng& r) {
    for (unsigned i = 0; i < BiroliSystem::size(); i++) {
      stepGrandCanonical(z, r);

      randomExchange(r);
    }
  }

  void sweepLocal(Rng& r) {
    for (unsigned i = 0; i < BiroliSystem::size() / 2; i++) {
      randomLocalExchange(r);
    }
  }

  void sweepSwap(Rng& r) {
    for (unsigned i = 0; i < BiroliSystem::size() / 2; i++) {
      randomExchange(r);
    }
  }

  Vertex<D, BiroliState>& apply(const Transformation<D>& R, const Vertex<D, BiroliState>& v) {
    return BiroliSystem::vertices[vectorToIndex(R.apply(v.position))];
  }

  bool eventChain(const Transformation<D>& R, Vertex<D, BiroliState>& v0) {
    unsigned n = 0;

    if (!v0.empty()) {
      Vertex<D, BiroliState>& v0New = apply(R, v0);;
      unsigned t0 = v0.maximumNeighbors;

      std::queue<std::pair<std::reference_wrapper<Vertex<D, BiroliState>>, unsigned>> q;

      q.push({v0New, t0});
      remove(v0);

      while (!q.empty()) {
        auto [vr, t] = q.front();
        Vertex<D, BiroliState>& v = vr;
        q.pop();

        if (!v.empty()) {
          q.push({apply(R, v), v.maximumNeighbors});
          remove(v);
        }

        insert(v, t, true);
        v.marked = true;

        for (Vertex<D, BiroliState>& vn : overlaps({v})) {
          if (!vn.marked) {
            q.push({apply(R, vn), vn.maximumNeighbors});
            remove(vn);
            vn.marked = true;
          } else {
            /* If a neighbor has already been moved but is still
             * frustrated, it may be due to one of its neighbors!
             */
            if (&vn != &v) {
              for (Vertex<D, BiroliState>& vnn : overlaps(vn)) {
                if (!vnn.marked) {
                  q.push({apply(R, vnn), vnn.maximumNeighbors});
                  remove(vnn);
                  vnn.marked = true;
                }
              }
            }
          }
        }
      }
    }

    for (Vertex<D, BiroliState>& v : BiroliSystem::vertices) {
      v.marked = false;
    }

    return true;
  }

  bool randomEventChain(Rng& r) {
    Transformation<D> R(BiroliSystem::L);
    R.v[r.uniform((unsigned)0, (unsigned)D-1)] = r.pick({-1, 1});
    return eventChain(R, r.pick(BiroliSystem::vertices));
  }

  unsigned flipCluster(const Transformation<D>& R, Vertex<D, BiroliState>& v0) {
    unsigned n = 0;

    Vertex<D, BiroliState>& v0New = BiroliSystem::vertices[vectorToIndex(R.apply(v0.position))];

    if (&v0New != &v0) {
      std::queue<std::reference_wrapper<Vertex<D, BiroliState>>> q;

      v0.marked = true;
      v0New.marked = true;

      swap(v0, v0New);

      for (Vertex<D, BiroliState>& vn : overlaps({v0, v0New})) {
        if (!vn.marked) {
          q.push(vn);
        }
      }

      while (!q.empty()) {
        Vertex<D, BiroliState>& v = q.front();
        q.pop();

        if (!v.marked && !overlaps(v).empty()) {
          v.marked = true;
          Vertex<D, BiroliState>& vNew = BiroliSystem::vertices[vectorToIndex(R.apply(v.position))];

          if (&vNew != &v) {
            vNew.marked = true;

            swap(v, vNew);

            for (Vertex<D, BiroliState>& vn : overlaps({v, vNew})) {
              if (!vn.marked) {
                q.push(vn);
              } else {
                /* If a neighbor has already been moved but is still
                 * frustrated, it may be due to one of its neighbors!
                 */
                if (&vn != &v && &vn != &vNew) {
                  for (Vertex<D, BiroliState>& vnn : overlaps(vn)) {
                    if (!vnn.marked) {
                      q.push(vnn);
                    }
                  }
                }
              }
            }
            n += 2;
          } else {
            n += 1;
          }
        }
      }
    } else {
      n = 1;
    }

    if (n > BiroliSystem::size() / 2) {
      BiroliSystem::orientation = R.apply(BiroliSystem::orientation);
    }

    for (Vertex<D, BiroliState>& v : BiroliSystem::vertices) {
      v.marked = false;
    }

    return n;
  }
};

void print(const BiroliSystem<2>& s) {
  for (const Vertex<2, BiroliState>& v : s.vertices) {
    if (!v.empty()) {
      std::cerr << v.state.maximumNeighbors;
    } else {
      std::cerr << " ";
    }

    if (v.position(0) == s.L - 1) {
      std::cerr << std::endl;
    }
  }
}

int main() {
  const unsigned D = 3;
  unsigned L = 15;
  unsigned Nmin = 2e2;
  unsigned Nmax = 2e5;
  double μmin = -3;
  double μmax = 10;
  double dμ = 0.05;

  BiroliSystem<D> s(L, 3);

  Rng r;

  for (double μ = μmin; μ <= μmax; μ += dμ) {
    double z = exp(μ);
    for (unsigned i = 0; i < 1e4; i++) {
      for (unsigned j = 0; j < s.BiroliSystem::size(); j++) {
        s.stepGrandCanonical(z, r);
//        s.randomEventChain(r);
        s.randomExchange(r);
      }

      if (!s.compatible()) {
        std::cerr << "Not compatible!" << std::endl;
        return 1;
      }

    }

    std::cout << μ << " " << s.density() << std::endl;
  }


  /*
  std::cerr << "Found state with appropriate density." << std::endl;

  System<D> s0 = s;

  std::vector<Matrix<D>> ms = generateTorusMatrices<D>();

  std::vector<unsigned> clusterDist(s.size() + 1);

  s.setInitial();

  unsigned n = 1;
  unsigned i = 0;
  double nC = 0;
  auto start = std::chrono::high_resolution_clock::now();
  while (nC / s.size() < 1e5) {
    if (n < 20 * log(i + 1)) {
      auto stop = std::chrono::high_resolution_clock::now();
      auto duration = duration_cast<std::chrono::microseconds>(stop - start);
      n++;
      std::cout << duration.count() << " " << s.overlap(s0) << std::endl;
    }
//    unsigned nn = s.flipCluster(Transformation<D>(L, ms, r), r.pick(s.vertices));
//    nC += nn;
//    clusterDist[nn]++;
    //s.sweepLocal(r);
    nC += s.size();
    s.sweepSwap(r);
    i++;
  }

  if (!s.compatible()) {
    std::cerr << "Not compatible!" << std::endl;
    return 1;
  }

  if (!s.valid()) {
    std::cerr << "Not valid!" << std::endl;
    return 1;
  }

  std::ofstream file("dist.dat");
  for (unsigned i : clusterDist) {
    file << i << " ";
  }
  file.close();
  */

  return 0;
}