#!/usr/bin/python


import numpy as np

from libs.consts import NOISE_FLOOR
from libs.consts import RANDOM_SEED
from libs.consts import GAUSSIAN_NOISE_STD
from libs.consts import GAUSSIAN_NOISE_MEAN

np.random.seed(RANDOM_SEED)


def log_gamma_loc(
    rx_loc: np.ndarray,  # dtype=float, avoid int
    tx_loc: np.ndarray,  # dtype=float, avoid int
    pwr: float,
    gamma: float,
    loss: float = 0.0,
    gaussian_noise: bool = False
):
    '''
    '''
    dist_squared = np.nansum((rx_loc - tx_loc) * (rx_loc - tx_loc), axis=1)
    dist_squared[dist_squared < 0.02] = 0.02
    noise = np.random.normal(GAUSSIAN_NOISE_MEAN, GAUSSIAN_NOISE_STD) if gaussian_noise else 0.0
    rss = pwr - 10.0 * gamma / 2 * np.log10(dist_squared) + noise
    rss[rss < NOISE_FLOOR] = NOISE_FLOOR
    rss[rss > pwr] = pwr + loss
    return rss


def log_gamma_dist(
    dist: np.ndarray,
    pwr: float,
    gamma: float,
    loss: float = 0.0,
    gaussian_noise: bool = False,
    is_squared: bool = False
):
    '''
    '''
    factor = 0.5 if is_squared else 1.0
    noise = np.random.normal(GAUSSIAN_NOISE_MEAN, GAUSSIAN_NOISE_STD) if gaussian_noise else 0.0
    rss = pwr - 10.0 * gamma * factor * np.log10(dist) + noise + loss
    rss[rss < NOISE_FLOOR] = NOISE_FLOOR
    rss[rss > pwr] = pwr
    return rss


def log_gamma_floorplan(
    floorplan,
    rx_loc,
    tx_loc,
    pwr: float,
    gaussian_noise: bool = False
):
    if isinstance(rx_loc, list) or isinstance(rx_loc, tuple):
        print("err: only support a single loc")
        return
    noise = np.random.normal(GAUSSIAN_NOISE_MEAN, GAUSSIAN_NOISE_STD) if gaussian_noise else 0.0
    floorplan.traceRays(pwr, tx_loc)
    rss, paths = floorplan.traceRay(rx_loc)
    return max(rss + noise, NOISE_FLOOR), paths