import numpy as np
import tensorflow as tf
from helpers import mappings
from helpers.utils import logsumexp10

class ThroughputModel():
    def __init__(self, resource_blocks=100, resource_block_bw=180e3, noise_level=-120.0, demand="medium"):
        self._noise_level = noise_level
        self._resource_blocks = resource_blocks
        self._resource_block_bw = resource_block_bw

        self._act_ues_mapping = mappings.Active_UEs_from_Connected(demand=demand)
        self._cell_load_mapping = mappings.Cell_Load_from_Active_UEs()
        self._cqi_mapping = mappings.CQI_from_SINR()
        self._spectral_efficiency_mapping =  mappings.Spectral_Efficiency_from_CQI()

    @tf.function
    def get_cell_assignment(self, R_matrix, softmax_temperature=3):
        R_sorted = tf.cast(tf.sort(R_matrix, axis=1), tf.float32)
        R_max = R_sorted[:, -1]
        delta_R = tf.reshape(R_max, (-1, 1)) - R_matrix
        soft_assignment = tf.math.softmax(-delta_R / softmax_temperature)

        return soft_assignment

    @tf.function
    def get_act_ues(self, number_ues_soft):
        return self._act_ues_mapping.map(number_ues_soft)

    @tf.function
    def get_cell_load(self, number_ues_soft):
        return self._cell_load_mapping.map(number_ues_soft)

    @tf.function
    def get_SINR_matrix(self, R_matrix, cell_load_vector):
        I_matrix = []
        for j in range(R_matrix.shape[1]):
            indices = list(range(R_matrix.shape[1]))
            indices.remove(j)

            relevant_interferers = tf.gather(R_matrix, indices, axis=1)
            relevant_cell_load = tf.gather(cell_load_vector, indices)

            noise_vector = tf.fill((relevant_interferers.shape[0], 1), tf.cast(self._noise_level, tf.float32))

            additive_terms = tf.concat((relevant_interferers, noise_vector), axis=1)
            scaling_factors = tf.concat((relevant_cell_load, tf.constant([1.], dtype=tf.float32)), axis=0)

            I_matrix.append(
                logsumexp10(additive_terms, alpha=scaling_factors)
            )

        SINR_matrix = R_matrix - tf.cast(tf.stack(I_matrix, axis=1), tf.float32)

        return SINR_matrix

    @tf.function
    def get_achievable_throughput(self, SINR_matrix):
        cqi_matrix = self._cqi_mapping.map(tf.reshape(SINR_matrix, (-1,)))
        achievable_throughput_matrix = tf.cast(
            self._resource_block_bw * self._spectral_efficiency_mapping.map(cqi_matrix), tf.float32
        )
        achievable_throughput_matrix = tf.reshape(achievable_throughput_matrix, SINR_matrix.shape)

        return achievable_throughput_matrix

    @tf.function
    def __call__(self, R_matrix, extended=False):
        soft_assignment_matrix = self.get_cell_assignment(R_matrix)
        connected_UEs_vector = tf.reduce_sum(soft_assignment_matrix, axis=0)

        #Compute the number of active UEs and the cell load
        active_UEs_vector = self.get_act_ues(connected_UEs_vector)
        cell_load_vector = self.get_cell_load(active_UEs_vector)

        #Compute the SINR and subsequently the RATE
        SINR_matrix = self.get_SINR_matrix(R_matrix, cell_load_vector)
        achievable_throughput_matrix = self.get_achievable_throughput(SINR_matrix)

        load_aware_throughput_vector = tf.reduce_sum(
            (self._resource_blocks / (active_UEs_vector + 1.0)) * soft_assignment_matrix * achievable_throughput_matrix, axis=1
        )

        #Compute the SINR and the theoretical empty cell rate per UE
        SINR_vector = tf.reduce_sum(soft_assignment_matrix * SINR_matrix, axis=1)
        load_unaware_throughput_vector = tf.reduce_sum(
            (self._resource_blocks / (1 + self.get_act_ues(R_matrix.shape[0] / R_matrix.shape[1])) ) * soft_assignment_matrix * achievable_throughput_matrix, axis=1
        )
        if extended:
            return load_aware_throughput_vector, load_unaware_throughput_vector, SINR_vector, connected_UEs_vector
        else:
            return load_aware_throughput_vector