#include <iostream>
#include <casa/Utilities/GenSort.h>
#include <casa/Arrays/Vector.h>
#include <casa/Arrays/Matrix.h>
#include <casa/Arrays/ArrayMath.h>
#include <casa/Arrays/ArrayIO.h>
#include <tables/Tables/ScalarColumn.h>
#include "STIdxIter.h"

namespace asap {

IndexIterator::IndexIterator( IPosition &shape ) 
  : niter_m( 0 )
{
  nfield_m = shape.nelements() ;
  prod_m.resize( nfield_m ) ;
  idx_m.resize( nfield_m ) ;
  prod_m[0] = 1 ;
  idx_m[0] = 0 ;
  for ( uInt i = 1 ; i < nfield_m ; i++ ) {
    prod_m[i] = shape[i-1] * prod_m[i-1] ;
    idx_m[i] = 0 ;
  }
  maxiter_m = prod_m[nfield_m-1] * shape[nfield_m-1] ;
//   cout << "maxiter_m=" << maxiter_m << endl ;
}

Bool IndexIterator::pastEnd() 
{
  return niter_m >= maxiter_m ;
}

void IndexIterator::next()
{
  niter_m++ ;
  uInt x = niter_m ;
  for ( Int i = nfield_m-1 ; i >= 0 ; i-- ) {
    idx_m[i] = x / prod_m[i] ;
    //cout << "i=" << i << ": prod=" << prod_m[i] 
    //     << ", idx=" << idx_m[i] << endl ;
    x %= prod_m[i] ;
  }
}

ArrayIndexIterator::ArrayIndexIterator( Matrix<uInt> &arr,
                                        vector< vector<uInt> > idlist )
  : arr_m( arr ),
    pos_m( 1 )
{
  ncol_m = arr_m.ncolumn() ;
  nrow_m = arr_m.nrow() ;
  vector< vector<uInt> > l = idlist ;
  if ( l.size() != ncol_m ) {
    l.resize( ncol_m ) ;
    for ( uInt i = 0 ; i < ncol_m ; i++ ) {
      Vector<uInt> a( arr_m.column( i ).copy() ) ;
      uInt n = genSort( a, Sort::Ascending, Sort::QuickSort|Sort::NoDuplicates ) ;
      a.resize(n,True) ;
      a.tovector( l[i] ) ;
    }
  }
  idxlist_m = l ;
  IPosition shape( ncol_m ) ;
  for ( uInt i = 0 ; i < ncol_m ; i++ ) {
    shape[i] = idxlist_m[i].size() ;
  }
//   cout << "shape=" << shape << endl ;
  iter_m = new IndexIterator( shape ) ;
  current_m.resize( ncol_m ) ;
}

ArrayIndexIterator::~ArrayIndexIterator()
{
  delete iter_m ;
}

Vector<uInt> ArrayIndexIterator::current() 
{
  Block<uInt> idx = iter_m->current() ;
  for ( uInt i = 0 ; i < ncol_m ; i++ )
    current_m[i] = idxlist_m[i][idx[i]] ;
  return current_m ;
}

Bool ArrayIndexIterator::pastEnd()
{
  return iter_m->pastEnd() ;
}

ArrayIndexIteratorNormal::ArrayIndexIteratorNormal( Matrix<uInt> &arr,
                                                    vector< vector<uInt> > idlist )
  : ArrayIndexIterator( arr, idlist )
{
  storage_m.resize( nrow_m ) ;
}

void ArrayIndexIteratorNormal::next() 
{
  iter_m->next() ;
}

Vector<uInt> ArrayIndexIteratorNormal::getRows( StorageInitPolicy policy )
{
  Vector<uInt> v = current() ;
  uInt len = 0 ;
  uInt *p = storage_m.storage() ;
  for ( uInt i = 0 ; i < nrow_m ; i++ ) {
    if ( allEQ( v, arr_m.row( i ) ) ) {
      *p = i ;
      len++ ;
      p++ ;
    }
  }
  pos_m[0] = len ;
  p = storage_m.storage() ;
  return Vector<uInt>( pos_m, p, policy ) ;
}

ArrayIndexIteratorAcc::ArrayIndexIteratorAcc( Matrix<uInt> &arr,
                                              vector< vector<uInt> > idlist )
  : ArrayIndexIterator( arr, idlist )
{
  storage_m.resize( arr_m.nelements()+nrow_m ) ;
  uInt *p = storage_m.storage() + arr_m.nelements() ;
  for ( uInt i = 0 ; i < nrow_m ; i++ ) {
    *p = i ;
    p++ ;
  }
  len_m.resize( ncol_m+1 ) ;
  p = len_m.storage() ;
  for ( uInt i = 0 ; i < ncol_m+1 ; i++ ) {
    *p = nrow_m ;
    p++ ;
//     cout << "len[" << i << "]=" << len_m[i] << endl ;
  }
  prev_m = iter_m->current() ;
  p = prev_m.storage() ;
  for ( uInt i = 0 ; i < ncol_m ; i++ ) {
    *p = *p - 1 ;
    p++ ;
  }
//   cout << "prev_m=" << Vector<uInt>(IPosition(1,ncol_m),prev_m.storage()) << endl ;
  skip_m.resize( ncol_m ) ;
  for ( uInt i = 0 ; i < ncol_m ; i++ ) {
    skip_m[i] = (Bool)(idxlist_m[i].size()==1) ;
//     cout << "skip_m[" << i << "]=" << skip_m[i] << endl ;
  }
}

void ArrayIndexIteratorAcc::next() 
{
  prev_m = iter_m->current() ;
  iter_m->next() ;
}

Vector<uInt> ArrayIndexIteratorAcc::getRows( StorageInitPolicy policy )
{
  Block<uInt> v = iter_m->current() ;
  Int c = isChanged( v ) ;
//   cout << "c=" << c << endl ;
  if ( c < 0 ) {
    pos_m[0] = len_m[0] ;
    uInt offset = 0 ;
    while( offset <= ncol_m && !(skip_m[offset]) )
      offset++ ;
//     cout << "offset=" << offset << endl ;
    return Vector<uInt>( pos_m, storage_m.storage()+offset*nrow_m, policy ) ;
  }
  uInt *base = storage_m.storage() + (c+1) * nrow_m ;
//   cout << "len_m[c+1]=" << len_m[c+1] << endl ;
//   cout << "base=" << Vector<uInt>(IPosition(1,len_m[c+1]),base,SHARE) << endl ;
  for ( Int i = c ; i >= 0 ; i-- ) {
    base = updateStorage( i, base, idxlist_m[i][v[i]] ) ;
//     cout << "len_m[" << i << "]=" << len_m[i] << endl ;
//     cout << "base=" << Vector<uInt>(IPosition(1,len_m[i]),base,SHARE) << endl ;
  }
  pos_m[0] = len_m[0] ;
  return Vector<uInt>( pos_m, base, policy ) ;
}

Int ArrayIndexIteratorAcc::isChanged( Block<uInt> &idx )
{
  Int i = ncol_m - 1 ;
  while( i >= 0 && idx[i] == prev_m[i] )
    i-- ;
  return i ;
}

uInt *ArrayIndexIteratorAcc::updateStorage( Int &icol,
                                            uInt *base,
                                            uInt &v )
{
  uInt *p ;
  if ( skip_m[icol] ) {
    // skip update, just update len_m[icol] and pass appropriate pointer
//     cout << "skip " << icol << endl ;
    p = base ;
    len_m[icol] = len_m[icol+1] ;
  }
  else {
    uInt len = 0 ;
    p = storage_m.storage() + icol * nrow_m ;
    uInt *work = p ;
    Bool b ;
    const uInt *arr_p = arr_m.getStorage( b ) ;
    long offset = 0 ;
    // warr_p points a first element of (icol)-th column
    const uInt *warr_p = arr_p + icol * nrow_m ;
    for ( uInt i = 0 ; i < len_m[icol+1] ; i++ ) {
      // increment warr_p by (*(base)-*(base-1))
      warr_p += *base - offset ;
      // check if target element is equal to value specified
      if ( *warr_p == v ) {
        // then, add current index to storage_m
        //       cout << "add " << *base << endl ;
        *work = *base ;
        len++ ;
        work++ ;
      }
      // update offset 
      offset = *base ;
      // next index
      base++ ;
    }
    arr_m.freeStorage( arr_p, b ) ;
    len_m[icol] = len ;
  }
  return p ;
}

STIdxIter::STIdxIter()
{
  iter_m = 0 ;
}

STIdxIter::STIdxIter( const string &name,
                      const vector<string> &cols ) 
{
}

STIdxIter::STIdxIter( const CountedPtr<Scantable> &s,
                      const vector<string> &cols )
{
}

STIdxIter::~STIdxIter()
{
  if ( iter_m != 0 )
    delete iter_m ;
}

vector<uInt> STIdxIter::tovector( Vector<uInt> v )
{
  vector<uInt> ret ;
  v.tovector( ret ) ;
  return ret ;
}

Vector<uInt> STIdxIter::getRows( StorageInitPolicy policy )
{ 
  return iter_m->getRows( policy ) ; 
} 

STIdxIterNormal::STIdxIterNormal()
  : STIdxIter()
{ 
}

STIdxIterNormal::STIdxIterNormal( const string &name,
                                  const vector<string> &cols ) 
  : STIdxIter( name, cols )
{
  Table t( name, Table::Old ) ;
  init( t, cols ) ;
}

STIdxIterNormal::STIdxIterNormal( const CountedPtr<Scantable> &s,
                                  const vector<string> &cols )
  : STIdxIter( s, cols )
{
  init( s->table(), cols ) ;
}

STIdxIterNormal::~STIdxIterNormal()
{
}

void STIdxIterNormal::init( Table &t, 
                            const vector<string> &cols )
{
  uInt ncol = cols.size() ;
  uInt nrow = t.nrow() ;
  Matrix<uInt> arr( nrow, ncol ) ;
  ROScalarColumn<uInt> col ;
  Vector<uInt> v ;
  for ( uInt i = 0 ; i < ncol ; i++ ) {
    col.attach( t, cols[i] ) ;
    v.reference( arr.column( i ) ) ;
    col.getColumn( v ) ;
  }
  iter_m = new ArrayIndexIteratorNormal( arr ) ;
}

STIdxIterAcc::STIdxIterAcc()
  : STIdxIter()
{ 
}

STIdxIterAcc::STIdxIterAcc( const string &name,
                            const vector<string> &cols ) 
  : STIdxIter( name, cols )
{
  Table t( name, Table::Old ) ;
  init( t, cols ) ;
}

STIdxIterAcc::STIdxIterAcc( const CountedPtr<Scantable> &s,
                            const vector<string> &cols )
  : STIdxIter( s, cols )
{
  init( s->table(), cols ) ;
}

STIdxIterAcc::~STIdxIterAcc()
{
}

void STIdxIterAcc::init( Table &t, 
                         const vector<string> &cols )
{
  uInt ncol = cols.size() ;
  uInt nrow = t.nrow() ;
  // array shape here is as follows if cols=["BEAMNO","POLNO","IFNO"]:
  // [[B0,B1,B2,...,BN],
  //  [P0,P1,P2,...,PN],
  //  [I0,I1,I2,...,IN]]
  // order of internal storage is
  // [B0,B1,B2,..,BN,P0,P1,P2,...,PN,I0,I1,I2,...,IN] 
  Matrix<uInt> arr( nrow, ncol ) ;
  Vector<uInt> v ;
  ROScalarColumn<uInt> col ;
  for ( uInt i = 0 ; i < ncol ; i++ ) {
    col.attach( t, cols[i] ) ;
    v.reference( arr.column( i ) ) ;
    col.getColumn( v ) ;
  }
  iter_m = new ArrayIndexIteratorAcc( arr ) ;
}

} // namespace