"""$ rio blocks"""


import json
import logging
import os.path

import click
import cligj

import rasterio
from rasterio.rio import options
from rasterio.rio.helpers import write_features
from rasterio.warp import transform_bounds


logger = logging.getLogger(__name__)


class _Collection(object):

    """For use with `rasterio.rio.helpers.write_features()`."""

    def __init__(self, dataset, bidx, precision=6, geographic=True):

        """Export raster dataset windows to GeoJSON polygon features.

        Parameters
        ----------
        dataset : a dataset object opened in 'r' mode
            Source dataset
        bidx : int
            Extract windows from this band
        precision : int, optional
            Coordinate precision
        geographic : bool, optional
            Reproject geometries to ``EPSG:4326`` if ``True``

        Yields
        ------
        dict
            GeoJSON polygon feature
        """

        self._src = dataset
        self._bidx = bidx
        self._precision = precision
        self._geographic = geographic

    def _normalize_bounds(self, bounds):
        if self._geographic:
            bounds = transform_bounds(self._src.crs, 'EPSG:4326', *bounds)
        if self._precision >= 0:
            bounds = (round(v, self._precision) for v in bounds)
        return bounds

    @property
    def bbox(self):
        return tuple(self._normalize_bounds(self._src.bounds))

    def __call__(self):
        gen = self._src.block_windows(bidx=self._bidx)
        for idx, (block, window) in enumerate(gen):
            bounds = self._normalize_bounds(self._src.window_bounds(window))
            xmin, ymin, xmax, ymax = bounds
            yield {
                'type': 'Feature',
                'id': '{0}:{1}'.format(os.path.basename(self._src.name), idx),
                'properties': {
                    'block': json.dumps(block),
                    'window': window.todict(),
                },
                'geometry': {
                    'type': 'Polygon',
                    'coordinates': [[
                        (xmin, ymin),
                        (xmin, ymax),
                        (xmax, ymax),
                        (xmax, ymin)
                    ]]
                }
            }


@click.command()
@options.file_in_arg
@options.output_opt
@cligj.precision_opt
@cligj.indent_opt
@cligj.compact_opt
@cligj.projection_projected_opt
@cligj.sequence_opt
@cligj.use_rs_opt
@click.option(
    '--bidx', type=click.INT, default=0,
    help="Index of the band that is the source of shapes.")
@click.pass_context
def blocks(
        ctx, input, output, precision, indent, compact, projection, sequence,
        use_rs, bidx):

    """Write dataset blocks as GeoJSON features.

    This command writes features describing a raster's internal blocks, which
    are used directly for raster I/O.  These features can be used to visualize
    how a windowed operation would operate using those blocks.

    Output features have two JSON encoded properties: block and window.  Block
    is a two element array like '[0, 0]' describing the window's position
    in the input band's window layout.  Window is a two element array
    containing two more two element arrays like '[[0, 256], [0, 256]' and
    describes the range of pixels the window covers in the input band.  Values
    are JSON encoded for better interoperability.

    Block windows are extracted from the dataset (all bands must have matching
    block windows) by default, or from the band specified using the '--bidx
    option:
    \b

        $ rio shapes --bidx 3 tests/data/RGB.byte.tif

    By default a GeoJSON 'FeatureCollection' is written, but the --sequence'
    option produces a GeoJSON feature stream instead.
    \b

        $ rio shapes tests/data/RGB.byte.tif --sequence

    Output features are reprojected to 'WGS84' unless the '--projected' flag is
    provided, which causes the output to be kept in the input datasource's
    coordinate reference system.

    For more information on exactly what blocks and windows represent, see
    'dataset.block_windows()'.
    """

    dump_kwds = {'sort_keys': True}

    if indent:
        dump_kwds['indent'] = indent
    if compact:
        dump_kwds['separators'] = (',', ':')

    stdout = click.open_file(
        output, 'w') if output else click.get_text_stream('stdout')

    with ctx.obj['env'], rasterio.open(input) as src:

        if bidx and bidx not in src.indexes:
            raise click.BadParameter("Not a valid band index")

        collection = _Collection(
            dataset=src,
            bidx=bidx,
            precision=precision,
            geographic=projection != 'projected')

        write_features(
            stdout, collection,
            sequence=sequence,
            geojson_type='feature' if sequence else 'collection',
            use_rs=use_rs,
            **dump_kwds)