#include <stdio.h>
#include <stdlib.h>
#include "mt19937ar.h" // a better RNG than the built-in one

// This is Dave's quick & dirty cs 490 assignment allocation program
// (c) 2012, Dave Tompkins

#define luckySeed 4901121

// assumes constant # people per assignment
#define numAssignments 11
#define studentsPerAssn 4
#define numStudents (numAssignments * studentsPerAssn)

int curAllocation[numStudents];
int bestAllocation[numStudents];
int studentLookup[numStudents];

// student preference data is stored in an external file
extern double pairScore[numStudents][numStudents];
extern double assignmentFraction[numStudents];
extern double assignmentScore[numStudents][numAssignments];

// determine assignment # for student s
#define getStudentAssn(s) (studentLookup[s] / studentsPerAssn)

/////////////////////////////////////

// generate a random number 0..n (n is small)
#define genRandN(n) (genrand_int32() % (n))

// swap two ints
#define swapInt(a,b) { int tmpSwapInt = a; a = b; b = tmpSwapInt; }

// print an array
void printArray (int *arr, int n) {
  int i;
  for (i=0; i<n; i++) {
    printf("%d",arr[i]);
    if (i==(n-1)) {
      printf("\n");
    } else {
      printf(",");
    }
  }
}

// generate an array of random numbers 0..n
void randomArray (int *arr, int n) {
  int i,j;
  for (i=0; i<n; i++) {
    arr[i] = i;
  }
  for (i=n-1; i>0; i--) {
    j = genRandN(i);
    swapInt(arr[i],arr[j])
  }
}

// generate reverse lookup of an array 0..n
void genLookup (int *arrOrig, int *arrLookup, int n) {
  int i;
  for (i=0; i<n; i++) {
    arrLookup[arrOrig[i]] = i;
  }
}

/////////////////////////////////////

// backup the current assignment
void saveSolution() {
  int i;
  for (i=0; i < numStudents; i++) {
    bestAllocation[i] = curAllocation[i];
  }
}

// return happiness for a student based on students in the group
double avgPartnerHappiness(int s) {
  int start = getStudentAssn(s) * studentsPerAssn;
  int end = start + studentsPerAssn - 1;
  int i;
  double score = 0;
  for (i=start; i<end; i++) {
    if (curAllocation[i] != s) {
      score += pairScore[s][curAllocation[i]];
    }
  }
  score /= (studentsPerAssn - 1);
  return score;
}

// return happiness for a student s
double studentHappiness(int s) {
  return assignmentScore[s][getStudentAssn(s)] * assignmentFraction[s] // score based on assignment
         + avgPartnerHappiness(s) * (1 - assignmentFraction[s]); // score based on partnerS
}

// return happiness for the class
double totalHappiness() {
  int i;
  double score = 0;
  for (i=0; i<numStudents; i++) {
    score += studentHappiness(i);
  }
  return score;
}

// swap the assignments for two students
void swapStudents(int s1, int s2) {
  int pos1 = studentLookup[s1];
  int pos2 = studentLookup[s2];
  curAllocation[pos1] = s2;
  curAllocation[pos2] = s1;
  studentLookup[s1] = pos2;
  studentLookup[s2] = pos1;
}

// check to see if swapping two students is better, and keep the swap if better
int goodStudentSwap(double curTotal, int s1, int s2) {
  double origScore = studentHappiness(s1) + studentHappiness(s2);
  double newScore;
  swapStudents(s1,s2);
  newScore = studentHappiness(s1) + studentHappiness(s2);
  if (newScore > origScore) {
    if (totalHappiness() > curTotal) {
      return 1;
    }
  }
  swapStudents(s1,s2);
  return 0;
}

// swap the assignments for 2 groups
void swapAssignments(int a1, int a2) {
  int i;
  for (i=0; i<studentsPerAssn; i++) {
    swapStudents(a1*studentsPerAssn+i, a2*studentsPerAssn+i);
  }
}

// check to see if swapping two groups is better, and keep the swap if better
int goodAssignmentSwap(double curTotal, int a1, int a2) {
  swapAssignments(a1,a2);
  if (totalHappiness() > curTotal) {
    return 1;
  }
  swapAssignments(a1,a2);
  return 0;
}

int main(int argc, char **argv) {
  int i,j;

  double bestSoFar = 0;
  double curScore;
  int curStudent;
  int doSwap;
  int studentOrder[numStudents];

  int numTries = 0;

  init_genrand(luckySeed);

  while (1) {
    numTries++;

    // assign the sutdents randomly
    randomArray(curAllocation, numStudents);
    genLookup(curAllocation, studentLookup, numStudents);
    curScore = totalHappiness();

    doSwap = 1;
    while (doSwap) {
      doSwap = 0;
      randomArray(studentOrder, numStudents);
      // try swapping all pairs of students, looking for improvement
      for (i=0; i < (numStudents - 1); i++) {
        curStudent = studentOrder[i];
        for (j=i+1; j < numStudents; j++) {
          if (getStudentAssn(curStudent) != getStudentAssn(studentOrder[j])) {
            doSwap = goodStudentSwap(curScore, curStudent, studentOrder[j]);
            if (doSwap) break;
          }
        }
        if (doSwap) break;
      }
      // if no students can be swapped, try swapping teams
      if (!doSwap) {
        for (i=0; i < (numAssignments - 1); i++) {
          for (j=i+1; j < numAssignments; j++) {
            doSwap = goodAssignmentSwap(curScore, i, j);
            if (doSwap) break;
          }
          if (doSwap) break;
        }
      }
      curScore = totalHappiness();
    }
    // no more swapping - stuck in a local minimum
    if (curScore > bestSoFar) {
      printf("# *** new best happiness *** = %g -> %g\n", bestSoFar, curScore);
      bestSoFar = curScore;
      saveSolution();
      printArray(bestAllocation, numStudents);
    }
    if (numTries % 1000 == 0) {
      printf("# tries: %d, best so far %g\n", numTries, bestSoFar);
      printArray(bestAllocation, numStudents);
    }
  }
}