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

#include <cstdint>

#include "atlas/grid/StructuredGrid.h"
#include "atlas/grid/detail/distribution/BandsDistribution.h"
#include "atlas/runtime/Exception.h"

namespace atlas {
namespace grid {
namespace detail {
namespace partitioner {

size_t BandsPartitioner::blocksize( const Grid& grid ) const {
    if ( blocksize_ == BLOCKSIZE_NX ) {
        if ( auto regular = RegularGrid( grid ) ) {
            return regular.nx();
        }
        return 1;
    }
    ATLAS_ASSERT( blocksize_ > 0 );
    return size_t( blocksize_ );
}

BandsPartitioner::BandsPartitioner() : Partitioner(), blocksize_{1} {}

BandsPartitioner::BandsPartitioner( int N, int blocksize ) : Partitioner( N ), blocksize_( blocksize ) {}

Distribution BandsPartitioner::partition( const Partitioner::Grid& grid ) const {
    if ( not distribution::BandsDistribution<int32_t>::detectOverflow( grid.size(), nb_partitions(),
                                                                       blocksize( ( grid ) ) ) ) {
        return Distribution{
            new distribution::BandsDistribution<int32_t>{grid, nb_partitions(), type(), blocksize( grid )}};
    }
    else {
        return Distribution{
            new distribution::BandsDistribution<int64_t>{grid, nb_partitions(), type(), blocksize( grid )}};
    }
}

void BandsPartitioner::partition( const Partitioner::Grid& grid, int part[] ) const {
    gidx_t gridsize = grid.size();
    if ( not distribution::BandsDistribution<int32_t>::detectOverflow( grid.size(), nb_partitions(),
                                                                       blocksize( ( grid ) ) ) ) {
        distribution::BandsDistribution<int32_t> distribution{grid, nb_partitions(), type(), blocksize( grid )};
        for ( gidx_t n = 0; n < gridsize; ++n ) {
            part[n] = distribution.function( n );
        }
    }
    else {
        distribution::BandsDistribution<int64_t> distribution{grid, nb_partitions(), type(), blocksize( grid )};
        for ( gidx_t n = 0; n < gridsize; ++n ) {
            part[n] = distribution.function( n );
        }
    }
}

}  // namespace partitioner
}  // namespace detail
}  // namespace grid
}  // namespace atlas

namespace {
atlas::grid::detail::partitioner::PartitionerBuilder<atlas::grid::detail::partitioner::BandsPartitioner> __Bands(
    atlas::grid::detail::partitioner::BandsPartitioner::static_type() );
}
