/*
 * (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.
 */

#include <algorithm>

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

namespace atlas {
namespace array {

namespace {
size_t compute_aligned_size(size_t size, size_t alignment) {
    size_t div           = size / alignment;
    size_t mod           = size % alignment;
    size_t _aligned_size = div * alignment;
    if (mod > 0) {
        _aligned_size += alignment;
    }
    return _aligned_size;
}
}  // namespace

ArraySpec::ArraySpec(): size_(), rank_(), allocated_size_(), contiguous_(true), default_layout_(true) {}

ArraySpec::ArraySpec(const ArrayShape& shape): ArraySpec(shape, ArrayAlignment()) {}

ArraySpec::ArraySpec(const ArrayShape& shape, const ArrayAlignment& alignment) {
    ArrayShape aligned_shape = shape;
    aligned_shape.back()     = compute_aligned_size(aligned_shape.back(), alignment);

    rank_           = static_cast<int>(shape.size());
    size_           = 1;
    allocated_size_ = 1;
    shape_.resize(rank_);
    strides_.resize(rank_);
    layout_.resize(rank_);
    for (int j = rank_ - 1; j >= 0; --j) {
        shape_[j]   = shape[j];
        strides_[j] = allocated_size_;
        layout_[j]  = j;
        size_ *= size_t(shape_[j]);
        allocated_size_ *= size_t(aligned_shape[j]);
    }
    ATLAS_ASSERT(allocated_size_ == compute_aligned_size(size_t(shape_[0]) * size_t(strides_[0]), alignment));
    contiguous_     = (size_ == allocated_size_);
    default_layout_ = true;
    alignment_      = alignment;

#ifdef ATLAS_HAVE_FORTRAN
    allocate_fortran_specs();
#endif
};

ArraySpec::ArraySpec(const ArrayShape& shape, const ArrayStrides& strides):
    ArraySpec(shape, strides, ArrayAlignment()) {}

ArraySpec::ArraySpec(const ArrayShape& shape, const ArrayStrides& strides, const ArrayAlignment& alignment) {
    ATLAS_ASSERT(shape.size() == strides.size(), "dimensions of shape and stride don't match");

    rank_ = static_cast<int>(shape.size());
    size_ = 1;
    shape_.resize(rank_);
    strides_.resize(rank_);
    layout_.resize(rank_);
    for (int j = rank_ - 1; j >= 0; --j) {
        shape_[j]   = shape[j];
        strides_[j] = strides[j];
        layout_[j]  = j;
        size_ *= shape_[j];
    }
    allocated_size_ = compute_aligned_size(shape_[0] * strides_[0], alignment);
    contiguous_     = (size_ == allocated_size_);
    default_layout_ = true;

#ifdef ATLAS_HAVE_FORTRAN
    allocate_fortran_specs();
#endif
}

ArraySpec::ArraySpec(const ArrayShape& shape, const ArrayStrides& strides, const ArrayLayout& layout):
    ArraySpec(shape, strides, layout, ArrayAlignment()) {}

ArraySpec::ArraySpec(const ArrayShape& shape, const ArrayStrides& strides, const ArrayLayout& layout,
                     const ArrayAlignment& alignment) {
    ATLAS_ASSERT(shape.size() == strides.size(), "dimensions of shape and stride don't match");

    rank_ = static_cast<int>(shape.size());
    size_ = 1;
    shape_.resize(rank_);
    strides_.resize(rank_);
    layout_.resize(rank_);
    default_layout_ = true;
    for (int j = rank_ - 1; j >= 0; --j) {
        shape_[j]   = shape[j];
        strides_[j] = strides[j];
        layout_[j]  = layout[j];
        size_ *= shape_[j];
        if (layout_[j] != idx_t(j)) {
            default_layout_ = false;
        }
    }
    allocated_size_ = compute_aligned_size(shape_[layout_[0]] * strides_[layout_[0]], alignment);
    contiguous_     = (size_ == allocated_size_);

#ifdef ATLAS_HAVE_FORTRAN
    allocate_fortran_specs();
#endif
}

const std::vector<int>& ArraySpec::shapef() const {
    return shapef_;
}

const std::vector<int>& ArraySpec::stridesf() const {
    return stridesf_;
}

void ArraySpec::allocate_fortran_specs() {
    shapef_.resize(rank_);
    for (idx_t j = 0; j < rank_; ++j) {
        shapef_[j] = shape_[rank_ - 1 - layout_[j]];
    }
    stridesf_.resize(strides_.size());
    std::reverse_copy(strides_.begin(), strides_.end(), stridesf_.begin());
}

}  // namespace array
}  // namespace atlas
