#!/usr/bin/python

import sys
sys.path.append(".")
sys.path.append("..")
sys.path.append("../..") # Adds higher directory to python modules path

import os
import time
import numpy as np
from multiprocessing import Process
from multiprocessing import Manager

from libs.util import convert_vector_to_mat
from libs.models import log_gamma_floorplan
from libs.spacemap import SpaceBlock
from libs.spacemap import SpaceRay
from libs.spacemap import SpaceMap
from libs.plotting import plotSpace
from libs.plotting import plotRSS


def getLocs(loc_x_range, loc_y_range, step_size: float = 0.3):
    loc_xs = np.arange(loc_x_range[0], loc_x_range[1] + step_size, step_size)
    loc_ys = np.arange(loc_y_range[0], loc_y_range[1] + step_size, step_size)
    return np.stack(np.meshgrid(loc_xs, loc_ys), -1).reshape(-1, 2)


def log_gamma_floorplan_multi(floormap, rx_loc, tx_loc, power, gaussian_noise, results, kk):
    rss, __ = log_gamma_floorplan(
        floormap,
        rx_loc,
        tx_loc,
        power,
        gaussian_noise=gaussian_noise
    )
    # print(kk, rss)
    results[kk] = rss


def test():

    def addRandomObject(penetrations, reflections, orientations):
        '''
        '''
        loss_p = np.random.randint(-20, -5)
        loss_r = np.random.randint(-20, -5)
        i = np.random.randint(2, 62)
        j = np.random.randint(2, 62)
        for ii in range(i-2, i+2):
            penetrations[ii, j-2] = loss_p
            reflections[ii, j-2] = wall_reflection
            orientations[ii, j-2] = 0
            penetrations[ii, j+2] = loss_p
            reflections[ii, j+2] = wall_reflection
            orientations[ii, j+2] = 0
        for jj in range(j-2, j+2):
            penetrations[i-2, jj] = loss_p
            reflections[i-2, jj] = wall_reflection
            orientations[i-2, jj] = 90
            penetrations[i+2, jj] = loss_p
            reflections[i+2, jj] = wall_reflection
            orientations[i+2, jj] = 90

    def setWalls(i, j, orientation_deg):
        penetrations[i, j] = wall_penetration
        reflections[i, j] = wall_reflection
        orientations[i, j] = orientation_deg / 180 * np.pi

    def setGlasses(i, j, orientation_deg):
        penetrations[i, j] = glass_penetration
        reflections[i, j] = glass_reflection
        orientations[i, j] = orientation_deg / 180 * np.pi

    penetrations = np.zeros((64, 64), dtype=float)
    reflections = np.ones((64, 64), dtype=float) * -100.0
    orientations = np.empty((64, 64), dtype=float) * float('nan')

    floormap = SpaceMap(width=6.4, length=6.4, block_size=0.1)

    # add walls
    wall_penetration = -20.0
    wall_reflection = -15.0
    glass_penetration = -3.0
    glass_reflection = -3.0
    # walls
    for j in range(5, 61):  # left
        setWalls(7, j, 90)
    for i in range(7, 62):  # up
        setWalls(i, 60, 0)
    for j in range(29, 60):  # right
        setWalls(61, j, 90)
    for j in range(5, 29):   # right
        setGlasses(61, j, 90)
    for i in range(7, 62):  # down
        setWalls(i, 5, 0)

    np.random.seed(None)
    num_of_objects_to_add = np.random.randint(5, 20)
    for __ in range(num_of_objects_to_add):
        addRandomObject(penetrations, reflections, orientations)

    floormap.setLosses(penetrations, reflections)
    floormap.setOrientations(orientations)

    plotSpace(floormap, cminmax=(-50, 0.0))

    rx_locs = getLocs([0, 6.3], [0, 6.3], step_size=0.1)
    rss_vec = []
    power = -40.0
    starttime = int(time.time())
    tx_loc = SpaceBlock(3.2, 3.2)
    floormap.traceRays(power, tx_loc)
    for j in range(0, rx_locs.shape[0], 8):
        procs = []
        results = Manager().dict()
        actual_num = min(8, rx_locs.shape[0] - j)
        for kk in range(actual_num):
            rx_loc = SpaceBlock(rx_locs[j + kk, 0], rx_locs[j + kk, 1])
            proc = Process(
                target=log_gamma_floorplan_multi,
                args=(floormap, rx_loc, tx_loc, power, False, results, kk)
            )
            proc.start()
            procs.append(proc)
        for proc in procs:
            proc.join()
        rss_vec.extend([results[kk] for kk in range(actual_num)])
    rss_map = convert_vector_to_mat(rx_locs, np.array(rss_vec), (64, 64))
    plotRSS(rss_map, cminmax=(-85,-30))



if __name__ == "__main__":
    test()