#####################################################################
##
## PYTHON VERSION OF DAC_READ.PRO IDL PROCEDURE
## FOR READING DAC FORMAT DATA GENERATED BY CANS+
## WRITTEN BY YOSUKE MATSUMOTO, CHIBA UNIV. (ymatumot@chiba-u.jp)
##
## USAGE:
## $ setenv PYTHONPATH $CANSPLUS_DIR/python
## $ python
##
## >>> from dac_read import dac_read
##
## FOR 2D DATA
## >>> data, x, y = dac_read('0010_*.dac',dimension=2)
##
## FOR 3D DATA
## >>> data, x, y, z = dac_read('0?00_ro_rank=*.dac',dimension=3)
##
#####################################################################
import numpy as np
from scipy.io import FortranFile
import glob
import os

def dacgetparam(file,name):

  f = open(file)

  for line in f:
    if len(line) >= 4:
      if line[0:4] == '#dac':
        strs=line.split()
        if strs[2] == name:
          mtype=int(strs[1])
          if mtype == 2:
            value=strs[3]
          if mtype == 4:
            value=int(strs[3])
          if mtype == 5:
            value=float(strs[3])
          break

  return value



def dacget1d(file):

  ff = FortranFile(file,"r")
  unity = ff.read_ints()
  mver  = ff.read_ints()
  mtype = ff.read_ints()
  mndim = ff.read_ints()
  mdim  = ff.read_ints()

  if mtype == 4:
    data = ff.read_ints(dtype=np.int32)
  if mtype == 5:
    data = ff.read_reals(dtype=np.float32)
  if mtype == 6:
    data = ff.read_reals(dtype=np.float64)

  return data



def dacget2s(file):

  ff = FortranFile(file,"r")
  unity = ff.read_ints()
  mver  = ff.read_ints()
  mtype = ff.read_ints()
  mndim = ff.read_ints()
  mdim  = ff.read_ints()

  if mtype == 4:
    data = ff.read_ints(dtype=np.int32)
  if mtype == 5:
    data = ff.read_reals(dtype=np.float32)
  if mtype == 6:
    data = ff.read_reals(dtype=np.float64)

  ff.close()

  return data.reshape(mdim[1],mdim[0]), mdim[0], mdim[1]



def dacget3s(file):

  ff = FortranFile(file,"r")
  unity = ff.read_ints()
  mver  = ff.read_ints()
  mtype = ff.read_ints()
  mndim = ff.read_ints()
  mdim  = ff.read_ints()

  if mtype == 4:
    data = ff.read_ints(dtype=np.int32)
  if mtype == 5:
    data = ff.read_reals(dtype=np.float32)
  if mtype == 6:
    data = ff.read_reals(dtype=np.float64)

  ff.close()

  return data.reshape(mdim[3],mdim[2],mdim[1],mdim[0]), mdim[0], mdim[1], mdim[2]



def dac_read2d(filename):

  flist = sorted(glob.glob(filename))
  if len(flist) == 0:
    print('No such file')
    return None, None, None

  dir = os.path.dirname(os.path.abspath(flist[0]))+'/'
  gm  = dacgetparam(dir+'params_rank=0000.txt','gm')
  ix  = dacgetparam(dir+'params_rank=0000.txt','ix')
  jx  = dacgetparam(dir+'params_rank=0000.txt','jx')
  npe = dacgetparam(dir+'params_rank=0000.txt','mpisize')
  margin = dacgetparam(dir+'params_rank=0000.txt','margin')
  dt = dacgetparam(dir+'params_rank=0000.txt','dtout')
  mpix = dacgetparam(dir+'params_rank=0000.txt','mpix')
  mpiy = dacgetparam(dir+'params_rank=0000.txt','mpiy')

  igx = ix*mpix-(mpix-1)*margin*2
  jgx = jx*mpiy-(mpiy-1)*margin*2

  x = np.zeros(igx)
  y = np.zeros(jgx)

# read x,y
  mpirank=0
  for my in range(1,mpiy+1):
    for mx in range(1,mpix+1):
      file = dir+'x_rank='+'{:04d}'.format(mpirank)+'.dac'
      xd = dacget1d(file)
      file = dir+'y_rank='+'{:04d}'.format(mpirank)+'.dac'
      yd = dacget1d(file)
      ixs = (mx-1)*(ix-2*margin)+1-1
      ixe = ixs+ix-1
      jxs = (my-1)*(jx-2*margin)+1-1
      jxe = jxs+jx-1
      x[ixs:ixe+1] = xd
      y[jxs:jxe+1] = yd
      mpirank = mpirank+1

  if len(flist) % mpirank != 0:
    print('Number of files and procs. is mismatch.')
    return None, None, None

  nproc = mpirank
  nd = int(len(flist)/nproc)
  data = np.zeros([nd,jgx,igx])

  for l in range(0,nd):
    mpirank=0
    for my in range(1,mpiy+1):
      for mx in range(1,mpix+1):
        print('Reading...'+flist[l*nproc+mpirank])
        temp, ix, jx = dacget2s(flist[l*nproc+mpirank])
        ixs = (mx-1)*(ix-2*margin)+1-1
        ixe = ixs+ix-1
        jxs = (my-1)*(jx-2*margin)+1-1
        jxe = jxs+jx-1
        data[l,jxs:jxe+1,ixs:ixe+1] = temp
        mpirank = mpirank+1

  return (data[0:nd,margin:jgx-margin,margin:igx-margin], 
          x[margin:igx-margin],
          y[margin:jgx-margin]) 



def dac_read3d(filename):

  flist = sorted(glob.glob(filename))
  if len(flist) == 0:
    print('No such file')
    return None, None, None, None

  dir = os.path.dirname(os.path.abspath(flist[0]))+'/'
  gm  = dacgetparam(dir+'params_rank=0000.txt','gm')
  ix  = dacgetparam(dir+'params_rank=0000.txt','ix')
  jx  = dacgetparam(dir+'params_rank=0000.txt','jx')
  kx  = dacgetparam(dir+'params_rank=0000.txt','kx')
  npe = dacgetparam(dir+'params_rank=0000.txt','mpisize')
  margin = dacgetparam(dir+'params_rank=0000.txt','margin')
  dt = dacgetparam(dir+'params_rank=0000.txt','dtout')
  mpix = dacgetparam(dir+'params_rank=0000.txt','mpix')
  mpiy = dacgetparam(dir+'params_rank=0000.txt','mpiy')
  mpiz = dacgetparam(dir+'params_rank=0000.txt','mpiz')

  igx = ix*mpix-(mpix-1)*margin*2
  jgx = jx*mpiy-(mpiy-1)*margin*2
  kgx = kx*mpiz-(mpiz-1)*margin*2

  x = np.zeros(igx)
  y = np.zeros(jgx)
  z = np.zeros(kgx)

# read x,y,z axes
  mpirank=0
  for mz in range(1,mpiz+1):
    for my in range(1,mpiy+1):
      for mx in range(1,mpix+1):
        file = dir+'x_rank='+'{:04d}'.format(mpirank)+'.dac'
        xd = dacget1d(file)
        file = dir+'y_rank='+'{:04d}'.format(mpirank)+'.dac'
        yd = dacget1d(file)
        file = dir+'z_rank='+'{:04d}'.format(mpirank)+'.dac'
        zd = dacget1d(file)
        ixs = (mx-1)*(ix-2*margin)+1-1
        ixe = ixs+ix-1
        jxs = (my-1)*(jx-2*margin)+1-1
        jxe = jxs+jx-1
        kxs = (mz-1)*(kx-2*margin)+1-1
        kxe = kxs+kx-1
        x[ixs:ixe+1] = xd
        y[jxs:jxe+1] = yd
        z[kxs:kxe+1] = zd
        mpirank=mpirank+1

  if len(flist) % mpirank != 0:
    print('Number of files and procs. is mismatch.')
    return None, None, None, None

  nproc = mpirank
  nd = int(len(flist)/nproc)
  data = np.zeros([nd,kgx,jgx,igx])

  for l in range(0,nd):
    mpirank=0
    for mz in range(1,mpiz+1):
      for my in range(1,mpiy+1):
        for mx in range(1,mpix+1):
          print('Reading...'+flist[l*nproc+mpirank])
          temp, ix, jx, kx = dacget3s(flist[l*nproc+mpirank])
          ixs = (mx-1)*(ix-2*margin)+1-1
          ixe = ixs+ix-1
          jxs = (my-1)*(jx-2*margin)+1-1
          jxe = jxs+jx-1
          kxs=(mz-1)*(kx-2*margin)+1-1
          kxe=kxs+kx-1
          data[l,kxs:kxe+1,jxs:jxe+1,ixs:ixe+1] = temp
          mpirank = mpirank+1

  return (data[0:nd,margin:kgx-margin,margin:jgx-margin,margin:igx-margin], 
          x[margin:igx-margin],
          y[margin:jgx-margin],
          z[margin:kgx-margin])



def dac_read(filename,dimension=0):

  if dimension != 2 and dimension != 3:
    print("Specify Dimension as dimension=2 or dimension=3")
    return None, None, None, None
  else:
    if dimension == 2:
      data, x, y = dac_read2d(filename)
      return data, x, y
    if dimension == 3:
      data, x, y, z = dac_read3d(filename)
      return data, x, y, z

