#!/usr/bin/env python3
import sys
import io
import os
import tarfile
import argparse
import numpy as np
import ucf
import configparser

parser = argparse.ArgumentParser(description='Unpacks an UCF tar file as MPIIO input file.')
parser.add_argument('infile', metavar='filein', help='input file')
parser.add_argument('iseq',  metavar='iseq', help='sequence number')
#parser.add_argument('base',  metavar='base', help='filebase to be archived: "uvwp" or "scal"')
parser.add_argument("-o", "--outdir", metavar='directory',nargs='?', default=None, help="name of the output directory", action="store")
parser.add_argument("-v", "--verbose", help="activate verbose output", action="store_true")
args = parser.parse_args()

file_in = args.infile
iseq    = int(args.iseq)
#base    = args.base
verbose = args.verbose

if args.outdir is None:
  dir_out = './'
else:
  dir_out = args.outdir

# Open tar file for input
ftar = tarfile.open(name=file_in,mode='r')

# Extract files
def transform_filename(filename,iseq):
  return filename.replace('.','_{:04d}.'.format(iseq))

files_all = ftar.getnames()
#files_aux = ('parameters.asc','proc.bin','grid.bin','particles.bin')
files_aux = [s for s in files_all if not s.startswith('uvwp') or s.startswith('scal')]
for files in files_aux:
  print(files)
  ucf_bytes = ftar.extractfile(files).read()
  file_out = dir_out+transform_filename(files,iseq)
  fidw = open(file_out,'wb')
  fidw.write(ucf_bytes)
  fidw.close()

files_data = [s for s in files_all if s.startswith('uvwp') or s.startswith('scal')]
base = files_data[0].split('.')[0]
iproc = 0
file_out = dir_out+'/'+base+'_{:04d}.bin'.format(iseq)
fidw = open(file_out,'wb')
for files in files_data:
  print(files)
  sproc = files.split('.')[1]
  if iproc!=int(sproc):
    raise ValueError('Invalid file order in tar: {}, {}'.format(iproc,files))
  ucf_bytes = ftar.extractfile(files).read()
  if iproc==0:
    fidw.write(ucf_bytes)
  else:
    fidw.write(ucf_bytes[64:])
  iproc += 1
fidw.close()

# Close tar file
ftar.close()