# This file is part of libkd.
# Licensed under a 3-clause BSD style license - see LICENSE
from __future__ import print_function
from . import spherematch_c

from ..util.starutil_numpy import radectoxyz, deg2dist, dist2deg, distsq2deg

import numpy as np

def match_xy(x1,y1, x2,y2, R, **kwargs):
    '''
    Like match_radec, except for plain old 2-D points.
    '''
    I,d = match(np.vstack((x1,y1)).T, np.vstack((x2,y2)).T, R, **kwargs)
    return I[:,0],I[:,1],d
    
# Copied from "celestial.py" by Sjoert van Velzen.
def match_radec(ra1, dec1, ra2, dec2, radius_in_deg, notself=False,
                nearest=False, indexlist=False, count=False):
    '''
    Cross-matches numpy arrays of RA,Dec points.

    Behaves like spherematch.pro of IDL.

    Parameters
    ----------
    ra1, dec1, ra2, dec2 : numpy arrays, or scalars.
        RA,Dec in degrees of points to match.

    radius_in_deg : float
        Search radius in degrees.

    notself : boolean
        If True, avoids returning 'identity' matches;
        ASSUMES that ra1,dec1 == ra2,dec2.

    nearest : boolean
        If True, returns only the nearest match in *(ra2,dec2)*
        for each point in *(ra1,dec1)*.

    indexlist : boolean
        If True, returns a list of length *len(ra1)*, containing *None*
        or a list of ints of matched points in *ra2,dec2*.
        

    Returns
    -------
    m1 : numpy array of integers
        Indices into the *ra1,dec1* arrays of matching points.
    m2 : numpy array of integers
        Same, but for *ra2,dec2*.
    d12 : numpy array, float
        Distance, in degrees, between the matching points.
    '''
    # Convert to coordinates on the unit sphere
    xyz1 = radectoxyz(ra1, dec1)
    #if all(ra1 == ra2) and all(dec1 == dec2):
    if ra1 is ra2 and dec1 is dec2:
        xyz2 = xyz1
    else:
        xyz2 = radectoxyz(ra2, dec2)
    r = deg2dist(radius_in_deg)

    extra = ()
    if nearest:
        X = _nearest_func(xyz2, xyz1, r, notself=notself, count=count)
        if not count:
            (inds,dists2) = X
            I = np.flatnonzero(inds >= 0)
            J = inds[I]
            d = distsq2deg(dists2[I])
        else:
            #print 'X', X
            #(inds,dists2,counts) = X
            J,I,d,counts = X
            extra = (counts,)
            print('I', I.shape, I.dtype)
            print('J', J.shape, J.dtype)
            print('counts', counts.shape, counts.dtype)
    else:
        X = match(xyz1, xyz2, r, notself=notself, indexlist=indexlist)
        if indexlist:
            return X
        (inds,dists) = X
        dist_in_deg = dist2deg(dists)
        I,J = inds[:,0], inds[:,1]
        d = dist_in_deg[:,0]
        
    return (I, J, d) + extra


def cluster_radec(ra, dec, R, singles=False):
    '''
    Finds connected groups of objects in RA,Dec space.

    Returns a list of lists of indices that are connected,
    EXCLUDING singletons.

    If *singles* is *True*, also returns the indices of singletons.
    '''
    I,J,d = match_radec(ra, dec, ra, dec, R, notself=True)

    # 'mgroups' maps each index in a group to a list of the group members
    mgroups = {}
    # 'ugroups' is a list of the unique groups
    ugroups = []
    
    for i,j in zip(I,J):
        # Are both sources already in groups?
        if i in mgroups and j in mgroups:
            # Are they already in the same group?
            if mgroups[i] == mgroups[j]:
                continue
            # merge if they are different;
            # assert(they are disjoint)
            lsti = mgroups[i]
            lstj = mgroups[j]
            merge = lsti + lstj
            for k in merge:
                mgroups[k] = merge

            ugroups.remove(lsti)
            ugroups.remove(lstj)
            ugroups.append(merge)

        elif i in mgroups:
            # Add j to i's group
            lst = mgroups[i]
            lst.append(j)
            mgroups[j] = lst
        elif j in mgroups:
            # Add i to j's group
            lst = mgroups[j]
            lst.append(i)
            mgroups[i] = lst
        else:
            # Create a new group
            lst = [i,j]
            mgroups[i] = lst
            mgroups[j] = lst

            ugroups.append(lst)


    if singles:
        S = np.ones(len(ra), bool)
        for g in ugroups:
            S[np.array(g)] = False
        S = np.flatnonzero(S)
        return ugroups,S

    return ugroups





def _cleaninputs(x1, x2):
    fx1 = x1.astype(np.float64)
    if x2 is x1:
        fx2 = fx1
    else:
        fx2 = x2.astype(np.float64)
    (N1,D1) = fx1.shape
    (N2,D2) = fx2.shape
    if D1 != D2:
        raise ValueError('Arrays must have the same dimensionality')
    return (fx1,fx2)

def _buildtrees(x1, x2):
    (fx1, fx2) = _cleaninputs(x1, x2)
    kd1 = spherematch_c.KdTree(fx1)
    if fx2 is fx1:
        kd2 = kd1
    else:
        kd2 = spherematch_c.KdTree(fx2)
    return (kd1, kd2)

def match(x1, x2, radius, notself=False, permuted=True, indexlist=False):
    '''
    ::

        (indices,dists) = match(x1, x2, radius):

    Or::

        inds = match(x1, x2, radius, indexlist=True):

    Returns the indices (Nx2 int array) and distances (Nx1 float
    array) between points in *x1* and *x2* that are within *radius*
    Euclidean distance of each other.

    *x1* is N1xD and *x2* is N2xD.  *x1* and *x2* can be the same
    array.  Dimensions D above 5-10 will probably not run faster than
    naive.

    Despite the name of this package, the arrays x1 and x2 need not be
    celestial positions; in particular, there is no RA wrapping at 0,
    and no special handling at the poles.  If you want to match
    celestial coordinates like RA,Dec, see the match_radec function.

    If *indexlist* is True, the return value is a python list with one
    element per data point in the first tree; that element is a python
    list containing the indices of points matched in the second tree.

    The *indices* return value has a row for each match; the matched
    points are:
    x1[indices[:,0],:]
    and
    x2[indices[:,1],:]

    This function doesn\'t know about spherical coordinates -- it just
    searches for matches in n-dimensional space.

    >>> from astrometry.util.starutil_numpy import *   
    >>> from astrometry.libkd import spherematch
    >>> # RA,Dec in degrees
    >>> ra1  = array([  0,  1, 2, 3, 4, 359,360])
    >>> dec1 = array([-90,-89,-1, 0, 1,  89, 90])
    >>> # xyz: N x 3 array: unit vectors
    >>> xyz1 = radectoxyz(ra1, dec1)
    >>> ra2  = array([ 45,   1,  4, 4, 4,  0,  1])
    >>> dec2 = array([-89, -88, -1, 0, 2, 89, 89])
    >>> xyz2 = radectoxyz(ra2, dec2)
    >>> # The \'radius\' is now distance between points on the unit sphere --
    >>> # for small angles, this is ~ angular distance in radians.  You can use
    >>> # the function:
    >>> radius_in_deg = 2.
    >>> r = sqrt(deg2distsq(radius_in_deg))
    >>> (inds,dists) = spherematch.match(xyz1, xyz2, r)
    >>> # Now *inds* is an Mx2 array of the matching indices,
    >>> # and *dists* the distances between them:
    >>> #  eg,  sqrt(sum((xyz1[inds[:,0],:] - xyz2[inds[:,1],:])**2, axis=1)) = dists
    >>> print inds
    [[0 0]
     [1 0]
     [1 1]
     [2 2]
     [3 2]
     [3 3]
     [4 3]
     [4 4]
     [5 5]
     [6 5]
     [5 6]
     [6 6]]
    >>> print sqrt(sum((xyz1[inds[:,0],:] - xyz2[inds[:,1],:])**2, axis=1))
    [ 0.01745307  0.01307557  0.01745307  0.0348995   0.02468143  0.01745307
      0.01745307  0.01745307  0.0003046   0.01745307  0.00060917  0.01745307]
    >>> print dists[:,0]
    [ 0.01745307  0.01307557  0.01745307  0.0348995   0.02468143  0.01745307
      0.01745307  0.01745307  0.0003046   0.01745307  0.00060917  0.01745307]
    >>> print vstack((ra1[inds[:,0]], dec1[inds[:,0]], ra2[inds[:,1]], dec2[inds[:,1]])).T
    [[  0 -90  45 -89]
     [  1 -89  45 -89]
     [  1 -89   1 -88]
     [  2  -1   4  -1]
     [  3   0   4  -1]
     [  3   0   4   0]
     [  4   1   4   0]
     [  4   1   4   2]
     [359  89   0  89]
     [360  90   0  89]
     [359  89   1  89]
     [360  90   1  89]]

    Parameters
    ----------
    x1 : numpy array, float, shape N1 x D
        First array of points to match

    x2 : numpy array, float, shape N2 x D
        Second array of points to match

    radius : float
        Scalar Euclidean distance to match
        
    Returns
    -------
    indices : numpy array, integers, shape M x 2, for M matches
        The array of matching indices; *indices[:,0]* are indices in *x1*,
        *indices[:,1]* are indices in *x2*.

    dists : numpy array, floats, length M, for M matches
        The distances between matched points.
        
    If *indexlist* is *True*:

    indices : list of ints of integers
        The list of matching indices.  One list element per *x1* element,
        containing a list of matching indices in *x2*.
        
    '''
    (kd1,kd2) = _buildtrees(x1, x2)
    if indexlist:
        inds = spherematch_c.match2(kd1, kd2, radius, notself, permuted)
    else:
        (inds,dists) = spherematch_c.match(kd1, kd2, radius, notself, permuted)
    if indexlist:
        return inds
    return (inds,dists)

def match_naive(x1, x2, radius, notself=False):
    ''' Does the same thing as match(), but the straight-forward slow
    way.  (Not necessarily the way you\'d do it in python either).
    Not very fair as a speed comparison, but useful to convince
    yourself that match() does the right thing.
    '''
    (fx1, fx2) = _cleaninputs(x1, x2)
    (N1,D1) = x1.shape
    (N2,D2) = x2.shape
    inds = []
    dists = []
    for i1 in range(N1):
        for i2 in range(N2):
            if notself and i1 == i2:
                continue
            d2 = sum((x1[i1,:] - x2[i2,:])**2)
            if d2 < radius**2:
                inds.append((i1,i2))
                dists.append(sqrt(d2))
    inds = array(inds)
    dists = array(dists)
    return (inds,dists)

def nearest(x1, x2, maxradius, notself=False, count=False):
    '''
    For each point in x2, returns the index of the nearest point in x1,
    if there is a point within 'maxradius'.

    (Note, this may be backward from what you want/expect!)
    '''
    (kd1,kd2) = _buildtrees(x1, x2)
    if count:
        X = spherematch_c.nearest2(kd1, kd2, maxradius, notself, count)
    else:
        X = spherematch_c.nearest(kd1, kd2, maxradius, notself)
    return X
_nearest_func = nearest

def tree_build_radec(ra=None, dec=None, xyz=None):
    '''
    Builds a kd-tree given *RA,Dec* or unit-sphere *xyz* coordinates.
    '''
    if ra is not None:
        (N,) = ra.shape
        xyz = np.zeros((N,3)).astype(float)
        xyz[:,2] = np.sin(np.deg2rad(dec))
        cosd = np.cos(np.deg2rad(dec))
        xyz[:,0] = cosd * np.cos(np.deg2rad(ra))
        xyz[:,1] = cosd * np.sin(np.deg2rad(ra))
    kd = spherematch_c.KdTree(xyz)
    return kd

def tree_build(X, nleaf=16, bbox=True, split=False):
    '''
    Builds a kd-tree given a numpy array of Euclidean points.
    
    Parameters
    ----------
    X: numpy array of shape (N,D)
        The points to index.
        
    Returns
    -------
    kd: integer
        kd-tree identifier (address).
    '''
    return spherematch_c.KdTree(X, nleaf=nleaf, bbox=bbox, split=split)

def tree_free(kd):
    '''
    Frees a kd-tree previously created with *tree_build*.
    '''
    print('No need for tree_free')
    pass

def tree_save(kd, fn):
    '''
    Writes a kd-tree to the given filename.
    '''
    print('Deprecated tree_save()')
    return kd.write(fn)
#rtn = spherematch_c.kdtree_write(kd, fn)
#return rtn

def tree_open(fn, treename=None):
    '''
    Reads a kd-tree from the given filename.
    '''
    if treename is None:
        return spherematch_c.KdTree(fn)
    else:
        return spherematch_c.KdTree(fn, treename)

def tree_close(kd):
    '''
    Closes a kd-tree previously opened with *tree_open*.
    '''
    print('No need for tree_close')
    pass

def tree_search(kd, pos, radius, getdists=False, sortdists=False):
    '''
    Searches the given kd-tree for points within *radius* of the given
    position *pos*.
    '''
    #print('Unnecessary call to tree_search(kd, ...); use kd.search(...)')
    return kd.search(pos, radius, int(getdists), int(sortdists))

def tree_search_radec(kd, ra, dec, radius, getdists=False, sortdists=False):
    '''
    ra,dec in degrees
    radius in degrees
    '''
    dec = np.deg2rad(dec)
    cosd = np.cos(dec)
    ra = np.deg2rad(ra)
    pos = np.array([cosd * np.cos(ra), cosd * np.sin(ra), np.sin(dec)])
    rad = deg2dist(radius)
    return tree_search(kd, pos, rad, getdists=getdists, sortdists=sortdists)

def trees_match(kd1, kd2, radius, nearest=False, notself=False,
                permuted=True, count=False):
    '''
    Runs rangesearch or nearest-neighbour matching on given kdtrees.

    'radius' is Euclidean distance.

    If 'nearest'=True, returns the nearest neighbour of each point in "kd1";
    ie, "I" will NOT contain duplicates, but "J" may.

    If 'count'=True, also counts the number of objects within range
    as well as returning the nearest neighbor of each point in "kd1";
    the return value becomes I,J,d,counts , counts a numpy array of ints.

    Returns (I, J, d), where
      I are indices into kd1
      J are indices into kd2
      d are distances-squared
      [counts is number of sources in range]

    >>> import numpy as np
    >>> X = np.array([[1, 2, 3, 6]]).T.astype(float)
    >>> Y = np.array([[1, 4, 4]]).T.astype(float)
    >>> kd1 = tree_build(X)
    >>> kd2 = tree_build(Y)
    >>> I,J,d = trees_match(kd1, kd2, 1.1, nearest=True)
    >>> print I
    [0 1 2]
    >>> print J
    [0 0 2]
    >>> print d
    [  0.  60.  60.]
    >>> I,J,d,count = trees_match(kd1, kd2, 1.1, nearest=True, count=True)
    >>> print I
    [0 1 2]
    >>> print J
    [0 0 2]
    >>> print d
    [  0.  60.  60.]
    >>> print count
    [1 1 2]
    '''
    rtn = None
    if nearest:
        rtn = spherematch_c.nearest2(kd2, kd1, radius, notself, count)
        # J,I,d,[count]
        rtn = (rtn[1], rtn[0], np.sqrt(rtn[2])) + rtn[3:]
        #distsq2deg(rtn[2]),
    else:
        (inds,dists) = spherematch_c.match(kd1, kd2, radius, notself, permuted)
        #d = dist2deg(dists[:,0])
        d = dists[:,0]
        I,J = inds[:,0], inds[:,1]
        rtn = (I,J,d)
    return rtn

def tree_permute(kd, I):
    print('Unnecessary call to tree_permute(kd, I): use kd.permute(I)')
    return kd.permute(I)

def tree_bbox(kd):
    print('Unnecessary call to tree_bbox(kd): use kd.bbox')
    return kd.bbox

def tree_n(kd):
    print('Unnecessary call to tree_n(kd): use kd.n')
    return kd.n

def tree_print(kd):
    print('Unnecessary call to tree_print(kd): use kd.print()')
    kd.print()
    
def tree_data(kd, I):
    print('Unnecessary call to tree_data(kd, I): use kd.get_data(I)')
    return kd.get_data(I)

if __name__ == '__main__':
    import doctest
    doctest.testmod()