/*
 * (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 "atlas/functionspace/StructuredColumns.h"

#include <fstream>
#include <iomanip>
#include <mutex>
#include <sstream>
#include <string>

#include "eckit/utils/MD5.h"

#include "atlas/array/Array.h"
#include "atlas/array/MakeView.h"
#include "atlas/domain.h"
#include "atlas/field/FieldSet.h"
#include "atlas/grid/Distribution.h"
#include "atlas/grid/Partitioner.h"
#include "atlas/grid/StructuredGrid.h"
#include "atlas/grid/StructuredPartitionPolygon.h"
#include "atlas/library/Library.h"
#include "atlas/mesh/Mesh.h"
#include "atlas/parallel/Checksum.h"
#include "atlas/parallel/GatherScatter.h"
#include "atlas/parallel/HaloExchange.h"
#include "atlas/parallel/mpi/mpi.h"
#include "atlas/parallel/omp/fill.h"
#include "atlas/parallel/omp/omp.h"
#include "atlas/runtime/Exception.h"
#include "atlas/runtime/Trace.h"
#include "atlas/util/Checksum.h"
#include "atlas/util/CoordinateEnums.h"
#include "atlas/util/detail/Cache.h"

#define REMOTE_IDX_BASE 0

namespace atlas {
namespace functionspace {
namespace detail {

namespace {


template <typename T, typename Field>
array::LocalView<T, 3> make_leveled_view( Field& field ) {
    using namespace array;
    if ( field.levels() ) {
        if ( field.variables() ) {
            return make_view<T, 3>( field ).slice( Range::all(), Range::all(), Range::all() );
        }
        else {
            return make_view<T, 2>( field ).slice( Range::all(), Range::all(), Range::dummy() );
        }
    }
    else {
        if ( field.variables() ) {
            return make_view<T, 2>( field ).slice( Range::all(), Range::dummy(), Range::all() );
        }
        else {
            return make_view<T, 1>( field ).slice( Range::all(), Range::dummy(), Range::dummy() );
        }
    }
}

template <typename T>
std::string checksum_3d_field( const parallel::Checksum& checksum, const Field& field ) {
    auto values = make_leveled_view<const T>( field );
    array::ArrayT<T> surface_field( values.shape( 0 ), values.shape( 2 ) );
    auto surface     = array::make_view<T, 2>( surface_field );
    const idx_t npts = values.shape( 0 );
    atlas_omp_for( idx_t n = 0; n < npts; ++n ) {
        for ( idx_t j = 0; j < surface.shape( 1 ); ++j ) {
            surface( n, j ) = 0.;
            for ( idx_t l = 0; l < values.shape( 1 ); ++l ) {
                surface( n, j ) += values( n, l, j );
            }
        }
    }
    return checksum.execute( surface.data(), surface_field.stride( 0 ) );
}

}  // namespace


class StructuredColumnsHaloExchangeCache : public util::Cache<std::string, parallel::HaloExchange>,
                                           public grid::detail::grid::GridObserver {
private:
    using Base = util::Cache<std::string, parallel::HaloExchange>;
    StructuredColumnsHaloExchangeCache() : Base( "StructuredColumnsHaloExchangeCache" ) {}

public:
    static StructuredColumnsHaloExchangeCache& instance() {
        static StructuredColumnsHaloExchangeCache inst;
        return inst;
    }
    util::ObjectHandle<value_type> get_or_create( const detail::StructuredColumns& funcspace ) {
        registerGrid( *funcspace.grid().get() );

        creator_type creator = std::bind( &StructuredColumnsHaloExchangeCache::create, &funcspace );
        return Base::get_or_create( key( funcspace ), remove_key( funcspace ), creator );
    }
    void onGridDestruction( grid::detail::grid::Grid& grid ) override { remove( remove_key( grid ) ); }

private:
    static Base::key_type key( const detail::StructuredColumns& funcspace ) {
        std::ostringstream key;
        key << "grid[address=" << funcspace.grid().get() << ",halo=" << funcspace.halo()
            << ",periodic_points=" << std::boolalpha << funcspace.periodic_points_
            << ",distribution=" << funcspace.distribution() << "]";
        return key.str();
    }

    static Base::key_type remove_key( const detail::StructuredColumns& funcspace ) {
        return remove_key( *funcspace.grid().get() );
    }

    static Base::key_type remove_key( const grid::detail::grid::Grid& grid ) {
        std::ostringstream key;
        key << "grid[address=" << &grid << "]";
        return key.str();
    }

    static value_type* create( const detail::StructuredColumns* funcspace ) {
        value_type* value = new value_type();

        value->setup( array::make_view<int, 1>( funcspace->partition() ).data(),
                      array::make_view<idx_t, 1>( funcspace->remote_index() ).data(), REMOTE_IDX_BASE,
                      funcspace->sizeHalo(), funcspace->sizeOwned() );
        return value;
    }
    ~StructuredColumnsHaloExchangeCache() override = default;
};

class StructuredColumnsGatherScatterCache : public util::Cache<std::string, parallel::GatherScatter>,
                                            public grid::detail::grid::GridObserver {
private:
    using Base = util::Cache<std::string, parallel::GatherScatter>;
    StructuredColumnsGatherScatterCache() : Base( "StructuredColumnsGatherScatterCache" ) {}

public:
    static StructuredColumnsGatherScatterCache& instance() {
        static StructuredColumnsGatherScatterCache inst;
        return inst;
    }
    util::ObjectHandle<value_type> get_or_create( const detail::StructuredColumns& funcspace ) {
        registerGrid( *funcspace.grid().get() );
        creator_type creator = std::bind( &StructuredColumnsGatherScatterCache::create, &funcspace );
        return Base::get_or_create( key( funcspace ), remove_key( funcspace ), creator );
    }
    void onGridDestruction( grid::detail::grid::Grid& grid ) override { remove( remove_key( grid ) ); }

private:
    static Base::key_type key( const detail::StructuredColumns& funcspace ) {
        std::ostringstream key;
        key << "grid[address=" << funcspace.grid().get() << ",halo=" << funcspace.halo()
            << ",periodic_points=" << std::boolalpha << funcspace.periodic_points_
            << ",distribution=" << funcspace.distribution() << "]";
        return key.str();
    }

    static Base::key_type remove_key( const detail::StructuredColumns& funcspace ) {
        return remove_key( *funcspace.grid().get() );
    }

    static Base::key_type remove_key( const grid::detail::grid::Grid& grid ) {
        std::ostringstream key;
        key << "grid[address=" << &grid << "]";
        return key.str();
    }

    static value_type* create( const detail::StructuredColumns* funcspace ) {
        value_type* value = new value_type();

        value->setup( array::make_view<int, 1>( funcspace->partition() ).data(),
                      array::make_view<idx_t, 1>( funcspace->remote_index() ).data(), REMOTE_IDX_BASE,
                      array::make_view<gidx_t, 1>( funcspace->global_index() ).data(), funcspace->sizeOwned() );
        return value;
    }
    ~StructuredColumnsGatherScatterCache() override = default;
};

class StructuredColumnsChecksumCache : public util::Cache<std::string, parallel::Checksum>,
                                       public grid::detail::grid::GridObserver {
private:
    using Base = util::Cache<std::string, parallel::Checksum>;
    StructuredColumnsChecksumCache() : Base( "StructuredColumnsChecksumCache" ) {}

public:
    static StructuredColumnsChecksumCache& instance() {
        static StructuredColumnsChecksumCache inst;
        return inst;
    }
    util::ObjectHandle<value_type> get_or_create( const detail::StructuredColumns& funcspace ) {
        registerGrid( *funcspace.grid().get() );
        creator_type creator = std::bind( &StructuredColumnsChecksumCache::create, &funcspace );
        return Base::get_or_create( key( funcspace ), remove_key( funcspace ), creator );
    }
    void onGridDestruction( grid::detail::grid::Grid& grid ) override { remove( remove_key( grid ) ); }

private:
    static Base::key_type key( const detail::StructuredColumns& funcspace ) {
        std::ostringstream key;
        key << "grid[address=" << funcspace.grid().get() << ",halo=" << funcspace.halo()
            << ",periodic_points=" << std::boolalpha << funcspace.periodic_points_
            << ",distribution=" << funcspace.distribution() << "]";
        return key.str();
    }

    static Base::key_type remove_key( const detail::StructuredColumns& funcspace ) {
        return remove_key( *funcspace.grid().get() );
    }

    static Base::key_type remove_key( const grid::detail::grid::Grid& grid ) {
        std::ostringstream key;
        key << "grid[address=" << &grid << "]";
        return key.str();
    }

    static value_type* create( const detail::StructuredColumns* funcspace ) {
        value_type* value = new value_type();
        util::ObjectHandle<parallel::GatherScatter> gather(
            StructuredColumnsGatherScatterCache::instance().get_or_create( *funcspace ) );
        value->setup( gather );
        return value;
    }
    ~StructuredColumnsChecksumCache() override = default;
};


const parallel::GatherScatter& StructuredColumns::gather() const {
    if ( gather_scatter_ ) {
        return *gather_scatter_;
    }
    gather_scatter_ = StructuredColumnsGatherScatterCache::instance().get_or_create( *this );
    return *gather_scatter_;
}

const parallel::GatherScatter& StructuredColumns::scatter() const {
    if ( gather_scatter_ ) {
        return *gather_scatter_;
    }
    gather_scatter_ = StructuredColumnsGatherScatterCache::instance().get_or_create( *this );
    return *gather_scatter_;
}

const parallel::Checksum& StructuredColumns::checksum() const {
    if ( checksum_ ) {
        return *checksum_;
    }
    checksum_ = StructuredColumnsChecksumCache::instance().get_or_create( *this );
    return *checksum_;
}

const parallel::HaloExchange& StructuredColumns::halo_exchange() const {
    if ( halo_exchange_ ) {
        return *halo_exchange_;
    }
    halo_exchange_ = StructuredColumnsHaloExchangeCache::instance().get_or_create( *this );
    return *halo_exchange_;
}

void StructuredColumns::set_field_metadata( const eckit::Configuration& config, Field& field ) const {
    field.set_functionspace( this );

    bool global( false );
    if ( config.get( "global", global ) ) {
        if ( global ) {
            idx_t owner( 0 );
            config.get( "owner", owner );
            field.metadata().set( "owner", owner );
        }
    }
    field.metadata().set( "global", global );

    idx_t levels( nb_levels_ );
    config.get( "levels", levels );
    field.set_levels( levels );

    idx_t variables( 0 );
    config.get( "variables", variables );
    field.set_variables( variables );

    if ( config.has( "type" ) ) {
        field.metadata().set( "type", config.getString( "type" ) );
    }
}

array::DataType StructuredColumns::config_datatype( const eckit::Configuration& config ) const {
    array::DataType::kind_t kind;
    if ( !config.get( "datatype", kind ) ) {
        throw_Exception( "datatype missing", Here() );
    }
    return array::DataType( kind );
}

std::string StructuredColumns::config_name( const eckit::Configuration& config ) const {
    std::string name;
    config.get( "name", name );
    return name;
}

idx_t StructuredColumns::config_levels( const eckit::Configuration& config ) const {
    idx_t levels( nb_levels_ );
    config.get( "levels", levels );
    return levels;
}

array::ArrayShape StructuredColumns::config_shape( const eckit::Configuration& config ) const {
    array::ArrayShape shape;

    shape.emplace_back( config_size( config ) );

    idx_t levels( nb_levels_ );
    config.get( "levels", levels );
    if ( levels > 0 ) {
        shape.emplace_back( levels );
    }

    idx_t variables( 0 );
    config.get( "variables", variables );
    if ( variables > 0 ) {
        shape.emplace_back( variables );
    }

    return shape;
}

array::ArrayAlignment StructuredColumns::config_alignment( const eckit::Configuration& config ) const {
    int alignment( 1 );
    config.get( "alignment", alignment );
    return alignment;
}

array::ArraySpec StructuredColumns::config_spec( const eckit::Configuration& config ) const {
    return array::ArraySpec( config_shape( config ), config_alignment( config ) );
}

size_t StructuredColumns::Map2to1::footprint() const {
    size_t size = sizeof( *this );
    size += data_.size() * sizeof( decltype( data_ )::value_type );
    return size;
}

void StructuredColumns::Map2to1::print( std::ostream& out ) const {
    for ( idx_t j = j_min_; j <= j_max_; ++j ) {
        out << std::setw( 4 ) << j << " : ";
        for ( idx_t i = i_min_; i <= i_max_; ++i ) {
            idx_t v = operator()( i, j );
            if ( v == missing() ) {
                out << std::setw( 6 ) << "X";
            }
            else {
                out << std::setw( 6 ) << v;
            }
        }
        out << '\n';
    }
}

void StructuredColumns::IndexRange::print( std::ostream& out ) const {
    for ( idx_t i = min_; i <= max_; ++i ) {
        idx_t v = operator()( i );
        if ( v == missing() ) {
            out << std::setw( 4 ) << "X";
        }
        else {
            out << std::setw( 4 ) << v;
        }
    }
    out << '\n';
}

idx_t StructuredColumns::config_size( const eckit::Configuration& config ) const {
    idx_t size = size_halo_;
    bool global( false );
    if ( config.get( "global", global ) ) {
        if ( global ) {
            idx_t owner( 0 );
            config.get( "owner", owner );
            size = ( mpi::rank() == owner ? grid_->size() : 0 );
        }
    }
    return size;
}

std::string StructuredColumns::distribution() const {
    return distribution_;
}

void StructuredColumns::throw_outofbounds( idx_t i, idx_t j ) const {
    std::stringstream ss;
    if ( j < j_begin_halo() || j >= j_end_halo() ) {
        ss << "OutofBounds: j out of range! : (i,j) = (" << i << "," << j << ") --- Expected: " << j_begin_halo()
           << " <= j < " << j_end_halo();
    }
    if ( i < i_begin_halo( j ) || i >= i_end_halo( j ) ) {
        ss << "OutofBounds: i out of range! : (i,j) = (" << i << "," << j << ") --- Expected: " << i_begin_halo( j )
           << " <= i < " << i_end_halo( j );
    }
    throw_Exception( ss.str(), Here() );
}

// ----------------------------------------------------------------------------
// Constructor
// ----------------------------------------------------------------------------
StructuredColumns::StructuredColumns( const Grid& grid, const eckit::Configuration& config ) :
    StructuredColumns::StructuredColumns( grid, grid::Partitioner(), config ) {}

StructuredColumns::StructuredColumns( const Grid& grid, const grid::Partitioner& p,
                                      const eckit::Configuration& config ) :
    StructuredColumns( grid, Vertical( config ), p, config ) {}

StructuredColumns::StructuredColumns( const Grid& grid, const grid::Distribution& distribution,
                                      const eckit::Configuration& config ) :
    StructuredColumns( grid, distribution, Vertical( config ), config ) {}

StructuredColumns::StructuredColumns( const Grid& grid, const grid::Distribution& distribution,
                                      const Vertical& vertical, const eckit::Configuration& config ) :
    vertical_( vertical ), nb_levels_( vertical_.size() ), grid_( new StructuredGrid( grid ) ) {
    setup( distribution, config );
}

StructuredColumns::StructuredColumns( const Grid& grid, const Vertical& vertical, const eckit::Configuration& config ) :
    StructuredColumns( grid, vertical, grid::Partitioner(), config ) {}

StructuredColumns::StructuredColumns( const Grid& grid, const Vertical& vertical, const grid::Partitioner& p,
                                      const eckit::Configuration& config ) :
    vertical_( vertical ), nb_levels_( vertical_.size() ), grid_( new StructuredGrid( grid ) ) {
    ATLAS_TRACE( "StructuredColumns constructor" );

    grid::Partitioner partitioner( p );
    if ( not partitioner ) {
        if ( config.has( "partitioner" ) ) {
            partitioner = grid::Partitioner( config.getSubConfiguration( "partitioner" ) );
        }
        else {
            if ( grid_->domain().global() ) {
                partitioner = grid::Partitioner( "equal_regions" );
            }
            else {
                partitioner = grid::Partitioner( "checkerboard" );
            }
        }
    }

    grid::Distribution distribution;
    ATLAS_TRACE_SCOPE( "Partitioning grid" ) { distribution = grid::Distribution( grid, partitioner ); }

    setup( distribution, config );
}

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

void StructuredColumns::compute_xy( idx_t i, idx_t j, PointXY& xy ) const {
    idx_t jj;
    if ( j < 0 ) {
        jj     = -j - 1 + north_pole_included_;
        xy[YY] = 180. - grid_->y( jj );
    }
    else if ( j >= ny_ ) {
        jj     = 2 * ny_ - j - 1 - south_pole_included_;
        xy[YY] = -180. - grid_->y( jj );
    }
    else {
        jj     = j;
        xy[YY] = grid_->y( jj );
    }
    xy[XX] = grid_->x( i, jj );
}


void StructuredColumns::Map2to1::resize( std::array<idx_t, 2> i_range, std::array<idx_t, 2> j_range ) {
    i_min_    = i_range[0];
    i_max_    = i_range[1];
    j_min_    = j_range[0];
    j_max_    = j_range[1];
    j_stride_ = ( i_max_ - i_min_ + 1 );
    data_.resize( ( i_max_ - i_min_ + 1 ) * ( j_max_ - j_min_ + 1 ) );
    atlas::omp::fill( data_.begin(), data_.end(), missing() + 1 );
}

const util::PartitionPolygon& StructuredColumns::polygon( idx_t halo ) const {
    if ( not polygon_ ) {
        polygon_.reset( new grid::StructuredPartitionPolygon( *this, halo ) );
    }
    return *polygon_;
}

const util::PartitionPolygons& StructuredColumns::polygons() const {
    if ( polygons_.size() ) {
        return polygons_;
    }
    polygon().allGather( polygons_ );
    return polygons_;
}

// ----------------------------------------------------------------------------
// Destructor
// ----------------------------------------------------------------------------
StructuredColumns::~StructuredColumns() {
    delete grid_;
}
// ----------------------------------------------------------------------------

// ----------------------------------------------------------------------------
// Create Field
// ----------------------------------------------------------------------------
Field StructuredColumns::createField( const eckit::Configuration& options ) const {
    Field field( config_name( options ), config_datatype( options ), config_spec( options ) );
    set_field_metadata( options, field );
    return field;
}

Field StructuredColumns::createField( const Field& other, const eckit::Configuration& config ) const {
    return createField( option::datatype( other.datatype() ) | option::levels( other.levels() ) |
                        option::variables( other.variables() ) |
                        option::type( other.metadata().getString( "type", "scalar" ) ) | config );
}
// ----------------------------------------------------------------------------

// ----------------------------------------------------------------------------
// Gather FieldSet
// ----------------------------------------------------------------------------
void StructuredColumns::gather( const FieldSet& local_fieldset, FieldSet& global_fieldset ) const {
    ATLAS_ASSERT( local_fieldset.size() == global_fieldset.size() );

    for ( idx_t f = 0; f < local_fieldset.size(); ++f ) {
        const Field& loc      = local_fieldset[f];
        Field& glb            = global_fieldset[f];
        const idx_t nb_fields = 1;
        idx_t root( 0 );
        glb.metadata().get( "owner", root );

        if ( loc.datatype() == array::DataType::kind<int>() ) {
            parallel::Field<int const> loc_field( make_leveled_view<const int>( loc ) );
            parallel::Field<int> glb_field( make_leveled_view<int>( glb ) );
            gather().gather( &loc_field, &glb_field, nb_fields, root );
        }
        else if ( loc.datatype() == array::DataType::kind<long>() ) {
            parallel::Field<long const> loc_field( make_leveled_view<const long>( loc ) );
            parallel::Field<long> glb_field( make_leveled_view<long>( glb ) );
            gather().gather( &loc_field, &glb_field, nb_fields, root );
        }
        else if ( loc.datatype() == array::DataType::kind<float>() ) {
            parallel::Field<float const> loc_field( make_leveled_view<const float>( loc ) );
            parallel::Field<float> glb_field( make_leveled_view<float>( glb ) );
            gather().gather( &loc_field, &glb_field, nb_fields, root );
        }
        else if ( loc.datatype() == array::DataType::kind<double>() ) {
            parallel::Field<double const> loc_field( make_leveled_view<const double>( loc ) );
            parallel::Field<double> glb_field( make_leveled_view<double>( glb ) );
            gather().gather( &loc_field, &glb_field, nb_fields, root );
        }
        else {
            throw_Exception( "datatype not supported", Here() );
        }
    }
}
// ----------------------------------------------------------------------------

// ----------------------------------------------------------------------------
// Gather Field
// ----------------------------------------------------------------------------
void StructuredColumns::gather( const Field& local, Field& global ) const {
    FieldSet local_fields;
    FieldSet global_fields;
    local_fields.add( local );
    global_fields.add( global );
    gather( local_fields, global_fields );
}
// ----------------------------------------------------------------------------

// ----------------------------------------------------------------------------
// Scatter FieldSet
// ----------------------------------------------------------------------------
void StructuredColumns::scatter( const FieldSet& global_fieldset, FieldSet& local_fieldset ) const {
    ATLAS_ASSERT( local_fieldset.size() == global_fieldset.size() );

    for ( idx_t f = 0; f < local_fieldset.size(); ++f ) {
        const Field& glb      = global_fieldset[f];
        Field& loc            = local_fieldset[f];
        const idx_t nb_fields = 1;
        idx_t root( 0 );
        glb.metadata().get( "owner", root );

        if ( loc.datatype() == array::DataType::kind<int>() ) {
            parallel::Field<int const> glb_field( make_leveled_view<const int>( glb ) );
            parallel::Field<int> loc_field( make_leveled_view<int>( loc ) );
            scatter().scatter( &glb_field, &loc_field, nb_fields, root );
        }
        else if ( loc.datatype() == array::DataType::kind<long>() ) {
            parallel::Field<long const> glb_field( make_leveled_view<const long>( glb ) );
            parallel::Field<long> loc_field( make_leveled_view<long>( loc ) );
            scatter().scatter( &glb_field, &loc_field, nb_fields, root );
        }
        else if ( loc.datatype() == array::DataType::kind<float>() ) {
            parallel::Field<float const> glb_field( make_leveled_view<const float>( glb ) );
            parallel::Field<float> loc_field( make_leveled_view<float>( loc ) );
            scatter().scatter( &glb_field, &loc_field, nb_fields, root );
        }
        else if ( loc.datatype() == array::DataType::kind<double>() ) {
            parallel::Field<double const> glb_field( make_leveled_view<const double>( glb ) );
            parallel::Field<double> loc_field( make_leveled_view<double>( loc ) );
            scatter().scatter( &glb_field, &loc_field, nb_fields, root );
        }
        else {
            throw_Exception( "datatype not supported", Here() );
        }

        glb.metadata().broadcast( loc.metadata(), root );
        loc.metadata().set( "global", false );
    }
}
// ----------------------------------------------------------------------------

// ----------------------------------------------------------------------------
// Scatter Field
// ----------------------------------------------------------------------------
void StructuredColumns::scatter( const Field& global, Field& local ) const {
    FieldSet global_fields;
    FieldSet local_fields;
    global_fields.add( global );
    local_fields.add( local );
    scatter( global_fields, local_fields );
}
// ----------------------------------------------------------------------------

std::string StructuredColumns::checksum( const FieldSet& fieldset ) const {
    eckit::MD5 md5;
    for ( idx_t f = 0; f < fieldset.size(); ++f ) {
        const Field& field = fieldset[f];
        if ( field.datatype() == array::DataType::kind<int>() ) {
            md5 << checksum_3d_field<int>( checksum(), field );
        }
        else if ( field.datatype() == array::DataType::kind<long>() ) {
            md5 << checksum_3d_field<long>( checksum(), field );
        }
        else if ( field.datatype() == array::DataType::kind<float>() ) {
            md5 << checksum_3d_field<float>( checksum(), field );
        }
        else if ( field.datatype() == array::DataType::kind<double>() ) {
            md5 << checksum_3d_field<double>( checksum(), field );
        }
        else {
            throw_Exception( "datatype not supported", Here() );
        }
    }
    return md5;
}
std::string StructuredColumns::checksum( const Field& field ) const {
    FieldSet fieldset;
    fieldset.add( field );
    return checksum( fieldset );
}

const StructuredGrid& StructuredColumns::grid() const {
    return *grid_;
}

namespace {


template <int RANK>
struct FixupHaloForVectors {
    FixupHaloForVectors( const StructuredColumns& ) {}
    template <typename DATATYPE>
    void apply( Field& field ) {
        std::string type = field.metadata().getString( "type", "scalar" );
        if ( type == "vector " ) {
            ATLAS_NOTIMPLEMENTED;
        }
    }
};

template <>
struct FixupHaloForVectors<2> {
    static constexpr int RANK = 2;
    const StructuredColumns& fs;
    FixupHaloForVectors( const StructuredColumns& _fs ) : fs( _fs ) {}

    template <typename DATATYPE>
    void apply( Field& field ) {
        std::string type = field.metadata().getString( "type", "scalar" );
        if ( type == "vector" ) {
            auto array = array::make_view<DATATYPE, RANK>( field );
            for ( idx_t j = fs.j_begin_halo(); j < 0; ++j ) {
                for ( idx_t i = fs.i_begin_halo( j ); i < fs.i_end_halo( j ); ++i ) {
                    idx_t n        = fs.index( i, j );
                    array( n, XX ) = -array( n, XX );
                    array( n, YY ) = -array( n, YY );
                }
            }
            for ( idx_t j = fs.grid().ny(); j < fs.j_end_halo(); ++j ) {
                for ( idx_t i = fs.i_begin_halo( j ); i < fs.i_end_halo( j ); ++i ) {
                    idx_t n        = fs.index( i, j );
                    array( n, XX ) = -array( n, XX );
                    array( n, YY ) = -array( n, YY );
                }
            }
        }
    }
};

template <>
struct FixupHaloForVectors<3> {
    static constexpr int RANK = 3;
    const StructuredColumns& fs;
    FixupHaloForVectors( const StructuredColumns& _fs ) : fs( _fs ) {}

    template <typename DATATYPE>
    void apply( Field& field ) {
        std::string type = field.metadata().getString( "type", "scalar" );
        if ( type == "vector" ) {
            auto array = array::make_view<DATATYPE, RANK>( field );
            for ( idx_t j = fs.j_begin_halo(); j < 0; ++j ) {
                for ( idx_t i = fs.i_begin_halo( j ); i < fs.i_end_halo( j ); ++i ) {
                    idx_t n = fs.index( i, j );
                    for ( idx_t k = fs.k_begin(); k < fs.k_end(); ++k ) {
                        array( n, k, XX ) = -array( n, k, XX );
                        array( n, k, YY ) = -array( n, k, YY );
                    }
                }
            }
            for ( idx_t j = fs.grid().ny(); j < fs.j_end_halo(); ++j ) {
                for ( idx_t i = fs.i_begin_halo( j ); i < fs.i_end_halo( j ); ++i ) {
                    idx_t n = fs.index( i, j );
                    for ( idx_t k = fs.k_begin(); k < fs.k_end(); ++k ) {
                        array( n, k, XX ) = -array( n, k, XX );
                        array( n, k, YY ) = -array( n, k, YY );
                    }
                }
            }
        }
    }
};


template <int RANK>
void dispatch_haloExchange( Field& field, const parallel::HaloExchange& halo_exchange, const StructuredColumns& fs ) {
    FixupHaloForVectors<RANK> fixup_halos( fs );
    if ( field.datatype() == array::DataType::kind<int>() ) {
        halo_exchange.template execute<int, RANK>( field.array(), false );
        fixup_halos.template apply<int>( field );
    }
    else if ( field.datatype() == array::DataType::kind<long>() ) {
        halo_exchange.template execute<long, RANK>( field.array(), false );
        fixup_halos.template apply<long>( field );
    }
    else if ( field.datatype() == array::DataType::kind<float>() ) {
        halo_exchange.template execute<float, RANK>( field.array(), false );
        fixup_halos.template apply<float>( field );
    }
    else if ( field.datatype() == array::DataType::kind<double>() ) {
        halo_exchange.template execute<double, RANK>( field.array(), false );
        fixup_halos.template apply<double>( field );
    }
    else {
        throw_Exception( "datatype not supported", Here() );
    }
    field.set_dirty( false );
}


template <int RANK>
void dispatch_adjointHaloExchange( Field& field, const parallel::HaloExchange& halo_exchange,
                                   const StructuredColumns& fs ) {
    FixupHaloForVectors<RANK> fixup_halos( fs );
    if ( field.datatype() == array::DataType::kind<int>() ) {
        halo_exchange.template execute_adjoint<int, RANK>( field.array(), false );
        fixup_halos.template apply<int>( field );
    }
    else if ( field.datatype() == array::DataType::kind<long>() ) {
        halo_exchange.template execute_adjoint<long, RANK>( field.array(), false );
        fixup_halos.template apply<long>( field );
    }
    else if ( field.datatype() == array::DataType::kind<float>() ) {
        halo_exchange.template execute_adjoint<float, RANK>( field.array(), false );
        fixup_halos.template apply<float>( field );
    }
    else if ( field.datatype() == array::DataType::kind<double>() ) {
        halo_exchange.template execute_adjoint<double, RANK>( field.array(), false );
        fixup_halos.template apply<double>( field );
    }
    else {
        throw_Exception( "datatype not supported", Here() );
    }
    field.set_dirty( false );
}
}  // namespace

void StructuredColumns::haloExchange( const FieldSet& fieldset, bool ) const {
    for ( idx_t f = 0; f < fieldset.size(); ++f ) {
        Field& field = const_cast<FieldSet&>( fieldset )[f];
        switch ( field.rank() ) {
            case 1:
                dispatch_haloExchange<1>( field, halo_exchange(), *this );
                break;
            case 2:
                dispatch_haloExchange<2>( field, halo_exchange(), *this );
                break;
            case 3:
                dispatch_haloExchange<3>( field, halo_exchange(), *this );
                break;
            case 4:
                dispatch_haloExchange<4>( field, halo_exchange(), *this );
                break;
            default:
                throw_Exception( "Rank not supported", Here() );
        }
    }
}

void StructuredColumns::adjointHaloExchange( const FieldSet& fieldset, bool ) const {
    for ( idx_t f = 0; f < fieldset.size(); ++f ) {
        Field& field = const_cast<FieldSet&>( fieldset )[f];
        switch ( field.rank() ) {
            case 1:
                dispatch_adjointHaloExchange<1>( field, halo_exchange(), *this );
                break;
            case 2:
                dispatch_adjointHaloExchange<2>( field, halo_exchange(), *this );
                break;
            case 3:
                dispatch_adjointHaloExchange<3>( field, halo_exchange(), *this );
                break;
            case 4:
                dispatch_adjointHaloExchange<4>( field, halo_exchange(), *this );
                break;
            default:
                throw_Exception( "Rank not supported", Here() );
        }
    }
}

void StructuredColumns::haloExchange( const Field& field, bool ) const {
    FieldSet fieldset;
    fieldset.add( field );
    haloExchange( fieldset );
}

void StructuredColumns::adjointHaloExchange( const Field& field, bool ) const {
    FieldSet fieldset;
    fieldset.add( field );
    adjointHaloExchange( fieldset );
}

size_t StructuredColumns::footprint() const {
    size_t size = sizeof( *this );
    size += ij2gp_.footprint();
    if ( field_xy_ ) {
        size += field_xy_.footprint();
    }
    if ( field_partition_ ) {
        size += field_partition_.footprint();
    }
    if ( field_global_index_ ) {
        size += field_global_index_.footprint();
    }
    if ( field_remote_index_ ) {
        size += field_remote_index_.footprint();
    }
    if ( field_index_i_ ) {
        size += field_index_i_.footprint();
    }
    if ( field_index_j_ ) {
        size += field_index_j_.footprint();
    }
    return size;
}

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


}  // namespace detail
}  // namespace functionspace
}  // namespace atlas
