#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 = Block<uInt>( nfield_m, 1 ) ;
  idx_m = Block<uInt>( nfield_m, 0 ) ;
  for ( uInt i = 1 ; i < nfield_m ; i++ ) {
    prod_m[i] = shape[i-1] * prod_m[i-1] ;
  }
  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( IPosition(1,0) )
{
  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() ;
  }
  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 ) ;
  for ( uInt i = 0 ; i < nrow_m ; i++ ) {
    storage_m[i+arr_m.nelements()] = i ;
  }
  len_m.resize( ncol_m+1 ) ;
  for ( uInt i = 0 ; i < ncol_m+1 ; i++ ) {
    len_m[i] = nrow_m ;
//     cout << "len[" << i << "]=" << len_m[i] << endl ;
  }
  prev_m = current() - (uInt)1 ;
}

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

Vector<uInt> ArrayIndexIteratorAcc::getRows( StorageInitPolicy policy )
{
  Vector<uInt> v = current() ;
  Int c = isChanged( v ) ;
//   cout << "c=" << c << endl ;
  if ( c < 0 ) {
    pos_m[0] = len_m[0] ;
    return Vector<uInt>( pos_m, storage_m.storage(), 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, 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, storage_m.storage(), policy ) ;
}

Int ArrayIndexIteratorAcc::isChanged( Vector<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 len = 0 ;
  uInt *p = base - nrow_m ;
  uInt *work = p ;
  for ( uInt i = 0 ; i < len_m[icol+1] ; i++ ) {
    if ( arr_m( *base, icol ) == v ) {
//       cout << "add " << *base << endl ;
      *work = *base ;
      len++ ;
      work++ ;
    }
    base++ ;
  }
  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 ;
  for ( uInt i = 0 ; i < ncol ; i++ ) {
    col.attach( t, cols[i] ) ;
    Vector<uInt> v( 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() ;
  Matrix<uInt> arr( nrow, ncol ) ;
  ROScalarColumn<uInt> col ;
  for ( uInt i = 0 ; i < ncol ; i++ ) {
    col.attach( t, cols[i] ) ;
    Vector<uInt> v( arr.column( i ) ) ;
    col.getColumn( v ) ;
  }
  iter_m = new ArrayIndexIteratorAcc( arr ) ;
}

} // namespace