#!/usr/bin/python

import os
import sys
import pickle
import argparse
import numpy as np

from libs.util import convert_vector_to_mat
from libs.consts import RANDOM_SEED_TEST
from libs.consts import RANDOM_SEED_TRAIN
from libs.models import log_gamma_loc


def getRandomTXLocs(
    num: int,
    width: float,
    length: float,
    offset_w: float = 0.0,
    offset_l: float = 0.0
):
    '''
    '''
    xys = np.random.rand(num, 2)
    xys[:, 0] = xys[:, 0] * width + offset_w
    xys[:, 1] = xys[:, 1] * length + offset_l
    return xys


def getLocs(loc_x_range, loc_y_range, step_size: float = 0.3):
    loc_xs = np.arange(loc_x_range[0], loc_x_range[1] + step_size, step_size)
    loc_ys = np.arange(loc_y_range[0], loc_y_range[1] + step_size, step_size)
    return np.stack(np.meshgrid(loc_xs, loc_ys), -1).reshape(-1, 2)


def getPowers(power_range, step_size: float = 2):
    return np.arange(power_range[0], power_range[1] + step_size, step_size)


def generateData(tx_locs, args):
    '''
    '''
    if args.single:
        tag = "single"
    elif args.train:
        tag = "input_synthetic_train"
    elif args.train_real:
        tag = "input_real_emu_train"
    elif args.test:
        tag = "testing_real_emu"
    tag = "{0}_gamma_{1:.1f}".format(tag, args.gamma)
    if args.withnoise:
        tag = "{}_noise_10dBvar".format(tag)

    # create folder
    folderp = "{}/{}".format(args.outfolder, tag)
    if not os.path.isdir(folderp):
        try:
            os.makedirs(folderp)
        except BaseException as e:
            print("err: {}".format(e))
            sys.exit(-1)

    # generate
    rx_locs = getLocs([-3.2, 3.1], [-3.2, 3.1], step_size=0.1)
    powers = getPowers([-40, 30])
    for i in range(tx_locs.shape[0]):
        fp_base = "{0}/img_{1:.2f}_{2:.2f}_{3:.1f}".format(folderp, tx_locs[i, 0] - 3.2, tx_locs[i, 1] - 3.2, args.gamma)
        for power in powers:
            rss_vec = log_gamma_loc(rx_locs, tx_locs[i, :], power, args.gamma, gaussian_noise=args.withnoise)
            rss_map = convert_vector_to_mat(rx_locs, rss_vec, (64, 64))
            with open("{}_{}.pickle".format(fp_base, int(power)), 'wb') as f:
                pickle.dump(rss_map, f, pickle.HIGHEST_PROTOCOL)


def main(args):
    tx_locs = None

    if args.single:
        if args.txloc is None:
            print("did not specify tx location")
            return
        locations = [float(x) for x in args.txloc.split(' ')]
        tx_locs = np.array(locations).reshape(-1, 2)
    elif args.train:
        np.random.seed(RANDOM_SEED_TRAIN)
        tx_locs = getLocs([-3, 3], [-3, 3], step_size=0.3)
    elif args.train_real:
        np.random.seed(RANDOM_SEED_TRAIN)
        tx_locs = getRandomTXLocs(400, 6.4, 6.4, offset_w=-3.2, offset_l=-3.2)
    elif args.test:
        np.random.seed(RANDOM_SEED_TEST)
        tx_locs = getRandomTXLocs(400, 6.4, 6.4, offset_w=-3.2, offset_l=-3.2)
    else:
        print("nothing specified")
        return

    generateData(tx_locs, args)


if __name__ == '__main__':
    p = argparse.ArgumentParser(description='Data Generator with Log Gamma Distance Model')
    p.add_argument(
        'outfolder',
        help='output generated data to folder'
    )
    p.add_argument(
        '--single',
        dest='single',
        action='store_true',
        default=False,
        help='generate data with tx location'
    )
    p.add_argument(
        '--txloc',
        dest='txloc',
        default=None,
        help='tx location in `x y x y ...` format, in meters'
    )
    p.add_argument(
        '--training',
        dest='train',
        action='store_true',
        default=False,
        help='generate training data (synthetic part)'
    )
    p.add_argument(
        '--training-real',
        dest='train_real',
        action='store_true',
        default=False,
        help='emulating training data (real data part)'
    )
    p.add_argument(
        '--testing',
        dest='test',
        action='store_true',
        default=False,
        help='emulating testing data (real data part)'
    )
    p.add_argument(
        '--gamma',
        dest='gamma',
        type=float,
        default=2.0,
        help='env gamma'
    )
    p.add_argument(
        '--with-noise',
        dest='withnoise',
        action='store_true',
        default=False,
        help='add 10dB noise to generated data'
    )
    try:
        args = p.parse_args()
    except BaseException as e:
        print(e)
        sys.exit()

    main(args)