#ifndef MPQC_ARRAY_CORE_HPP
#define MPQC_ARRAY_CORE_HPP

#include "mpqc/array/forward.hpp"

#include "mpqc_config.h"
#include "mpqc/mpi.hpp"
#ifdef HAVE_MPI
#include "mpqc/array/parallel.hpp"
#ifdef HAVE_ARMCI
extern "C" {
#include <armci.h>
}
#endif
#endif // HAVE_MPI

#include "mpqc/range.hpp"
#include <boost/foreach.hpp>
#include "mpqc/utility/mutex.hpp"
#include "mpqc/utility/exception.hpp"

#include <boost/noncopyable.hpp>

namespace mpqc {
namespace detail {

    struct array_core_driver {
        array_core_driver() {}
    };

    template<typename T>
    struct array_impl<T, array_core_driver>
	: ArrayBase, boost::noncopyable
    {

        template<typename Extent>
	array_impl(const std::string &name,
                   const std::vector<Extent> &extents)
            : ArrayBase(name, extents)
        {
            data_.resize(this->size());
	}

        void sync() {}

    protected:

	void _put(const std::vector<range> &r, const void *buffer) {
            apply(putv, r, this->dims_, &this->data_[0], (const T*)buffer);
	}

	void _get(const std::vector<range> &r, void *buffer) const {
            apply(getv, r, this->dims_, &this->data_[0], (T*)buffer);
	}

    private:

        std::vector<T> data_;

        static size_t putv(size_t size, T *data, const T *buffer) {
            std::copy(buffer, buffer+size, data);
            return size;
        }

        static size_t getv(size_t size, const T *data, T *buffer) {
            std::copy(data, data+size, buffer);
            return size;
        }

        template<class F, typename data_ptr, typename buffer_ptr>
        static size_t apply(F f,
                            const std::vector<range> &r,
                            const std::vector<size_t> &dims,
                            data_ptr data, buffer_ptr buffer) {
            size_t begin = 0;
            std::vector<size_t> counts;
            std::vector<size_t> strides;
            counts.push_back(1);
            strides.push_back(1);
            for (size_t i = 0, stride = 1; i < dims.size(); ++i) {
                counts.back() *= r[i].size();
                begin += *r[i].begin()*stride;
                stride *= dims[i];
                // std::cout << "counts = " << counts.back() << std::endl;
                // std::cout << "strides = " << strides.back() << std::endl;
                if (dims[i] != r[i].size()) {
                    counts.push_back(1);
                    strides.push_back(stride);
                }
            }
            return apply(f, counts, strides, data+begin, buffer, strides.size());
        }

        template<class F, typename data_ptr, typename buffer_ptr>
        static size_t apply(F f,
                            const std::vector<size_t> &counts,
                            const std::vector<size_t> &strides,
                            data_ptr data, buffer_ptr buffer,
                            int level) {
            if (level == 0) 
                return 0;
            if (level == 1) {
                //std::cout << "count1 = " << counts[0] << std::endl;
                return f(counts[0], data, buffer);
            }
            size_t size = 0;
            --level;
            for (size_t i = 0; i < counts[level]; ++i) {
                //std::cout << size << " " << strides[level] << std::endl;
                size_t n = apply(f, counts, strides, data, buffer+size, level);
                data += strides[level];
                size += n;
            }
            return size;
        }


    };

#ifdef HAVE_ARMCI

    template<typename T>
    struct array_parallel_impl<T, array_core_driver>
	: ArrayBase, boost::noncopyable
    {

	array_parallel_impl(const std::string &name,
			    const std::vector<size_t> &dims,
			    const MPI::Comm &comm)
            : ArrayBase(name, dims)
        {
            initialize(ArrayBase::dims_, comm);
        }

    private:

        void initialize(const std::vector<size_t> &dims, const MPI::Comm &comm) {

            // dont have code to use ARMCI groups yet
	    MPQC_CHECK(comm == MPI::Comm::World());

	    size_t size = 1;
	    std::vector<range> extents;
	    for (int i = 0; i < int(dims.size())-1; ++i) {
		size *= dims[i];
		extents.push_back(range(0, dims[i]));
	    }
	    size_t block = (dims.back() + comm.size() - 1)/comm.size();
	    size *= block;

	    check(ARMCI_Init(), "ARMCI_Init", comm);
	    
	    std::vector<void*> data(comm.size(), NULL);
	    check(ARMCI_Malloc(&data[0], size*sizeof(T)), "ARMCI_Malloc", comm);
            data_ = data[comm.rank()];

	    for (size_t i = 0; i < comm.size(); ++i) {
		Tile tile;
		tile.data = data[i];
		tile.proc = i;
		tile.local = (i == comm.rank());
		size_t begin = std::min<size_t>(dims.back(), i*block);
		size_t end = std::min<size_t>(dims.back(), begin+block);
		tile.extents = extents;
		tile.extents.push_back(range(begin, end));
		tiles_.push_back(tile);
		//printf("tile[%i].ptr=%p\n", i, tile.ptr);
	    }

	}

	~array_parallel_impl() {
            ARMCI_Free(this->data_);
	}

	void sync() {
	    ARMCI_Barrier();
	}

        const MPI::Comm& comm() const {
            return this->comm_;
        }

    protected:

	void _put(const std::vector<range> &r, const void *buffer) {
	    apply<PUT>(this->tiles_, r, (void*)buffer);
	}

	void _get(const std::vector<range> &r, void *buffer) const {
	    apply<GET>(this->tiles_, r, buffer);
	}

    private:

	enum OP { PUT, GET };

        typedef array_tile<void*> Tile;
        void *data_;
	std::vector<Tile> tiles_;

	static void check(int err, const std::string &func,
			  MPI::Comm comm = MPI::Comm::Self()) {
	    if (comm.any((err != 0))) {
		throw std::runtime_error(func + " failed");
	    }
	}

	template<OP op>
	static void apply(const std::vector<Tile> &tiles, 
			  const std::vector<range> &r,
			  void *buffer) {
	    
	    T *local = (T*)buffer;

	    BOOST_FOREACH (Tile t, tiles) {

		std::vector<range> x = t.subset(r);
		//std::cout << t.extents << " ^ " << r << " = " << x << std::endl;
		if (x.empty()) continue;

		size_t total = 1;
		size_t N = x.size();
		int remote_strides[N];
		int local_strides[N];
		int count[N];

		for (size_t i = 0; i < N; ++i) {
		    total *= x[i].size();
		    count[i] = x[i].size();
		    local_strides[i] =
			x[i].size()*((i > 0) ? local_strides[i-1] : 1);
		    remote_strides[i] =
			t.extents[i].size()*((i > 0) ? remote_strides[i - 1] : 1);
		    local_strides[i] *= sizeof(T);
		    remote_strides[i] *= sizeof(T);
		}
                count[0] *= sizeof(T);

                T *remote = (T*)t.data;
                for (size_t i = 0, stride = 1; i < N; ++i) {
                    remote += (x[i].begin() - t.extents[i].begin())*stride;
                    stride *= t.extents[i].size();
                }

		// for (size_t i = 0; i < N; ++i) {
		//     printf("%i count=%i, remote_strides=%i, local_strides=%i, total=%lu\n",
		// 	   i, count[i], remote_strides[i], local_strides[i], total);
		// }
		// printf("tile.ptr=%p\n", t.ptr);

                mutex::global::lock();
		if (op == GET) {
		    int err;
		    err = ARMCI_GetS(remote, remote_strides,
				     local, local_strides,
				     count, N-1, t.proc);
		    check(err, "ARMCI_GetS failed");
		}
		if (op == PUT) {
		    int err;
		    err = ARMCI_PutS(local, local_strides,
				     remote, remote_strides,
				     count, N-1, t.proc);
		    check(err, "ARMCI_PutS failed");
		}
                mutex::global::unlock();

		local += total;

	    }
	}

    };
#endif // HAVE_ARMCI

} // namespace detail
} // namespace mpqc


namespace mpqc {
    static const detail::array_core_driver ARRAY_CORE;
}


#endif /* MPQC_ARRAY_CORE_HPP */
