/*
* (C) Copyright 2013 ECMWF.
*
* This software is licensed under the terms of the Apache Licence Version 2.0
* which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
* In applying this licence, ECMWF does not waive the privileges and immunities
* granted to it by virtue of its status as an intergovernmental organisation nor
* does it submit to any jurisdiction.
*/

#pragma once

#include <cassert>
#include <cstddef>

#include "atlas/library/config.h"
#include "atlas/runtime/Exception.h"

#if ATLAS_GRIDTOOLS_STORAGE_BACKEND_CUDA
#include <cuda_runtime.h>
#endif

namespace atlas {
namespace array {

//------------------------------------------------------------------------------

template <typename T>
class Vector {
public:
    Vector(idx_t N = 0): data_(N ? new T[N] : nullptr), data_gpu_(nullptr), size_(N) {}

    void resize_impl(idx_t N) {
        ATLAS_ASSERT(not data_gpu_, "we can not resize a vector after has been cloned to device");
        ATLAS_ASSERT(N >= size_);
        if (N == size_)
            return;

        T* d_ = new T[N];
        for (idx_t c = 0; c < size_; ++c) {
            d_[c] = data_[c];
        }
        if (data_)
            delete[] data_;
        data_ = d_;
    }

    void resize(idx_t N) {
        resize_impl(N);
        size_ = N;
    }

    void resize(idx_t N, T val) {
        resize_impl(N);
        for (idx_t c = size_; c < N; ++c) {
            data_[c] = val;
        }

        size_ = N;
    }

    void updateDevice() {
        if (!data_gpu_) {
#if ATLAS_GRIDTOOLS_STORAGE_BACKEND_CUDA
            ::cudaMalloc((void**)(&data_gpu_), sizeof(T*) * size_);

            T* buff = new T[size_];

            for (idx_t i = 0; i < size(); ++i) {
                data_[i]->updateDevice();
                buff[i] = data_[i]->gpu_object_ptr();
            }
            ::cudaMemcpy(data_gpu_, buff, sizeof(T*) * size_, cudaMemcpyHostToDevice);
            delete buff;
#else
            data_gpu_ = data_;
#endif
            size_gpu_ = size_;
        }
        else {
            ATLAS_ASSERT(size_gpu_ == size_);
#if ATLAS_GRIDTOOLS_STORAGE_BACKEND_CUDA
            for (idx_t i = 0; i < size(); ++i) {
                data_[i]->updateDevice();
                assert(data_gpu_[i] == data_[i]->gpu_object_ptr());
            }
#endif
        }
    }
    void updateHost() {
        ATLAS_ASSERT(data_gpu_ != nullptr);

#if ATLAS_GRIDTOOLS_STORAGE_BACKEND_CUDA

        for (idx_t i = 0; i < size(); ++i) {
            data_[i]->updateHost();
        }
#endif
    }

    T* gpu_object_ptr() { return data_gpu_; }

    T* data() { return data_; }
    T* data_gpu() { return data_gpu_; }

    idx_t size() const { return size_; }

private:
    T* data_;
    T* data_gpu_;
    idx_t size_;
    idx_t size_gpu_;
};

template <typename T>
class VectorView {
public:
    VectorView(): vector_(nullptr), data_(nullptr), size_(0) {}
    VectorView(Vector<T> const& vector, T* data): vector_(&vector), data_(data), size_(vector.size()) {}

    ATLAS_HOST_DEVICE
    T& operator[](idx_t idx) {
        assert(idx < size_);

        return data_[idx];
    }

    ATLAS_HOST_DEVICE
    T const& operator[](idx_t idx) const {
        assert(idx < size_);
        return data_[idx];
    }
    ATLAS_HOST_DEVICE
    T base() { return *data_; }

    ATLAS_HOST_DEVICE
    idx_t size() const { return size_; }

    bool is_valid(Vector<T>& vector) { return (&vector) == vector_ && (data_ != nullptr); }

public:
    Vector<T> const* vector_;
    T* data_;
    idx_t size_;
};

template <typename T>
VectorView<T> make_host_vector_view(Vector<T> vector_) {
    return VectorView<T>(vector_, vector_.data());
}

template <typename T>
VectorView<T> make_device_vector_view(Vector<T> vector_) {
    return VectorView<T>(vector_, vector_.data_gpu());
}

//------------------------------------------------------------------------------

}  // namespace array
}  // namespace atlas
