import sys
import os
import numpy as np
from itertoolset import *
from tools import timeit
from hashlib import md5
import cPickle as pickle
import gzip

# all material/surface properties are interpolated at these
# wavelengths when they are sent to the gpu
standard_wavelengths = np.arange(200, 810, 20).astype(np.float32)

class Mesh(object):
    def __init__(self, vertices, triangles, remove_duplicate_vertices=False):
        vertices = np.asarray(vertices, dtype=np.float32)
        triangles = np.asarray(triangles, dtype=np.int32)

        if len(vertices.shape) != 2 or vertices.shape[1] != 3:
            raise ValueError('shape mismatch')

        if len(triangles.shape) != 2 or triangles.shape[1] != 3:
            raise ValueError('shape mismatch')

        if (triangles < 0).any():
            raise ValueError('indices in `triangles` must be positive.')

        if (triangles >= len(vertices)).any():
            raise ValueError('indices in `triangles` must be less than the '
                             'length of the vertex array.')

        self.vertices = vertices
        self.triangles = triangles

        if remove_duplicate_vertices:
            self.remove_duplicate_vertices()

    def get_bounds(self):
        return np.min(self.vertices, axis=0), np.max(self.vertices, axis=0)

    def remove_duplicate_vertices(self):
        unique_vertices, inverse = np.unique(self.vertices.view([('', self.vertices.dtype)]*self.vertices.shape[1]), return_inverse=True)
        self.vertices = unique_vertices.view(self.vertices.dtype).reshape((unique_vertices.shape[0], self.vertices.shape[1]))
        self.triangles = np.vectorize(lambda i: inverse[i])(self.triangles)

    def assemble(self, key=slice(None), group=True):
        if group:
            vertex_indices = self.triangles[key]
        else:
            vertex_indices = self.triangles[key].flatten()

        return self.vertices[vertex_indices]

    def __len__(self):
        return len(self.triangles)

    def __add__(self, other):
        return Mesh(np.concatenate((self.vertices, other.vertices)), np.concatenate((self.triangles, other.triangles + len(self.vertices))))

class Solid(object):
    def __init__(self, mesh, material1=None, material2=None, surface=None, color=0xffffff):
        self.mesh = mesh

        if np.iterable(material1):
            if len(material1) != len(mesh):
                raise ValueError('shape mismatch')
            self.material1 = np.array(material1, dtype=np.object)
        else:
            self.material1 = np.tile(material1, len(self.mesh))

        if np.iterable(material2):
            if len(material2) != len(mesh):
                raise ValueError('shape mismatch')
            self.material2 = np.array(material2, dtype=np.object)
        else:
            self.material2 = np.tile(material2, len(self.mesh))

        if np.iterable(surface):
            if len(surface) != len(mesh):
                raise ValueError('shape mismatch')
            self.surface = np.array(surface, dtype=np.object)
        else:
            self.surface = np.tile(surface, len(self.mesh))

        if np.iterable(color):
            if len(color) != len(mesh):
                raise ValueError('shape mismatch')
            self.color = np.array(color, dtype=np.uint32)
        else:
            self.color = np.tile(color, len(self.mesh)).astype(np.uint32)

        self.unique_materials = \
            np.unique(np.concatenate([self.material1, self.material2]))

        self.unique_surfaces = np.unique(self.surface)

    def __len__(self):
        return len(self.mesh)

    def __add__(self, other):
        return Solid(self.mesh + other.mesh, np.concatenate((self.material1, other.material1)), np.concatenate((self.material2, other.material2)), np.concatenate((self.surface, other.surface)), np.concatenate((self.color, other.color)))

class Material(object):
    """Material optical properties."""
    def __init__(self, name='none'):
        self.name = name

        self.refractive_index = None
        self.absorption_length = None
        self.scattering_length = None
        self.density = 0.0 # g/cm^3
        self.composition = {} # by mass

    def set(self, name, value, wavelengths=standard_wavelengths):
        if np.iterable(value):
            if len(value) != len(wavelengths):
                raise ValueError('shape mismatch')
        else:
            value = np.tile(value, len(wavelengths))

        self.__dict__[name] = np.array(zip(wavelengths, value), dtype=np.float32)

class Surface(object):
    """Surface optical properties."""
    def __init__(self, name='none'):
        self.name = name

        self.set('detect', 0)
        self.set('absorb', 0)
        self.set('reflect_diffuse', 0)
        self.set('reflect_specular', 0)

    def set(self, name, value, wavelengths=standard_wavelengths):
        if np.iterable(value):
            if len(value) != len(wavelengths):
                raise ValueError('shape mismatch')
        else:
            value = np.tile(value, len(wavelengths))

        if (np.asarray(value) < 0.0).any():
            raise Exception('all probabilities must be >= 0.0')

        self.__dict__[name] = np.array(zip(wavelengths, value), dtype=np.float32)

def interleave(arr, bits):
    """
    Interleave the bits of quantized three-dimensional points in space.

    Example
        >>> interleave(np.identity(3, dtype=np.int))
        array([4, 2, 1], dtype=uint64)
    """
    if len(arr.shape) != 2 or arr.shape[1] != 3:
        raise Exception('shape mismatch')

    z = np.zeros(arr.shape[0], dtype=np.uint64)
    for i in range(bits):
        z |= (arr[:,2] & 1 << i) << (2*i) | \
             (arr[:,1] & 1 << i) << (2*i+1) | \
             (arr[:,0] & 1 << i) << (2*i+2)
    return z

def morton_order(mesh, bits):
    """
    Return a list of zvalues for triangles in `mesh` by interleaving the
    bits of the quantized center coordinates of each triangle. Each coordinate
    axis is quantized into 2**bits bins.
    """
    lower_bound, upper_bound = mesh.get_bounds()

    if bits <= 0 or bits > 21:
        raise Exception('number of bits must be in the range (0,21].')

    max_value = 2**bits - 1

    def quantize(x):
        return np.uint64((x-lower_bound)*max_value/(upper_bound-lower_bound))

    mean_positions = quantize(np.mean(mesh.assemble(), axis=1))

    return interleave(mean_positions, bits)

class Geometry(object):
    def __init__(self):
        self.solids = []
        self.solid_rotations = []
        self.solid_displacements = []

    def add_solid(self, solid, rotation=np.identity(3), displacement=(0,0,0)):
        rotation = np.asarray(rotation, dtype=np.float32)

        if rotation.shape != (3,3):
            raise ValueError('shape mismatch')

        self.solid_rotations.append(rotation.astype(np.float32))

        displacement = np.asarray(displacement, dtype=np.float32)

        if displacement.shape != (3,):
            raise ValueError('shape mismatch')

        self.solid_displacements.append(displacement)

        self.solids.append(solid)

        return len(self.solids)-1

    @timeit
    def build(self, bits=8, shift=3):
        offsets = [ (0,0) ]
        for solid in self.solids:
            offsets.append( (offsets[-1][0] + len(solid.mesh.vertices),
                             offsets[-1][1] + len(solid.mesh.triangles)) )

        vertices = np.zeros(shape=(offsets[-1][0], 3), dtype=np.float32)
        triangles = np.zeros(shape=(offsets[-1][1],3), dtype=np.int32)

        for i, (solid, (vertex_offset, triangle_offset)) in \
                enumerate(zip(self.solids, offsets[:-1])):
            triangles[triangle_offset:triangle_offset+len(solid.mesh.triangles),:] = \
                solid.mesh.triangles + vertex_offset
            vertices[vertex_offset:vertex_offset+len(solid.mesh.vertices),:] = \
                np.inner(solid.mesh.vertices, self.solid_rotations[i]) + self.solid_displacements[i]
            
        self.mesh = Mesh(vertices, triangles)

        self.colors = np.concatenate([solid.color for solid in self.solids])

        self.solid_id = np.concatenate([np.fromiter(repeat(i, len(solid.mesh.triangles)), dtype=np.uint32) for i, solid in enumerate(self.solids)])

        self.unique_materials = list(np.unique(np.concatenate([solid.unique_materials for solid in self.solids])))

        material_lookup = dict(zip(self.unique_materials, range(len(self.unique_materials))))

        self.material1_index = \
            np.fromiter(imap(material_lookup.get, chain(*[solid.material1 for solid in self.solids])), dtype=np.int32)

        self.material2_index = \
            np.fromiter(imap(material_lookup.get, chain(*[solid.material2 for solid in self.solids])), dtype=np.int32)

        self.unique_surfaces = list(np.unique(np.concatenate([solid.unique_surfaces for solid in self.solids])))

        surface_lookup = dict(zip(self.unique_surfaces, range(len(self.unique_surfaces))))

        self.surface_index = \
            np.fromiter(imap(surface_lookup.get, chain(*[solid.surface for solid in self.solids])), dtype=np.int32)

        try:
            self.surface_index[self.surface_index == self.unique_surfaces.index(None)] = -1
        except ValueError:
            pass

        checksum = md5(str(bits))
        checksum.update(str(shift))
        checksum.update(self.mesh.vertices)
        checksum.update(self.mesh.triangles)

        cache_dir = os.path.expanduser('~/.chroma')
        cache_file = checksum.hexdigest()
        cache_path = os.path.join(cache_dir, cache_file)

        try:
            f = gzip.GzipFile(cache_path, 'rb')
        except IOError:
            pass
        else:
            print 'loading cache.'
            data = pickle.load(f)

            reorder = data.pop('reorder')
            self.mesh.triangles = self.mesh.triangles[reorder]
            self.material1_index = self.material1_index[reorder]
            self.material2_index = self.material2_index[reorder]
            self.surface_index = self.surface_index[reorder]
            self.colors = self.colors[reorder]
            self.solid_id = self.solid_id[reorder]

            for key, value in data.iteritems():
                setattr(self, key, value)
            f.close()
            return

        zvalues_mesh = morton_order(self.mesh, bits)
        reorder = np.argsort(zvalues_mesh)
        zvalues_mesh = zvalues_mesh[reorder]

        if (np.diff(zvalues_mesh) < 0).any():
            raise Exception('zvalues_mesh out of order.')

        self.mesh.triangles = self.mesh.triangles[reorder]

        self.material1_index = self.material1_index[reorder]
        self.material2_index = self.material2_index[reorder]
        self.surface_index = self.surface_index[reorder]
        self.colors = self.colors[reorder]
        self.solid_id = self.solid_id[reorder]

        unique_zvalues = np.unique(zvalues_mesh)

        self.lower_bounds = np.empty((unique_zvalues.size,3), dtype=np.float32)
        self.upper_bounds = np.empty((unique_zvalues.size,3), dtype=np.float32)

        assembled_mesh = self.mesh.assemble(group=False)
        self.node_map = np.searchsorted(zvalues_mesh, unique_zvalues)
        self.node_map_end = np.searchsorted(zvalues_mesh, unique_zvalues, side='right')

        for i, (zi1, zi2) in enumerate(izip(self.node_map, self.node_map_end)):
            self.lower_bounds[i] = assembled_mesh[zi1*3:zi2*3].min(axis=0)
            self.upper_bounds[i] = assembled_mesh[zi1*3:zi2*3].max(axis=0)

        self.layers = np.zeros(unique_zvalues.size, dtype=np.uint32)
        self.first_node = unique_zvalues.size

        begin_last_layer = 0

        for layer in count(1):
            bit_shifted_zvalues = unique_zvalues >> shift
            unique_zvalues = np.unique(bit_shifted_zvalues)

            i0 = begin_last_layer + bit_shifted_zvalues.size

            self.node_map.resize(self.node_map.size+unique_zvalues.size)
            self.node_map[i0:] = np.searchsorted(bit_shifted_zvalues, unique_zvalues) + begin_last_layer
            self.node_map_end.resize(self.node_map_end.size+unique_zvalues.size)
            self.node_map_end[i0:] = np.searchsorted(bit_shifted_zvalues, unique_zvalues, side='right') + begin_last_layer

            self.layers.resize(self.layers.size+unique_zvalues.size)
            self.layers[i0:] = layer

            self.lower_bounds.resize((self.lower_bounds.shape[0]+unique_zvalues.size,3))
            self.upper_bounds.resize((self.upper_bounds.shape[0]+unique_zvalues.size,3))

            for i, zi1, zi2 in izip(count(i0), self.node_map[i0:], self.node_map_end[i0:]):
                self.lower_bounds[i] = self.lower_bounds[zi1:zi2].min(axis=0)
                self.upper_bounds[i] = self.upper_bounds[zi1:zi2].max(axis=0)

            begin_last_layer += bit_shifted_zvalues.size

            if unique_zvalues.size == 1:
                break

        print >>sys.stderr, 'Writing BVH to cache directory...'

        if not os.path.exists(cache_dir):
            os.makedirs(cache_dir)

        with gzip.GzipFile(cache_path, 'wb', compresslevel=1) as f:
            data = {}
            for key in ['lower_bounds', 'upper_bounds', 'node_map', 'node_map_end', 'layers', 'first_node']:
                data[key] = getattr(self, key)
            data['reorder'] = reorder
            pickle.dump(data, f, -1)