#!/usr/bin/python


import pickle
import argparse


def loadData(filepath):
    try:
        return pickle.load(open(filepath, "rb"))
    except BaseException as e:
        print("err: {}".format(e))
        pass
    return None


def trainedModelFromGAN(data):
    '''
    '''
    print("not implemented yet")
    raise


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

    tx_loc, tx_pwr, est_rss_map, metrics = trainedModelFromGAN(data)
    print("|- final:")
    print("|--- tx loc: {:.3f}, {:.3f}".format(tx_loc[0], tx_loc[1]))
    print("|--- tx pwr: {:.2f}".format(tx_pwr))
    print("|--- metrics: {}".format(metrics))

    if args.outputfile:
        filename = os.path.basename(args.outputfile)
        with open(args.outputfile, 'a') as f:
            f.write(
                "{},".format(filename) +
                "{:.3f},{:.3f},".format(tx_loc[0], tx_loc[1]) +
                "{:.2f},".format(tx_pwr) +
                ",".join(["{:.6f}".format(x) for x in metrics]) +
                "\n"
            )

    if args.visualize:
        plotRSS(
            data,
            est_rss_map,
            est_tx_loc=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\n")


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


if __name__ == '__main__':
    p = argparse.ArgumentParser(description='NN from GAN')
    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(
        '--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)