#include <assert.h>
#include <stdlib.h>
#include <string.h>

#include <array>
#include <bitset>
#include <functional>
#include <iostream>
#include <list>
#include <unordered_map>
#include <unordered_set>
#include <thread>
#include <boost/integer.hpp>

#include "cachepad.hpp"
#include "mmalloc.hpp"
#include "timer.hpp"

using namespace std;
using boost::optional;

#ifndef BIT_WIDTH
#define BIT_WIDTH 16
#endif

const uint8_t logState = BIT_WIDTH;
typedef typename boost::uint_t<logState    >::least State;
typedef typename boost::uint_t<logState + 1>::least StateIter;
const StateIter numState  = (StateIter) 1 << logState;
const StateIter nullState = ~((StateIter) 0);
const State maxState = ~((State) 0);

bitset<8> rule(110);

State update(State s) {
  State r(0);
  bitset<logState> b(s);
  for (unsigned i=0; i<logState; i++)
    r |= rule[1 * b[(i + logState - 1) % logState]
	    + 2 * b[i]                
	    + 4 * b[(i + 1) % logState]] << i;
  return r;
}

typedef array<State,   numState> Trans;
typedef array<uint8_t, numState> pbitset;

bool iterStateP(function<void(State, bool&)> f, optional<string> msg = optional<string>(), bool parallel = false, bool skipWorkTest = false) {
  PerfPrinter perfPrinter(msg);
  int numThreads = 1;
  if (parallel)
    numThreads = min<uint64_t>(thread::hardware_concurrency(), numState);
  cache_pad<bool> *perThreadWorked = new cache_pad<bool>[numThreads];
  list<thread*> tasks;
  for (int t=0; t<numThreads; t++) {
    perThreadWorked[t]() = false;
    tasks.push_front(new thread([=]{
	  for (StateIter s = numState / numThreads * t;
	       s < ((t == numThreads - 1) ? numState : (numState / numThreads * (t+1)));
	       s++)
	    f(s, perThreadWorked[t]());
	}));
  }
  for (; !tasks.empty(); tasks.front()->join(), delete tasks.front(), tasks.pop_front());
  bool worked = skipWorkTest;
  for (int t=0; t<numThreads; t++)
    worked |= perThreadWorked[t];
  delete perThreadWorked;
  return worked;
}

bool iterState(function<void(State)> f, optional<string> msg = optional<string>(), bool parallel = false) {
  return iterStateP([=](State s, bool &) { f(s); }, msg, parallel, true);
}

void iterTransP(int times, function<void(State, bool&)> f, optional<string> msg = optional<string>(), bool parallel = false, bool skipWorkTest = false) {
  PerfPrinter perfPrinter(msg);
  auto msg2 = [=,&msg] (int i) {
    return msg ? (*msg + string(" ") + to_string(times-i) + string("/") + to_string(times)) : msg; };
  while (times--)
    if (not iterStateP(f, msg2(times), parallel, skipWorkTest)) return;
}

void iterTrans(int times, function<void(State)> f, optional<string> msg = optional<string>(), bool parallel = false) {
  iterTransP(times, [=](State s, bool &) { f(s); }, msg, parallel, true);
}

void init(Trans &t, Trans &c, pbitset &reachable) {
  iterState([&](State s) {
      t[s] = update(s);
      c[s] = s;
      reachable[t[s]] = 1;
    }, (string) "single step transition table", true);
}

void findCycle(Trans &t, Trans &c) {
  // forward to t=numState; now every state is in a cycle
  iterTransP(logState, [&](State s, bool &worked) {
      State n = t[s];
      t[s] = t[n];
      c[s] = min<State>(c[s], c[n]);
      if (n != t[s]) worked = true;
    }, (string) "fwd time", true);
  // Transients may have a cycle id (minimum state) that is on the
  // transient (not in the cycle) and thus different from the cycle id
  // in the cycle. Thus we copy the cycle id from a state in the cycle
  // to all states in the basin.
  iterState([&](State s) {
      c[s] = c[t[s]];
    }, (string) "fix transient cycle id", true);
}

State canonize(State s) {
  // TODO: test if a variant base on clz is faster
  State cs = s;
  for (int i=0; i<logState; i++)
    cs = min<State>(cs, (((s<<i) | (s>>(logState-i))) & maxState));
  return cs;
};

void cycleStat(Trans &t, Trans &c, pbitset &reachable) {
  struct Stat {
    StateIter basin, len, eden, totalBasin;
    explicit Stat() : basin(0), len(0), eden(0), totalBasin(0) {}
  };
  unordered_map<State, Stat> cycStat;
  unordered_set<State> redCycleCounted;

  // Is this cycle the canonical one?
  // How big is the basin of attraction?
  // How many garden of eden states does it contain?
  { PerfPrinter perfPrinter((string) "cycle count");
  iterState([&](State s) {
    auto cyc = c[s];
    auto canonCyc = canonize(cyc);
    auto &stat = cycStat[canonCyc];
    stat.totalBasin++;
    if (cyc == canonCyc) {
      stat.basin++;
      if (!reachable[s])
	stat.eden++;
    }
  }, (string) "basin & eden"); }

  // how long is the cycle, what is the actual minimal state
  { PerfPrinter perfPrinter((string) "cycle length");
  for (auto &i : cycStat) {
    Stat &stat = i.second;
    State cur, start;
    cur = start = t[i.first];
    do {
      stat.len++;
      cur = update(cur);
    } while (cur != start);
    // TODO: parallelize loop
  }}

  // print it
  for (auto &i : cycStat) {
    auto &s = i.second;
    assert(s.totalBasin % s.basin == 0);
    cout << bitset<logState>(i.first) << "\t"
	 << (s.totalBasin / s.basin) << "\t"
	 << s.len << "\t"
	 << s.basin << "\t"
	 << s.eden  << "\t"
	 << i.first << endl;
  }
}

void print(Trans &t) {
  for (auto s : t)
    cout << bitset<logState>(s) << endl;
}

void printTraj(State s, int count) {
  while (count--) {
    cout << bitset<logState>(s) << endl;
    s = update(s);
  }
}

int main(int argc, char **argv) {
  assert(argc >= 2);
  rule = atoi(argv[1]);
  if (!strcmp(argv[2], "traj")) {
    assert(argc == 5);
    printTraj(atoi(argv[3]), atoi(argv[4]));
  }
  if (!strcmp(argv[2], "cycle")) {
    Trans *t, *c;
    pbitset *r;
    { PerfPrinter perfPrinter((string) "allocating memory");
      t = mmalloc<Trans>();
      c = mmalloc<Trans>();
      r = mmalloc<pbitset>();
    }
    init(*t, *c, *r);
    findCycle(*t, *c);
    cycleStat(*t, *c, *r);
  }
  return 0;
}