//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Base/Type/Field2D.h
//! @brief     Defines two-dimensional data arrays.
//!
//! @homepage  http://www.bornagainproject.org
//! @license   GNU General Public License v3 or higher (see COPYING)
//! @copyright Forschungszentrum Jülich GmbH 2024
//! @authors   Scientific Computing Group at MLZ (see CITATION, AUTHORS)
//
//  ************************************************************************************************

#ifndef BORNAGAIN_BASE_TYPE_FIELD2D_H
#define BORNAGAIN_BASE_TYPE_FIELD2D_H

#include <functional>
#include <heinz/Complex.h> // complex_t
#include <vector>

template <DoubleOrComplex T> using Field2D = std::vector<std::vector<T>>;

using double2d_t = Field2D<double>;
using complex2d_t = Field2D<complex_t>;

namespace FieldUtil {

template <DoubleOrComplex T> std::vector<T> flatten(const Field2D<T>& src)
{
    std::vector<T> result;
    for (const auto& row : src)
        result.insert(result.end(), row.begin(), row.end());
    return result;
}

template <DoubleOrComplex T>
Field2D<T> make(size_t n_rows, size_t n_cols, std::function<T(size_t, size_t)> setter)
{
    Field2D<T> result(n_rows, std::vector<T>(n_cols));
    for (size_t i = 0; i < n_rows; i++)
        for (size_t j = 0; j < n_cols; j++)
            result[i][j] = setter(i, j);
    return result;
}

template <DoubleOrComplex T> Field2D<T> reshapeTo2D(const std::vector<T>& src, size_t n_rows)
{
    size_t n_cols = src.size() / n_rows;
    return make<T>(n_rows, n_cols,
                   [&src, n_cols](size_t i, size_t j) -> T { return src[i * n_cols + j]; });
}

} // namespace FieldUtil

#endif // BORNAGAIN_BASE_TYPE_FIELD2D_H
