#include "charm.h"
#include <stdlib.h>

int idx_CountMsg;
int idx_Counter;
int idx_createCounter;
int idx_childCount;
int idx_sendCounts;

typedef struct {
  int count;
} CountMsg;

typedef struct {
  CkGroupID thisgroup;
  int myCount;
  int totalCount;
  int waitFor;
  CthThread threadID;
} Counter;

void createCounter(void *msg, Counter *ctr)
{
  CkFreeMsg(msg);
  ctr->thisgroup = CkGetGroupID();
  ctr->myCount = ctr->totalCount = 0;
  ctr->waitFor = CkNumPes();
  ctr->threadID = 0;
}

void childCount(CountMsg *msg, Counter *ctr)
{
  ctr->totalCount += msg->count;
  CkFreeMsg(msg);
  ctr->waitFor--;
  if(ctr->waitFor == 0 && ctr->threadID) {
    CthAwaken(ctr->threadID);
  }
}

void increment(Counter *ctr)
{
  ctr->myCount++;
}

void sendCounts(void *msg, Counter *ctr)
{
  CountMsg *m = (CountMsg *) CkAllocMsg(idx_CountMsg, sizeof(CountMsg), 0);
  CkFreeMsg(msg);
  m->count = ctr->myCount;
  CkSendMsgBranch(idx_childCount, m, 0, ctr->thisgroup);
}

int getTotalCount(Counter *ctr)
{
  void *msg = CkAllocSysMsg();
  CkBroadcastMsgBranch(idx_sendCounts, msg, ctr->thisgroup);
  ctr->threadID = CthSelf();
  while(ctr->waitFor != 0) CthSuspend();
  return ctr->totalCount;
}

CkGroupID counterInit(void)
{
  void *msg = CkAllocSysMsg();
  return CkCreateGroup(idx_Counter, idx_createCounter, msg, 0, 0);
}

void registerCounter(void)
{
  idx_CountMsg = CkRegisterMsg("CountMsg", (CkPackFnPtr) 0, (CkUnpackFnPtr) 0,
                               (CkCoerceFnPtr) 0, sizeof(CountMsg));
  idx_Counter = CkRegisterChare("Counter", sizeof(Counter));
  idx_createCounter = CkRegisterEp("createCounter", (CkCallFnPtr) createCounter,
                                   0, idx_Counter);
  idx_childCount = CkRegisterEp("childCount", (CkCallFnPtr) childCount,
                                idx_CountMsg, idx_Counter);
  idx_sendCounts = CkRegisterEp("sendCounts", (CkCallFnPtr) sendCounts,
                                0, idx_Counter);
}

#define Nmax 20

/* readonly variables */

int grain, N;
CkGroupID counterGroup;
CkChareID mainhandle;

int idx_PartialBoard;
int idx_Main;
int idx_createMain;
int idx_onQD;
int idx_queens;
int idx_createQueens;


typedef struct {
  int nextRow;
  int Queens[Nmax];
} PartialBoard;

typedef struct {
  double t0;
} Main;

void createMain(CkArgMsg *msg, Main *main)
{
  PartialBoard *pMsg; 
  CkGetChareID(&mainhandle);

  if(msg->argc < 3)
    CkAbort("Usage: pgm N grain\n");
  N = atoi(msg->argv[1]);
  grain = atoi(msg->argv[2]);
  CkFreeMsg(msg);
  pMsg = CkAllocMsg(idx_PartialBoard, sizeof(PartialBoard), 0);
  pMsg->nextRow = 0;
  CkCreateChare(idx_queens, idx_createQueens, pMsg, 0, 0);
  main->t0 = CkTimer();
  counterGroup = counterInit();
  CkStartQD(idx_onQD, &mainhandle);
}

void thr_onQD(Main *main)
{
  int nSolutions = getTotalCount(CkLocalBranch(counterGroup));
  CkPrintf("There are %d solutions to %d queens. Finish time=%f\n",
          nSolutions, N, CkTimer()-main->t0);
  CkExit();
}

void onQD(void *msg, Main *main)
{
  CthThread thr = CthCreate((CthVoidFn)thr_onQD, main, 0);
  CthAwaken(thr);
  CkFreeMsg(msg);
}

typedef struct {
  int dummy;
} queens;

void solutionFound(int board[])
{
  increment((Counter*) CkLocalBranch(counterGroup));
}

int consistent(queens *q, int board[], int lastRow, int col)
{
  int x,y;

  for (x=0; x<lastRow; x++)
  {
    y = board[x];
    if ((y==col)  || ((lastRow-x) == (col-y)) || ((lastRow-x) == (y-col)) ) {
      return(0);
    }
  }
  return(1);
}

void seqQueens(queens *q, int board[], int nextRow)
{
  int col;

  if (nextRow == N) { 
    solutionFound(board);
    return;
  }
  for (col = 0; col<N; col++) {
    if (consistent(q, board, nextRow, col))
    {
      board[nextRow] = col;
      seqQueens(q, board, nextRow+1);
    }
  }
}


void createQueens(PartialBoard *m, queens *q)
{
  int col,i, row;
  PartialBoard *newMsg;

  row = m->nextRow;
  if (row == N)
    solutionFound(m->Queens);
  else if ( (N-row) < grain)
    seqQueens(q, m->Queens, row);
  else
  {
    for (col = 0; col<N; col++)
      if (consistent(q, m->Queens, row, col))
      {
        newMsg = CkAllocMsg(idx_PartialBoard, sizeof(PartialBoard), 0);
        newMsg->nextRow = row + 1;
        for (i=0; i<N; i++)
          newMsg->Queens[i] = m->Queens[i];
        newMsg->Queens[row] = col;
        CkCreateChare(idx_queens, idx_createQueens, newMsg, 0, 0);
      }
  }
  CkFreeMsg(m);
  free(q);
}

void CkRegisterMainModule(void)
{
  registerCounter();
  idx_PartialBoard = CkRegisterMsg("PartialBoard", (CkPackFnPtr)0, 
                                   (CkUnpackFnPtr)0, (CkCoerceFnPtr)0, 
                                   sizeof(PartialBoard));
  idx_Main = CkRegisterChare("Main", sizeof(Main));
  idx_createMain = CkRegisterEp("createMain", (CkCallFnPtr) createMain,
                                0, idx_Main);
  CkRegisterMainChare(idx_Main, idx_createMain);
  idx_onQD = CkRegisterEp("onQD", (CkCallFnPtr) onQD,
                                0, idx_Main);
  idx_queens = CkRegisterChare("queens", sizeof(queens));
  idx_createQueens = CkRegisterEp("createQueens", (CkCallFnPtr) createQueens,
                                  idx_PartialBoard, idx_queens);
  CkRegisterReadonly(sizeof(int), &grain);
  CkRegisterReadonly(sizeof(int), &N);
  CkRegisterReadonly(sizeof(CkGroupID), &counterGroup);
  CkRegisterReadonly(sizeof(CkChareID), &mainhandle);
}
