#!/usr/bin/env python3

import sys, os
sys.path.append(os.getcwd())
import argparse
import subprocess

from pipeline.se3solver import solve, merge_forest, presample, merge_pdsc, merge_blooming_forest

def main():
    # subprocess.call(['/usr/bin/env'])
    main_parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    subparsers = main_parser.add_subparsers(dest='command')
    # Subcommand 'solve'
    parser = subparsers.add_parser("solve", help='Solve a puzzle with the specified planner.', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('puzzle', help='Configure file generated by OMPL GUI')
    parser.add_argument('planner_id', help='Choose a planner', choices=range(20), type=int)
    parser.add_argument('days', help='Time limit in day(s)', type=float)
    parser.add_argument('--ec_budget', help='''Budget of edge connections.''', type=int, default=-1)
    parser.add_argument('--out', help='Output complete planning data', default='')
    parser.add_argument('--trajectory_out', help='Output complete planning data', default='')
    parser.add_argument('--bloom_out', help='Output bloom results', default=None)
    parser.add_argument('--sampler_id', help='Valid state sampler', type=int, default=0)
    parser.add_argument('--saminj', help='Sample injection file', type=str, default='')
    parser.add_argument('--samset', help='Predefined sample set (PDS) file', type=str, default='')
    parser.add_argument('--use_roots_from',
                        help='Use predicated roots from file. Commonly used as baseline comparison',
                        type=str, nargs='*', default=[])
    parser.add_argument('--use_blooming_tree',
                        help='Use blooming tree in Phase 2. '
                        'If current RDT tree directly connected to a PDS sample, '
                        'the whole tree from blooming would be added to the current RDT tree.',
                        action='store_true')
    parser.add_argument('--samset2', help='Predefined sample set, format: current total prefix', type=str, nargs=3, default=[])
    parser.add_argument('--skip_existing', help='Quit if the output file already exists', action='store_true')
    parser.add_argument('--rdt_k', help='K Nearest in RDT algorithm', type=int, default=1)
    parser.add_argument('--cdres', help='Collision detection resolution, set zero/negative to enable continuouse collision detection', type=float, default=0.005)
    parser.add_argument('--replace_istate', help='''
Replace Initial State. Syntax: file=<path to npz>,key=<npz key>,offset=<number>,size=<number>,out=<dir>,
in which key=, size= are optional.
key is default to the first array in NPZ,
and size is default to 1.
A -1 size selects all states in the file.
''',
            type=str, default=None)
    parser.add_argument('--replace_gstate', help='''Same syntax with replace_istate, but replaces goal state''', type=str, default=None)
    parser.add_argument('--bvresize', help='''add this number to the bounding volume defined by the puzzle''', type=float, default=0.0)
    parser.add_argument('solver_option_vector', help='Option vectors passed over to OMPL planner', nargs=argparse.REMAINDER, type=str, default=[])
    # Subcommand 'merge_forest'
    parser = subparsers.add_parser("merge_forest", help='Merge the planning data from multiple planners')
    parser.add_argument('puzzle', help='Configure file generated by OMPL GUI')
    parser.add_argument('days', help='Time limit in day(s)', type=float)
    parser.add_argument('trees', help='''Tree/Graph files/directories created by individual planner.
If a directory is provided, all .mat files in this directory will be loaded''', nargs='+')
    parser.add_argument('--out', help='Output file of the merged graph, in .npz format', default=None)
    parser.add_argument('--cdres', help='Collision detection resolution', type=float, default=0.005)
    # Subcommand 'presample'
    parser = subparsers.add_parser("presample", help='Presample a set of samples.', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('puzzle', help='Configure file generated by OMPL GUI')
    parser.add_argument('nsamples', help='Total Number of samples', type=int)
    parser.add_argument('out', help='Output file for samples')
    parser.add_argument('--sampler', help='Valid state sampler', type=int, default=0)
    parser.add_argument('--cdres', help='Collision detection resolution', type=float, default=0.005)
    # Subcommand 'merge_pdsc'
    parser = subparsers.add_parser("merge_pdsc", help='Merge connectivity matrix created from PreDefined set of samples.', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('pdsf', help='Pre-Defined Sample set (PDS) File')
    parser.add_argument('dir', help='''Directory that stores the connectivity sparse matrices''')
    parser.add_argument('out', help='''output file''')

    # Subcommand 'merge_blooming_forest'
    parser = subparsers.add_parser("merge_blooming_forest", help='Merge the forest from blooming algorithm')
    parser.add_argument('puzzle', help='Configure file generated by OMPL GUI')
    parser.add_argument('bloom_dir', help='Directory that stores the blooming dir')
    parser.add_argument('--out', help='Output file of inter-tree edges, in .npz format', default=None)
    parser.add_argument('--cdres', help='Override collision detection resolution defined in the puzzle file', type=float, default=None)
    parser.add_argument('--knn', help='K-nearest', type=int, default=8)
    parser.add_argument('--algo_version',
                        help='Reversion of Algorithm. 0: single KNN DS, 1: KNN for each DS',
                        type=int, default=0, choices=[0,1,2,3])
    parser.add_argument('--subset',
                        type=int, default=None, nargs='*')

    args = main_parser.parse_args()
    # Empty dic 
    args.istate_dic = {}
    args.gstate_dic = {}
    if hasattr(args, 'replace_istate') and args.replace_istate is not None:
        dic = dict(item.split("=") for item in args.replace_istate.split(","))
        print(dic)
        assert ('file' in dic) and ('offset' in dic) and ('out' in dic), '--replace_istate need file, offset and out'
        if 'size' in dic and int(dic['size']) > 1:
            assert 'out' in dic, 'size= requires out= for storage'
        args.istate_dic = dic
    if hasattr(args, 'replace_gstate') and args.replace_gstate is not None:
        dic = dict(item.split("=") for item in args.replace_gstate.split(","))
        print(dic)
        assert ('file' in dic) and ('offset' in dic) and ('out' in dic), '--replace_gstate need file, offset and out'
        if 'size' in dic and int(dic['size']) > 1:
            assert 'out' in dic, 'size= requires out= for storage'
        args.gstate_dic = dic

    globals()[args.command](args)

if __name__ == '__main__':
    main()
