
#include "antsUtilities.h"
#include "antsAllocImage.h"
#include <algorithm>

#include <algorithm>
#include <vector>
#include <iostream>
#include <fstream>
#include <cstdlib>
#include <ctime>
#include <string>
#include <math.h>
#include <float.h>
#include <assert.h>
#include "ReadWriteData.h"

#include <vnl/algo/vnl_symmetric_eigensystem.h>

#include <itkArray.h>
#include <itkMatrix.h>
#include <itkImageRegionIteratorWithIndex.h>

#include "itkTDistribution.h"
#include "vnl/vnl_math.h"
#include "vnl/vnl_erf.h"

namespace ants
{
// for computing F distribution
extern "C" double dbetai_(double *x, double *pin, double *qin);

extern "C" double dgamma_(double *x);

typedef struct
  {
  double statVal;
  int index;
  } StatElement;

int smallerStatElem(StatElement * elem1, StatElement * elem2)
// comparison function for sorting
{
  if( elem1->statVal > elem2->statVal )
    {
    return 1;
    }
  else if( elem1->statVal < elem2->statVal )
    {
    return -1;
    }
  else
    {
    return 0;
    }
}

template <class TImageType>
void ReadImage(itk::SmartPointer<TImageType> & target, const char *file, bool copy)
{
  //  std::cout << " reading b " << std::string(file) << std::endl;
  typedef itk::ImageFileReader<TImageType> readertype;
  typename readertype::Pointer reader = readertype::New();
  reader->SetFileName(file);
  reader->Update();
  if( !copy )
    {
    target = (reader->GetOutput() );
    }
  else
    {
    typedef itk::ImageRegionIteratorWithIndex<TImageType> Iterator;
    Iterator vfIter2( target,  target->GetLargestPossibleRegion() );
    for(  vfIter2.GoToBegin(); !vfIter2.IsAtEnd(); ++vfIter2 )
      {
      vfIter2.Set( reader->GetOutput()->GetPixel(vfIter2.GetIndex() ) );
      }
    }
}

#define GROUPALABEL   0
#define GROUPBLABEL   1

typedef struct
  {
  int randNum;
  int index;
  } PermElement;

void generatePerm(int length, int * genPerm);

// generate a permutation of numbers from 0 to length-1
void generatePermGroup(int * groupID, int lengthGroupA, int lengthGroupB, int * genGroupID);

// generate a permutation of group assignments

int smallerPermElem(PermElement * elem1, PermElement * elem2);

static int first = 0;

void generatePermGroup(int * groupID, int lengthGroupA, int lengthGroupB,
                       int * genGroupID)
// generate a permutation of group assignments
{
  int numSubjects = lengthGroupA + lengthGroupB;

  if( !first )
    {
    first = 1; srand(time(NULL) );
    // cout << "generatePermGroup called" << endl;
    }
  int * newPerm = new int[numSubjects];
  generatePerm(numSubjects, newPerm);
  for( int i = 0; i < numSubjects; i++ )
    {
    genGroupID[i] = groupID[newPerm[i]];
    }

  delete [] newPerm;
}

void generatePerm(int length, int * genPerm)
{
  if( !first )
    {
    first = 1; srand(time(NULL) );
    // cout << "generatePerm called" << endl;
    }
  PermElement * newPerm = new PermElement[length];
  int           cnt;
  for( cnt = 0; cnt < length; cnt++ )
    {
    newPerm[cnt].randNum = rand();
    newPerm[cnt].index = cnt;
    }
  qsort(newPerm, length, sizeof(PermElement),
        (int (*)(const void *, const void *) )smallerPermElem);
  for( cnt = 0; cnt < length; cnt++ )
    {
    genPerm[cnt] = newPerm[cnt].index;
    }
  delete [] newPerm;
}

int smallerPermElem(PermElement * elem1, PermElement * elem2)
{
  if( elem1->randNum > elem2->randNum )
    {
    return 1;
    }
  else if( elem1->randNum < elem2->randNum )
    {
    return -1;
    }
  else
    {
    return 0;
    }
}

double computeQuantile(int numObs, double * stat, double quantile)
// computes the value in the provided statistic for the given quantile value
//
// numObs = Number of Observations
// stat = double array[numObs ] contains the test statisticss
{
  static int first = 0;

  if( !first )
    {
    first = 1;
    srand(time(NULL) );
    }

  StatElement * sortStat = new StatElement[numObs];
  for( int perm = 0; perm < numObs; perm++ )
    {
    sortStat[perm].statVal = stat[perm];
    sortStat[perm].index = perm;
    }

  // sort, smallest first
  qsort(sortStat, numObs, sizeof(StatElement),
        (int (*)(const void *, const void *) )smallerStatElem);

  // index at value
  double quantindex = (double) numObs * quantile;
  if( quantile == 1.0 )
    {
    quantindex = numObs;
    }

  double retval = stat[sortStat[(int) quantindex].index];

  delete [] sortStat;

  return retval;
}

void computePermStatPval(int numFeatures, int numPerms,
                         double * permStat, double * permStatPval)
// computes the Pval for all permutation statistics
// the p-val is computed as the percentual ordered rank over all permutations
//
// numFeatures = Number of scalar features per Subject
// numPerms = Number of Permutations
// permStat = double array[numPerms * numFeatures] contains the test statistics
// permStatPval = double array [numPerms * numFeatures] returns the p-val of the statistics
{
  static int first = 0;

  if( !first )
    {
    first = 1;
    srand(time(NULL) );
    }

  int feat;
  int perm;

  StatElement * sortPermStat = new StatElement[numPerms];
  for( feat = 0; feat < numFeatures; feat++ )
    {
    for( perm = 0; perm < numPerms; perm++ )
      {
      sortPermStat[perm].statVal = permStat[perm * numFeatures + feat];
      sortPermStat[perm].index = perm;
      }

    // sort, smallest first
    qsort(sortPermStat, numPerms, sizeof(StatElement),
          (int (*)(const void *, const void *) )smallerStatElem);

    double curPval = 0;
    for( perm = 0; perm < numPerms; perm++ )
      {
      // percentual rank 0..1 -> cumulative probability -> p-val
      double nextPval = 1.0 - (double) (perm + 1) / (double) numPerms;

      int curIndex = sortPermStat[perm].index;

      if( (perm == 0) || (sortPermStat[perm].statVal != sortPermStat[perm - 1].statVal) )
        {
        // current value is different from previous value (or first value),
        // thus step up p-value
        curPval = nextPval;
        }

      permStatPval[curIndex * numFeatures + feat] = curPval;
      }
    }
  delete [] sortPermStat;
}

double decode_ieee_single( unsigned char *v, int natural_order)
{
  unsigned char *data = v;
  int            s, e;
  unsigned long  src;
  long           f;
  double         value = 0.0;

  if( natural_order )
    {
    src = ( (unsigned long)data[0] << 24)
      | ( (unsigned long)data[1] << 16)
      | ( (unsigned long)data[2] << 8)
      | ( (unsigned long)data[3]);
    }
  else
    {
    src = ( (unsigned long)data[3] << 24)
      | ( (unsigned long)data[2] << 16)
      | ( (unsigned long)data[1] << 8)
      | ( (unsigned long)data[0]);
    }

  s = (src & 0x80000000UL) >> 31;
  e = (src & 0x7F800000UL) >> 23;
  f = (src & 0x007FFFFFUL);

  if( e == 255 && f != 0 )
    {
/* NaN (Not a Number) */
    value = DBL_MAX;
    }
  else if( e == 255 && f == 0 && s == 1 )
    {
/* Negative infinity */
    value = -DBL_MAX;
    }
  else if( e == 255 && f == 0 && s == 0 )
    {
/* Positive infinity */
    value = DBL_MAX;
    }
  else if( e > 0 && e < 255 )
    {
/* Normal number */
    f += 0x00800000UL;
    if( s )
      {
      f = -f;
      }
    value = ldexp(f, e - 150);
    }
  else if( e == 0 && f != 0 )
    {
/* Denormal number */
    if( s )
      {
      f = -f;
      }
    value = ldexp(f, -149);
    }
  else if( e == 0 && f == 0 && s == 1 )
    {
/* Negative zero */
    value = 0;
    }
  else if( e == 0 && f == 0 && s == 0 )
    {
/* Positive zero */
    value = 0;
    }
  else
    {
/* Never happens */
    printf("s = %d, e = %d, f = %lu\n", s, e, f);
    assert(!"Woops, unhandled case in decode_ieee_single()");
    }

  return value;
}

double factorial( double x)
{
  if( x <= 1 )
    {
    return 1;
    }
  double fac = x;
  double n = fac - 1;
  while(  n >= 1 )
    {
    fac *= n;
    n = n - 1;
    }

  return fac;
}

double betadist(  double a, double b )
{
  double numer = factorial(  a - 1) * factorial(b - 1);
  double denom = factorial( a + b - 1);

  return numer / denom;
}

double TTest(int numSubjects,   int* groupLabel, double * featureValue )
{
  int    numSubjA = 0;
  int    numSubjB = 0;
  double meanA = 0, meanB = 0;

  //  unsigned int GROUPALABEL=0;
  // unsigned int GROUPBLABEL=1;
  for( int subj = 0; subj < numSubjects; subj++ )
    {
    if( groupLabel[subj] == GROUPALABEL )
      {
      numSubjA++;
      meanA += featureValue[subj];
      }
    else if( groupLabel[subj] == GROUPBLABEL )
      {
      numSubjB++;
      meanB += featureValue[subj];
      }
    else
      {
      std::cout << " group label " << groupLabel[subj] << " does not exist" << std::endl;
      }
    }
  meanA /= (float)numSubjA;
  meanB /= (float)numSubjB;

  double varA = 0, varB = 0;
  for( int subj = 0; subj < numSubjects; subj++ )
    {
    if( groupLabel[subj] == GROUPALABEL )
      {
      varA += (featureValue[subj] - meanA) * (featureValue[subj] - meanA);
      }
    else if( groupLabel[subj] == GROUPBLABEL )
      {
      varB += (featureValue[subj] - meanB) * (featureValue[subj] - meanB);
      }
    }

  float n1 = (float) numSubjA;
  float n2 = (float) numSubjB;
  varA /= (n1); // use n1 -1  for unbiased estimator, here assume normal distribution
  varB /= (n2); // use n2 - 1 " ... "
//  float sdA=sqrt(varA);
//  float sdB=sqrt(varB);

//  float df  =  n1 + n2 - 2;
// unequal vars
  float denom = varA / n1 + varB / n2;
  // for equal vars
  //    float var =  ( (n1-1.)*newvar1 + (n2-1.)*newvar2 ) / df;
  //    denom = var*(1.0/n1+1.0/n2);
  double tt = 0;
  if( denom > 0 )
    {
    tt =  (meanA - meanB) / sqrt(denom);
    }

  return tt;
}

template <unsigned int ImageDimension>
int StudentsTestOnImages(int argc, char *argv[])
{
  typedef float                                 PixelType;
  typedef itk::Image<PixelType, ImageDimension> ImageType;

  typename ImageType::Pointer mask = NULL;
//  ReadImage<ImageType>(mask, argv[1], false);

  unsigned int numSubjectsA = atoi(argv[3]);
  unsigned int numSubjectsB = atoi(argv[4]);
  unsigned int numSubjects = numSubjectsA + numSubjectsB;
  std::string  outname = std::string(argv[2]);
  unsigned int numvals = numSubjects;
  int*         groupLabel = new int[numSubjects];
  for( unsigned int i = 0; i < numSubjectsA; i++ )
    {
    groupLabel[i] = 0;
    }
  for( unsigned int i = numSubjectsA; i < numSubjects; i++ )
    {
    groupLabel[i] = 1;
    }
  double* feature = new double[numvals];
  for( unsigned int i = 0; i < numvals; i++ )
    {
    feature[i] = 0;
    }

  std::cout << " Numvals " << numvals << std::endl;
  // Get the image dimension
  std::string               fn = std::string(argv[5]);
  itk::ImageIOBase::Pointer imageIO =
    itk::ImageIOFactory::CreateImageIO(fn.c_str(), itk::ImageIOFactory::ReadMode);
  imageIO->SetFileName(fn.c_str() );
  imageIO->ReadImageInformation();
  typename ImageType::SizeType size;
  typename ImageType::SpacingType spacing;
  typename ImageType::PointType origin;
  typename ImageType::DirectionType direction;
  std::vector<double> axis;
  for( unsigned int i = 0; i < ImageDimension; i++ )
    {
    size[i] = imageIO->GetDimensions(i);
//       if (size[i] !=  mask->GetLargestPossibleRegion().GetSize()[i])
    // {
    //  std::cout <<  " mask not same size as data !! " << std::endl;
    //  throw std::exception();
    // }

    spacing[i] = imageIO->GetSpacing(i);
    origin[i]  = imageIO->GetOrigin(i);
    axis = imageIO->GetDirection(i);
    for( unsigned j = 0; j < ImageDimension; j++ )
      {
      if( j < imageIO->GetNumberOfDimensions() )
        {
        direction[j][i] = axis[j];
        }
      else
        {
        direction[j][i] = 0.0;
        }
      }
    }
  std::cout << " size " << size << std::endl;
  typename ImageType::RegionType region;
  region.SetSize(size );
  // ORIENTATION ALERT. the code this replaced originally didn't
  // bother setting the origins even though the directions were
  // grabbed from the ImageIO.  I'm assuming that was supposed to
  // happen, and was left out as an oversight.
  typename ImageType::Pointer StatImage = AllocImage<ImageType>(region,
                                                                spacing,
                                                                origin,
                                                                direction,
                                                                0);
  typename ImageType::Pointer PImage = AllocImage<ImageType>(region,
                                                             spacing,
                                                             origin,
                                                             direction,
                                                             0);

//   unsigned int sizeofpixel=sizeof(PixelType);

  std::vector<typename ImageType::Pointer> imagestack;
  imagestack.resize(numvals);
  for( unsigned int j = 0; j < numvals; j++ )
    {
    std::string ifn = std::string(argv[5 + j]);
    std::cout << "reading " << ifn << std::endl;
    ReadImage<ImageType>(imagestack[j], ifn.c_str(), false);
    }

  typedef itk::ImageRegionIteratorWithIndex<ImageType> Iterator;
  Iterator      vfIter(PImage, PImage->GetLargestPossibleRegion() );
  unsigned long nvox = 1;
  for( unsigned int i = 0; i < ImageDimension; i++ )
    {
    nvox *= PImage->GetLargestPossibleRegion().GetSize()[i];
    }

  unsigned long ct = 0;
  unsigned long prog = nvox / 20;
  std::cout << " NVals " << numvals << " NSub " << numSubjects <<  std::endl;
  for(  vfIter.GoToBegin(); !vfIter.IsAtEnd(); ++vfIter )
    {
    typename ImageType::IndexType index = vfIter.GetIndex();
    for( unsigned int subj = 0; subj < numSubjects; subj++ )
      {
      feature[subj] = imagestack[subj]->GetPixel(index);
      }
    if( ct % prog == 0 )
      {
      std::cout << " % " << (float) ct / (float) nvox << std::endl;
      }

    double stat = 0;
//      if (mask->GetPixel(index) >= 0.5)
    if( true )
      {
      stat = TTest(numSubjects, groupLabel, feature);
      }
    ct++;
    StatImage->SetPixel(  index,   stat);
    }

  typedef itk::Statistics::TDistribution DistributionType;
  typename DistributionType::Pointer distributionFunction = DistributionType::New();

  WriteImage(StatImage, outname.c_str() );

  delete [] feature;
  delete [] groupLabel;

  return 1;
}

// entry point for the library; parameter 'args' is equivalent to 'argv' in (argc,argv) of commandline parameters to
// 'main()'
int StudentsTestOnImages( std::vector<std::string> args, std::ostream* out_stream = NULL )
{
  // put the arguments coming in as 'args' into standard (argc,argv) format;
  // 'args' doesn't have the command name as first, argument, so add it manually;
  // 'args' may have adjacent arguments concatenated into one argument,
  // which the parser should handle
  args.insert( args.begin(), "StudentsTestOnImages" );

  int     argc = args.size();
  char* * argv = new char *[args.size() + 1];
  for( unsigned int i = 0; i < args.size(); ++i )
    {
    // allocate space for the string plus a null character
    argv[i] = new char[args[i].length() + 1];
    std::strncpy( argv[i], args[i].c_str(), args[i].length() );
    // place the null character in the end
    argv[i][args[i].length()] = '\0';
    }
  argv[argc] = 0;
  // class to automatically cleanup argv upon destruction
  class Cleanup_argv
  {
public:
    Cleanup_argv( char* * argv_, int argc_plus_one_ ) : argv( argv_ ), argc_plus_one( argc_plus_one_ )
    {
    }

    ~Cleanup_argv()
    {
      for( unsigned int i = 0; i < argc_plus_one; ++i )
        {
        delete[] argv[i];
        }
      delete[] argv;
    }

private:
    char* *      argv;
    unsigned int argc_plus_one;
  };
  Cleanup_argv cleanup_argv( argv, argc + 1 );

  // antscout->set_stream( out_stream );

  std::cout <<  " df     P = 0.05  P = 0.01   P = 0.001  " << std::endl;
  std::cout << " 1             12.71     63.66     636.61  " << std::endl;
  std::cout << " 2    4.30     9.92     31.60    " << std::endl;
  std::cout << " 3    3.18     5.84     12.92" << std::endl;
  std::cout << " 4    2.78     4.60     8.61" << std::endl;
  std::cout << " 5    2.57     4.03     6.87" << std::endl;
  std::cout << " 6     2.45     3.71     5.96" << std::endl;
  std::cout << " 7    2.36     3.50     5.41" << std::endl;
  std::cout << " 8    2.31     3.36     5.04" << std::endl;
  std::cout << " 9    2.26     3.25     4.78" << std::endl;
  std::cout << " 10    2.23     3.17     4.59" << std::endl;
  std::cout << " 15    2.13     2.95     4.07" << std::endl;
  std::cout << " 20    2.09     2.85     3.85" << std::endl;
  std::cout << " 30    2.04     2.75     3.65" << std::endl;
  std::cout << " 50    2.01     2.68     3.50" << std::endl;
  std::cout << " 100    1.98     2.63     3.39  " << std::endl;

  if( argc < 6 )
    {
    std::cout << "Usage: " << argv[0] <<  " ImageDimension  OutName NGroup1 NGroup2 ControlV1*   SubjectV1*   "
             << std::endl;
    std::cout << " Assume all images the same size " << std::endl;
    std::cout << " Writes out an F-Statistic image " << std::endl;
    std::cout <<  " \n example call \n  \n ";
    std::cout << argv[0] << "  2  TEST.nii.gz 4 8 FawtJandADCcon/*SUB.nii  FawtJandADCsub/*SUB.nii  \n ";
    return 1;
    }

  switch( atoi(argv[1]) )
    {
    case 2:
      {
      StudentsTestOnImages<2>(argc, argv);
      }
      break;
    case 3:
      {
      StudentsTestOnImages<3>(argc, argv);
      }
      break;
    default:
      std::cout << "Unsupported dimension" << std::endl;
      return EXIT_FAILURE;
    }

  return 0;
}
} // namespace ants
