/*
 * (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 <array>
#include <type_traits>
#include <utility>
#include <vector>

#include "atlas/array.h"
#include "atlas/array/ArrayUtil.h"
#include "atlas/array/DataType.h"
#include "atlas/array/gridtools/GridToolsTraits.h"
#include "atlas/array_fwd.h"
#include "atlas/library/config.h"
#include "atlas/runtime/Exception.h"
#include "atlas/runtime/Log.h"

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

namespace atlas {
namespace array {
namespace gridtools {

template <unsigned int Rank>
std::array<idx_t, Rank> get_array_from_vector(std::vector<idx_t> const& values) {
    std::array<idx_t, Rank> array;
    std::copy(values.begin(), values.end(), array.begin());
    return array;
}

template <unsigned int TotalDims, unsigned int Dim, typename = void>
struct check_dimension_lengths_impl {
    template <typename FirstDim, typename... Dims>
    static void apply(ArrayShape const& shape, FirstDim first_dim, Dims... d) {
        if (first_dim < shape[Dim]) {
            std::stringstream err;
            err << "Attempt to resize array with original size for dimension " << Dim << " of " << shape[Dim] << " by "
                << first_dim << std::endl;
            throw_Exception(err.str(), Here());
        }
        check_dimension_lengths_impl<TotalDims, Dim + 1>::apply(shape, d...);
    }
};

template <unsigned int TotalDims, unsigned int Dim>
struct check_dimension_lengths_impl<TotalDims, Dim,
                                    typename std::enable_if<(Dim == TotalDims - 1) && (TotalDims > 1)>::type> {
    template <typename FirstDim>
    static void apply(ArrayShape const& shape, FirstDim first_dim) {
        if (first_dim < shape[Dim]) {
            std::stringstream err;
            err << "Attempt to resize array with original size for dimension " << Dim - 1 << " of " << shape[Dim - 1]
                << " by " << first_dim << std::endl;
            throw_Exception(err.str(), Here());
        }
    }
};

template <unsigned int TotalDims, unsigned int Dim>
struct check_dimension_lengths_impl<TotalDims, Dim,
                                    typename std::enable_if<(Dim == TotalDims - 1) && (TotalDims == 1)>::type> {
    template <typename FirstDim>
    static void apply(ArrayShape const& shape, FirstDim first_dim) {
        if (first_dim < shape[Dim]) {
            std::stringstream err;
            err << "Attempt to resize array with original size for dimension " << Dim << " of " << shape[Dim] << " by "
                << first_dim << std::endl;
            throw_Exception(err.str(), Here());
        }
    }
};

template <typename... Dims>
void check_dimension_lengths(ArrayShape const& shape, Dims... d) {
    check_dimension_lengths_impl<sizeof...(d), 0>::apply(shape, d...);
}

template <unsigned int Rank>
struct default_layout_t {
    template <typename T>
    struct get_layout;

    template <typename UInt, UInt... Indices>
    struct get_layout<std::integer_sequence<UInt, Indices...>> {
        using type = ::gridtools::layout_map<Indices...>;
    };

    using type = typename get_layout<std::make_integer_sequence<::gridtools::uint_t, Rank>>::type;
};


template <typename Value, typename LayoutMap>
struct get_layout_map_component {
    // TODO: static_assert( ::gridtools::is_layout_map<LayoutMap>(), "Error: not a
    // layout_map" );
    template <int Idx>
    struct get_component {
        ATLAS_HOST_DEVICE
        constexpr get_component() {}

        ATLAS_HOST_DEVICE constexpr static Value apply() { return LayoutMap::template at<Idx>(); }
    };
};

template <typename Value>
struct get_stride_component {
    template <int Idx>
    struct get_component {
        ATLAS_HOST_DEVICE
        constexpr get_component() {}

        template <typename StorageInfoPtr>
        ATLAS_HOST_DEVICE constexpr static Value apply(StorageInfoPtr a) {
            static_assert((::gridtools::is_storage_info<typename std::remove_pointer<StorageInfoPtr>::type>::value),
                          "Error: not a storage_info");
            return a->template stride<Idx>();
        }
    };
};

template <typename Value>
struct get_shape_component {
    template <int Idx>
    struct get_component {
        ATLAS_HOST_DEVICE
        constexpr get_component() {}

        template <typename StorageInfoPtr>
        ATLAS_HOST_DEVICE constexpr static Value apply(StorageInfoPtr a) {
            static_assert((::gridtools::is_storage_info<typename std::remove_pointer<StorageInfoPtr>::type>::value),
                          "Error: not a storage_info");
            return a->template total_length<Idx>();
        }
    };
};

// indirection around C++11 sizeof... since it is buggy for nvcc and cray
template <typename... T>
struct get_pack_size {
    using type = ::gridtools::static_uint<sizeof...(T)>;
};

template <typename Value, typename LayoutMap, typename Alignment, size_t Rank>
struct gt_storage_t {
    using type =
        gridtools::storage_traits::data_store_t<Value, gridtools::storage_traits::custom_layout_storage_info_align_t<
                                                           0, LayoutMap, ::gridtools::zero_halo<Rank>, Alignment>>;
};

template <typename Value, typename LayoutMap, typename Alignment, typename... UInts>
typename gt_storage_t<Value, LayoutMap, Alignment, get_pack_size<UInts...>::type::value>::type* create_gt_storage(
    UInts... dims) {
    static_assert((sizeof...(dims) > 0), "Error: can not create storages without any dimension");
    constexpr static unsigned int rank = get_pack_size<UInts...>::type::value;
    typedef gridtools::storage_traits::custom_layout_storage_info_align_t<0, LayoutMap, ::gridtools::zero_halo<rank>,
                                                                          Alignment>
        storage_info_ty;
    typedef gridtools::storage_traits::data_store_t<Value, storage_info_ty> data_store_t;

    data_store_t* ds;
    if (::gridtools::accumulate(::gridtools::multiplies(), dims...) == 0) {
        ds = new data_store_t();
    }
    else {
        storage_info_ty si(dims...);
        ds = new data_store_t(si);
    }
    return ds;
}

template <typename Value, typename LayoutMap, typename... UInts>
typename gt_storage_t<Value, LayoutMap, ::gridtools::alignment<1>, get_pack_size<UInts...>::type::value>::type*
create_gt_storage(UInts... dims) {
    return create_gt_storage<Value, LayoutMap, ::gridtools::alignment<1>>(dims...);
}


template <typename Value, size_t Rank>
struct gt_wrap_storage_t {
    using type = gridtools::storage_traits::data_store_t<Value, gridtools::storage_traits::storage_info_t<0, Rank>>;
};

template <typename Value, unsigned int Rank>
static typename gt_wrap_storage_t<Value, Rank>::type* wrap_gt_storage(Value* data, std::array<idx_t, Rank>&& shape,
                                                                      std::array<idx_t, Rank>&& strides) {
    static_assert((Rank > 0), "Error: can not create storages without any dimension");
    typedef gridtools::storage_traits::storage_info_t<0, Rank, ::gridtools::zero_halo<Rank>> storage_info_ty;
    typedef gridtools::storage_traits::data_store_t<Value, storage_info_ty> data_store_t;
    ::gridtools::array<::gridtools::uint_t, Rank> _shape;
    ::gridtools::array<::gridtools::uint_t, Rank> _strides;
    for (unsigned int i = 0; i < Rank; ++i) {
        _shape[i]   = shape[i];
        _strides[i] = strides[i];
    }
    storage_info_ty si(_shape, _strides);
    data_store_t* ds = new data_store_t(si, data);

    return ds;
}

constexpr idx_t zero(idx_t) {
    return 0;
}

template <idx_t... Is>
ArrayStrides make_null_strides(std::integer_sequence<idx_t, Is...>) {
    return make_strides(zero(Is)...);
}

template <typename UInt>
struct my_apply_gt_integer_sequence {
    template <typename Container, template <UInt T> class Lambda, typename... ExtraTypes>
    ATLAS_HOST_DEVICE static constexpr Container apply(ExtraTypes const&... args_) {
        static_assert((std::is_same<Container, Container>::value),
                      "ERROR: my_apply_gt_integer_sequence only accepts a "
                      "std::integer_sequence type. Check the call");
        return Container(args_...);
    }
};

template <typename UInt, UInt... Indices>
struct my_apply_gt_integer_sequence<std::integer_sequence<UInt, Indices...>> {
    /**
         @brief duplicated interface for the case in which the container is an
     aggregator
       */
    template <typename Container, template <UInt T> class Lambda, typename... ExtraTypes>
    ATLAS_HOST_DEVICE static constexpr Container apply(ExtraTypes const&... args_) {
        return Container{Lambda<Indices>::apply(args_...)...};
    }
};

#ifndef __CUDACC__

template <typename DataStore, typename... Dims>
ArraySpec ATLAS_HOST make_spec(DataStore* gt_data_store_ptr, Dims... dims) {
    static_assert((::gridtools::is_data_store<DataStore>::value), "Internal Error: passing a non GT data store");

    if (gt_data_store_ptr->valid()) {
        auto storage_info_ptr = gt_data_store_ptr->get_storage_info_ptr().get();
        using Layout          = typename DataStore::storage_info_t::layout_t;
        using Alignment       = typename DataStore::storage_info_t::alignment_t;

        using seq = my_apply_gt_integer_sequence<std::make_integer_sequence<int, sizeof...(dims)>>;

        ArraySpec spec(
            ArrayShape{(idx_t)dims...},
            seq::template apply<ArrayStrides, get_stride_component<idx_t>::template get_component>(storage_info_ptr),
            seq::template apply<ArrayLayout, get_layout_map_component<idx_t, Layout>::template get_component>(),
            ArrayAlignment(Alignment::value));
        ATLAS_ASSERT(spec.allocatedSize() == storage_info_ptr->padded_total_length());
        return spec;
    }
    else {
        return ArraySpec(make_shape({dims...}),
                         make_null_strides(std::make_integer_sequence<idx_t, sizeof...(dims)>()));
    }
}
#endif

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

}  // namespace gridtools
}  // namespace array
}  // namespace atlas
