#!/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 non-overlap 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, indicating that the two

    bounding boxes 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 range(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 range(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 range(numPoly):

                if polyId not in edgeCache:

                    iEdges = {}

                    iVerts = self.collection[polyId].vertices

                    nv = len(iVerts)

                    ne = nv - 1

                    for i in range(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 range(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(list(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))

