import numpy as np
from scipy.constants import speed_of_light
import pymc3 as pm
import matplotlib.pyplot as plt
import os

class ShadowFading():
    def __init__(self, sigma_f=6, correlation_distance=50):
        self._sigma_f = sigma_f
        self._correlation_distance = correlation_distance
        self._covariance = None

    @staticmethod
    def correlation_function(distance, correlation_distance):
        return np.exp(-distance/correlation_distance)

    def _generate_covariance(self, user_locations):
        covariance = np.zeros((len(user_locations), len(user_locations)), dtype="float64")
        for i in range(len(user_locations)):
            for j in range(len(user_locations)):
                covariance[i, j] = ShadowFading.correlation_function(
                    np.linalg.norm(
                        user_locations[i] - user_locations[j]
                    ),
                    self._correlation_distance
                )

        self._covariance = self._sigma_f**2 *  covariance

        while not np.all(np.linalg.eigvals(self._covariance) > 0):
            print("Not positive semidefinite -- add to diagonal")
            self._covariance = np.diag(np.full(len(self._covariance), 1e-1)) + self._covariance

    def covariance(self, user_locations):
        self._generate_covariance(user_locations)

        return self._covariance

    def __call__(self, user_locations):
        self._generate_covariance(user_locations)

        return np.random.multivariate_normal(
            np.full(len(user_locations), 0),
            self._covariance,
            check_valid="warn"
        )

class PathlossModel():
    def __init__(self, frequency=1.8):
        self._frequency = frequency

    @staticmethod
    def getLOSpathlos(d3d, d2d, dbp, fc, h_b, h_t):
        PL1 = 28+22*np.log10(d3d)+20*np.log10(fc)
        PL2 = 28+40*np.log10(d3d)+20*np.log10(fc) - 9*np.log10(dbp**2+(h_b - h_t)**2)
        PL = np.zeros((d3d.shape))
        PL = PL2
        PL[(np.greater_equal(d2d,10) & np.less_equal(d2d,dbp))] = PL1[(np.greater_equal(d2d,10) & np.less_equal(d2d,dbp))]
        return PL

    @staticmethod
    def getNLOSpathloss(street_width, average_building_height, bs_height, ue_height, d3d, f_ghz):
        pathlossdB = 161.04 - 7.1*np.log10(street_width) + 7.5*np.log10(average_building_height) - \
            (24.37 - 3.7 * (average_building_height /bs_height)**2) * np.log10(bs_height) + \
            (43.42 - 3.1*np.log10(bs_height)) * (np.log10(d3d) - 3) + \
            20*np.log10(f_ghz) - \
            (3.2*(np.log10(17.625))**2 - 4.97) - 0.6 * (ue_height - 1.5)
        return pathlossdB

    @staticmethod
    def pathlos_36873(h_b, h_ue, f_ghz, street_w, building_h, d2d):
        '''
        d2d ... Distance 2D in Meters
        '''
        h_e = h_b - h_ue
        d3d = np.sqrt(d2d**2+h_e**2)

        dbp =  4*h_b*h_ue*f_ghz*10e8/speed_of_light

        nlos_pathloss = PathlossModel.getNLOSpathloss(
            street_w,
            building_h,
            h_b, h_ue,
            d3d, f_ghz
        )
        los_pathloss = PathlossModel.getLOSpathlos(
            d3d, d2d,
            dbp, f_ghz,
            h_b, h_ue
        )

        return np.maximum(nlos_pathloss, los_pathloss)

    def __call__(self, measurement_locations, return_distance=False):
        d2d = np.linalg.norm(measurement_locations, axis=1)
        return self.pathlos_36873(15, 5, self._frequency, 10, 15, d2d)

class SystemModel():
    def __init__(self, frequency=1.8, attenuation_max=30, three_db_angle=65):
        self._frequency = frequency
        self._attenuation_max = attenuation_max
        self._three_db_angle = three_db_angle

    def apply_pathloss_model(self, locations, pathloss):
        y = PathlossModel(self._frequency)(locations) - pathloss
        return y

    def plot(self, locations, pathloss, azimuths, save_path=None, name=None):
        y = self.apply_pathloss_model(locations, pathloss)

        plt.figure()
        plt.scatter(azimuths, y)
        plt.ylabel("Pathloss Offset over Angle [dB]")
        plt.xlabel("Azimuth [Degrees]")
        if save_path is not None:
            plt.savefig(os.path.join(save_path, f"pathloss_offset_over_angle_{name}.jpeg"), bbox_inches="tight")
        plt.close()

    def theano_system_model(self, azimuths, sector_orienation):
        bs_azimuth = sector_orienation * 180 / np.pi
        difference = pm.math.abs_((bs_azimuth - azimuths))
        theta = pm.math.minimum(difference, pm.math.abs_(difference - 360))

        pattern = 12 * (theta/self._three_db_angle)**2
        pattern = -pm.math.minimum(
            pattern,
            self._attenuation_max
        )

        return pattern