#!/usr/bin/python

import os
import sys
import pickle
import argparse

from libs.util import touch
from libs.util import stats
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 loadData(filepath):
    try:
        return pickle.load(open(filepath, "rb"))
    except BaseException as e:
        print("err: {}".format(e))
        pass
    return None


def fittingSingle(filepath, args):
    '''
    '''
    # load data from file
    data = loadData(filepath)
    if data is None:
        print("err: failed to load file {}".format(filepath))
        return

    # convert data matrix to vectors
    rx_locs, rx_rsses = convert_mat_to_vector(data, block_size=0.1)

    # start fitting
    min_fit_mse = float('inf')
    best_tx_loc = None
    best_tx_pwr = None
    best_env_gamma = None
    best_rsses = None
    for i in range(args.iternum):
        result = modelfit_log_gamma(
            rx_locs,
            rx_rsses,
            bounds_pwr=(-60, 0),
            bounds_gamma=(2, 6),
            bounds_loc_x=(0, 6.4),
            bounds_loc_y=(0, 6.4),
            monte_carlo_sampling=args.montecarlo,
            monte_carlo_sampling_rate=args.montecarlo_rate
        )
        # unpack
        fit_mse, est_tx_loc, est_tx_pwr, est_env_gamma, est_rsses = result
        if fit_mse < min_fit_mse:
            print("|----- found better:", fit_mse, est_tx_loc, est_tx_pwr, est_env_gamma)
            min_fit_mse = fit_mse
            best_tx_loc = est_tx_loc
            best_tx_pwr = est_tx_pwr
            best_env_gamma = est_env_gamma
            best_rsses = est_rsses

    # calculate diff
    data_best_est = convert_vector_to_mat(rx_locs, best_rsses, data.shape, block_size=0.1)
    data_diff = data - data_best_est
    stat = stats(data_diff)

    print("|- final:")
    print("|--- tx loc: {:.3f}, {:.3f}".format(best_tx_loc[0], best_tx_loc[1]))
    print("|--- tx pwr: {:.2f}".format(best_tx_pwr))
    print("|--- env gamma: {:.2f}".format(best_env_gamma))
    print("|--- fitting mse: {:.6f}".format(min_fit_mse))
    print("|--- rss diff: 5perc/95perc/median/mean: {:.4f}/{:.4f}/{:.4f}/{:.4f} m"
          .format(stat['5perc'], stat['95perc'], stat['median'], stat['mean']))

    if args.outputfile:
        filename = os.path.basename(args.outputfile)
        with open(args.outputfile, 'a') as f:
            f.write(
                "{},".format(filename) +
                "{:.3f},{:.3f},".format(best_tx_loc[0], best_tx_loc[1]) +
                "{:.2f},{:.2f},".format(best_tx_pwr, best_env_gamma) +
                "{:.6f},".format(min_fit_mse) +
                "{:.4f},{:.4f},{:.4f},{:.4f}".format(stat['5perc'], stat['95perc'], stat['median'], stat['mean']) +
                "\n"
            )

    if args.visualize:
        plotRSS(
            data,
            data_best_est,
            est_tx_loc=best_tx_loc / 0.1
        )


def main(args):

    # finding files
    filepaths = []
    if args.filepath:
        filepaths.append(args.filepath)
    if args.folderpath:
        files = os.listdir(args.folderpath)
        filepaths.extend(["{}/{}".format(args.folderpath, file) for file in files if '.pickle' in file])

    # prepare
    if args.outputfile and not os.path.isfile(args.outputfile):
        with open(args.outputfile, 'w') as f:
            f.write("filename,tx_x,tx_y,tx_pwr,env_gamma,fit_mse,rss_diff_5perc,rss_diff_95perc,rss_diff_med,rss_diff_avg\n")


    # loop through and fit
    for i in range(len(filepaths)):
        print("- model fitting on file {}".format(filepaths[i]))
        fittingSingle(filepaths[i], args)


if __name__ == '__main__':
    p = argparse.ArgumentParser(description='Traditional Model Fitting')
    p.add_argument(
        '--filepath', '-f',
        dest='filepath',
        default=None,
        help='input filepath for a pickle'
    )
    p.add_argument(
        '--folderpath', '-fd',
        dest='folderpath',
        default=None,
        help='input folderpath for many pickles'
    )
    p.add_argument(
        '--visualize', '-v',
        action='store_true',
        default=False,
        help='enable visualize'
    )
    p.add_argument(
        '--montecarlo', '-m',
        dest='montecarlo',
        action='store_true',
        default=False,
        help='enable monte carlo sampling'
    )
    p.add_argument(
        '--montecarlo-rate', '-mr',
        dest='montecarlo_rate',
        type=float,
        default=0.8,
        help='if monte carlo sampling enabled, specify the rate (0-1, default 0.8)'
    )
    p.add_argument(
        '--iternum', '-n',
        dest='iternum',
        type=int,
        default=1000,
        help='num of iterations'
    )
    p.add_argument(
        '--output', '-o',
        dest='outputfile',
        default=None,
        help='output results to filepath'
    )
    try:
        args = p.parse_args()
    except BaseException as e:
        print(e)
        sys.exit()

    if args.filepath is None and args.folderpath is None:
        print("at least specify file `-f` or folder `-fd`")
        sys.exit()
    elif args.filepath is None and not os.path.isdir(args.folderpath):
        print("folder {} not exit".format(args.folderpath))
        sys.exit()
    elif args.folderpath is None and not os.path.isfile(args.filepath):
        print("file {} not exit".format(args.filepath))
        sys.exit()

    if args.outputfile and not touch(args.outputfile):
        print("cannot create file {}".format(args.outputfile))
        sys.exit()

    main(args)