import os
import random
import argparse

def readTxt(fname):
    data = []
    with open(fname, 'rb') as fin:
        for line in fin:
            data.append(line.decode('utf-8').strip())
    print("Reading {} example from {}".format(len(data), fname))
    return data

def saveTxt(data, fname):
    with open(fname, 'w') as fout:
        for d in data:
            fout.write('{}\n'.format(d))
    print('Save {} example to {}'.format(len(data), fname))

def sampleByBatch(args):
    datas1 = readTxt(args.i1)
    datas2 = readTxt(args.i2)
    batch_size = args.b
    ratio = len(datas1) / (len(datas1) + len(datas2)) * batch_size 
    left_bound = int(ratio)
    right_bound = int(ratio + 1)
    outputs = []
    i1 = 0
    i2 = 0
    while i1 < len(datas1) and i2 < len(datas2):
        p = random.random() + left_bound
        stride = min(left_bound if p < ratio else right_bound, len(datas1) - i1)
        outputs.extend(datas1[i1:i1+stride])
        i1 += stride
        outputs.extend(datas2[i2:i2+batch_size-stride])
        i2 += (batch_size-stride)

    print("data1: sample {}/{}, data2: sample: {}/{}".format(i1, len(datas1), i2, len(datas2)))
    if i1 < len(datas1):
        outputs.extend(datas1[i1:])

    if i2 < len(datas2):
        outputs.extend(datas2[i2:])
    
    saveTxt(outputs, args.o)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('-i1', help='input1', type=str, default='input.txt')
    parser.add_argument('-i2', help='input2', type=str, default='input.txt')
    parser.add_argument('-o', help='output', type=str, default='output.txt')
    parser.add_argument('-b', help='batch_size', type=float, default=16)
    parser.add_argument('-m', help='mode', type=str)
    args = parser.parse_args()

    eval("{}(args)".format(args.m))
