#!/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='''Unveil Hi-C Anchors and Peaks.''',
                                     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='''Calculate the total number of intra-chromosomal chromatin contacts
            and select the most appropriate pre-trained model for you.''')
    subdepth.set_defaults(func=calculate_depth.main)
    subpool = subparsers.add_parser('pool',
            help='''Print centroid loci from score_genome/score_chromosome output''')
    subpool.set_defaults(func=call_loops.main)
    subs=[subdepth,subtrain,subchrom,subgen,subpool]
    
    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.''')

    subchrom.add_argument('-C', '--chrom', help='''Chromosome label. Only contact data within the
                          specified chromosome will be considered.''')
    subgen.add_argument('-C', '--chroms', nargs = '*', default = ['#', 'X'],
                        help = 'List of chromosome labels. Only contact data within the specified '
                        'chromosomes will be included. Specially, "#" stands for chromosomes '
                        'with numerical labels. "--chroms" with zero argument will include '
                        'all chromosome data. (default "#" X)')

    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=6, help = '''Lower bound of distance between loci in bins (default 6).''')
        i.add_argument('-u', '--upper', type = int, default=300, help = '''Upper bound of distance between loci in bins (default 300).''')
        i.add_argument('--minimum-prob', type = float, default=0.5, help = '''Only output pixels with probability score greater than this value (default 0.5)''')
        i.add_argument('-O', '--output', help='Output file name.')

    subdepth.add_argument('--min-dis', default=0, type=int,
                         help = '''Only count reads with genomic distance (in base pairs) greater than this value. (default 0)''')
    
    subtrain.add_argument('-b', '--bedpe',
                        help = '''Path to the bedpe file containing positive training set.''')
    subtrain.add_argument('-w', '--width', type = int, default=5,
                        help = '''Number of bins added to center of window. 
                        default width=5 corresponds to 11x11 windows''')
    subtrain.add_argument('--nproc', type=int, default=4, help='''Number of worker processes that
                          will be allocated for training. (default 4)''')
    subtrain.add_argument('-O', '--output', help='Folder path to store trained models.')

    subpool.add_argument('-i', '--infile',
                        help = '''Path to the bedpe file outputted from score_chromosome or score_genome''')
    subpool.add_argument('-o', '--outfile', help='Output file name.')
    subpool.add_argument('-t', '--threshold', type = float, default = 0.9,
                        help = 'Probability threshold applied before peak calling (default 0.9)')
   
    ## 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()
