#!/usr/bin/env python
# -*- coding: utf-8 -*-
import io
import sys
import click
import subprocess

import numpy as np

from . import _fileio, _pairsam_format, cli, _headerops

UTIL_NAME = 'pairsam_restrict'

@cli.command()

@click.argument(
    'pairsam_path', 
    type=str,
    required=False)

@click.option(
    '-f', '--frags',
    type=str,
    required=True,
    help='a tab-separated BED file with the positions of restriction fragments '
         '(chrom, start, end). Can be generated using cooler digest.')

@click.option(
    '-o', "--output", 
    type=str, 
    default="", 
    help='output pairsam file.'
        ' If the path ends with .gz, the output is bgzip-compressed.'
        ' By default, the output is printed into stdout.')


def restrict(pairsam_path, frags, output):
    '''appends positions of restriction fragments to a pairs/pairsam file. 

    PAIRSAM_PATH : input .pairsam file. If the path ends with .gz, the input is
    gzip-decompressed. By default, the input is read from stdin.
    '''
    restrict_py(pairsam_path, frags, output)

def restrict_py(pairsam_path, frags, output):
    instream = (_fileio.auto_open(pairsam_path, mode='r') 
                if pairsam_path else sys.stdin)
    outstream = (_fileio.auto_open(output, mode='w') 
                 if output else sys.stdout)


    header, body_stream = _headerops.get_header(instream)
    header = _headerops.append_new_pg(header, ID=UTIL_NAME, PN=UTIL_NAME)
    outstream.writelines((l+'\n' for l in header))
    rfrags = np.genfromtxt(
        frags, delimiter='\t', comments='#', dtype=None,
        names=['chrom', 'start', 'end'])


    rfrags.sort(order=['chrom', 'start','end'])
    rfrags.sort(order=['chrom', 'start', 'end'])
    chrom_borders = np.r_[0,
                          1+np.where(rfrags['chrom'][:-1] != rfrags['chrom'][1:])[0],
                          rfrags.shape[0]]
    rfrags = {rfrags['chrom'][i]:rfrags['end'][i:j] +1
              for i, j in zip(chrom_borders[:-1], chrom_borders[1:])}


    for line in body_stream:
        cols = line[:-1].split(_pairsam_format.PAIRSAM_SEP)
        chrom1, pos1 = cols[_pairsam_format.COL_C1], int(cols[_pairsam_format.COL_P1])
        rfrag_idx1, rfrag_start1, rfrag_end1 = find_rfrag(rfrags, chrom1, pos1)
        chrom2, pos2 = cols[_pairsam_format.COL_C2], int(cols[_pairsam_format.COL_P2])
        rfrag_idx2, rfrag_start2, rfrag_end2 = find_rfrag(rfrags, chrom2, pos2)
        cols += [str(rfrag_idx1), str(rfrag_start1), str(rfrag_end1)]
        cols += [str(rfrag_idx2), str(rfrag_start2), str(rfrag_end2)]
        outstream.write(_pairsam_format.PAIRSAM_SEP.join(cols))
        outstream.write('\n')

    if instream != sys.stdin:
        instream.close()
    if outstream != sys.stdout:
        outstream.close()


def find_rfrag(rfrags, chrom, pos):
    rsites_chrom = rfrags[chrom.encode('ascii')]
    idx = min(max(0,rsites_chrom.searchsorted(pos, 'right')-1), len(rsites_chrom)-2)
    return idx, rsites_chrom[idx], rsites_chrom[idx+1]

if __name__ == '__main__':
    restrict()

@@ -167,16 +167,35 @@
 eaderops
+, common_io_options
 %0A%0AUTIL_N
@@ -664,24 +664,29 @@
 nds with .gz
+/.lz4
 , the output
@@ -693,14 +693,8 @@
  is 
-bgzip-
 comp
@@ -699,16 +699,31 @@
 mpressed
+ by pbgzip/lz4c
 .'%0A     
@@ -777,16 +777,35 @@
 out.')%0A%0A
+@common_io_options%0A
 %0Adef res
@@ -833,24 +833,34 @@
 rags, output
+, **kwargs
 ):%0A    '''ap
@@ -992,30 +992,40 @@
 with .gz
+/.lz4
 , the 
+%0A    
 input is
 %0A    gzi
@@ -1020,18 +1020,9 @@
 t is
-%0A    gzip-
+ 
 deco
@@ -1029,16 +1029,31 @@
 mpressed
+ by pbgzip/lz4c
 . By def
@@ -1139,16 +1139,26 @@
 , output
+, **kwargs
 )%0A%0Adef r
@@ -1195,16 +1195,26 @@
 , output
+, **kwargs
 ):%0A    i
@@ -1264,16 +1264,151 @@
 mode='r'
+, %0A                                  nproc=kwargs.get('nproc_in'),%0A                                  command=kwargs.get('cmd_in', None)
 ) %0A     
@@ -1450,16 +1450,17 @@
 .stdin)%0A
+%0A
     outs
@@ -1502,16 +1502,155 @@
 mode='w'
+, %0A                                   nproc=kwargs.get('nproc_out'),%0A                                   command=kwargs.get('cmd_out', None)
 ) %0A     
@@ -2431,21 +2431,25 @@
 s = line
-%5B:-1%5D
+.rstrip()
 .split(_

add lz4 and custom compression