"""
Author(s): Matthew Loper

See LICENCE.txt for licensing and contact information.
"""

__author__ = 'matt'

__all__ = ['get_vertices_per_edge', 'get_faces_per_edge', 'get_vert_connectivity']

import zlib
import numpy as np
import scipy.sparse as sp
import cPickle as pickle
from chumpy.utils import row, col


def get_vert_connectivity(mesh_v, mesh_f):
    """Returns a sparse matrix (of size #verts x #verts) where each nonzero
    element indicates a neighborhood relation. For example, if there is a
    nonzero element in position (15,12), that means vertex 15 is connected
    by an edge to vertex 12."""

    vpv = sp.csc_matrix((len(mesh_v),len(mesh_v)))

    # for each column in the faces...
    for i in range(3):
        IS = mesh_f[:,i]
        JS = mesh_f[:,(i+1)%3]
        data = np.ones(len(IS))
        ij = np.vstack((row(IS.flatten()), row(JS.flatten())))
        mtx = sp.csc_matrix((data, ij), shape=vpv.shape)
        vpv = vpv + mtx + mtx.T

    return vpv


def get_vertices_per_edge(mesh_v, mesh_f):
    """Returns an Ex2 array of adjacencies between vertices, where
    each element in the array is a vertex index. Each edge is included
    only once. If output of get_faces_per_edge is provided, this is used to
    avoid call to get_vert_connectivity()"""

    faces = mesh_f
    cache_fname = '/tmp/verts_per_edge_cache_' + str(zlib.crc32(faces.flatten())) + '.pkl'
    try:
        with open(cache_fname, 'rb') as fp:
            return(pickle.load(fp))
    except:
        vc = sp.coo_matrix(get_vert_connectivity(mesh_v, mesh_f))
        result = np.hstack((col(vc.row), col(vc.col)))
        result = result[result[:,0] < result[:,1]] # for uniqueness

        with open(cache_fname, 'wb') as fp:
            pickle.dump(result, fp, -1)
        return result


def get_faces_per_edge(mesh_v, mesh_f, verts_per_edge=None):
    faces = mesh_f
    suffix = str(zlib.crc32(verts_per_edge.flatten())) if verts_per_edge is not None else ''
    cache_fname = '/tmp/edgecache_new_' + str(zlib.crc32(faces.flatten())) + '_' + suffix + '.pkl'

    try:
        with open(cache_fname, 'rb') as fp:
            return(pickle.load(fp))
    except:
        if verts_per_edge is None:
            verts_per_edge = get_vertices_per_edge(mesh_v, mesh_f)

        v2f = {i: set([]) for i in range(len(mesh_v))}
        # TODO: cythonize?
        for idx, f in enumerate(mesh_f):
            v2f[f[0]].add(idx)
            v2f[f[1]].add(idx)
            v2f[f[2]].add(idx)

        fpe = -np.ones_like(verts_per_edge)
        for idx, edge in enumerate(verts_per_edge):
            faces = v2f[edge[0]].intersection(v2f[edge[1]])
            faces = list(faces)[:2]
            for i, f in enumerate(faces):
                fpe[idx,i] = f

        result = fpe
        # f = faces
        # IS = np.repeat(np.arange(len(f)), 3)
        # JS = f.ravel()
        # data = np.ones(IS.size)
        # f2v = sp.csc_matrix((data, (IS, JS)), shape=(len(f), np.max(f.ravel())+1))
        # f2f = f2v.dot(f2v.T)
        # f2f = f2f.tocoo()
        # f2f = np.hstack((col(f2f.row), col(f2f.col), col(f2f.data)))
        # which = (f2f[:,0] < f2f[:,1]) & (f2f[:,2] >= 2)
        # result = np.asarray(f2f[which, :2], np.uint32)

        with open(cache_fname, 'wb') as fp:
            pickle.dump(result, fp, -1)
        return result



def loop_subdivider(mesh_v, mesh_f):

    IS = []
    JS = []
    data = []

    vc = get_vert_connectivity(mesh_v, mesh_f)
    ve = get_vertices_per_edge(mesh_v, mesh_f)
    vo = get_vert_opposites_per_edge(mesh_v, mesh_f)

    if True:
        # New values for each vertex
        for idx in xrange(len(mesh_v)):

            # find neighboring vertices
            nbrs = np.nonzero(vc[:,idx])[0]

            nn = len(nbrs)

            if nn ==3:
                wt = 3./16.
            elif nn > 3:
                wt = 3. / (8. * nn)
            else:
                raise Exception('nn should be 3 or more')
            for nbr in nbrs:
                IS.append(idx)
                JS.append(nbr)
                data.append(wt)

            JS.append(idx)
            IS.append(idx)
            data.append(1. - (wt * nn))

    start = len(mesh_v)
    edge_to_midpoint = {}

    if True:
        # New values for each edge:
        # new edge verts depend on the verts they span
        for idx, vs in enumerate(ve):

            vsl = list(vs)
            vsl.sort()
            IS.append(start + idx)
            IS.append(start + idx)
            JS.append(vsl[0])
            JS.append(vsl[1])
            data.append(3./8)
            data.append(3./8)

            opposites = vo[(vsl[0], vsl[1])]
            for opp in opposites:
                IS.append(start + idx)
                JS.append(opp)
                data.append(2./8./len(opposites))

            edge_to_midpoint[(vsl[0], vsl[1])] = start + idx
            edge_to_midpoint[(vsl[1], vsl[0])] = start + idx

    f = []

    for f_i, old_f in enumerate(mesh_f):
        ff = np.concatenate((old_f, old_f))

        for i in range(3):
            v0 = edge_to_midpoint[(ff[i], ff[i+1])]
            v1 = ff[i+1]
            v2 = edge_to_midpoint[(ff[i+1], ff[i+2])]
            f.append(row(np.array([v0,v1,v2])))

        v0 = edge_to_midpoint[(ff[0], ff[1])]
        v1 = edge_to_midpoint[(ff[1], ff[2])]
        v2 = edge_to_midpoint[(ff[2], ff[3])]
        f.append(row(np.array([v0,v1,v2])))

    f = np.vstack(f)

    IS = np.array(IS, dtype=np.uint32)
    JS = np.array(JS, dtype=np.uint32)

    if True: # for x,y,z coords
        IS = np.concatenate((IS*3, IS*3+1, IS*3+2))
        JS = np.concatenate((JS*3, JS*3+1, JS*3+2))
        data = np.concatenate ((data,data,data))

    ij = np.vstack((IS.flatten(), JS.flatten()))
    mtx = sp.csc_matrix((data, ij))

    return mtx, f


def get_vert_opposites_per_edge(mesh_v, mesh_f):
    """Returns a dictionary from vertidx-pairs to opposites.
    For example, a key consist of [4,5)] meaning the edge between
    vertices 4 and 5, and a value might be [10,11] which are the indices
    of the vertices opposing this edge."""
    result = {}
    for f in mesh_f:
        for i in range(3):
            key = [f[i], f[(i+1)%3]]
            key.sort()
            key = tuple(key)
            val = f[(i+2)%3]

            if key in result:
                result[key].append(val)
            else:
                result[key] = [val]
    return result


