#include <assert.h>
#include <algorithm>
#include <utility>
#include <iterator>
#include <blitz/array.h>
#include <random/uniform.h>

// compile and link with g++ -ggdb3 blitzsort.cpp -o blitzsort -lblitz

using namespace std;
using namespace blitz;
using namespace ranlib;

class randuni
{
public:
    randuni()  
    {
    }
    
    double operator()(double v)
    {
        return u.random();
    }
    BZ_DECLARE_FUNCTOR(randuni);
    Uniform<double> u;
};

struct mypair
{
    double v;
    int    idx;
};

class mycomp : public binary_function<mypair, mypair, bool>
{
    bool operator()(mypair &a, mypair &b)
    {
        return a.v < b.v;
    }
};

template <class RandomAccessIterator>
class strided_iterator 
{
  typedef strided_iterator self;
  
public:
  
  typedef typename std::iterator_traits<RandomAccessIterator>::difference_type   difference_type; 
  typedef typename std::iterator_traits<RandomAccessIterator>::value_type        value_type;
  typedef typename std::iterator_traits<RandomAccessIterator>::iterator_category iterator_category;
  typedef typename std::iterator_traits<RandomAccessIterator>::reference         reference;
  typedef typename std::iterator_traits<RandomAccessIterator>::pointer           pointer;

  typedef difference_type      Distance;
  typedef RandomAccessIterator iterator_type;

  inline strided_iterator() : stride(0), pos(0) { }

  inline strided_iterator(const RandomAccessIterator& x, int s, int p) 
    : iter(x), stride(s), pos(p) { }

  inline strided_iterator(const self& x)
    : iter(x.iter), stride(x.stride), pos(x.pos) { }

  inline self& operator=(const self& x) 
  {
    iter = x.iter; stride = x.stride; pos = x.pos; return *this;
  }

  inline int index() const { return pos; }
  
  inline operator RandomAccessIterator () const { return iter; }

  inline RandomAccessIterator base() const { return iter; }

  inline reference operator*() const { return *iter; }

  inline self& operator++ ()
  {
    ++pos; iter += stride; return *this;
  }

  inline self operator++ (int)
  { 
    self tmp = *this; ++pos; iter += stride; return tmp; 
  }

  inline self& operator-- ()
  { 
    --pos; iter -= stride; return *this; 
  }
  
  inline self operator-- (int)
  { 
    self tmp = *this; --pos; iter -= stride; return tmp; 
  }
 
  inline self operator+ (Distance n) const 
  {
    return self (iter + n*stride, stride, pos + n);
  }
  
  inline self& operator+= (Distance n) 
  { 
    iter += n*stride; pos += n; return *this; 
  }
  
  inline self operator- (Distance n) const 
  {
    return self (iter - n*stride, stride, pos - n);
  }
  
  inline self& operator-= (Distance n) 
  {
    iter -= n*stride; pos -= n; return *this; 
  }

  inline self operator+ (const self& x) const 
  {
    return self(iter + x.iter, stride, pos + x.pos); 
  }
  
  inline Distance operator- (const self& x) const { return pos - x.pos; }
  
  inline reference operator[] (Distance n) const { return *(*this+n); }
 
  inline bool operator==(const self& x) const { return pos == x.pos; }
  
  inline bool operator!=(const self& x) const { return pos != x.pos; }
  
  inline bool operator<(const self& x) const { return pos < x.pos; }  

protected:
  
  RandomAccessIterator iter;
  
  int stride;
  int pos;
};

template<typename T, int rank_N>
void partial_sort1(Array<T, rank_N> &array, int num)
{
    TinyVector<int, rank_N> lower(0);
    TinyVector<int, rank_N> upper = array.ubound();
    upper[rank_N - 1] = 0;
    
    RectDomain<rank_N> domain(lower, upper);
    Array<T, rank_N> top = array(domain); 
       
    typename Array<T, rank_N>::iterator it = top.begin();
    typename Array<T, rank_N>::iterator et = top.end();
    
    // only handles arrays with last dimension stored with stride 1
    assert(array.stride(rank_N-1) == 1);
    
    for(; it != et; ++it)
    {
        T *start = &(*it);
        T *end = start + array.extent(rank_N - 1);
        
        make_heap(start, end);
        for (int j = 0; j < num; j++)
        {
            pop_heap(start, end - j);
        }
    }
}

template<typename T, int rank_N>
void partial_sort(Array<T, rank_N> &array, int sortDim, int num)
{
    assert(0 <= sortDim && sortDim < rank_N);
    
    TinyVector<int, rank_N> lower(0);
    TinyVector<int, rank_N> upper = array.ubound();
    upper[sortDim] = 0;
    
    RectDomain<rank_N> domain(lower, upper);
    Array<T, rank_N> top = array(domain); 
    
    typename Array<T, rank_N>::iterator it = top.begin();
    typename Array<T, rank_N>::iterator et = top.end();
    
    assert(array.stride(rank_N-1) == 1);
        
    int stride = array.stride(sortDim);
        
    for(; it != et; ++it)
    {
        strided_iterator<T*> start(&(*it), stride, 0);
        strided_iterator<T*> end = start + array.extent(sortDim);
        
        make_heap(start, end);
        for (typename strided_iterator<T*>::difference_type j = 0; j < num; j++)
        {
            pop_heap(start, end - j);
        }
    }
}

template<typename T, int rank_N>
void print(Array<T, rank_N> &array, int sortDim)
{
    assert(0 <= sortDim && sortDim < rank_N);
    
    TinyVector<int, rank_N> lower(0);
    TinyVector<int, rank_N> upper = array.ubound();
    upper[sortDim] = 0;
    
    RectDomain<rank_N> domain(lower, upper);
    Array<T, rank_N> top = array(domain); 
       
    typename Array<T, rank_N>::iterator it = top.begin();
    typename Array<T, rank_N>::iterator et = top.end();
    
    assert(array.stride(rank_N-1) == 1);
    
    int stride = array.stride(sortDim);
        
    for(; it != et; ++it)
    {    
        cout << "position = " << it.position() << " : ";
    
        strided_iterator<T*> start(&(*it), stride, 0);
        strided_iterator<T*> end = start + array.extent(sortDim);
        
        if (start != end) 
        {   
            cout << *start; 
            ++start;
        }
        for (; start != end; ++start)
        {
            cout << ", " << *start;
        }
        cout << endl;
    }
}

int main()
{
    Array<double, 3> A3(4,5,7); // choose 7 as last dimension for better visualization of results
    
    // random functor
    randuni r;
    
    A3 = r(A3);    
    cout << "unsorted A3 = " << A3 << endl;            
    partial_sort1(A3, 4);
    cout << "partial sorted A3 = " << A3 << endl; 
    
    A3 = r(A3);    
    cout << "unsorted A3 = " << A3 << endl;            
    partial_sort(A3, 2, 4);
    cout << "partial sorted A3 = " << A3 << endl; 
    print(A3, 2);
    
    A3 = r(A3);    
    cout << "unsorted A3 = " << A3 << endl;            
    partial_sort(A3, 1, 3);
    cout << "partial sorted A3 = " << A3 << endl; 
    print(A3, 1);
}