#include "neededroutines.h"

const int MAXPERSON = 20000;
const int MAXGROUPS = 100;
const int NSIMUL = 1000;

int readdata(int familyno[], int indivno[], int edssv[], int durationv[], int groupv[], int &ngroups);

void calcglobal(float msssv[], int edssv[], int durationv[], int nperson);

void calclocal(float msssv[], int edssv[], int durationv[], int nperson);

void kruskal_test(float msssv[], int durationv[], int groupv[], int nperson, int permp);

void kruskal_perm(float ranksv[], int groupv[], int nperson, float TKWv, int numgroupv[]);

inline void permute(int n, int labels[]);

void dump_msss(int familynov[], int indivnov[], int edssv[], int durationv[], int groupv[], float msssv[], int nperson);


int main(int argc, char *argv[])
{
  // arguments to main() are 

  int edss[MAXPERSON], duration[MAXPERSON];
  // edss actually means edss score * 2 (in order to be an integer)
  int familyno[MAXPERSON], indivno[MAXPERSON];
  float msss[MAXPERSON];
  int group[MAXPERSON], ngroups;
  int edssfreq[20];
  int seed, flag, key, nperson, tabletype, permp, i, j;

  seed = 10;
  unsigned long iseed=seed;
  sdrni(&iseed);

  cout << "Borland C++ 5.5.1 for Win32 Copyright (c) 1993, 2000 Borland\n";
  cout << "MSSStest version 2.0 September 2004\n";
  nperson = readdata(familyno, indivno, edss, duration, group, ngroups);
  
  ofstream fout("msss.out", ios::app);
  if (!fout)
    {cout << "Cannot open `msss.out' file to store results.\n";
    cout << "Type 0 and return to end program\n";
    cin >> key;
    exit(0);
    }

  flag = 0;
  while (!flag)
    {cout << "Do you wish to use global MSSS or local MSSS?\n(We strongly recommend global MSSS as the default.)\n";
    cout << "Type 1 for global MSSS or 2 for local MSSS.\nThen press return.\n";
    cin >> key;
    if (key==1 || key==2)
      {flag = 1;
      if (key==1)
        {cout << "You selected global MSSS.\n";
	fout << "Using global MSSS.\n";
        }
      if (key==2)
        {cout << "You selected local MSSS.\n";
	fout << "Using local MSSS.\n";
	}
      }
    }
  tabletype = key;
  fout.close();
  
  if (ngroups>1)
    {flag = 0;
    while (!flag)
      {cout << "Do you also wish to calculate permutation p-value?\n";
      cout << "Type 1 for permutation p-value or 0 for no permutation p-value.\nThen press return.\n";
      cin >> key;
      if (key==1 || key==0)
        {flag = 1;
	if (key==1)
	  cout << "You chose to calculate permutation p-value.\n";
	if (key==0)
	  cout << "You selected no permutation p-value.\n";
	}
      }
    permp = key;
    }  // end of if (ngroups>1) section

  if (tabletype==1)
    calcglobal(msss, edss, duration, nperson);
  else
    calclocal(msss, edss, duration, nperson);

  if (ngroups==1)
    permp = 0;
  kruskal_test(msss, duration, group, nperson, permp);
  dump_msss(familyno, indivno, edss, duration, group, msss, nperson);

  if (ngroups>1)
    cout << "Mean MSSS and test result (if any) saved in file `msss.out'.\n";
  cout << "Type 0 and return to end program\n";
  cin >> key;
}



int readdata(int familynov[], int indivnov[], int edssv[], int durationv[], int groupv[], int &ngroups)
{
  int i, usegroup[MAXGROUPS], key, tempi, nper;
  float temp;
  int errorflag1, errorflag2;

  cout << "Reading data file `edss.txt'.\n";

  for (i=0; i<MAXGROUPS; i++)
    usegroup[i] = 0;

  i = 0;

  ifstream fin("edss.txt");
  if (!fin)
    {cout << "Cannot find file `edss.txt'." << endl;
    cout << "Type 0 and return to end program\n";
    cin >> key;
    exit(0);
    }
  while (!fin.eof())
    {fin >> tempi;
    familynov[i] = tempi;
    if (fin.eof())
      break;
    fin >> indivnov[i];
    fin >> temp;
    edssv[i] = (int) floor(temp * 2 + 0.0001);
    fin >> temp;
    durationv[i] = (int) floor(temp + 0.0001);
    fin >> groupv[i];
    if (groupv[i] < 0 || groupv[i] >= MAXGROUPS)
      {cout << "Error in data.  Groups' designations should be >=0 and <" << MAXGROUPS << ".\n";
      cout << "Type 0 and return to end program\n";
      cin >> key;
      exit(0);
      }
    usegroup[groupv[i]] = 1;
    i++;
    }
  fin.close();
  nper = i;

  ngroups = 0;
  for (i=0; i<MAXGROUPS; i++)
    if (usegroup[i]>0)
      ngroups++;

  cout << "Found data on " << nper << " individuals in " << ngroups << " groups.\n";
  cout << "Family number, individual number, EDSS, duration and group for first and last individuals are:\n";
  cout << familynov[0] << " " << indivnov[0] << " " << edssv[0] * 1.0 / 2 << " " << durationv[0] << " " << groupv[0] << endl;
  cout << familynov[nper-1] << " " << indivnov[nper-1] << " " << edssv[nper-1] * 1.0 / 2 << " " << durationv[nper-1] << " " << groupv[nper-1] << endl;

/*
  if (ngroups<2)
    {cout << "Error in data.  There must be at least two groups.\n";
    cout << "Type 0 and return to end program\n";
    cin >> key;
    exit(0);
    }
*/

  errorflag1 = 0;
  errorflag2 = 0;
  for (i=0; i<nper; i++)
    {if (edssv[i]<0 || edssv[i]>20 || edssv[i]==1)
      errorflag1 = 1;
    if (edssv[i]==20)
      errorflag2 = 1;
    }

  if (errorflag1)
    {cout << "File contains illegal EDSS values.  Legal values are 0, 1, 1.5, 2, ..., 9.5.\n";
    cout << "Type 0 and return to end program\n";
    cin >> key;
    exit(0);
    }
  if (errorflag2)
    cout << "WARNING: File contains EDSS values of 10.  These will be treated as though 9.5.\n";

  ofstream fout("msss.out", ios::trunc);
  if (!fout)
    {cout << "Cannot create new `msss.out' file to store output.\n";
    cout << "Type 0 and return to end program\n";
    cin >> key;
    exit(0);
    }
  fout << nper << " individuals.\n";
  fout << ngroups << " groups.\n";
  fout.close();

  return nper;
}


void calcglobal(float msssv[], int edssv[], int durationv[], int nperson)
{
  int i, j, key, durtemp, edsstemp;
  float mssstable[20][31];
  
  cout << "Reading global MSSS file `global.dat'.\n";
  
  ifstream fin("global.dat");
  if (!fin)
    {cout << "Cannot find file `global.dat'." << endl;
    cout << "Type 0 and return to end program\n";
    cin >> key;
    exit(0);
    }

  for (i=0; i<20; i++)
    for (j=0; j<31; j++)
      fin >> mssstable[i][j];

  fin.close();

  for (i=0; i<nperson; i++)
    {if (durationv[i]>30)
      durtemp = 30;
    else
      durtemp = durationv[i];
    if (edssv[i]==20)
      edsstemp = 19;
    else
      edsstemp = edssv[i];
    msssv[i] = mssstable[edsstemp][durtemp];
    }
}


void calclocal(float msssv[], int edssv[], int durationv[], int nperson)
{
  int i, j, mindur, maxdur, durtemp, edsstemp;
  int counts[20][31], cumcounts[20][31];
  float mssstable[20][31];

  for (i=0; i<20; i++)
    for (j=0; j<31; j++)
      counts[i][j] = 0;

  for (i=0; i<nperson; i++)
    {if (durationv[i]<=1)
      {mindur = durationv[i];
      maxdur = durationv[i]+2;
      }
    if (durationv[i]>=2 && durationv[i]<=3)
      {mindur = durationv[i] - 1;
      maxdur = durationv[i]+2;
      }
    if (durationv[i]>3)
      {mindur = durationv[i] - 2;
      maxdur = durationv[i] + 2;
      if (maxdur>30)
	maxdur = 30;
      }
    for (j=mindur; j<=maxdur; j++)
      counts[edssv[i]][j]++;
    }

  for (i=0; i<31; i++)
    {cumcounts[0][i] = counts[0][i];
    cumcounts[1][i] = counts[0][i];
    // That is for an EDSS of 0.5 (which does not exist)
    for (j=2; j<20; j++)
      cumcounts[j][i] = cumcounts[j-1][i] + counts[j][i];
    }

  for (i=0; i<31; i++)
    {mssstable[0][i] = (cumcounts[0][i] + 1) / 2;
    mssstable[0][i] /= (cumcounts[19][i] + 1);
    for (j=1; j<20; j++)
      {mssstable[j][i] = (cumcounts[j-1][i] + 1 + cumcounts[j][i]) / 2;
      mssstable[j][i] /= (cumcounts[19][i] + 1);
      }
    }

  for (i=0; i<nperson; i++)
    {if (durationv[i]>30)
      durtemp = 30;
    else
      durtemp = durationv[i];
    if (edssv[i]==20)
      edsstemp = 19;
    else
      edsstemp = edssv[i];
    msssv[i] = mssstable[edsstemp][durtemp] * 10;
    }
}


void kruskal_test(float msssv[], int durationv[], int groupv[], int npersonv, int permp)
{
  int i, numgroup[MAXGROUPS], ngroups, key, count, newgroupv[MAXPERSON+1];
  float msssvec1[MAXPERSON+1], msssvec2[MAXPERSON+1], ranks[MAXPERSON+1], sumranks[MAXGROUPS], summsss[MAXGROUPS];
  float TKW, CKW, Tstat, pval;

  count = 0;
  for (i=0; i<npersonv ; i++)
      if (durationv[i]>0)
        {msssvec1[count] = msssv[i];
	newgroupv[count] = groupv[i];
	count++;
	}
  // Just take people with duration>0

    ofstream fout("msss.out", ios::app);
    if (!fout)
      {cout << "Cannot open `msss.out' file to store results.\n";
      cout << "Type 0 and return to end program\n";
      cin >> key;
      exit(0);
      }

  if (count<npersonv)
    {cout << "Only individuals with duration >0 are used.\n" << npersonv - count << " individuals with duration=0 have been excluded.\n";
    fout << "Only individuals with duration >0 are used.\n" << npersonv - count << " individuals with duration=0 have been excluded.\n";
    npersonv = count;
    }
      
  for (i=0; i<npersonv; i++)
    msssvec2[i+1] = msssvec1[i];
  // The rankties() function works with a vector v[1:n+1], rather than v[0:n]

  CKW = rankties(npersonv, msssvec2, ranks);

  for (i=0; i<MAXGROUPS; i++)
    {sumranks[i] = 0;
    numgroup[i] = 0;
    summsss[i] = 0;
    }

  for (i=0; i<npersonv; i++)
    {sumranks[newgroupv[i]] += ranks[i+1];
    // i+1 because ranks[] is indexed from 1 up, whereas newgroupv is from 0 up
    summsss[newgroupv[i]] += msssvec1[i];
    numgroup[newgroupv[i]]++;
    }

  ngroups = 0;
  TKW = 0;

  for (i=0; i<MAXGROUPS; i++)
    if (numgroup[i]>0)
      {ngroups++;
      TKW += pow(sumranks[i], 2) / numgroup[i];
      }

  TKW *= 12.0 / (npersonv * (npersonv+1));
  Tstat = TKW;
  TKW -= 3 * (npersonv+1);

  // Correction for ties
  TKW /= (1 - CKW);
  // If H_0 true, TKW is asymptotic chisq with (ngroups-1) df

  cout << "Mean MSSS score in each group is:\n";
  fout << "Mean MSSS score in each group is:\n";
  for (i=0; i<MAXGROUPS; i++)
    if (numgroup[i]>0)
      {cout << "Group " << i << ": " << summsss[i] / numgroup[i] << endl;
      fout << "Group " << i << ": " << summsss[i] / numgroup[i] << endl;
      }

  if (ngroups>1)
    {cout << "Kruskal-Wallis test statistic is " << TKW << ".\nAsymptotically this is distributed chi-square on " << ngroups-1 << " degrees of freedom.\n";
    fout << "Kruskal-Wallis test statistic is " << TKW << ".\nAsymptotically this is distributed chi-square on " << ngroups-1 << " degrees of freedom.\n";

    pval = pchisq(ngroups-1, TKW);
    if (pval>0.000001)
      {cout << "This gives a p-value of " << pval << ".\n";
      fout << "This gives a p-value of " << pval << ".\n";
      }
    else
      {cout << "This gives a p-value of <10e-6.\n";
      fout << "This gives a p-value of <10e-6.\n";
      }
    }

  fout.close();

  if (ngroups>1 & permp)
    kruskal_perm(ranks, newgroupv, npersonv, Tstat, numgroup);
}


void kruskal_perm(float ranksv[], int groupv[], int nperson, float Tstatv, int numgroupv[])
{
  int sim, i, greater, key;
  int groupperm[MAXPERSON+1];
  float sumranksperm[MAXGROUPS], Tperm;

  greater = 0;
  for (i=0; i<nperson; i++)
    groupperm[i] = groupv[i];
  
  cout << "Now performing permutations.  Please wait.\n";

  for (sim=0; sim<NSIMUL; sim++)
    {Tperm = 0; 
    permute(nperson, groupperm);
    
    for (i=0; i<MAXGROUPS; i++)
      sumranksperm[i] = 0;
    for (i=0; i<nperson; i++)
      sumranksperm[groupperm[i]] += ranksv[i+1];
    
    for (i=0; i<MAXGROUPS; i++)
      if (numgroupv[i]>0)
	Tperm += pow(sumranksperm[i], 2) / numgroupv[i];

    Tperm *= 12.0 / (nperson * (nperson+1));

    if (Tperm>Tstatv)
      greater++;
    }

  cout << "Based on " << NSIMUL << " permutations, the permutation p-value is " << (greater + 1.0) / (NSIMUL+1) << ".\n";

  ofstream fout("msss.out", ios::app);
  if (!fout)
    {cout << "Cannot open `msss.out' file to store results.\n";
    cout << "Type 0 and return to end program\n";
    cin >> key;
    exit(0);
    }
  fout << "Based on " << NSIMUL << " permutations, the permutation p-value is " << (greater + 1.0) / (NSIMUL+1) << ".\n";
  fout.close();
}


inline void permute(int n, int labels[])
{
  // This takes vector of labels labels[] and randomly permutes it

  int i, choose, temp;

  for (i=0; i<n; i++)
    {choose = i + (int) floor( randunif() * (n-i) );
    temp = labels[i];
    labels[i] = labels[choose];
    labels[choose] = temp;
    }
}


void dump_msss(int familynov[], int indivnov[], int edssv[], int durationv[], int groupv[], float msssv[], int nperson)
{
  int i, key;

  cout << "Saving MSSS scores to file `indivmsss.out'.\n";

  ofstream fout("indivmsss.out", ios::trunc);
  if (!fout)
    {cout << "Cannot create new `indivmsss.out' file.\n";
    cout << "Type 0 and return to end program\n";
    cin >> key;
    exit(0);
    }

  fout << "Family\tIndiv\tEDSS\tDur\tGroup\tMSSS\n";
  for (i=0; i<nperson; i++)
      fout << familynov[i] << "\t" << indivnov[i] << "\t" << edssv[i]*0.5 << "\t" << durationv[i] << "\t" << groupv[i] << "\t" << setprecision(4) << msssv[i] << endl;

  fout.close();
}