from .coordinate_helpers import CoordinateHelper
from .transforms import WCSWorld2PixelTransform
from .utils import coord_type_from_ctype
from .frame import RectangularFrame

from . import six


class CoordinatesMap(object):

    def __init__(self, axes, wcs, transform=None):

        # Keep track of parent axes and WCS
        self._axes = axes
        self._wcs = wcs

        # Set up transform
        if transform is None:
            self._transform = WCSWorld2PixelTransform(self._wcs)
        else:
            self._transform = transform

        self.frame = RectangularFrame(axes, self._transform)

        # Set up coordinates
        self._coords = {}
        for coord_index in [0, 1]:
            coord_type = coord_type_from_ctype(wcs.wcs.ctype[coord_index])
            self._coords[coord_index] = CoordinateHelper(parent_axes=axes,
                                                         transform=self._transform,
                                                         coord_index=coord_index,
                                                         coord_type=coord_type,
                                                         frame=self.frame)


        # Set up aliases for coordinates
        name_1 = self._wcs.wcs.ctype[0][:4].replace('-', '')
        self._coords[name_1.lower()] = self._coords[0]
        name_2 = self._wcs.wcs.ctype[1][:4].replace('-', '')
        self._coords[name_2.lower()] = self._coords[1]

        # Set up default labels if possible

    def __getitem__(self, item):
        if isinstance(item, six.string_types):
            return self._coords[item.lower()]
        else:
            return self._coords[item]

    def set_visible(self, visibility):
        raise NotImplementedError()

    def enable_offset_mode(self, reference_coordinates):
        raise NotImplementedError()

    def disable_offset_mode(self):
        raise NotImplementedError()

    def grid(self, draw_grid=True, **kwargs):
        """
        Plot gridlines for both coordinates.

        Standard matplotlib appearance options (color, alpha, etc.) can be
        passed as keyword arguments.

        Parameters
        ----------
        draw_grid : bool
            Whether to show the gridlines
        """
        self[0].grid(draw_grid=draw_grid, **kwargs)
        self[1].grid(draw_grid=draw_grid, **kwargs)

@@ -65,35 +65,35 @@
 s import WCS
-World2Pixel
+Pixel2World
 Transform%0Afr
@@ -267,16 +267,28 @@
 orm=None
+, slice=None
 ):%0A%0A    
@@ -472,19 +472,19 @@
  WCS
-World2Pixel
+Pixel2World
 Tran
@@ -498,16 +498,29 @@
 elf._wcs
+, slice=slice
 )%0A      
@@ -716,14 +716,34 @@
  in 
-%5B0, 1%5D
+range(self._wcs.wcs.naxis)
 :%0A  
@@ -1945,32 +1945,144 @@
 ementedError()%0A%0A
+    def __iter__(self):%0A        for coord in range(self._wcs.wcs.naxis):%0A            yield self._coords%5Bcoord%5D%0A%0A
     def grid(sel
@@ -2411,67 +2411,44 @@
     
-self%5B0%5D.grid(draw_grid=draw_grid, **kwargs)%0A        self%5B1%5D
+for coord in self:%0A            coord
 .gri

Generalize CoordinatesMap to arbitrary number of coordinates, with slicing