#!/usr/bin/env python3
import random, math, subprocess, numpy as np, util
from h5py import File
import os.path

def copy_dataset(fromdset, todset):
	for l, r in util.lrange(0, fromdset.shape[0], progress_bar=True):
		todset[l:r] = fromdset[l:r]

def create_dataset(fromdset, f, dsetname):
	return f.create_dataset(dsetname, fromdset.shape, fromdset.dtype)

def line_generator(dset, start, end):
	for l, r in util.lrange(start, end):
		yield from dset[l:r]

def shuffle_line_generator(fromdset):
	start, mid, end = 0, fromdset.shape[0] // 2, fromdset.shape[0]
	gens = [line_generator(fromdset, start, mid), line_generator(fromdset, mid, end)]
	sizes = [mid - start, end - mid]
	while sizes[0] + sizes[1] > 0:
		logit = int(random.random() * (sizes[0] + sizes[1]) < sizes[1])
		yield next(gens[logit])
		sizes[logit] -= 1
	gens[0].close()
	gens[1].close()

def shuffle_loop(fromdset, todset, buffer_size=100000):
	start, end = 0, fromdset.shape[0]
	buf = np.zeros([buffer_size, *fromdset.shape[1:]])
	gen = shuffle_line_generator(fromdset)
	for l, r in util.lrange(start, end, buffer_size, progress_bar=True):
		for i in range(r - l):
			buf[i] = next(gen)
		todset[l:r] = buf[:r-l]
	gen.close()

def shuffle(fromdset, todset, tmpfile="/tmp/shuffle_tmp.h5"):
	ftmp = File(tmpfile, "w")
	dsets = [create_dataset(fromdset, ftmp, "left"), create_dataset(fromdset, ftmp, "right")]
	print("Copying fromdset to tmpdset ...")
	copy_dataset(fromdset, dsets[0])
	for i in range(int(math.log(fromdset.shape[0], 2)) + 2):
		print("Shuffle loop {} ...".format(i + 1))
		shuffle_loop(dsets[i % 2], dsets[(i + 1) % 2])
	print("Copying tmpdset to dest todset ...")
	copy_dataset(dsets[(i + 1) % 2], todset)
	ftmp.close()
	subprocess.run(["rm", tmpfile])

def main(args):
	fin, fout = File(args.source, "r"), File(args.target, "w")
	for k, v in fin.attrs.items():
		print("Copying attribute {} ...".format(k))
		fout.attrs[k] = v
	for k, fromdset in fin.items():
		todset = create_dataset(fromdset, fout, k)
		if k in args.dsets:
			print("Shuffling dataset {} ...".format(k))
			shuffle(fromdset, todset)
		else:
			print("Copying dataset {} ...".format(k))
			copy_dataset(fromdset, todset)
	fin.close(); fout.close()

import argparse
parser = argparse.ArgumentParser()
parser.add_argument("source")
parser.add_argument("--dset", dest="dsets", action="append")
parser.add_argument("target")
args = parser.parse_args()

try:
	main(args)
except KeyboardInterrupt:
	print("Stop.")
	if os.path.exists(args.outfile):
		subprocess.run(["rm", args.outfile])

