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 = [i]
        jmat = [j]

        # latitude
        if i:
            imat.append(i-1)
        if i < self.n[0]-1:
            imat.append(i+1)

        # longitude: wraps around
        if j:
            jmat.append(j-1)
        else:
            jmat.append(self.n[-1] - 1)
        if j == self.n[-1] - 1:
            jmat.append(0)
        else:
            jmat.append(j+1)

        return [(_i, _j) for _i in imat for _j in jmat
                if (_i,_j) != (i,j)]

    def neighbor_points(self, i, j):
        """return all points in a lat-lon box"""
        geoids = set()
        geoids.update(self[(i,j)])
        for ii, jj in self.neighbors(i,j):
            geoids.update(self[ii,jj])
        return geoids


class GriddedLocations(object):

    def __init__(self, locations):
        raise NotImplementedError('TODO')