#!/usr/bin/python

import sys
sys.path.append(".")
sys.path.append("..")
sys.path.append("../..") # Adds higher directory to python modules path

import os
import time
import pickle
import numpy as np
from libs.consts import NOISE_FLOOR
from libs.util import convert_mat_to_vector
from libs.util import convert_vector_to_mat
from libs.fitting import modelfit_log_gamma
from libs.plotting import plotRSS


def test():
    data_ori = pickle.load(open(os.path.join(os.path.dirname(__file__), "dev_map.pickle"), "rb"))
    rx_locs, rx_rsses = convert_mat_to_vector(data_ori)

    min_pmse = float('inf')
    best_tx_loc = None
    best_tx_pwr = None
    best_env_gamma = None
    best_rsses = None
    for i in range(10):
        pmse, est_tx_loc, est_tx_pwr, est_env_gamma, est_rsses = modelfit_log_gamma(rx_locs, rx_rsses)
        if pmse < min_pmse:
            print("found better:", pmse, est_tx_loc, est_tx_pwr, est_env_gamma)
            data_est = convert_vector_to_mat(rx_locs, est_rsses, data_ori.shape)
            min_pmse = pmse
            best_tx_loc = est_tx_loc
            best_tx_pwr = est_tx_pwr
            best_env_gamma = est_env_gamma
            best_rsses = est_rsses
    print("final:", min_pmse, best_tx_loc, best_tx_pwr, best_env_gamma)
    data_best_est = convert_vector_to_mat(rx_locs, best_rsses, data_ori.shape)
    plotRSS(
        data_ori,
        data_best_est,
        est_tx_loc=best_tx_loc / 0.1
    )


if __name__ == "__main__":
    test()