#!/usr/bin/env python
import argparse, sys
from peakachu import train_models,score_genome,score_chromosome,call_loops,calculate_depth
def getargs():
    ## Construct an ArgumentParser object for command-line arguments
    parser = argparse.ArgumentParser(description='''Train Random Forest with Hi-C data and training peaklist.
                                     ''',
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    subparsers = parser.add_subparsers(dest='subcommands')
    subtrain = subparsers.add_parser('train',
            help='''Train RandomForest model per chromosome''')
    subtrain.set_defaults(func=train_models.main)
    subchrom = subparsers.add_parser('score_chromosome',
            help='''Calculate interaction probability per pixel for a chromosome''')
    subchrom.set_defaults(func=score_chromosome.main)
    subgen = subparsers.add_parser('score_genome',
            help='''Calculate interaction probability per pixel for the whole genome''')
    subgen.set_defaults(func=score_genome.main)
    subdepth=subparsers.add_parser('depth',
            help='''Print total intra-chromosomal read count''')
    subdepth.set_defaults(func=calculate_depth.main)
    subpool = subparsers.add_parser('pool',
            help='''Print centroid loci from score_x output''')
    subpool.set_defaults(func=call_loops.main)
    subs=[subdepth,subtrain,subchrom,subgen,subpool]
    #subpool.add_argument('-l','--lower',
    #                     help = 'Lower bound of report in bp (20000)',
    #                     type = int, default=20000)
    #subpool.add_argument('-u','--upper',
    #                     help = 'Upper bound of report in bp (300000)',
    #                     type = int, default=300000)
    for i in subs[1:]:
        i.add_argument('-r', '--resolution',
                         help = 'Resolution in bp, default 10000',
                         type = int, default=10000)
    for i in subs[:-1]:
        i.add_argument('-p', '--path',
                         help = 'Path to a .cool URI string or a .hic file.')
    for i in subs[1:-1]:
        i.add_argument('--balance',action='store_true', help='Whether or not using the ICE/KR-balanced matrix.')
        i.add_argument('-O', '--output', help='Folder path to store results.')
        i.add_argument('-w', '--width', type = int, default=5,
                        help = '''Number of bins added to center of window. 
                        default width=5 corresponds to 11x11 windows''')
    for i in subs[2:-1]:
        i.add_argument('-m', '--model', type = str, help = '''Path to pickled model file.''')
        i.add_argument('-l', '--lower', type = int, default=2, help = '''Lower bound of distance between loci in bins (default 2).''')
        i.add_argument('-u', '--upper', type = int, default=300, help = '''Upper bound of distance between loci in bins (default 300).''')
    subpool.add_argument('-i', '--infile',
                        help = '''Path to the bedpe file containing model output''')

    subpool.add_argument('-t', '--threshold', type = float, default = .9,
                        help = 'Probability threshold applied before peak calling, default .9')
    subtrain.add_argument('-b', '--bedpe',
                        help = '''Path to the bedpe file containing positive training set.''')
   
    ## Parse the command-line arguments
    commands = sys.argv[1:]
    if ((not commands) or ((commands[0] in ['train', 'score_chromosome', 'score_genome', 'depth', 'pool'])
        and len(commands) == 1)):
        commands.append('-h')
    args = parser.parse_args(commands)

    return args, commands


def run():
    # Parse Arguments
    args, commands = getargs()
    # Improve the performance if you don't want to run it
    if commands[0] not in ['-h','--help']:
        args.func(args)
        

if __name__ == '__main__':
    run()
