import pysal as ps
import json

# compatibility for Python 2 & Python 3:
try:
    from urllib.request import urlopen
except ImportError:
    from urllib2 import urlopen

import copy
import numpy as np

"""
Weights Meta Data

Prototyping meta data functions and classes for weights provenance


Based on Anselin, L., S.J. Rey and W. Li (2014) "Metadata and provenance for
spatial analysis: the case of spatial weights." International Journal of
Geographical Information Science.  DOI:10.1080/13658816.2014.917313



TODO
----

- Document each public function with working doctest
- Abstract the test files as they currently assume location is source
  directory
- have wmd_reader take either a wmd file or a wmd dictionary/object
"""
__author__ = "Sergio J. Rey <srey@asu.edu>, Wenwen Li <wenwen@asu.edu>"


def wmd_reader(fileName):
    """

    Examples
    --------
    >>> import wmd
    >>> wr = wmd.wmd_reader('w1rook.wmd')
    wmd_reader failed:  w1rook.wmd
    >>> wr = wmd.wmd_reader('wrook1.wmd')
    >>> wr.neighbors[2]
    [0, 1, 3, 4]
    """

    try:
        meta_data = _uri_reader(fileName)
    except:
        try:
            with open(fileName, 'r') as fp:
                meta_data = json.load(fp)
                global fullmeta
                fullmeta = {}
                fullmeta['root'] = copy.deepcopy(meta_data)
                w = _wmd_parser(meta_data)
                return w
        except:
            print 'wmd_reader failed: ', fileName


class WMD(ps.W):
    """Weights Meta Data Class"""
    def __init__(self, neighbors=None,  weights=None, id_order=None):
        self.meta_data = {}
        super(WMD, self).__init__(neighbors, weights, id_order)

    # override transform property to record any post-instantiation
    # transformations in meta data

    @ps.W.transform.setter
    def transform(self, value):
        super(WMD, WMD).transform.__set__(self, value)
        self.meta_data['transform'] = self._transform

    def write(self, fileName, data=False):
        """

        Examples
        --------
        >>> import wmd
        >>> wr = wmd.wmd_reader('w1rook.wmd')
        wmd_reader failed:  w1rook.wmd
        >>> wr = wmd.wmd_reader('wrook1.wmd')
        >>> wr.write('wr1.wmd')
        >>> wr1 = wmd.wmd_reader('wr1.wmd')
        >>> wr.neighbors[2]
        [0, 1, 3, 4]
        >>> wr1.neighbors[2]
        [0, 1, 3, 4]
        >>>

        """
        _wmd_writer(self, fileName, data=data)

######################### Private functions #########################


def _wmd_writer(wmd_object, fileName, data=False):
    try:
        with open(fileName, 'w') as f:
            if data:
                wmd_object.meta_data['data'] = {}
                wmd_object.meta_data['data']['weights'] = wmd_object.weights
                wmd_object.meta_data['data']['neighbors'] = wmd_object.neighbors
            json.dump(wmd_object.meta_data,
                      f,
                      indent=4,
                      separators=(',', ': '))
    except:
        print 'wmd_writer failed.'


def _block(arg_dict):
    """
    General handler for block weights


    Examples
    --------
    >>> w = wmd_reader('taz_block.wmd')
    >>> w.n
    4109
    >>> w.meta_data
    {'root': {u'input1': {u'data1': {u'type': u'dbf',
                                     u'uri': u'http://toae.org/pub/taz.dbf'}},
                          u'weight_type': u'block', u'transform': u'O',
                          u'parameters': {u'block_variable': u'CNTY'}}}

    """
    input1 = arg_dict['input1']
    for key in input1:
        input1 = input1[key]
        break
    uri = input1['uri']
    weight_type = arg_dict['weight_type'].lower()

    var_name = arg_dict['parameters']['block_variable']
    dbf = ps.open(uri)
    block = np.array(dbf.by_col(var_name))
    dbf.close()
    w = ps.weights.util.block_weights(block)
    w = WMD(w.neighbors, w.weights)
    w.meta_data = {}
    w.meta_data['input1'] = {"type": 'dbf', 'uri': uri}
    w.meta_data['transform'] = w.transform
    w.meta_data['weight_type'] = weight_type
    w.meta_data['parameters'] = {'block_variable': var_name}

    return w


def _contiguity(arg_dict):
    """
    General handler for building contiguity weights from shapefiles

    Examples
    --------

    >>> w = wmd_reader('wrook1.wmd')
    >>> w.n
    49
    >>> w.meta_data
    {'root': {u'input1': {u'data1': {u'type': u'shp',
                                     u'uri': u'http://toae.org/pub/columbus.shp'}},
                          u'weight_type': u'rook', u'transform': u'O'}}
    """
    input1 = arg_dict['input1']
    for key in input1:
        input1 = input1[key]
        break
    uri = input1['uri']
    weight_type = arg_dict['weight_type']
    weight_type = weight_type.lower()
    if weight_type == 'rook':
        w = ps.rook_from_shapefile(uri)
    elif weight_type == 'queen':
        w = ps.queen_from_shapefile(uri)
    else:
        print "Unsupported contiguity criterion: ", weight_type
        return None
    if 'parameters' in arg_dict:
        order = arg_dict['parameters'].get('order', 1)  # default to 1st order
        lower = arg_dict['parameters'].get('lower', 0)  # default to exclude lower orders
        if order > 1:
            w_orig = w
            w = ps.higher_order(w, order)
            if lower:
                for o in xrange(order-1, 1, -1):
                    w = ps.weights.w_union(ps.higher_order(w_orig, o), w)
                w = ps.weights.w_union(w, w_orig)
        parameters = arg_dict['parameters']
    else:
        parameters = {'lower': 0, 'order': 1}
    w = WMD(w.neighbors, w.weights)
    w.meta_data = {}
    w.meta_data["input1"] = {"type": 'shp', 'uri': uri}
    w.meta_data["transform"] = w.transform
    w.meta_data["weight_type"] = weight_type
    w.meta_data['parameters'] = parameters
    return w


def _kernel(arg_dict):
    """
    General handler for building kernel based weights from shapefiles

    Examples
    --------

    >>> w = wmd_reader('kernel.wmd')
    >>> w.n
    49
    >>> w.meta_data
    {'root': {u'input1': {u'data1': {u'type': u'shp',
                                     u'uri': u'../examples/columbus.shp'}},
              u'weight_type': u'kernel', u'transform': u'O',
              u'parameters': {u'function': u'triangular',
              u'bandwidths': None, u'k': 2}}}



    """
    input1 = arg_dict['input1']['data1']
    uri = input1['uri']
    weight_type = arg_dict['weight_type']
    weight_type = weight_type.lower()
    k = 2
    bandwidths = None
    function = 'triangular'
    if 'parameters' in arg_dict:
        k = arg_dict['parameters'].get('k', k)  # set default to 2
        bandwidths = arg_dict['parameters'].get('bandwidths', bandwidths)
        function = arg_dict['parameters'].get('function', function)
    else:
        parameters = {}
        parameters['k'] = k
        parameters['bandwidths'] = bandwidths
        parameters['function'] = function
        arg_dict['parameters'] = parameters


    if weight_type == 'akernel':
        # adaptive kernel
        w = ps.adaptive_kernelW_from_shapefile(uri, bandwidths = bandwidths,
                k=k, function = function)
    elif weight_type == 'kernel':
        w = ps.kernelW_from_shapefile(uri, k=k, function = function)
    else:
        print "Unsupported kernel: ",weight_type
        return None
    w = WMD(w.neighbors, w.weights)
    w.meta_data = {}
    w.meta_data["input1"] = {"type": 'shp', 'uri':uri}
    w.meta_data["transform"] = w.transform
    w.meta_data["weight_type"] =  weight_type
    w.meta_data['parameters'] = arg_dict['parameters']
    return w


def _distance(arg_dict):
    """
    General handler for distance based weights obtained from shapefiles
    """
    input1 = arg_dict['input1']
    uri = input1['uri']
    weight_type = arg_dict['weight_type']
    weight_type = weight_type.lower()
    k = 2
    id_variable = None
    p = 2
    radius = None
    if 'parameters' in arg_dict:
        k = arg_dict['parameters'].get('k',k) # set default to 2
        id_variable = arg_dict['parameters'].get('id_variable', id_variable)
        p = arg_dict['parameters'].get('p',p)
        radius = arg_dict['parameters'].get('radius', radius)

    else:
        parameters = {}
        parameters['k'] = 2
        parameters['id_variable'] = None
        parameters['radius'] = None
        parameters['p'] = 2
        arg_dict['parameters'] = parameters

    if weight_type == 'knn':
        w = ps.knnW_from_shapefile(uri,k=k,p=p,idVariable=id_variable,
                radius=radius)
        w = WMD(w.neighbors, w.weights)
        w.meta_data = {}
        w.meta_data["input1"] = {"type": 'shp', 'uri':uri}
        w.meta_data["weight_type"] =  'knn'
        w.meta_data["transform"] = w.transform
        w.meta_data['parameters'] = arg_dict['parameters']
        return w

def _higher_order(arg_dict):
    wmd = arg_dict['wmd']
    order = 2
    if 'parameters' in arg_dict:
        order = arg_dict['parameters'].get('order', order)
    else:
        parameters = {}
        parameters['order'] = order
        arg_dict['parameters'] = parameters


    w = ps.higher_order(wmd, order)
    w = WMD(w.neighbors, w.weights)
    w.meta_data = {}
    w.meta_data['input1'] = arg_dict['input1']
    w.meta_data['parameters'] = arg_dict['parameters']

    return w

def _intersection(arg_dict):
    #wmd = arg_dict['wmd']
    w1 = arg_dict['input1']['data1']['uri']
    w2 = arg_dict['input1']['data2']['uri']
    w = ps.w_intersection(w1,w2)
    w = WMD(w.neighbors, w.weights)
    return w

def _geojsonf(arg_dict):
    """
    Handler for local geojson files
    """
    input1 = arg_dict['input1']
    uri = input1['uri']
    weight_type = arg_dict['weight_type']
    weight_type = weight_type.lower()
    id_variable = None

    if weight_type == 'queen_geojsonf':
        w = ps.weights.user.queen_from_geojsonf(uri)
        w.meta_data = {}
        w.meta_data["input1"] = {"type": 'geojsonf', 'uri':uri}
        w.meta_data["weight_type"] =  'queen'
        w.meta_data["transform"] = w.transform
        return w

# wrapper dict that maps specific weights types to a handler function that
# builds the specific weights instance
WEIGHT_TYPES = {}
WEIGHT_TYPES['rook'] = _contiguity
WEIGHT_TYPES['queen'] = _contiguity
WEIGHT_TYPES['akernel'] = _kernel
WEIGHT_TYPES['kernel'] = _kernel
WEIGHT_TYPES['knn'] = _distance
WEIGHT_TYPES['higher_order'] = _higher_order
WEIGHT_TYPES['block'] = _block
WEIGHT_TYPES['intersection'] = _intersection
#WEIGHT_TYPES['queen_geojsonf'] = geojsonf
#WEIGHT_TYPES['geojsons'] = geojsons


def _uri_reader(uri):
    j = json.load(urlopen(uri))
    return j

def _wmd_read_only(fileName):
    try:
        meta_data = _uri_reader(fileName)
    except:
        try:
            with open(fileName, 'r') as fp:
                meta_data = json.load(fp)
                return meta_data
        except:
            print '_wmd_read_only failed: ', fileName

def _wmd_parser(wmd_object):
    if 'root' in wmd_object:
        wmd_object = wmd_object['root']
    weight_type = wmd_object['weight_type'].lower()
    for key in wmd_object['input1']:
        #print key
        if wmd_object['input1'][key]['type'] == 'prov':
            #      call wmd_reader
            uri = wmd_object['input1'][key]['uri']

            meta_data = _wmd_read_only(uri)
            fullmeta[uri] = copy.deepcopy(meta_data) #add full metadata
            wmd = _wmd_parser(meta_data)
            wmd_object['input1'][key]['uri'] = wmd
        else:
            # handle distributed files
            uri = wmd_object['input1'][key]['uri']
            try:
                tmp = open(uri)
                #print ' tmp: ', tmp
                wmd_object['input1'][key]['uri'] = uri
            except:
                _download_shapefiles(uri)
                uri = uri.split("/")[-1]
                wmd_object['input1'][key]['uri'] = uri # use local copy

    if weight_type in WEIGHT_TYPES:
        #print weight_type
        wmd  = WEIGHT_TYPES[weight_type](wmd_object)
        wmd.meta_data = fullmeta
    else:
        print 'Unsupported weight type: ', weight_type

    return wmd

def _download_shapefiles(file_name):
    file_parts = file_name.split("/")
    file_prefix = file_parts[-1].split(".")[0]
    exts = [ ".shp", ".dbf", ".shx" ]
    for ext in exts:
        # rebuild url
        file_name = file_prefix + ext
        file_parts[-1] = file_name
        new_url = "/".join(file_parts)
        #print file_name, new_url
        u = urlopen(new_url)
        f = open(file_name, 'wb')
        meta = u.info()
        file_size = int(meta.getheaders("Content-Length")[0])
        #print "Downloading: %s Bytes: %s" % (file_name, file_size)
        file_size_dl = 0
        block_sz = 8192
        while True:
            bf = u.read(block_sz)
            if not bf:
                break
            file_size_dl += len(bf)
            f.write(bf)
            status = r"%10d [%3.2f%%]" % (file_size_dl, file_size_dl * 100. /
                    file_size)
            status = status + chr(8)* (len(status)+1)
        #print status, f.close()

if __name__ == '__main__':

    # distributed file
    w1 = wmd_reader("wrook1.wmd")

##    # order
##    w1o = wmd_reader('wrooko1.wmd')
##    w2o = wmd_reader('wrooko2.wmd')
##    w2ol = wmd_reader('wrooko2l.wmd')
##
##    # kernels
    ak1 = wmd_reader('akern1.wmd')
    kern = wmd_reader('kernel.wmd')
##
##    # knn
##    knn = wmd_reader('knn.wmd')
##
##
##
##    # moran workflow
##    import pysal as ps


    # geojson
    #wj = wmd_reader("wgeojson.wmd")


    # here we test chaining
#    r1 = wmd_reader('chain2inputs.wmd')
#    print "full metadata is listed below: \n", fullmeta
    # r2 = wmd_reader('chain2.wmd')

    taz_int = wmd_reader("taz_intersection.wmd")



    ## intersection between queen and block weights
    #import numpy as np
    #w = ps.lat2W(4,4)
    #block_variable = np.ones((w.n,1))
    #block_variable[:8] = 0
    #w_block = ps.weights.util.block_weights(block_variable)

    #w_intersection = ps.w_intersection(w, w_block)


    ## with Columbus example using EW as the block and queen
    #dbf = ps.open("columbus.dbf")
    #ew = np.array(dbf.by_col("EW"))
    #dbf.close()
    #w_ew = ps.weights.util.block_weights(ew)
    #wr = ps.rook_from_shapefile("columbus.shp")
    #w_int = ps.w_intersection(w_ew, wr)


    #blk = wmd_reader('block2.wmd')

    #taz_int = wmd_reader("http://spatial.csf.asu.edu/taz_intersection.wmd")

