#!/usr/bin/python

import os
import sys
import pickle
import argparse

from libs.plotting import plotSingleImage


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


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

    print('--- plotting for {}...'.format(filepath))

    plotSingleImage(
        data,
        fp=filepath,
        display=args.showimg,
        dump=args.dump,
        cminmax=(-80, -40),
        title=os.path.basename(filepath).replace('.pickle', '')
    )


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])

    # loop through and plot
    for i in range(len(filepaths)):
        print("- processing file {}".format(filepaths[i]))
        plotSingle(filepaths[i], args)


if __name__ == '__main__':
    p = argparse.ArgumentParser(description='Visualization')
    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(
        '--dump',
        dest='dump',
        action='store_true',
        default=False,
        help='input folderpath for many pickles'
    )
    p.add_argument(
        '--no-disp',
        dest='showimg',
        action='store_false',
        default=True,
        help='input folderpath for many pickles'
    )
    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()

    main(args)