import math

class LatLonGrid(object):

    lat_range = (-90., 90)
    lon_range = (-180., 180)  # however, this wraps around

    def __init__(self, nlat, nlon):
        self.n = (nlat, nlon)
        self.range = (self.lat_range, self.lon_range)
        self.d = ((self.lat_range[-1] - self.lat_range[0])/self.n[0],
                  (self.lon_range[-1] - self.lon_range[0])/self.n[1])
        self.create_grid()

    def create_grid(self):
        self.grid = []
        for _ in xrange(self.n[0]):
            self.grid.append([set() for _ in xrange(self.n[-1])])

    def add(self, geoid, lat, lon):
        latlon = (lat, lon)
        self[self.index(lat, lon)].add(geoid)

    def __getitem__(self, index):
        """
        index -- 2-tuple or list of i and j indices
        """
        return self.grid[index[0]][index[1]]

    def index(self, lat, lon):
        return [int(math.floor((val-self.range[i][0])/self.d[i]))
                for i, val in enumerate((lat, lon))]

    def neighbors(self, i, j):
        """
        return neighbors of points i, j
        """
        imat = []
        jmat = []

        # latitude
        if i:
            imat.append(i-1)
            raise NotImplementedError('TODO')