import numpy as np
from scipy import constants
import matplotlib.pyplot as plt

def to_decibel(x):
    return 10*np.log10(x)

### BEGIN SOLUTION
def length_los(x_tx, x_rx, y_tx, y_rx, h_tx, h_rx, **kwargs):
    distance_sq = (x_tx-x_rx)**2 + (y_tx-y_rx)**2
    return np.sqrt(distance_sq + (h_tx-h_rx)**2)

def length_ground(x_tx, x_rx, y_tx, y_rx, h_tx, h_rx, **kwargs):
    distance_sq = (x_tx-x_rx)**2 + (y_tx-y_rx)**2
    return np.sqrt(distance_sq + (h_tx+h_rx)**2)

def length_left(x_tx, x_rx, y_tx, y_rx, h_tx, h_rx, **kwargs):
    return np.sqrt((x_tx+x_rx)**2 + (y_tx-y_rx)**2 + (h_tx-h_rx)**2)

def length_right(x_tx, x_rx, y_tx, y_rx, h_tx, h_rx, width, **kwargs):
    return np.sqrt((2*width-x_tx-x_rx)**2 + (y_tx-y_rx)**2 + (h_tx-h_rx)**2)

def length_top(x_tx, x_rx, y_tx, y_rx, h_tx, h_rx, height, **kwargs):
    return np.sqrt((x_tx-x_rx)**2 + (y_tx-y_rx)**2 + (2*height-h_tx-h_rx)**2)

def receive_power(path_len, refl_coeff, freq, pow_tx=1, c=constants.c):
    omega = 2*np.pi*freq
    _factor = pow_tx * (c/(2*omega))**2
    _sum_components = refl_coeff * np.exp(-1j*omega/c*path_len)
    path_sum = np.sum(_sum_components, axis=-1)
    rec_power = _factor * np.abs(path_sum)**2
    return rec_power
### END SOLUTION

def main():
    y_rx = np.logspace(1, 3, 3000)
    y_tx = np.zeros_like(y_rx)
    h_tx = 3
    h_rx = 1.5
    x_tx = 5
    x_rx = 5
    width = 10
    height = 10
    freq=2.4e9

    ### BEGIN SOLUTION
    func_lengths = {
                    "All components": [length_los, length_ground, length_left, length_right, length_top],
                    "Two-Ray Ground": [length_los, length_ground],
                    "LOS only": [length_los],
                   }
    for name, _func_lengths in func_lengths.items():
        path_len = np.vstack([_length(x_tx, x_rx, y_tx, y_rx, h_tx, h_rx,
                                      width=width, height=height)
                              for _length in _func_lengths]).T
        refl_coeff = -1./path_len
        refl_coeff[:, 0] = -refl_coeff[:, 0]
        rec_power = receive_power(path_len, refl_coeff, freq)
        plt.semilogx(y_rx, to_decibel(rec_power), label=name)
    plt.legend()
    ### END SOLUTION
    return

if __name__ == "__main__":
    main()
    plt.show()
