import os
import re
import sys
import timeit
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path, PurePosixPath

import h5py
import numpy as np
import psutil
from mpi4py import MPI
from tqdm import tqdm

from .h5 import h5_print_contents
from .signal import ccor
from .utils import even_print, format_time_string

# ======================================================================

def _calc_ccor_xpln(self, **kwargs):
    '''
    calculate cross-correlation in [z,t] at every [x,y]
    - designed for analyzing unsteady, thin planes in [x]
    - Multithreaded with ThreadPoolExecutor()
        - scipy.signal.csd() automatically tries to run multithreaded
        - set OMP_NUM_THREADS=1 and pass 'n_threads' to as kwarg manually
    '''
    
    if (self.rank==0):
        verbose = True
    else:
        verbose = False
    
    if verbose: print('\n'+'rgd.calc_ccor_xpln()'+'\n'+72*'-')
    t_start_func = timeit.default_timer()
    
    ## assert that the opened RGD has fsubtype 'unsteady' (i.e. is NOT a prime file)
    if (self.fsubtype!='unsteady'):
        raise ValueError
    if not self.usingmpi:
        raise NotImplementedError('function is not implemented for non-MPI usage')
    
    h5py_is_mpi_build = h5py.h5.get_config().mpi
    if not h5py_is_mpi_build:
        if verbose: print('h5py was not compiled for parallel usage! exiting.')
        sys.exit(1)
    
    rx = kwargs.get('rx',1)
    ry = kwargs.get('ry',1)
    rz = kwargs.get('rz',1)
    rt = kwargs.get('rt',1)
    
    # cy = kwargs.get('cy',1) ## number of subdivisions per rank [y] range
    # if not isinstance(cy,int):
    #     raise TypeError('cy should be an int')
    # if (cy<1):
    #     raise TypeError('cy should be an int')
    
    sy = kwargs.get('sy',1) ## number of [y] layers to read at a time
    if not isinstance(sy,int) or (sy<1):
        raise TypeError('sy should be a positive non-zero int')
    
    n_threads = kwargs.get('n_threads',1)
    #try:
    #    n_threads = int(os.environ.get('OMP_NUM_THREADS'))
    #except TypeError: ## not set
    #    n_threads = os.cpu_count()
    
    fn_h5_out = kwargs.get('fn_h5_out',None) ## filename for output HDF5 (.h5) file
    
    ## Debug Rank:Proc Affinity
    #pp = psutil.Process()
    #print(f"[Rank {self.rank}] sees CPUs: {pp.cpu_affinity()}  |  n_threads={n_threads}  |  OMP_NUM_THREADS={os.environ.get('OMP_NUM_THREADS')}")
    
    ## only distribute data across [y]
    if (rx!=1):
        raise AssertionError('rx!=1')
    if (rz!=1):
        raise AssertionError('rz!=1')
    if (rt!=1):
        raise AssertionError('rt!=1')
    
    if not isinstance(ry,int) or (ry<1):
        raise ValueError('ry should be a positive non-zero int')
    
    ## check the choice of ranks per dimension
    if (rx*ry*rz*rt != self.n_ranks):
        raise AssertionError('rx*ry*rz*rt != self.n_ranks')
    if (rx>self.nx):
        raise AssertionError('rx>self.nx')
    if (ry>self.ny):
        raise AssertionError('ry>self.ny')
    if (rz>self.nz):
        raise AssertionError('rz>self.nz')
    if (rt>self.nt):
        raise AssertionError('rt>self.nt')
    
    if (self.ny%ry!=0):
        raise ValueError('ny not divisible by ry')
    
    ## distribute 4D data over ranks --> here only in [y]
    ryl_ = np.array_split(np.arange(self.ny,dtype=np.int64),min(ry,self.ny))
    ryl = [[b[0],b[-1]+1] for b in ryl_ ]
    ry1,ry2 = ryl[self.rank]
    nyr = ry2 - ry1
    
    ## check all [y] ranges have same size
    for ryl_ in ryl:
        if not (ryl_[1]-ryl_[0]==nyr):
            raise ValueError('[y] chunks are not even in size')
    
    # ## [y] sub chunk range --> cyl = list of ranges in ry1:ry2
    # ## cy is the NUMBER of chunks for the rank sub-range
    # cyl_ = np.array_split( np.arange(ry1,ry2) , min(cy,nyr) )
    # cyl  = [[b[0],b[-1]+1] for b in cyl_ ]
    # 
    # for nyc_ in [ cyl_[1]-cyl_[0] for cyl_ in cyl ]:
    #     if (nyc_ < 1):
    #         #raise ValueError
    #         print(f'rank {self.rank:d}: sub-range is <1')
    #         self.comm.Abort(1)
    # 
    # if 1: ## assert that [y] sub-chunk ranges are correct
    #     
    #     yi = np.arange(self.ny, dtype=np.int32)
    #     
    #     local_indices = []
    #     for cyl_ in cyl:
    #         cy1, cy2 = cyl_
    #         local_indices += [ yi_ for yi_ in yi[cy1:cy2] ]
    #     
    #     G = self.comm.gather([ self.rank , local_indices ], root=0)
    #     G = self.comm.bcast(G, root=0)
    #     
    #     all_indices = []
    #     for G_ in G:
    #         all_indices += G_[1]
    #     all_indices = np.array( sorted(all_indices), dtype=np.int32 )
    #     
    #     if not np.array_equal( all_indices , yi ):
    #         raise AssertionError
    
    if (nyr%sy!=0):
        raise ValueError('nyr not divisible by sy')
    
    ## output filename : HDF5 (.h5)
    if (fn_h5_out is None): ## automatically determine name
        fname_path = os.path.dirname(self.fname)
        fname_base = os.path.basename(self.fname)
        fname_root, fname_ext = os.path.splitext(fname_base)
        fname_root = re.findall(r'io\S+_mpi_[0-9]+', fname_root)[0]
        fn_h5_out_base = fname_root+'_ccor.h5'
        fn_h5_out = str(PurePosixPath(fname_path, fn_h5_out_base))
    if (Path(fn_h5_out).suffix != '.h5'):
        raise ValueError(f"fn_h5_out='{str(fn_h5_out)}' must end in .h5")
    if os.path.isfile(fn_h5_out):
        #if (os.path.getsize(fn_h5_out) > 8*1024**3):
        #    raise ValueError(f"fn_h5_out='{str(fn_h5_out)}' exists and is >8 [GB]. exiting for your own safety.")
        if (fn_h5_out == self.fname):
            raise ValueError(f"fn_h5_out='{str(fn_h5_out)}' cannot be same as input filename.")
    
    if verbose: even_print( 'fn_h5'      , self.fname )
    if verbose: even_print( 'fn_h5_out'  , fn_h5_out  )
    if verbose: print(72*'-')
    self.comm.Barrier()
    
    ## the data dictionary to be pickled later
    data = {}
    
    ## infile
    fsize = os.path.getsize(self.fname)/1024**3
    if verbose: even_print(os.path.basename(self.fname),'%0.1f [GB]'%fsize)
    if verbose: even_print('nx',f'{self.nx:d}')
    if verbose: even_print('ny',f'{self.ny:d}')
    if verbose: even_print('nz',f'{self.nz:d}')
    if verbose: even_print('nt',f'{self.nt:d}')
    if verbose: even_print('ngp',f'{self.ngp/1e6:0.1f} [M]')
    #if verbose: even_print('cy',f'{cy:d}')
    if verbose: even_print('sy',f'{sy:d}')
    if verbose: even_print('n_ranks',f'{self.n_ranks:d}')
    if verbose: even_print('n_threads',f'{n_threads:d}')
    if verbose: print(72*'-')
    
    ## 0D freestream scalars
    lchar   = self.lchar   ; data['lchar']   = lchar
    U_inf   = self.U_inf   ; data['U_inf']   = U_inf
    rho_inf = self.rho_inf ; data['rho_inf'] = rho_inf
    T_inf   = self.T_inf   ; data['T_inf']   = T_inf
    
    #data['M_inf'] = self.M_inf
    data['Ma'] = self.Ma
    data['Pr'] = self.Pr
    
    ## read in 1D coordinate arrays & re-dimensionalize
    x = np.copy( self['dims/x'][()] * self.lchar )
    y = np.copy( self['dims/y'][()] * self.lchar )
    z = np.copy( self['dims/z'][()] * self.lchar )
    t = np.copy( self['dims/t'][()] * self.tchar )
    
    nx = self.nx ; data['nx'] = nx
    ny = self.ny ; data['ny'] = ny
    nz = self.nz ; data['nz'] = nz
    nt = self.nt ; data['nt'] = nt
    
    ## assert constant Δz
    dz0 = np.diff(z)[0]
    if not np.all(np.isclose(np.diff(z), dz0, rtol=1e-6)):
        raise NotImplementedError('Δz not constant')
    dz = np.diff(z)[0]
    
    ## dimensional [s]
    dt = self.dt * self.tchar
    np.testing.assert_allclose(dt, t[1]-t[0], rtol=1e-12, atol=1e-12)
    
    t_meas = self.duration * self.tchar
    np.testing.assert_allclose(t_meas, t.max()-t.min(), rtol=1e-12, atol=1e-12)
    
    zrange = z.max() - z.min()
    
    data['x'] = x
    data['y'] = y
    data['z'] = z
    
    data['t']      = t
    data['t_meas'] = t_meas
    data['dt']     = dt
    data['dz']     = dz
    data['zrange'] = zrange
    
    if verbose: even_print( 'Δt/tchar'       , f'{dt/self.tchar:0.8f}' )
    if verbose: even_print( 'Δt'             , f'{dt:0.3e} [s]'        )
    if verbose: even_print( 'duration/tchar' , f'{self.duration:0.1f}' )
    if verbose: even_print( 'duration'       , f'{self.duration*self.tchar:0.3e} [s]' )
    if verbose: print(72*'-')
    
    ## report
    if verbose:
        even_print('Δt'     , f'{dt    :0.5e} [s]' )
        even_print('t_meas' , f'{t_meas:0.5e} [s]' )
        even_print('Δz'     , f'{dz0   :0.5e} [m]' )
        even_print('zrange' , f'{zrange:0.5e} [m]' )
        print(72*'-')
    
    ## get lags [t]
    #lags_t,_  = ccor( np.ones(win_len,dtype=np.float32) , np.ones(win_len,dtype=np.float32), get_lags=True )
    #n_lags_t_ = win_len*2-1
    lags_t,_  = ccor( np.ones(nt,dtype=np.float32) , np.ones(nt,dtype=np.float32), get_lags=True )
    n_lags_t_ = nt*2-1
    n_lags_t  = lags_t.shape[0]
    if (n_lags_t!=n_lags_t_):
        raise AssertionError('check lags [t]')
    
    data['lags_t']   = lags_t
    data['n_lags_t'] = n_lags_t
    
    if verbose:
        even_print('n lags (Δt)' , '%i'%(n_lags_t,))
    
    ## get lags [z]
    lags_z,_  = ccor( np.ones(nz,dtype=np.float32) , np.ones(nz,dtype=np.float32), get_lags=True )
    n_lags_z_ = nz*2-1
    n_lags_z  = lags_z.shape[0]
    if (n_lags_z!=n_lags_z_):
        raise AssertionError('check lags [z]')
    
    data['lags_z']   = lags_z
    data['n_lags_z'] = n_lags_z
    
    if verbose:
        even_print('n lags (Δz)' , '%i'%(n_lags_z,))
    
    # ===
    
    ## cross-correlation pairs
    ## [ str:var1, str:var2, bool:do_density_weighting]
    ccor_combis = [
    [ 'u' , 'v' , True  ], ## [ ρ·u″ , ρ·v″ ]
    [ 'u' , 'u' , True  ], ## [ ρ·u″ , ρ·u″ ]
    [ 'v' , 'v' , True  ], ## [ ρ·v″ , ρ·v″ ]
    [ 'w' , 'w' , True  ], ## [ ρ·w″ , ρ·w″ ]
    [ 'u' , 'v' , False ], ## [ u′   , v′   ]
    [ 'u' , 'u' , False ], ## [ u′   , u′   ]
    [ 'v' , 'v' , False ], ## [ v′   , v′   ]
    [ 'w' , 'w' , False ], ## [ w′   , w′   ]
    ]
    
    ## generate cross-correlation scalar names
    scalars = []
    for ccor_combi in ccor_combis:
        s1,s2,do_density_weighting = ccor_combi
        if do_density_weighting:
            scalars.append(f'r{s1}II_r{s2}II')
        else:
            scalars.append(f'{s1}I_{s2}I')
    
    ## generate avg scalar names
    scalars_Re_avg = []
    scalars_Fv_avg = []
    for ccor_combi in ccor_combis:
        s1,s2,do_density_weighting = ccor_combi
        if do_density_weighting and ('rho' not in scalars_Re_avg):
            scalars_Re_avg.append('rho')
        if do_density_weighting:
            if (s1 not in scalars_Fv_avg):
                scalars_Fv_avg.append(s1)
            if (s2 not in scalars_Fv_avg):
                scalars_Fv_avg.append(s2)
        else:
            if (s1 not in scalars_Re_avg):
                scalars_Re_avg.append(s1)
            if (s2 not in scalars_Re_avg):
                scalars_Re_avg.append(s2)
    
    ## numpy formatted arrays: buffers for PSD & other data (rank-local)
    Rz         = np.zeros(shape=(nyr,n_lags_z ) , dtype={'names':scalars        , 'formats':[ np.dtype(np.float64) for s in scalars        ]})
    Rt         = np.zeros(shape=(nyr,n_lags_t ) , dtype={'names':scalars        , 'formats':[ np.dtype(np.float64) for s in scalars        ]})
    covariance = np.zeros(shape=(nyr,         ) , dtype={'names':scalars        , 'formats':[ np.dtype(np.float64) for s in scalars        ]})
    avg_Re     = np.zeros(shape=(nyr,         ) , dtype={'names':scalars_Re_avg , 'formats':[ np.dtype(np.float64) for s in scalars_Re_avg ]})
    avg_Fv     = np.zeros(shape=(nyr,         ) , dtype={'names':scalars_Fv_avg , 'formats':[ np.dtype(np.float64) for s in scalars_Fv_avg ]})
    
    if verbose:
        even_print('n cross-correlation scalar combinations' , f'{len(ccor_combis):d}')
    
    # ==============================================================
    # check memory
    # ==============================================================
    
    hostname = MPI.Get_processor_name()
    mem_free_gb = psutil.virtual_memory().free / 1024**3
    G = self.comm.gather([ self.rank , hostname , mem_free_gb ], root=0)
    G = self.comm.bcast(G, root=0)
    
    host_mem = {}
    for rank, host, mem in G:
        if host not in host_mem or mem < host_mem[host]:
            host_mem[host] = mem
    total_free = sum(host_mem.values())
    
    if verbose:
        print(72*'-')
        for key,value in host_mem.items():
            even_print(f'RAM free {key}', f'{int(np.floor(value)):d} [GB]')
        even_print('RAM free (local,min)', f'{int(np.floor(min(host_mem.values()))):d} [GB]')
        even_print('RAM free (global)', f'{int(np.floor(total_free)):d} [GB]')
    
    shape_read = (nx,sy,nz,nt) ## local
    if verbose: even_print('read shape (local)', f'[{nx:d},{sy:d},{nz:d},{nt:d}]')
    data_gb = np.dtype(np.float64).itemsize * np.prod(shape_read) / 1024**3
    if verbose: even_print('read size (global)', f'{int(np.ceil(data_gb*ry)):d} [GB]')
    
    if verbose: even_print('read size (global) ×6', f'{int(np.ceil(data_gb*ry*6)):d} [GB]')
    ram_usage_est = data_gb*ry*6/total_free
    if verbose: even_print('RAM usage estimate', f'{100*ram_usage_est:0.1f} [%]')
    
    self.comm.Barrier()
    if (ram_usage_est>0.80):
        print('RAM consumption might be too high. exiting.')
        self.comm.Abort(1)
    
    # ==============================================================
    # main loop
    # ==============================================================
    
    if verbose:
        progress_bar = tqdm(
            #total=len(ccor_combis)*cy,
            total=len(ccor_combis)*(nyr//sy),
            ncols=100,
            desc='ccor',
            leave=True,
            file=sys.stdout,
            mininterval=0.1,
            smoothing=0.,
            #bar_format="\033[B{l_bar}{bar}| {n}/{total} [{percentage:.1f}%] {elapsed}/{remaining}\033[A\n\b",
            bar_format="{l_bar}{bar}| {n}/{total} [{percentage:.1f}%] {elapsed}/{remaining}",
            ascii="░█",
            colour='#FF6600',
            )
    
    for cci,cc in enumerate(ccor_combis): ## ccor pairs
        
        if verbose: tqdm.write(72*'-')
        
        scalar_L, scalar_R, do_density_weighting = cc
        
        if do_density_weighting:
            msg = f'ccor[ρ·{scalar_L}″,ρ·{scalar_R}″]'
        else:
            msg = f'ccor[{scalar_L}′,{scalar_R}′]'
        if verbose:
            tqdm.write(even_print('computing',msg,s=True,))
        
        dset_L   = self[f'data/{scalar_L}']
        dset_R   = self[f'data/{scalar_R}']
        dset_rho = self['data/rho']
        
        scalar = scalars[cci]
        
        ## assert scalar name
        if do_density_weighting:
            if (f'r{scalar_L}II_r{scalar_R}II' != scalar ):
                raise ValueError
        else:
            if (f'{scalar_L}I_{scalar_R}I' != scalar ):
                raise ValueError
        
        # ## [y] loop outer (chunks within rank)
        # for cyl_ in cyl:
        #     cy1, cy2 = cyl_
        #     nyc = cy2 - cy1
        
        for ci in range(nyr//sy):
            cy1 = ry1 + ci*sy
            cy2 = cy1 + sy
            nyc = cy2 - cy1
            
            self.comm.Barrier()
            t_start = timeit.default_timer()
            
            ## read data L
            n_scalars_read = 1 ## initialize
            scalar_str = scalar_L
            with dset_L.collective:
                data_L = np.copy( dset_L[:,:,cy1:cy2,:].T ).astype(np.float64)
            
            ## read data R (if != data L)
            if (scalar_L==scalar_R):
                data_R = np.copy( data_L )
            else:
                n_scalars_read += 1
                scalar_str += f',{scalar_R}'
                with dset_R.collective:
                    data_R = np.copy( dset_R[:,:,cy1:cy2,:].T ).astype(np.float64)
            
            ## read ρ
            if do_density_weighting:
                n_scalars_read += 1
                scalar_str += ',ρ'
                with dset_rho.collective:
                    rho = np.copy( dset_rho[:,:,cy1:cy2,:].T ).astype(np.float64)
            else:
                rho = None
            
            self.comm.Barrier()
            t_delta = timeit.default_timer() - t_start
            data_gb = n_scalars_read * ( self.nx * ry * (cy2-cy1) * self.nz * self.nt * dset_L.dtype.itemsize ) / 1024**3
            if verbose:
                tqdm.write(even_print(f'read: {scalar_str}', '%0.3f [GB]  %0.3f [s]  %0.3f [GB/s]'%(data_gb,t_delta,(data_gb/t_delta)), s=True))
            
            ## data_L and data_R should be [nx,nyc,nz,nt] where nyc is the chunk [y] range
            if ( data_L.shape != (nx,nyc,nz,nt) ) or ( data_R.shape != (nx,nyc,nz,nt) ):
                print(f'rank {self.rank:d}: shape violation')
                self.comm.Abort(1)
            if (rho is not None) and ( rho.shape != (nx,nyc,nz,nt) ):
                print(f'rank {self.rank:d}: shape violation')
                self.comm.Abort(1)
            
            # === redimensionalize
            
            if scalar_L in ['u','v','w',]:
                data_L *= U_inf
            else:
                raise ValueError
            
            if scalar_R in ['u','v','w',]:
                data_R *= U_inf
            else:
                raise ValueError
            
            if (rho is not None): ## i.e. if do_density_weighting
                rho *= rho_inf
            
            # === compute mean-removed data
            
            ## avg(□) or avg(ρ·□)/avg(ρ) in [t]
            if do_density_weighting:
                rho_avg     = np.mean(        rho , axis=3, dtype=np.float64, keepdims=True) ## [x,y,z,1]
                data_L_avg  = np.mean( rho*data_L , axis=3, dtype=np.float64, keepdims=True) ## [x,y,z,1]
                data_L_avg /= rho_avg
                data_R_avg  = np.mean( rho*data_R , axis=3, dtype=np.float64, keepdims=True) ## [x,y,z,1]
                data_R_avg /= rho_avg
            else:
                data_L_avg = np.mean( data_L , axis=3, dtype=np.float64, keepdims=True) ## [x,y,z,1]
                data_R_avg = np.mean( data_R , axis=3, dtype=np.float64, keepdims=True) ## [x,y,z,1]
            
            ## Reynolds prime □′ or Favre prime □″
            data_L -= data_L_avg
            data_R -= data_R_avg
            
            ## assert stationarity / definition averaging
            ## avg(□′)==0 or avg(ρ·□″)==0
            if do_density_weighting:
                a_ = np.mean(rho*data_L, axis=3, dtype=np.float64, keepdims=True)
                b_ = np.mean(rho*data_R, axis=3, dtype=np.float64, keepdims=True)
            else:
                a_ = np.mean(data_L, axis=3, dtype=np.float64, keepdims=True)
                b_ = np.mean(data_R, axis=3, dtype=np.float64, keepdims=True)
            if not np.allclose( a_, np.zeros_like(a_), atol=1e-6 ) or not np.allclose( b_, np.zeros_like(b_), atol=1e-6 ):
                print(f'rank {self.rank:d}: avg(□′)!=0 or avg(ρ·□″)!=0')
                self.comm.Abort(1)
            
            ## covariance: <□′·□′> OR <ρ□″·ρ□″> --> note that this is NOT the typical Favre <ρ·□″□″>
            if do_density_weighting:
                covariance_ = np.mean( rho*data_L * rho*data_R , axis=3 , dtype=np.float64, keepdims=True)
            else:
                covariance_ = np.mean( data_L*data_R , axis=3 , dtype=np.float64, keepdims=True)
            
            ## write this chunk/scalar's covariance to covariance buffer
            ## avg over [x,z] : [x,y,z,1] --> [y]
            yiA = cy1 - ry1
            yiB = cy2 - ry1
            covariance[scalar][yiA:yiB] = np.squeeze( np.mean( covariance_ , axis=(0,2,3) , dtype=np.float64) )
            
            ## write (rank-local) 1D [y] averages
            if do_density_weighting:
                avg_Fv[scalar_L][yiA:yiB] = np.squeeze( np.mean( data_L_avg , axis=(0,2,3) , dtype=np.float64) )
                avg_Fv[scalar_R][yiA:yiB] = np.squeeze( np.mean( data_R_avg , axis=(0,2,3) , dtype=np.float64) )
                avg_Re['rho'][yiA:yiB]    = np.squeeze( np.mean( rho_avg    , axis=(0,2,3) , dtype=np.float64) )
            else:
                avg_Re[scalar_L][yiA:yiB] = np.squeeze( np.mean( data_L_avg , axis=(0,2,3) , dtype=np.float64) )
                avg_Re[scalar_R][yiA:yiB] = np.squeeze( np.mean( data_R_avg , axis=(0,2,3) , dtype=np.float64) )
            
            # ===============================================================================
            # At this point you have 4D [x,y,z,t] [□′,□′] or [ρ·□″,ρ·□″] data
            # ===============================================================================
            
            def __ccor_z_thread_kernel(xi,ti,yii,do_density_weighting):
                if do_density_weighting:
                    uL = rho[xi,yii,:,ti] * data_L[xi,yii,:,ti]
                    uR = rho[xi,yii,:,ti] * data_R[xi,yii,:,ti]
                else:
                    uL = data_L[xi,yii,:,ti]
                    uR = data_R[xi,yii,:,ti]
                return xi,ti,ccor(uL,uR)
            
            def __ccor_t_thread_kernel(xi,zi,yii,do_density_weighting):
                if do_density_weighting:
                    uL = rho[xi,yii,zi,:] * data_L[xi,yii,zi,:]
                    uR = rho[xi,yii,zi,:] * data_R[xi,yii,zi,:]
                else:
                    uL = data_L[xi,yii,zi,:]
                    uR = data_R[xi,yii,zi,:]
                return xi,zi,ccor(uL,uR)
            
            self.comm.Barrier()
            t_start = timeit.default_timer()
            
            ## [y] loop inner (indices within chunk)
            for yi in range(cy1,cy2):
                
                yii  = yi - cy1 ## chunk local
                yiii = yi - ry1 ## rank local
                
                ## ccor buffers for [y] loop inner
                R_xt = np.zeros((nx,nt,n_lags_z) , dtype=np.float64) ## [x,t] range for ccor(z)
                R_xz = np.zeros((nx,nz,n_lags_t) , dtype=np.float64) ## [x,z] range for ccor(t)
                
                # ===========================================================================
                # ccor(z) : loop over [x,t]
                # ===========================================================================
                
                ## concurrent/threaded execution for ccor(z)
                tasks = [(xi,ti,yii,do_density_weighting) for xi in range(nx) for ti in range(nt)]
                with ThreadPoolExecutor(max_workers=n_threads) as executor:
                    results = executor.map(lambda t: __ccor_z_thread_kernel(*t,), tasks)
                    for xi,ti,result in results:
                        R_xt[xi,ti,:] = result
                
                # for xi in range(nx):
                #     for ti in range(nt):
                #         
                #         ## 1D [z] □′ or ρ·□″ vectors
                #         if do_density_weighting:
                #             uL = np.copy( rho[xi,yii,:,ti] * data_L[xi,yii,:,ti] )
                #             uR = np.copy( rho[xi,yii,:,ti] * data_R[xi,yii,:,ti] )
                #         else:
                #             uL = np.copy( data_L[xi,yii,:,ti] )
                #             uR = np.copy( data_R[xi,yii,:,ti] )
                #         
                #         R_xt[xi,ti,:] = ccor( uL , uR )
                
                ## avg in [x,t] & write in rank context
                Rz[scalar][yiii,:] = np.mean(R_xt, axis=(0,1), dtype=np.float64)
                
                # ===========================================================================
                # ccor(t) : loop over [x,z]
                # ===========================================================================
                
                ## concurrent/threaded execution for ccor(t)
                tasks = [(xi,zi,yii,do_density_weighting) for xi in range(nx) for zi in range(nz)]
                with ThreadPoolExecutor(max_workers=n_threads) as executor:
                    results = executor.map(lambda t: __ccor_t_thread_kernel(*t,), tasks)
                    for xi,zi,result in results:
                        R_xz[xi,zi,:] = result
                
                # for xi in range(nx):
                #     for zi in range(nz):
                #         
                #         ## 1D [z] □′ or ρ·□″ vectors
                #         if do_density_weighting:
                #             uL = np.copy( rho[xi,yii,zi,:] * data_L[xi,yii,zi,:] )
                #             uR = np.copy( rho[xi,yii,zi,:] * data_R[xi,yii,zi,:] )
                #         else:
                #             uL = np.copy( data_L[xi,yii,zi,:] )
                #             uR = np.copy( data_R[xi,yii,zi,:] )
                #         
                #         R_xz[xi,zi,:] = ccor( uL , uR )
                
                ## avg in [x,z] & write in rank context
                Rt[scalar][yiii,:] = np.mean(R_xz, axis=(0,1), dtype=np.float64)
            
            self.comm.Barrier()
            t_delta = timeit.default_timer() - t_start
            if verbose: tqdm.write(even_print(msg, format_time_string(t_delta), s=True))
            if verbose: progress_bar.update() ## (scalar, [y] chunk) progress
            #break ## debug
        #break ## debug
    if verbose: progress_bar.close()
    self.comm.Barrier()
    if verbose: print(72*'-')
    
    # ==============================================================
    # write HDF5 (.h5) file
    # ==============================================================
    
    ## overwrite outfile!
    ## open on rank 0 and write attributes, dimensions, etc.
    if (self.rank==0):
        with h5py.File(fn_h5_out, 'w') as hfw:
            
            ## write floats,ints as top-level attributes
            for key,val in data.items():
                if isinstance(data[key], (int,np.int32,np.int64)):
                    hfw.attrs[key] = val
                elif isinstance(data[key], (float,np.float32,np.float64)):
                    hfw.attrs[key] = val
                elif isinstance(data[key], np.ndarray):
                    pass
                else:
                    print(f'key {key} is type {str(type(data[key]))}')
                    self.comm.Abort(1)
            
            ## write numpy arrays
            hfw.create_dataset( 'dims/x'      , data=x      ) ## [m]
            hfw.create_dataset( 'dims/y'      , data=y      ) ## [m]
            hfw.create_dataset( 'dims/z'      , data=z      ) ## [m]
            hfw.create_dataset( 'dims/t'      , data=t      ) ## [s]
            hfw.create_dataset( 'dims/lags_z' , data=lags_z ) ## [m]
            hfw.create_dataset( 'dims/lags_t' , data=lags_t ) ## [s]
            
            ## initialize datasets
            for scalar in scalars:
                hfw.create_dataset( f'covariance/{scalar}'  , shape=(ny,)         , dtype=np.float64, chunks=None         , data=np.full((ny,),0.,np.float64)         )
                hfw.create_dataset( f'Rz/{scalar}'          , shape=(ny,n_lags_z) , dtype=np.float64, chunks=(1,n_lags_z) , data=np.full((ny,n_lags_z),0.,np.float64) )
                hfw.create_dataset( f'Rt/{scalar}'          , shape=(ny,n_lags_t) , dtype=np.float64, chunks=(1,n_lags_t) , data=np.full((ny,n_lags_t),0.,np.float64) )
            
            ## initialize datasets 1D [y] mean
            for scalar in avg_Re.dtype.names:
                hfw.create_dataset( f'avg/Re/{scalar}', shape=(ny,), dtype=np.float64, chunks=None, data=np.full((ny,),0.,np.float64) )
            for scalar in avg_Fv.dtype.names:
                hfw.create_dataset( f'avg/Fv/{scalar}', shape=(ny,), dtype=np.float64, chunks=None, data=np.full((ny,),0.,np.float64) )
    
    self.comm.Barrier()
    
    with h5py.File(fn_h5_out, 'a', driver='mpio', comm=self.comm) as hfw:
        
        ## collectively write covariance,Rz,Rt
        for scalar in scalars:
            dset = hfw[f'covariance/{scalar}']
            with dset.collective:
                dset[ry1:ry2] = covariance[scalar][:]
            dset = hfw[f'Rz/{scalar}']
            with dset.collective:
                dset[ry1:ry2,:] = Rz[scalar][:,:]
            dset = hfw[f'Rt/{scalar}']
            with dset.collective:
                dset[ry1:ry2,:] = Rt[scalar][:,:]
        
        ## collectively write 1D [y] avgs (Reynolds,Favre)
        for scalar in avg_Re.dtype.names:
            dset = hfw[f'avg/Re/{scalar}']
            with dset.collective:
                dset[ry1:ry2] = avg_Re[scalar][:]
        for scalar in avg_Fv.dtype.names:
            dset = hfw[f'avg/Fv/{scalar}']
            with dset.collective:
                dset[ry1:ry2] = avg_Fv[scalar][:]
    
    ## report file contents
    self.comm.Barrier()
    if (self.rank==0):
        even_print( os.path.basename(fn_h5_out) , f'{(os.path.getsize(fn_h5_out)/1024**2):0.1f} [MB]' )
        print(72*'-')
        with h5py.File(fn_h5_out,'r') as hfr:
            h5_print_contents(hfr)
    self.comm.Barrier()
    
    if verbose: print(72*'-')
    if verbose: print('total time : rgd.calc_ccor_xpln() : %s'%format_time_string((timeit.default_timer() - t_start_func)))
    if verbose: print(72*'-')
    return
