/*
 * (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 <functional>
#include <type_traits>

#include "atlas/array/LocalView.h"
#include "atlas/field/Field.h"
#include "atlas/functionspace/FunctionSpace.h"
#include "atlas/functionspace/detail/FunctionSpaceImpl.h"
#include "atlas/library/config.h"
#include "atlas/parallel/mpi/mpi.h"
#include "atlas/parallel/omp/omp.h"
#include "atlas/runtime/Exception.h"
#include "atlas/util/Config.h"

namespace atlas {
class Field;
class FieldSet;
}  // namespace atlas

namespace atlas {
namespace trans {
class Trans;
}
}  // namespace atlas

namespace atlas {
namespace functionspace {
namespace detail {

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

class Spectral : public functionspace::FunctionSpaceImpl {
    /*
  Spectral data is organised as:
     m = zonal wavenumber
     n = total wavenumber
  
  const auto zonal_wavenumbers = Spectral::zonal_wavenumbers();
  idx_t jc=0;
  for( int jm=0; jm<zonal_wavenumbers.size(); ++jm ) {
    int m = zonal_wavenumbers(m);
    for( int n=m; m<=truncation; ++n ) {
      data( jc++, jfld ) = func_real_part(m,n);
      data( jc++, jfld ) = func_imag_part(m,n);
    }
  }
  
*/

public:
    Spectral( const eckit::Configuration& );

    Spectral( const int truncation, const eckit::Configuration& = util::NoConfig() );

    Spectral( const trans::Trans&, const eckit::Configuration& = util::NoConfig() );

    virtual ~Spectral() override;

    virtual std::string type() const override { return "Spectral"; }

    virtual std::string distribution() const override;

    using FunctionSpaceImpl::createField;
    virtual Field createField( const eckit::Configuration& ) const override;
    virtual Field createField( const Field&, const eckit::Configuration& ) const override;

    void gather( const FieldSet&, FieldSet& ) const;
    void gather( const Field&, Field& ) const;

    void scatter( const FieldSet&, FieldSet& ) const;
    void scatter( const Field&, Field& ) const;

    std::string checksum( const FieldSet& ) const;
    std::string checksum( const Field& ) const;

    void norm( const Field&, double& norm, int rank = 0 ) const;
    void norm( const Field&, double norm_per_level[], int rank = 0 ) const;
    void norm( const Field&, std::vector<double>& norm_per_level, int rank = 0 ) const;

    array::LocalView<const int, 1> zonal_wavenumbers() const;  // zero-based, OK

    idx_t levels() const { return nb_levels_; }

    class For {
    public:
        For( const Spectral& fs, const util::Config& config = util::NoConfig() ) :
            truncation{fs.truncation()},
            zonal_wavenumbers{fs.zonal_wavenumbers()},
            global{config.getBool( "global", false )},
            owner{config.getInt( "owner", 0 )} {}

    protected:
        using View = const array::LocalView<const int, 1>;
        int truncation;
        View zonal_wavenumbers;
        bool global;
        idx_t owner;

    public:
#define FunctorArgs( ... )                                                                                             \
    typename std::enable_if<std::is_convertible<Functor, std::function<void( __VA_ARGS__ )>>::value, Functor>::type* = \
        nullptr

        // Functor: void f(real,imag,n,m)
        template <typename Functor, FunctorArgs( idx_t, idx_t, int, int )>
        void operator()( const Functor& f ) const {
            idx_t index = 0;
            if ( global ) {
                if ( owner == mpi::rank() ) {
                    atlas_omp_parallel_for( int m = 0; m <= truncation; ++m ) {
                        for ( int n = m; n <= truncation; ++n ) {
                            f( index, index + 1, n, m );
                            index += 2;
                        }
                    }
                }
            }
            else {
                const int nb_zonal_wavenumbers{static_cast<int>( zonal_wavenumbers.size() )};
                atlas_omp_parallel_for( int jm = 0; jm < nb_zonal_wavenumbers; ++jm ) {
                    const int m = zonal_wavenumbers( jm );
                    for ( int n = m; n <= truncation; ++n ) {
                        f( index, index + 1, n, m );
                        index += 2;
                    }
                }
            }
        }

        // Functor: void f(real,imag,n)
        template <typename Functor, FunctorArgs( idx_t, idx_t, int )>
        void operator()( const Functor& f ) const {
            idx_t index = 0;
            if ( global ) {
                if ( owner == mpi::rank() ) {
                    atlas_omp_parallel_for( int m = 0; m <= truncation; ++m ) {
                        for ( int n = m; n <= truncation; ++n ) {
                            f( index, index + 1, n );
                            index += 2;
                        }
                    }
                }
            }
            else {
                const int nb_zonal_wavenumbers{static_cast<int>( zonal_wavenumbers.size() )};
                atlas_omp_parallel_for( int jm = 0; jm < nb_zonal_wavenumbers; ++jm ) {
                    const int m = zonal_wavenumbers( jm );
                    for ( int n = m; n <= truncation; ++n ) {
                        f( index, index + 1, n );
                        index += 2;
                    }
                }
            }
        }
#undef FunctorArgs
    };
    template <typename Functor>
    void parallel_for( const Functor& f ) const {
        For( *this, util::NoConfig() )( f );
    }

    template <typename Functor>
    void parallel_for( const util::Config& config, const Functor& f ) const {
        For( *this, config )( f );
    }

public:  // methods
    idx_t nb_spectral_coefficients() const;
    idx_t nb_spectral_coefficients_global() const;
    int truncation() const { return truncation_; }

    virtual idx_t size() const override { return nb_spectral_coefficients(); }

private:  // methods
    array::DataType config_datatype( const eckit::Configuration& ) const;
    std::string config_name( const eckit::Configuration& ) const;
    idx_t config_size( const eckit::Configuration& ) const;
    idx_t config_levels( const eckit::Configuration& ) const;
    void set_field_metadata( const eckit::Configuration&, Field& ) const;
    size_t footprint() const override;


private:  // Fortran access
    friend struct SpectralFortranAccess;
    int nump() const;                               // equivalent to nmyms().size()
    array::LocalView<const int, 1> nvalue() const;  // Return wave number n for a given index
    array::LocalView<const int, 1> nmyms() const;   // Return list of local zonal wavenumbers "m"
    array::LocalView<const int, 1> nasm0() const;

private:  // data
    idx_t nb_levels_;
    int truncation_;

    class Parallelisation;
    std::unique_ptr<Parallelisation> parallelisation_;
};

}  // namespace detail

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

class Spectral : public FunctionSpace {
public:
    Spectral();
    Spectral( const FunctionSpace& );
    Spectral( const eckit::Configuration& );
    Spectral( const int truncation, const eckit::Configuration& = util::NoConfig() );
    Spectral( const trans::Trans&, const eckit::Configuration& = util::NoConfig() );

    operator bool() const { return valid(); }
    bool valid() const { return functionspace_; }

    void gather( const FieldSet&, FieldSet& ) const;
    void gather( const Field&, Field& ) const;

    void scatter( const FieldSet&, FieldSet& ) const;
    void scatter( const Field&, Field& ) const;

    std::string checksum( const FieldSet& ) const;
    std::string checksum( const Field& ) const;

    void norm( const Field&, double& norm, int rank = 0 ) const;
    void norm( const Field&, double norm_per_level[], int rank = 0 ) const;
    void norm( const Field&, std::vector<double>& norm_per_level, int rank = 0 ) const;

    idx_t nb_spectral_coefficients() const;
    idx_t nb_spectral_coefficients_global() const;
    int truncation() const;
    idx_t levels() const { return functionspace_->levels(); }

    template <typename Functor>
    void parallel_for( const Functor& f ) const {
        functionspace_->parallel_for( f );
    }
    template <typename Functor>
    void parallel_for( const util::Config& config, const Functor& f ) const {
        functionspace_->parallel_for( config, f );
    }

private:
    const detail::Spectral* functionspace_;
};

}  // namespace functionspace
}  // namespace atlas
