module UnibenchInput;

use Structures: all except {transpose};
use StdIO: all;
use SysErr: all;

export {getDataCube, getLabels};

string FibreScanString( File &fd)
{
  /* Skip rest of line */
  skip =  fscans( fd, 1000);
  str = fscans( fd, 1000);
  return( str);
}

double[*], bool[*] getDataCube( int[.] dimensions)
{
  data, mask, innerDimLabels, indexLabels, numInnerDims = readSparseIntoDenseDouble( "data.out");
  printf( "numInnerDims = %d", (:int)numInnerDims);
  printf( "[ ");
  for( it=0 ; it<shape( innerDimLabels)[[0]] ; it++) {
    printf( " %s ", to_string( StringArray::sel( [it], innerDimLabels)));
  }
  printf( " ]\n");


  printf("v\n");
  /*
   * Add dimensions from data returned from observation
   * script to list of dimensions to keep
   */
  numOuterDims = dim( data) - numInnerDims;
  inner_dimensions = drop( [numOuterDims], iota( dim( data)));
  data_dimensions = dimensions ++ inner_dimensions;

  occ = contained( iota( dim(data)), data_dimensions);
  if( sum( occ) == shape( data_dimensions)[0] ) {
    if( all( occ == 1)) { 
      /* we have selected the entire cube */
      data = transpose( data_dimensions, data);
      mask = transpose( dimensions, mask);
    } else if( all( occ <= 1)) {
      /* remove some axes */
      data = transpose( data_dimensions, data);
      mask = transpose( dimensions, mask);
      data, mask = collapse( sum( toi( occ < 1)), data, mask);
    } else {
      RuntimeError::error( 1,
        "duplicate axis references in \"getDataCube\" call.");
    }
  } else {
    RuntimeError::error( 1,
        "non existing axis referenced in \"getDataCube\" call.");
  }
  return( data, mask);
}
   
stringArray getLabels( int[.] dimensions) 
{
  labels, mask = readSparseIntoDenseString( "labels.out");
  occ = contained( iota( shape(mask)[0]), dimensions);
  if( sum( occ) == shape( dimensions)[0] ) {
    if( all( occ <= 1)) {
      labels, mask = filter( labels, mask, dimensions);
    } else {
      RuntimeError::error( 1,
        "duplicate axis references in \"getLabels\" call.");
    }
  } else {
    RuntimeError::error( 1,
        "non existing axis referenced in \"getLabels\" call.");
  }

  return( labels);
}

/**
 * helper functions that might be better placed in a more general module:
 */

int[*] contained( int[*] val, int[*] a)
{
  return( { iv -> sum( toi( val[iv] == a))});
}

int[.] perm( int[.] vect, int[.] pv)
{
  res = with {
          ( . <= iv <= .) : vect[ pv[iv]];
        } : genarray( shape(vect), 0);
  return( res);
}

int[.] invPerm( int[.] pv)
{
  res = 0*pv;
  for( i=0; i<shape(pv)[0]; i++) {
    res[pv[i]] = i;
  }
  return( res);
}

double[*], bool[*] collapse( int dims, double[*] a, bool[*] mask)
{
  shp_data = take( [dim(a) - dims], shape( a));
  skip = dim(a) - dim(mask);
  res_a = with {
           ( . <= iv <= .) {
             hyper_a = a[iv];
             hyper_mask = mask[drop([-skip], iv)];
             if( sum( toi(hyper_mask)) <= 1) {
               val = sum( where( hyper_mask, hyper_a, 0d));
             } else {
               RuntimeError::error( 1, "cannot collapse axes!");
               val = 0d;
             }
           } : val;
         } : genarray( shp_data, 0d);

  shp_mask = take( [dim(mask) - dims], shape( mask));
  res_mask = with {
              ( . <= iv <= .) : any( mask[iv]);
            } : genarray( shp_mask, false);

  return(res_a, res_mask);
}

double[*] transpose( int[.] dims, double[*] a)
{
  occ = contained( iota( dim(a)), dims);
  pv_add = genarray( [sum( toi( occ==0))], 0);
  k = 0;
  for( i=0; i< dim(a); i++) {
    if( occ[i] == 0) {
      pv_add[k] = i;
      k++;
    }
  }
  pv = dims ++pv_add;

  res = with {
         ( . <= iv <= .) : a[ perm(iv, invPerm(pv))];
       } : genarray( perm( shape(a), pv), 0d);
  return(res);
}

bool[*] transpose( int[.] dims, bool[*] a)
{
  occ = contained( iota( dim(a)), dims);
  pv_add = genarray( [sum( toi( occ==0))], 0);
  k = 0;
  for( i=0; i< dim(a); i++) {
    if( occ[i] == 0) {
      pv_add[k] = i;
      k++;
    }
  }
  pv = dims ++pv_add;

  res = with {
         ( . <= iv <= .) : a[ perm(iv, invPerm(pv))];
       } : genarray( perm( shape(a), pv), false);
  return(res);
}


stringArray, bool[.,.] filter( stringArray arr, bool[.,.] mask, int[.] dimensions)
{
  /*
   * we assume here, that all entries in d_i in dimensions
   * we have 0 <= d_i < shape(arr)[0] 
   *         and (d_i == d_j) => (i == j)
   */
  new_shp = shape(dimensions) ++ drop( [1], shape( mask));
  res = genarray( new_shp, to_stringArray( ""));
  for( i=0; i<new_shp[0]; i++) {
    for( j=0; j<new_shp[1]; j++) {
      res[[i,j]] = arr[[dimensions[i], j]];
    }
  }

  mask = with {
           ([0] <= [i] < shape(dimensions)) : mask[dimensions[i]];
         } : genarray( shape(dimensions), genarray( drop( [1], shape( mask)),
                                                    false));

  return( res, mask);
}

/**
 * local helper functions:
 */

int[.] readIntVect( int dim, File &fd)
{
  int extent;

  shp = genarray( [dim], 0);
  for( i=0; i<dim; i++) {
    num_found, extent = fscanf( fd, "%d\n");
    if( num_found != 1 ){
      RuntimeError::error( 1, "failed to read integer from input file!");
    }
    shp[i] = extent;
  }
  return( shp);
}

double[*] readDoubleArraySingleLine( File &fd)
{
  double value;
  /**
   * TODO: currently only reads scalars; should read double[*]!
   */
  num_found, value = fscanf( fd, "%lf\n");
  if( num_found != 1 ){
    RuntimeError::error( 1,
      "failed to read scalar double value from input file!");
  }
  return( value);
}

string readString( File &fd)
{
  str =  fscans( fd, 1000) ;
  num = fscanf( fd, "\n");
  return( str);
}



/**
 * generic sparse array reading functions:
 */

double[*], bool[*], stringArray, stringArray, int readSparseIntoDenseDouble( string file)
{
  err, fd = fopen( file, "r");
  if( fail( err)) {
    RuntimeError::error( 1, "could not open file \"%s\": %s.",
                            file, strerror( err));
  }

  dataRows = (:int) FibreScanIntArray( fd);

  /**
   * read shape of cube:
   */
  shp = FibreScanIntArray( fd);

  innerDimLabels = FibreScanStringArray( fd);
  numIndexedInnerDims = shape( innerDimLabels)[0];



  /* Pass 1 needs to start here. Read all data and count the number of indicies
   * for each inner dimension */

  dimLabels = to_stringArray( "");
  innerDimCounts = genarray( [numIndexedInnerDims], 0);

  /* This is realy shit and hard-coded but until StringArray is
   * fixed I am desparate
   */
  innerDimSeen = genarray( [numIndexedInnerDims, 100], to_stringArray( ""));

  printf("a\n");
  innerDimBlankSeen = genarray( [numIndexedInnerDims], false);
  unlabelledDims = [:int];
  for( x=0 ; x < dataRows ; x++) { 
    printf("aa\n");
    idx = FibreScanIntArray( fd);
    result = FibreScanDoubleArray( fd);
    mask = FibreScanIntArray( fd);
    dimLabels = FibreScanStringArray( fd);
    for( i = 0 ; i < shape( innerDimLabels)[[0]] ; i++) {
      if( ( ! in( to_string( innerDimLabels[[i]]), dimLabels))) {
        innerDimBlankSeen[[i]] = true;
      }
    }
    numLabels = (:int) FibreScanIntArray( fd);
    for( i=0 ; i<numLabels ; i++) {
      labels = FibreScanStringArray( fd);
      indexIntoInnerDimLabels = 
        indexOf( to_string( StringArray::sel( [i],
                 dimLabels)), innerDimLabels);
      for( j=0 ; j < shape( labels)[[0]] ; j++) {
         /* Keep count of new index labels for dimenson 
          * innerDimLabels[ indexIntoInnerDimLabels] in 
          * innerDimCounts[ indexIntoInnerDimLabels]
          */
        newLabel = labels[[j]];
        if( ! in( to_string( newLabel), innerDimSeen[ indexIntoInnerDimLabels])) {
          innerDimSeen = 
            StringArray::modarray( 
              innerDimSeen,  
              indexIntoInnerDimLabels ++ [innerDimCounts[ indexIntoInnerDimLabels]],
              newLabel);
          innerDimCounts[ indexIntoInnerDimLabels] = 
            innerDimCounts[ indexIntoInnerDimLabels] + 1;
        }
      }
    }
    if( numLabels ==0) {
      resShape = shape( result);
      if( shape( resShape)[[0]] > shape( unlabelledDims)[[0]]) {
        unlabelledDims = with {
          ( [0] <= iv <= [shape( unlabelledDims)[[0]]]) : 
            (resShape[iv] > unlabelledDims[iv] ? resShape[iv] : unlabelledDims[iv]);
                         } : modarray( resShape);
      }
      else {
        unlabelledDims = with {
          ( [0] <= iv <= [shape( resShape)[[0]]]) : 
            (resShape[iv] > unlabelledDims[iv] ? resShape[iv] : unlabelledDims[iv]);
                         } : modarray( unlabelledDims);
      }
    }
  }
  fclose( fd);

  printf("a\n");
  for( i = 0 ; i < shape( innerDimLabels)[[0]] ; i++) {
    if( (shape( unlabelledDims)[[0]] != 0) || ( innerDimBlankSeen[ i])) {
      innerDimSeen = StringArray::modarray( 
          innerDimSeen, 
          [i] ++ [innerDimCounts[ i]],
          to_stringArray( "BLANK"));
      innerDimCounts[ i] = innerDimCounts[ i] + 1;
    }
  }

  oldInnerDimLabels = innerDimLabels;
  innerDimLabels = 
    genarray( [shape(innerDimLabels)[[0]] + shape(unlabelledDims)[[0]]],
              to_stringArray( ""));
  for( k=0 ; k < shape( oldInnerDimLabels)[[0]] ; k++) {
    innerDimLabels = StringArray::modarray( innerDimLabels, [k], 
        StringArray::sel( [k], oldInnerDimLabels));
  }
  for( l=0 ; l < shape( unlabelledDims)[[0]] ; l++) {
    innerDimLabels = StringArray::modarray( innerDimLabels, [k+l], 
        to_stringArray( sprintf( "Unamed Dimension %d", l)));
  }

  /* TODO: if any results have no dim labels then each given label
   * needs an extra unknown entry
   */

  innerDimCounts = innerDimCounts ++ unlabelledDims;

  printf("a\n");
  /* debugging start */
  /*
  indicies = getIndicies( shape( innerDimLabels));
  printf( "[ ");
  for( i = 0 ; i < shape( indicies)[0] ; i++) {
    printf( "%s ", to_string( innerDimLabels[ indicies[[i]]]));
  }
  printf( "]\n");
  indicies = getIndicies( shape( innerDimCounts));
  printf( "[ ");
  for( i = 0 ; i < shape( indicies)[0] ; i++) {
    print( innerDimCounts[ indicies[[i]]]);
  }
  printf( "]\n");
  */
  /* debugging end */



  /* Pass 2: Insert elements into datacube */
  err, fd = fopen( file, "r");
  if( fail( err)) {
    RuntimeError::error( 1, "could not open file \"%s\": %s.",
                            file, strerror( err));
  }
  
  dataRows = (:int) FibreScanIntArray( fd);
  shp = FibreScanIntArray( fd);

  rubbish = FibreScanStringArray( fd);
  numInnerDims = shape( innerDimLabels)[0];

  if( prod( shp ++ innerDimCounts) > 10000000) {
    RuntimeError::error( 1, "This cube has a size bigger than 10.000.000!"
      "Please select less dimensions or narrow down the query and try again.");
  }

  data = genarray( shp ++ innerDimCounts, 0d);
  mask = genarray( shp ++ innerDimCounts, false);

  printf("a\n");
  for( x=0 ; x < dataRows ; x++)  { 
    printf( "x=%d\n", (:int)x);
    idx = FibreScanIntArray( fd);
    if( dim( idx) != 1) {
      RuntimeError::error( 1, "index for data is not a vector!");
    }
    if( any( idx <0) || any( idx >shp)) {
      RuntimeError::error( 1, "index for data exceeds the specified shape!");
    }
    result = FibreScanDoubleArray( fd);
    /* we'll ignore the new mask for now */
    newMask = FibreScanDoubleArray( fd);
    newDimLabels = FibreScanStringArray( fd);
    numLabels = (:int) FibreScanIntArray( fd);
    labelArray = genarray( [ numIndexedInnerDims, maxInt( innerDimCounts)], 
                           to_stringArray(""));
    if( numLabels != 0) {
      for( i=0 ; i<numLabels ; i++) {
        labels = FibreScanStringArray( fd);
        for( li=0 ; li < shape(labels)[[0]] ; li++) {
          labelArray = 
            StringArray::modarray( labelArray, 
                                   [i,li],
                                   StringArray::sel( [li], labels));
        }
      }
      /* Iterate over the array just read in from the file */
      dimIndex = genarray( [numIndexedInnerDims], [0]);
      dimInUse = genarray( [numIndexedInnerDims], false);
      for( d=0 ; d <numIndexedInnerDims ; d++) {
        dimIndex_, dimInUse_ = 
          indexOf2( to_string( StringArray::sel( [d], innerDimLabels)), 
                    newDimLabels);
        dimIndex[[ d]] = dimIndex_;
        dimInUse[[ d]] = dimInUse_;
      }
      indicies = getIndicies( shape( result));
      for( i=0 ; i<shape(indicies)[[0]] ; i++) {
        /* calculate index into inner array */
        indexIntoInnerDim = genarray( [numIndexedInnerDims], 0);
        /* for each inner dimension in our result array */
        for( d=0 ; d < numIndexedInnerDims ; d++) {
          if( dimInUse[[d]]) {
            newDimString = to_string( 
                StringArray::sel( dimIndex[[d]] ++ [ indicies[[ i,d]]], 
                                  labelArray));
            indexIntoInnerDim[[d]] = 
              indexOfUpto( newDimString, 
                           innerDimSeen[[ d]], 
                           innerDimCounts[[d]])[[0]];
          }
          else {
            indexToUse = innerDimCounts[[ d]] - 1; /* last place */
            indexIntoInnerDim[[d]] = indexToUse;
            /*printf( "dim %d missing using index %d\n", d, indexToUse);*/
          }
        }

        /* place there */
        data[ idx ++ indexIntoInnerDim] = result[ indicies[[i]]];
        mask[ idx ++ indexIntoInnerDim] = true;
      }
    }
    else {
      innerTop = with {
        ( . <= [x] <= .) : innerDimCounts[x];
      } : genarray( [numIndexedInnerDims], 0);
     
      /* Where the dimensionallity of the array being inserted into
       * and the result match and the shape of the result is less 
       * than the shape of the array being insertedinto
       * this is equivenant to a common modarray call.
       * Where the dimensionslity of result is smaller this should replicated 
       * among the other dimensions.
       */
      /*
      printf("before insert data is\n");
      print( data);
      */
      if( dim( result) == 0) {
        insertdata = result;
      }
      else {
        insertdata = tod( 0 * shape( data[ idx ++ innerTop])) + result;
      }
      data[ idx ++ innerTop ] = insertdata;
      data = with {
        ( idx ++ (0 * innerTop) <= iv <= idx ++ innerTop) : insertdata;
             } : modarray( data);
      printf( "inserting into \n");
      print( idx++innerTop);
      /*
      printf( "with\n"); 
      print( insertdata);
      */
      mask[ idx ++ innerTop ] = genarray( shape( mask[ idx ++ innerTop ]) , true);
      mask = with {
        ( idx ++ (0 * innerTop) <= iv <= idx ++ innerTop) : 
          genarray( shape( mask[ iv]) , true);
             } : modarray( mask);

    }
  }

  fclose( fd);

  return( data, mask, innerDimLabels, innerDimSeen, numInnerDims);
}




stringArray, bool[*] readSparseIntoDenseString( string file)
{
  int elements;
  int in;
  stringArray first_result;
  stringArray result;
  err, fd = fopen( file, "r");
  if( fail( err)) {
    RuntimeError::error( 1, "could not open file \"%s\": %s.",
                            file, strerror( err));
  }

  count, elements =  fscanf( fd, "%d\n");

  /**
   * read shape of cube:
   */
  shp = FibreScanIntArray( fd);
  
  if( prod( shp) > 10000000) {
    RuntimeError::error( 1, "This cube has a size bigger than 10.000.000!"
      "Please select less dimensions or narrow down the query and try again.");
  }

  idx = FibreScanIntArray( fd);
  
  if( dim( idx) != 1) {
    RuntimeError::error( 1, "index for data is not a vector!");
  }
  if( any( idx <0) || any( idx >shp)) {
    RuntimeError::error( 1, "index for data exceeds the specified shape!");
  }
  first_result = FibreScanStringArray( fd);

  data = genarray( shp, first_result);
  mask = genarray( shp, false);
  mask[idx] = true;

  for( i=1 ; i < elements ; i++) {
    idx = FibreScanIntArray( fd);
    if( dim( idx) != 1) {
      RuntimeError::error( 1, "index for data is not a vector!");
    }
    if( any( idx <0) || any( idx >shp)) {
      RuntimeError::error( 1, "index for data exceeds the specified shape!");
    }
    result = FibreScanStringArray( fd);
    data = StringArray::modarray( data, idx, result);
    mask[idx] = true;
  }
  fclose( fd);

  return( data, mask);
}

inline
int[.,.] getIndicies( int[.] shp)
{
    iota_shp = Array::iota( shp);
    return(
       Array::reshape( [ Array::prod( Array::take( [ Array::dim( Array::iota( shp)) Array::- 1 ], 
                                                   Array::shape( iota_shp))),
                         Array::sel( [0], Array::reverse( Array::shape( shp))) ],
                       iota_shp));
}

inline
bool in( string needle, stringArray haystack)
{
  res = false;
  indicies = getIndicies( shape( haystack));
  for( i = 0 ; (i < shape( indicies)[0]) && (res == false) ; i++) {
    if( StringArray::to_string( 
          StringArray::sel( indicies[ i], haystack)) == needle) {
      res = true;
    }
  }
  return( res);
}


inline
int[.] indexOfUpto( string value, stringArray array, int max)
{
  indicies = getIndicies( shape( array));
  res = [];
  for( i = 0 ; (i < shape( indicies)[0]) && (i < max) ; i++) {
    if( StringArray::to_string( StringArray::sel( indicies[ i], array)) == value) {
      res = indicies[ i];
    }
  }
  return( res);
}

inline
int[.], bool indexOf2( string value, stringArray array)
{
  indicies = getIndicies( shape( array));
  res = [];
  for( i = 0 ; (i < shape( indicies)[0])  ; i++) {
    if( StringArray::to_string( StringArray::sel( indicies[ i], array)) == value) {
      res = indicies[ i];
    }
  }
  return( res, dim(res)!=0);
}

inline
int[.] indexOf( string value, stringArray array)
{
  indicies = getIndicies( shape( array));
  res = [];
  for( i = 0 ; (i < shape( indicies)[0])  ; i++) {
    if( StringArray::to_string( StringArray::sel( indicies[ i], array)) == value) {
      res = indicies[ i];
    }
  }
  return( res);
}



inline
int maxInt( int[.] vect)
{
  res = vect[[0]];
  for( i=1; i< shape( vect)[[0]] ; i++) {
    if( vect[[i]] > res) {
      res = vect[[i]];
    }
  }
  return( res);
}