#!/usr/bin/python


import numpy as np
import matplotlib.pyplot as plt

from libs.consts import COLORMAP_MIN
from libs.consts import COLORMAP_MAX
from libs.spacemap import SpaceMap


def plotSpace(space_map: SpaceMap):
    '''
    '''
    tx_locs = space_map.getTransmitterLocs()
    penetration_loss, reflection_loss = space_map.getLosses()

    fig = plt.figure(figsize=(8, 3))

    ax1 = fig.add_subplot(1, 2, 1)
    ax1.set_title("penetration loss")
    plt.imshow(
        np.transpose(penetration_loss),
        cmap='hot',
        origin='lower',
        interpolation='nearest',
        vmin=COLORMAP_MIN,
        vmax=COLORMAP_MAX
    )
    plt.colorbar()

    ax2 = fig.add_subplot(1, 2, 2)
    ax2.set_title("reflection loss")
    plt.imshow(
        np.transpose(reflection_loss),
        cmap='hot',
        origin='lower',
        interpolation='nearest',
        vmin=COLORMAP_MIN,
        vmax=COLORMAP_MAX
    )
    plt.colorbar()

    # plt.show()
    plt.draw()
    plt.pause(0.1)
    q = input("press enter..")
    plt.close()


def plotRSS(
    ori_rss: np.ndarray,
    est_rss: np.ndarray = None,
    est_tx_loc: np.ndarray = None,
    true_tx_loc: np.ndarray = None
):
    '''
    '''
    fig = plt.figure(figsize=(8, 3))
    ax1 = fig.add_subplot(1, 2, 1)
    ax1.set_title("original rss map")
    plt.imshow(
        np.transpose(ori_rss),
        cmap='hot',
        origin='lower',
        interpolation='nearest',
        vmin=COLORMAP_MIN,
        vmax=COLORMAP_MAX
    )
    plt.colorbar()

    ax2 = fig.add_subplot(1, 2, 2)
    ax2.set_title("est rss map")
    if est_rss is not None:
        plt.imshow(
            np.transpose(est_rss),
            cmap='hot',
            origin='lower',
            interpolation='nearest',
            vmin=COLORMAP_MIN,
            vmax=COLORMAP_MAX
        )
        plt.colorbar()
    if est_tx_loc is not None:
        ax2.plot(est_tx_loc[0], est_tx_loc[1], 'bo')
        ax2.text(est_tx_loc[0] - 5, est_tx_loc[1] + 2, 'est tx', color='b')
    if true_tx_loc is not None:
        ax2.plot(true_tx_loc[0], true_tx_loc[1], 'bs')
        ax2.text(true_tx_loc[0], true_tx_loc[1], 'true tx', color='b')

    # plt.show()
    plt.draw()
    plt.pause(0.1)
    q = input("press enter..")
    plt.close()