<filename>bin/sample_imgt_without_replacement.py
#! /usr/bin/env python
# Copyright 2014 <NAME>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""sample_vdjxml_without_replacement.py

Subsample an input vdjxml stream. Sample without replacement. If the sampling
level is larger than the number of chains in the file, it simply copies the
file and issues a warning. 
"""

import sys
import random
import optparse
import subprocess
import warnings

import vdj

warnings.simplefilter('always')

option_parser = optparse.OptionParser()
option_parser.add_option('-n','--num',type='int')
(options,args) = option_parser.parse_args()

if len(args) == 2:
    inhandle = open(args[0],'r')
    outhandle = open(args[1],'w')
elif len(args) == 1:
    inhandle = open(args[0],'r')
    outhandle = sys.stdout
elif len(args) == 0:
    raise ValueError, "must provide at least an input file"

# determine the total number of chains in the file (using unix grep and wc)
p = subprocess.Popen('cat %s | grep ^ID | wc -l' % args[0],shell=True,stdout=subprocess.PIPE)
total_chains = int(p.stdout.read().strip())

# check if subsampling should include the entire file
if options.num >= total_chains:
    warnings.warn("Subsampling level is greater than or equal to number of chains in file: printing whole file to output.")
    for chain in vdj.parse_imgt(inhandle):
        print >>outhandle, chain
else:
    # choose a random set of indices to select for the output file
    random.seed()
    idxs = sorted(random.sample(xrange(total_chains),options.num))
    for (i,chain) in enumerate(vdj.parse_imgt(inhandle)):
        if len(idxs) == 0:
            break
        if i == idxs[0]:
            print >>outhandle, chain
            idxs.pop(0)
