import numpy as np
from ..crs import CartesianCRS

def coordinate_transformer(crs1, crs2):
    if crs1 == crs2:
        def transformer(x, y):
            return x, y
    else:
        def transformer(x, y):
            return crs1.transform(crs2, x, y)
    return transformer

class CoordinateGenerator(object):
    """ Class that generates coordinates from grid indices in a specified
    coordinate system. """

    def __init__(self, transform, size, transform_crs=CartesianCRS, output_crs=CartesianCRS):
        """ Creates a CoordinateGenerator. A CoordinateGenerator generates
        world positions (geographical or projected) from indices or slices
        based on a RegularGrid transform tuple.

        Parameters
        ----------
        transform : tuple
            RegularGrid transform tuple in the form (x0, y0, dx, dy, sx, sy)
        size : tuple
            grid size for slicing
        transform_crs : CRS
            coordinate system used by the transform tuple
        output_crs : CRS
            coordinate system of the positions generated by indexing the
            CoordinateGenerator

        Usage
        -----
        # grid is a RegularGrid instance
        cg = grid.coordinates(crs=LonLatWGS84)

        # indexing returns the coordinates of cells in the grid, transformed to
        # the desired coordinate system
        x, y = cg[100, 250]

        # slicing is supported
        X, Y = cg[:, 200:32:-2]
        allX, allY = cg[:,:]
        """
        self.transform = transform
        self.size = size
        self._coordtransformer = coordinate_transformer(transform_crs, output_crs)

    def __getitem__(self, key):
        if len(key) != 2:
            raise KeyError("CoordinateGenerator indices must be length 2")
        ikey, jkey = key
        if isinstance(ikey, slice):
            ikey = np.arange(*ikey.indices(self.size[0]))
        if isinstance(jkey, slice):
            jkey = np.arange(*jkey.indices(self.size[1]))
        if hasattr(ikey, "__len__") and hasattr(jkey, "__len__") and \
                (len(ikey) > 1) and (len(jkey) > 1):
            jkey, ikey = np.meshgrid(jkey, ikey)
        T = self.transform
        x_ = T[0] + jkey*T[2] + ikey*T[4] + 0.5*T[2] + 0.5*T[4]
        y_ = T[1] + ikey*T[3] + jkey*T[5] + 0.5*T[3] + 0.5*T[5]
        return self._coordtransformer(x_, y_)

