#!/usr/bin/env python

"""
grid locations, for speed and fidelity
"""

import math
import sys
from .cli import CitiesParser
from .locations import locations

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)
        index = self.index(lat, lon)
        self[index].add(geoid)
        return index

    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, nlat=90, nlon=90):
        self.locations = _locations
        self.grid = LatLonGrid(nlat, nlon)
        self.gridloc = {}
        for geoid, (lat, lon) in locations.items():
            gridloc[geoid] = self.grid.add(geoid, lat, lon)


def main(args=sys.argv[1:]):
    """CLI"""

    # parse command line
    parser = CitiesParser(description=__doc__)
    options = parser.parse_args(args)

    # read locations
    city_locations = locations(parser.read_cities())

    # make a grid
    gridded_locations = GriddedLocations(city_locations)

if __name__ == '__main__':
    main()