#!/usr/bin/env python

# Created on Sat Oct 13 14:04:26 2018
# Author: XiaoTao Wang

## Required modules

import argparse, sys, os, neoloop

currentVersion = neoloop.__version__

def getargs():
    ## Construct an ArgumentParser object for command-line arguments
    parser = argparse.ArgumentParser(description='''Search matched SV regions by gene.''',
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    
    # Version
    parser.add_argument('-v', '--version', action='version',
                        version=' '.join(['%(prog)s',currentVersion]),
                        help='Print version number and exit.')
    
    # Input
    parser.add_argument('-L', '--loop-file',
                        help = '''Path to the loop file generated by neoloop-caller.''')
    parser.add_argument('-G', '--gene-name',
                        help = '''Gene name.''')
    parser.add_argument('--ensembl-release', default=93, type=int,
                        help='''Release number of the Ensembl database. The reference genome assembly should
                        be consistent with your Hi-C data used in neoloop-caller.''')
    parser.add_argument('--species', default='human', help='''Species name of your cell line.''')

    ## Parse the command-line arguments
    commands = sys.argv[1:]
    if not commands:
        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', '-v', '--help', '--version']:
        
        from pyensembl import EnsemblRelease
        from neoloop.visualize.loops import Loops

        # Create/Load Ensembl Release cache ...
        db = EnsemblRelease(args.ensembl_release, species=args.species)
        db.download()
        db.index()
        
        IDs = set()
        loops = Loops(args.loop_file).loops
        min_interval = 15000
        for k in loops:
            for loci1, loci2, n in loops[k]:
                if n != 1:
                    continue
                c1, s1, e1 = loci1
                c2, s2, e2 = loci2
                res = e1 - s1
                if res < min_interval:
                    half = (min_interval - res) // 2
                    s1 = max(0, s1 - half)
                    e1 = e1 + half
                    s2 = max(0, s2 - half)
                    e2 = e2 + half
                genes = db.gene_names_at_locus(c1.lstrip('chr'), s1, e1)
                for g in genes:
                    if g == args.gene_name:
                        IDs.add(k)
                        break
                
                genes = db.gene_names_at_locus(c2.lstrip('chr'), s2, e2)
                for g in genes:
                    if g == args.gene_name:
                        IDs.add(k)
                        break
        
        for k in IDs:
            print(k)

if __name__ == '__main__':
    run()

