#!/usr/bin/python
#import math
import pysal
from pysal.cg.standalone import get_shared_segments, get_bounding_box

__author__ = "Sergio J. Rey <srey@asu.edu> "
__all__ = ["QUEEN", "ROOK", "ContiguityWeights_binning",
           "ContiguityWeightsPolygons"]


import time

# delta to get buckets right
DELTA = 0.000001

QUEEN = 1
ROOK = 2

# constants for bucket sizes
BUCK_SM = 8
BUCK_LG = 80
SHP_SMALL = 1000


def bbcommon(bb, bbother):
    """
    Checks for overlaps of bounding boxes. First, east-west, then north-south.
    Element 0 is west, element 2 is east, element 3 is north, element 1 is
    south.
    All four checks must be false for chflag to be true, meaning the two
    bounding boxes do not overlap.
    """
    chflag = 0
    if not ((bbother[2] < bb[0]) or (bbother[0] > bb[2])):
        if not ((bbother[3] < bb[1]) or (bbother[1] > bb[3])):
            chflag = 1
    return chflag


class ContiguityWeightsPolygons:

    """
    Contiguity for a collection of polygons using a binning algorithm
    """

    def __init__(self, collection, wttype=1, bbox=None):
        """

        Parameters
        ==========

        collection: PySAL PolygonCollection 

        wttype: int
                1: Queen
                2: Rook
        """

        self.collection = collection
        self.wttype = wttype
        if bbox is None:
            try:
                bbox = collection.bbox
            except AttributeError:
                bbox = get_bounding_box(collection)
        self.shapebox = bbox
        self.numPoly = len(collection)
        self.do_weights()

    def do_weights(self):

        shapebox = self.shapebox      # bounding box

        numPoly = self.numPoly

        # bucket size
        if (numPoly < SHP_SMALL):
            bucketmin = numPoly // BUCK_SM + 2
        else:
            bucketmin = numPoly // BUCK_LG + 2
            # print 'bucketmin: ', bucketmin
        # bucket length
        lengthx = ((shapebox[2] + DELTA) - shapebox[0]) / bucketmin
        lengthy = ((shapebox[3] + DELTA) - shapebox[1]) / bucketmin

        # print lengthx, lengthy

        # initialize buckets
        columns = [set() for i in range(bucketmin)]
        rows = [set() for i in range(bucketmin)]

        minbox = shapebox[:2] * 2             # minx,miny,minx,miny
        binWidth = [lengthx, lengthy] * 2      # lenx,leny,lenx,leny
        bbcache = {}
        poly2Column = [set() for i in range(numPoly)]
        poly2Row = [set() for i in range(numPoly)]
        for i in range(numPoly):
            shpObj = self.collection[i]
            bbcache[i] = shpObj.bbox[:]
            projBBox = [int((shpObj.bbox[:][j] -
                             minbox[j]) / binWidth[j]) for j in xrange(4)]
            for j in range(projBBox[0], projBBox[2] + 1):
                columns[j].add(i)
                poly2Column[i].add(j)
            for j in range(projBBox[1], projBBox[3] + 1):
                rows[j].add(i)
                poly2Row[i].add(j)

        w = {}
        if self.wttype == QUEEN:
            # loop over polygons rather than bins
            vertCache = {}
            for polyId in xrange(numPoly):
                if polyId not in vertCache:
                    vertCache[polyId] = set(self.collection[polyId].vertices)
                idRows = poly2Row[polyId]
                idCols = poly2Column[polyId]
                rowPotentialNeighbors = set()
                colPotentialNeighbors = set()
                for row in idRows:
                    rowPotentialNeighbors = rowPotentialNeighbors.union(
                        rows[row])
                for col in idCols:
                    colPotentialNeighbors = colPotentialNeighbors.union(
                        columns[col])
                potentialNeighbors = rowPotentialNeighbors.intersection(
                    colPotentialNeighbors)
                if polyId not in w:
                    w[polyId] = set()
                for j in potentialNeighbors:
                    if polyId < j:
                        if j not in vertCache:
                            vertCache[j] = set(self.collection[j].vertices)
                        if bbcommon(bbcache[polyId], bbcache[j]):
                            vertCache[j] = set(self.collection[j].vertices)
                            common = vertCache[
                                polyId].intersection(vertCache[j])
                            if len(common) > 0:
                                w[polyId].add(j)
                                if j not in w:
                                    w[j] = set()
                                w[j].add(polyId)
        elif self.wttype == ROOK:
            # check for a shared edge
            edgeCache = {}
            # loop over polygons rather than bins
            for polyId in xrange(numPoly):
                if polyId not in edgeCache:
                    iEdges = {}
                    iVerts = self.collection[polyId].vertices
                    nv = len(iVerts)
                    ne = nv - 1
                    for i in xrange(ne):
                        l = iVerts[i]
                        r = iVerts[i + 1]
                        iEdges[(l, r)] = []
                        iEdges[(r, l)] = []
                    edgeCache[polyId] = iEdges
                iEdgeSet = set(edgeCache[polyId].keys())
                idRows = poly2Row[polyId]
                idCols = poly2Column[polyId]
                rowPotentialNeighbors = set()
                colPotentialNeighbors = set()
                for row in idRows:
                    rowPotentialNeighbors = rowPotentialNeighbors.union(
                        rows[row])
                for col in idCols:
                    colPotentialNeighbors = colPotentialNeighbors.union(
                        columns[col])
                potentialNeighbors = rowPotentialNeighbors.intersection(
                    colPotentialNeighbors)
                if polyId not in w:
                    w[polyId] = set()
                for j in potentialNeighbors:
                    if polyId < j:
                        if bbcommon(bbcache[polyId], bbcache[j]):
                            if j not in edgeCache:
                                jVerts = self.collection[j].vertices
                                jEdges = {}
                                nv = len(jVerts)
                                ne = nv - 1
                                for e in xrange(ne):
                                    l = jVerts[e]
                                    r = jVerts[e + 1]
                                    jEdges[(l, r)] = []
                                    jEdges[(r, l)] = []
                                edgeCache[j] = jEdges
                            # for edge in edgeCache[j]:
                            if iEdgeSet.intersection(edgeCache[j].keys()):
                                w[polyId].add(j)
                                if j not in w:
                                    w[j] = set()
                                w[j].add(polyId)
                                # break
        else:
            raise Exception("Unsupported weight type.")

        self.w = w

if __name__ == "__main__":
    import time
    fname = pysal.examples.get_path('NAT.shp')
    print('QUEEN binning')
    t0 = time.time()
    qb = ContiguityWeights_binning(pysal.open(fname), QUEEN)
    t1 = time.time()
    print("using " + str(fname))
    print("time elapsed for queen... using bins: " + str(t1 - t0))

    t0 = time.time()
    rb = ContiguityWeights_binning(pysal.open(fname), ROOK)
    t1 = time.time()
    print('Rook binning')
    print("using " + str(fname))
    print("time elapsed for rook... using bins: " + str(t1 - t0))

    from pysal.weights._contW_rtree import ContiguityWeights_rtree

    t0 = time.time()
    rt = ContiguityWeights_rtree(pysal.open(fname), ROOK)
    t1 = time.time()

    print("time elapsed for rook... using rtree: " + str(t1 - t0))
    print(rt.w == rb.w)

    print('QUEEN')
    t0 = time.time()
    qt = ContiguityWeights_rtree(pysal.open(fname), QUEEN)
    t1 = time.time()
    print("using " + str(fname))
    print("time elapsed for queen... using rtree: " + str(t1 - t0))
    print(qb.w == qt.w)

    print('knn4')
    t0 = time.time()
    knn = pysal.knnW_from_shapefile(fname, k=4)
    t1 = time.time()
    print(t1 - t0)

    print('rook from shapefile')
    t0 = time.time()
    knn = pysal.rook_from_shapefile(fname)
    t1 = time.time()
    print(t1 - t0)
