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

parser = argparse.ArgumentParser(description='Reads an ucf.tar archive, downsamples it and saves it to a new ucf.tar archive. Can be used as a pipe.')
parser.add_argument("-c", "--chunks", metavar='file',nargs='?', type=int, default=None, help="number of chunks [default: not chunked]", action="store")
parser.add_argument("-i", "--infile", metavar='file',nargs='?', default=None, help="name of the input file [default: stdin]", action="store")
parser.add_argument("-o", "--outfile", metavar='file',nargs='?', default=None, help="name of the output file [default: stdout]", action="store")
parser.add_argument("-n", "--nskip", metavar='N',nargs='?', type=int, default=2, help="keep every Nth grid point [default: 2]", action="store")
parser.add_argument("-sp", "--single-precision", help="output data in single-precision? [default: False]", action="store_true")
args = parser.parse_args()

nchunk   = args.chunks
nskip    = args.nskip
file_in  = args.infile
file_out = args.outfile
saveSinglePrecision = args.single_precision

if nchunk is None:
  if file_in is None:
    istream = tarfile.open(fileobj=sys.stdin.buffer,mode='r|',bufsize=512*1024**2,ignore_zeros=True)
  else:
    filehandle_in = open(file_in,'rb')
    istream = tarfile.open(fileobj=filehandle_in,mode='r')
else:
  if file_in is None:
    raise ValueError('Chunked mode only works with input file, not stream')
  else:
    ichunk = 0
    filehandle_in = open(file_in+'.{:d}'.format(ichunk),'rb')
    istream = tarfile.open(fileobj=filehandle_in,mode='r')

if file_out is None:
  ostream = tarfile.open(fileobj=sys.stdout.buffer,mode='w|',bufsize=512*1024**2,pax_headers=tarfile.USTAR_FORMAT)
else:
  filehandle_out = open(file_out,'wb')
  ostream = tarfile.open(fileobj=filehandle_out,mode='w',pax_headers=tarfile.USTAR_FORMAT)

while True:
    iinfo = istream.next()
    if iinfo is None:
      if nchunk is not None:
        if ichunk+1<nchunk:
          istream.close()
          filehandle_in.close()
          ichunk = ichunk+1
          filehandle_in = open(file_in+'.{:d}'.format(ichunk),'rb')
          istream = tarfile.open(fileobj=filehandle_in,mode='r')
          iinfo = istream.next()
        else:
          break
      else:
        break
    print(iinfo.name,file=sys.stderr)

    ucfbytes_in = istream.extractfile(iinfo).read()
    ucfbytes_out = b''

    if iinfo.name=='parameters.asc':
      ucfbytes_out += ucfbytes_in

    if iinfo.name=='particles.bin':
      ucfbytes_out += ucfbytes_in

    if iinfo.name=='grid.bin':
      ucfhandle = ucf.UCF(file=ucfbytes_in,verbosity=False)
      ucfhandle.copyFileHeaderToBuffer()
      ucfhandle.copyStepToBuffer(1,step_out=1,recursive=False)
      for iset in range(0,ucfhandle.NumDataset):
        (data,params) = ucfhandle.readSet(step=1,dset=iset+1)
        params = list(params)
        nx = params[0]
        ny = params[1]
        nz = params[2]
        x = data[0:nx:nskip]
        y = data[nx:nx+ny:nskip]
        z = data[nx+ny:nx+ny+nz:nskip]
        params[0] = len(x)
        params[1] = len(y)
        params[2] = len(z)
        data = np.concatenate((x,y,z))
        ucfhandle.addDatasetToBuffer(data,params=params,step=1,dset=iset+1)
      ucfbytes_out += ucfhandle.flushBuffer()
      ucfhandle.close()

    if iinfo.name=='proc.bin':
      ucfhandle = ucf.UCF(file=ucfbytes_in,verbosity=False)
      ucfhandle.copyFileHeaderToBuffer()
      ucfhandle.copyStepToBuffer(1,step_out=1,recursive=False)
      for iset in range(0,ucfhandle.NumDataset):
        (data,params) = ucfhandle.readSet(step=1,dset=iset+1)
        nxp = params[0]
        nyp = params[1]
        nzp = params[2]
        ibeg = np.copy(data[0:nxp]                            )
        iend = np.copy(data[nxp:2*nxp]                        )
        jbeg = np.copy(data[2*nxp:2*nxp+nyp]                  )
        jend = np.copy(data[2*nxp+nyp:2*nxp+2*nyp]            )
        kbeg = np.copy(data[2*nxp+2*nyp:2*nxp+2*nyp+nzp]      )
        kend = np.copy(data[2*nxp+2*nyp+nzp:2*nxp+2*nyp*2*nzp])
        for ixp in range(0,nxp):
          ibeg[ixp] = (ibeg[ixp]-1)//nskip+1
          iend[ixp] = (iend[ixp]-1)//nskip+1
        for iyp in range(0,nyp):
          jbeg[iyp] = (jbeg[iyp]-1)//nskip+1
          jend[iyp] = (jend[iyp]-1)//nskip+1 
        for izp in range(0,nzp):
          kbeg[izp] = (kbeg[izp]-1)//nskip+1
          kend[izp] = (kend[izp]-1)//nskip+1
        data = np.concatenate((ibeg,iend,jbeg,jend,kbeg,kend))
        ucfhandle.addDatasetToBuffer(data,params=params,step=1,dset=iset+1)
      ucfbytes_out += ucfhandle.flushBuffer()
      ucfhandle.close()

    if 'uvwp.' in iinfo.name:
      ucfhandle = ucf.UCF(file=ucfbytes_in,verbosity=False)
      ucfhandle.copyFileHeaderToBuffer()
      ucfhandle.copyStepToBuffer(1,step_out=1,recursive=False)
      for iset in range(0,4):
        (data,params_in) = ucfhandle.readSet(step=1,dset=iset+1)
        ighost  = params_in[0]
        (ibeg,jbeg,kbeg) = params_in[1:4]
        (nxl,nyl,nzl)    = params_in[4:7]
        (nxg,nyg,nzg)    = params_in[7:10]
        data    = data.reshape((nxl+2*ighost,nyl+2*ighost,nzl+2*ighost),order='F')

        if nskip>1:
          islice = [ii-ibeg+ighost for ii in range(ibeg,ibeg+nxl) if (ii-1)%nskip==0] 
          jslice = [ii-jbeg+ighost for ii in range(jbeg,jbeg+nyl) if (ii-1)%nskip==0] 
          kslice = [ii-kbeg+ighost for ii in range(kbeg,kbeg+nzl) if (ii-1)%nskip==0]

          data = data[np.ix_(islice,jslice,kslice)]
          ibeg = (islice[0]+ibeg-ighost)//nskip+1 
          jbeg = (jslice[0]+jbeg-ighost)//nskip+1 
          kbeg = (kslice[0]+kbeg-ighost)//nskip+1 
          (nxl,nyl,nzl) = data.shape
          nxg = (nxg-1)//nskip+1
          nyg = (nyg-1)//nskip+1
          nzg = (nzg-1)//nskip+1

        params_out = list(params_in)
        if nskip>1:
          params_out[0]    = 0
          params_out[1:4]  = (ibeg,jbeg,kbeg)
          params_out[4:7]  = (nxl,nyl,nzl)
          params_out[7:10] = (nxg,nyg,nzg)

        if saveSinglePrecision:
          data = data.astype(np.float32,casting='same_kind')
        ucfhandle.addDatasetToBuffer(data,params=params_out,step=1,dset=iset+1)
      ucfbytes_out += ucfhandle.flushBuffer()
      ucfhandle.close()
    
    if 'scal.' in iinfo.name:
      ucfhandle = ucf.UCF(file=ucfbytes_in,verbosity=False)
      ucfhandle.copyFileHeaderToBuffer()
      ucfhandle.copyStepToBuffer(1,step_out=1,recursive=False)
      for iset in range(0,ucfhandle.NumDataset):
        (data,params_in) = ucfhandle.readSet(step=1,dset=iset+1)
        ighost  = params_in[0]
        (ibeg,jbeg,kbeg) = params_in[1:4]
        (nxl,nyl,nzl)    = params_in[4:7]
        (nxg,nyg,nzg)    = params_in[7:10]
        data    = data.reshape((nxl+2*ighost,nyl+2*ighost,nzl+2*ighost),order='F')
 
        if nskip>1:
          islice = [ii-ibeg+ighost for ii in range(ibeg,ibeg+nxl) if (ii-1)%nskip==0] 
          jslice = [ii-jbeg+ighost for ii in range(jbeg,jbeg+nyl) if (ii-1)%nskip==0] 
          kslice = [ii-kbeg+ighost for ii in range(kbeg,kbeg+nzl) if (ii-1)%nskip==0]

          data = data[np.ix_(islice,jslice,kslice)]
          ibeg = (islice[0]+ibeg-ighost)//nskip+1 
          jbeg = (jslice[0]+jbeg-ighost)//nskip+1 
          kbeg = (kslice[0]+kbeg-ighost)//nskip+1 
          (nxl,nyl,nzl) = data.shape
          nxg = (nxg-1)//nskip+1
          nyg = (nyg-1)//nskip+1
          nzg = (nzg-1)//nskip+1

        params_out = list(params_in)
        if nskip>1:
          params_out[0]    = 0
          params_out[1:4]  = (ibeg,jbeg,kbeg)
          params_out[4:7]  = (nxl,nyl,nzl)
          params_out[7:10] = (nxg,nyg,nzg)
        
        if saveSinglePrecision:
          data = data.astype(np.float32,casting='same_kind')
        ucfhandle.addDatasetToBuffer(data,params=params_out,step=1,dset=iset+1)
      ucfbytes_out += ucfhandle.flushBuffer()
      ucfhandle.close()

    oinfo = tarfile.TarInfo(name=iinfo.name)
    oinfo.size = len(ucfbytes_out)
    ostream.addfile(oinfo,fileobj=io.BytesIO(ucfbytes_out))

istream.close()
ostream.close()
if file_in is not None:
  filehandle_in.close()
if file_out is not None:
  filehandle_out.close()