# -*- coding: utf-8 -*-

import time
import threading
from Queue import Queue, Empty
input_queue = Queue()
output_queue = Queue()

from libc.stdlib cimport malloc, free, calloc
cimport cython

cdef import from "stdlib.h":
    void *memset(void *str, int c, size_t n)
    void *memcpy(void *str1, void *str2, size_t n)
    void *memmove(void *str1, void *str2, size_t n)
    
import numpy as np
cimport numpy as np

import math
cdef import from "math.h":
    double cos(double)
    double sin(double)
    double log10(double)
    

cdef struct Ray:
    double dx
    double dy
    double x
    double y
    double power
    int dist
    int source_id
    int generation
    
#~ cdef double NONE = 10**9
#~ cdef int NONE_INT = 10**9

cdef inline int fast_10log10(double v):
    cdef long long int_repr_of_double
    cdef long long first_sign_bit = 2<<50
    cdef long long second_sign_bit = 2<<49
    cdef long long third_sign_bit = 2<<48
    cdef int res 
        
    # cast from double to int64
    int_repr_of_double = (<long long*>&v)[0]
    # estimate log10 (extract Bit 62 - 52, see IEEE-754)
    # assume no negative values (ignore bit 63)
    # divide by 3 to change from log2 to log10 (log10(2) == 0.301
    # and then multiplicate by 10 ( so we multiplicate by 3...)
    
    res = ((int_repr_of_double>>52) - 1023) * 3

    # analyze first 3 bits of significant
    # that will give us more accuracy (and we underestimate res by range(0..3) )
    if (int_repr_of_double & first_sign_bit) > 0:
        res += 1
    
    if (int_repr_of_double & second_sign_bit) > 0:
        res += 1
    
    if (int_repr_of_double & third_sign_bit) > 0:
        res += 1
    
    return res
    
cdef class Rays:
    cdef Ray *handle
    cdef public int count # number of used rays
    cdef int spawn_idx
    cdef public int size # number of allocated rays
        
    def __init__(Rays self, source=None, int numrays=0, int angle=0):
        self.count = 0
        if source is not None:
            self.initFromSource(source, numrays, angle)
        elif numrays > 0:
            self.alloc(numrays)
            self.spawn_idx = 0
        
        
    def alloc(Rays self, num):
        self.count = num
        self.size = num
        free(self.handle) 
        self.handle = <Ray*>calloc(self.count, sizeof(Ray))
            
    def initFromSource(Rays self, source, int numrays, int angle):
        self.alloc(numrays)
        
        cdef Ray *ray
        cdef double initial_power, x, y, dx, dy, alpha, pi = math.pi
        cdef int sourceidx, j
        
        cdef double increment = 360.0 / numrays
        
        sourceidx = source.idx
        initial_power = float(source.power)# / numrays
        x = source.x
        y = source.y
        
        for j in range(numrays):
            alpha = j * increment
            
            # increase alpha by one degree
            # to ensure very low number of rays are not orthogonal to walls in
            # rectangular maps
            alpha = alpha + angle
            
            dx = cos(pi / 180.0 * alpha)
            dy = sin(pi / 180.0 * alpha)
            
            ray = &self.handle[j]
            ray.x = x
            ray.y = y
            ray.dx = dx
            ray.dy = dy
            ray.power = initial_power
            ray.dist = 0
            ray.generation = 1
            ray.source_id = sourceidx
            
    
    cdef void _spawn(Rays self, Ray ray, double dampening, int dist, double dx, double dy, double x, double y) nogil:
        ray.generation += 1
        ray.power = ray.power * dampening
        ray.dist = dist
        #~ print 'spawn, %s, %s' % (dx, dy)
        
        ray.dy = dy
        ray.dx = dx
        ray.x = x
        ray.y = y
        
        self.handle[self.spawn_idx] = ray
        self.spawn_idx += 1
        
    
    cdef void endspawning(Rays self) nogil:
        ''' set's self.count to number of spawned rays '''
        self.count = self.spawn_idx
        
    def toNumpy(self):
        ''' generate a numpy representation of ray data 
        build one array for doubles and on for ints
        '''
        
        cdef np.ndarray[np.float32_t, ndim=2] doubles
        cdef np.ndarray[np.int_t, ndim=2] ints
        cdef int i
        cdef Ray *ray
        doubles = np.zeros([self.count, 5], dtype=np.float32)
        ints = np.zeros([self.count, 4], dtype=np.int)
        
        for i in range(self.count):
            ray = &self.handle[i]
            doubles[i, 0] = ray.dx
            doubles[i, 1] = ray.dy
            doubles[i, 2] = ray.x
            doubles[i, 3] = ray.y
            doubles[i, 4] = ray.power
            ints[i, 0] = ray.dist
            ints[i, 1] = ray.source_id
            ints[i, 2] = ray.generation
            #~ ints[i, 3] = ray.medium
        return doubles, ints
        
    def __dealloc__(Rays self):
        free(self.handle) 

cdef class Map:
    cdef double *map # hold signal strength for each x, y point
    cdef double *map_xs, *map_ys
    cdef public int width
    cdef public int height
    cdef public int numsources
    cdef public list sources
    
    cdef char *bgra_output
    cdef char *bgra_buffer
        
    cdef int mapsize
    
    def __init__(Map self, int width, int height, sources):
        self.mapsize = width * height
        self.numsources = len(sources)
        self.sources = sources
        
        self.map = <double*>calloc(self.mapsize * self.numsources, sizeof(double))
        self.map_xs = <double*>calloc(self.mapsize * self.numsources, sizeof(double))
        self.map_ys = <double*>calloc(self.mapsize * self.numsources, sizeof(double))
        
        
        self.bgra_buffer = <char*>calloc(self.mapsize * 4, sizeof(char))
            
        self.width = width
        self.height = height
    
    def __dealloc__(Map self):
        free(self.map)
        free(self.map_xs)
        free(self.map_ys)
    
    def clear(self):
        cdef int i
        for i in range(self.mapsize * self.numsources):
            self.map[i] = 0
    
    cdef inline int _translate(Map self, int x, int y, int source):
        return source * self.mapsize + y * self.width + x
    
    @cython.boundscheck(False)    
    cdef inline double _get(self, int x, int y, int source):
        return self.map[self._translate(x, y, source)]
    
    @cython.boundscheck(False)    
    cdef inline _set(self, int x, int y, int source, double v):
        self.map[self._translate(x, y, source)] = v
    
    def get(self, x, y, source):
        return self._get(x, y, source)
        
    
    @cython.boundscheck(False)
    @cython.cdivision(True)
    def writeBGRA(Map self, int bits, np.ndarray[np.int16_t] envmap, int alpha):
        '''write BGRA to pointer given by bits '''
        cdef char *ptr = <char*>bits
        self.bgra_output = ptr
        
        cdef double fact
        cdef int idx = 0, i, source_offset, source_idx, color
        cdef int b = 0, 
            
        memset(ptr, 0, self.mapsize * 4)
        for source_idx in range(self.numsources):
            color = source_idx % 3
            
            fact =  float(self.sources[source_idx].power)
            source_offset = source_idx * self.mapsize
            
            #BGRA
            for i in range(self.mapsize):
                v = self.map[source_offset + i]
                
                if v > 0:
                    # (v / fact) -> normalize to 1.0
                    b = fast_10log10(v / fact)
                    #~ b = int(10 * log10(v / fact))
                    
                    # normalize to color space 0-255
                    # we expect values from -1 to -100
                    b = (150 + b * 2) * 2
                    
                    if b > 255:
                        b = 255
                else:
                    b = 0
                
                if b > 0:
                    idx = i * 4
                    ptr[idx + color] = b
                    #~ ptr[idx + 1] = ptr[idx]
                    #~ ptr[idx + 2] = ptr[idx]
                    if envmap[i] == 0:
                        ptr[idx + 3] = alpha # alpha mask for no envmap layers
                    else:
                        ptr[idx + 3] = 20 # alpha mask
        
        #TODO:  only on demand
        memcpy(self.bgra_buffer, ptr, self.mapsize * 4)
        
        return
    
    def drawRect(Map self, int x1, int y1, int x2, int y2):
        cdef int i
        cdef char *ptr = self.bgra_output
        
        y1 = max(0, y1)
        y2 = min(y2, self.height) - 1
        x1 = max(0, x1)
        x2 = min(x2, self.width) - 1
        memcpy(ptr, self.bgra_buffer, self.mapsize * 4)
        
        horz1 = ((y1 * self.width + x1) * 4, (y1 * self.width + x2) * 4, 4)
        horz2 = ((y2 * self.width + x1) * 4, (y2 * self.width + x2) * 4, 4)
        vert1 = ((y1 * self.width + x1) * 4, (y2 * self.width + x1) * 4, self.width * 4)
        vert2 = ((y1 * self.width + x2) * 4, (y2 * self.width + x2) * 4, self.width * 4)
        
        for r in (horz1, horz2, vert1, vert2):
            for i in range(r[0], r[1], r[2]):
                ptr[i + 0] = 255
                ptr[i + 1] = 255
                ptr[i + 2] = 255
                ptr[i + 3] = 255
    
    @cython.boundscheck(False)
    def asArray(Map self, int source_idx, invertY=False):
        ''' convert map data to numpy array with shape width/height '''
        cdef np.ndarray[np.double_t, ndim=2] map = np.zeros([self.height, self.width], dtype=np.double)
        cdef int x, y, _y
        for x in range(self.width):
            for y in range(self.height):
                _y = self.height - 1 - y if invertY else y
                map[y, x] = self._get(x, _y, source_idx)
        return map
        
    @cython.boundscheck(False)
    def fromArray(Map self, np.ndarray[np.float32_t] a):
        cdef int i
        for i in range(self.mapsize * self.numsources):
            self.map[i] = a[i]
        
        return
                    
cdef class Simulation:
    cdef public list rays
    cdef public Map map
        
    cdef public int width, height
    cdef public int max_generations, enable_reflections, enable_refractions
    cdef public double pixel_per_meter
        
    cdef np.ndarray material_normals
        
    #~ cdef public list generations
    #~ cdef public list materials 
    
    #~ cdef double *raypower
    cdef double materials[100][2] # max 99 materials
        
    def __init__(Simulation self, int width, int height, 
            int max_generations, enable_reflections, enable_refractions, 
            list sources, list materials, double pixel_per_meter):
        self.width = width
        self.height = height
        self.map = Map(width, height, sources)
        self.pixel_per_meter = pixel_per_meter
        
        # use initRays() to initialize
        self.rays = []
        self.max_generations = max_generations
        self.enable_reflections = enable_reflections
        self.enable_refractions = enable_refractions
        
        # the [0] index is not used (reserved for NO material)
        for i, (refl, refr) in enumerate(materials[:99]):
            self.materials[i+1][0] = refl
            self.materials[i+1][1] = refr
            
    def __dealloc__(Map self):
        pass
    
    @cython.boundscheck(False)
    @cython.cdivision(True)
    def initRays(Simulation self, int numrays, int angle):
        ''' configure ray-specific aspects '''
        
        for source in self.map.sources:
            self.rays.append(Rays(source=source, numrays=numrays, angle=angle))

            
    @cython.boundscheck(False)
    @cython.cdivision(True)
    def loopRays(Simulation self, Rays rays, np.ndarray[np.int16_t, ndim=2] envmap):
            
        cdef int i, j, xi, yi, medium, currmedium, mapidx, dist_offset, material_normal
        cdef int width = self.width, height = self.height
        cdef double x, y, newpower, currpower, reflection, refraction
        cdef Ray *ray
        cdef Ray r
        cdef Rays newrays = Rays(numrays=rays.count*2)
        cdef int done, dist, map_source_offset, source_id
        
        cdef double clamp = 10**-10
        
        cdef np.ndarray[np.int16_t, ndim=2] material_normals = self.material_normals
        cdef int pixel_per_meter_log2 = int(round(math.log(self.pixel_per_meter, 2)))
       
        #~ with nogil:
        source_id = (&rays.handle[0]).source_id
        map_source_offset = source_id * self.map.mapsize
        
        for j in range(rays.count):
            ray = &rays.handle[j]
            
            x, y = ray.x, ray.y
            
            xi, yi = int(x), int(y)
            currmedium = envmap[xi, yi]
            #~ print xi, yi, currmedium
            #~ print currmedium
            dist = ray.dist
            while True:
                
                if x <= 0 or x >= width-1 or y <= 0 or y >= height:
                    break
                
                xi, yi = int(x), int(y)
                
                if dist > 32:
                    power = ray.power / (dist >> pixel_per_meter_log2)**2
                else:
                    power = ray.power / (dist / self.pixel_per_meter)**2
                #~ print '%.2f %.2f %i %s' % (x, y, envmap[xi, yi], power)
                if power < clamp:
                    break
                    
                medium = envmap[xi, yi]
                if medium != currmedium: 
                    if medium == 0:
                        currmedium = medium
                    else:
                        material_normal = material_normals[xi, yi]
                        
                        r = rays.handle[j] # copy ray struct
                        
                        if self.enable_reflections and currmedium == 0:
                            reflection = self.materials[medium][0]
                            if material_normal == 1:
                                # horicontal material
                                newrays._spawn(r, reflection, dist, ray.dx, -ray.dy, x-ray.dx, y-ray.dy)
                            elif material_normal == 0:
                                # vertical material
                                newrays._spawn(r, reflection, dist, -ray.dx, ray.dy, x-ray.dx, y-ray.dy)
                        
                        if self.enable_refractions:
                            refraction = self.materials[medium][1]
                            newrays._spawn(r, refraction, dist, ray.dx, ray.dy, x, y)
                        
                        break
                    
                 
                    
                mapidx = map_source_offset + yi * self.map.width + xi 
                currpower = self.map.map[mapidx]
                
                # handle interference?
                if power > currpower:
                    self.map.map[mapidx] = power
                    # TODO: prepare for bresenham 10% speed loss?
                    #~ self.map.map_xs[mapidx] = ray.x 
                    #~ self.map.map_ys[mapidx] = ray.y
                    
                x = x + ray.dx
                y = y + ray.dy
                dist += 1
                
        newrays.endspawning()
            
        return newrays
    
    @cython.boundscheck(False)
    def loadMaterialNormals(self, np.ndarray[np.int16_t, ndim=2] envmap):
        ''' we do not use real normals atm, only 0 and 1 for horizontally, vertically '''
        cdef np.ndarray[np.int16_t, ndim=2] material_normals = np.zeros([self.width, self.height], dtype=np.int16)
            
        
        cdef int x, y, medium
        for x in range(1, self.width-1): # ignore border pixels
            for y in range(1, self.height-1):
                medium = envmap[x, y]
                if medium > 0:
                    if envmap[x-1, y] == medium and envmap[x+1, y] == medium:
                        material_normals[x, y] = 1
                    elif envmap[x, y-1] == medium and envmap[x, y+1] == medium:
                        material_normals[x, y] = 0
                    elif envmap[x-1, y] == medium or envmap[x+1, y] == medium:
                        material_normals[x, y] = 1
                    elif envmap[x, y-1] == medium or envmap[x, y+1] == medium:
                        material_normals[x, y] = 0
                    #~ else:
                        # a single pixel?
                        #~ raise RuntimeError('unsupported normal handling')
                        #~ self.material_normals[x, y] = 1
        
        self.material_normals = material_normals
                        
    def build(Simulation self, np.ndarray[np.int16_t, ndim=2] envmap):
        self.map.clear()
        self.loadMaterialNormals(envmap)

        todo = self.rays
        gen = 0
        while len(todo) > 0 and gen < self.max_generations:
            for _rays in todo:
                input_queue.put((self, _rays, envmap))
            
            _todo = []
            for i in range(len(todo)):
                newRays = output_queue.get()
                if newRays.count > 0:
                    _todo.append(newRays)
            
            gen += 1
            todo[:] = _todo
        return
        
        
        # single threaded
        #~ t = time.time()
        #~ for rays in self.rays:
            #~ for i in range(self.max_generations):
                #~ rays = self.loopRays(rays, envmap)
                #~ if rays.count == 0:
                    #~ break
        

def builder(i):
    try:
        while True:
            try:
                sim, _rays, envmap = input_queue.get()
            except Empty:
                continue
            else:
                newrays = sim.loopRays(_rays, envmap)
                output_queue.put(newrays)
    except Exception, e:
        print e
        output_queue.put(Rays())

for i in range(4):
    from threading import Thread
    thread = Thread(target=builder, args=(i, ))
    thread.setDaemon(True)
    thread.start()