"""Apply a set of rolling filter to a data cubes

This copy read complex filter dataset.
"""
import argparse
import os
from glob import glob
from multiprocessing import Pool, Array, current_process

import numpy as np
import xarray as xr

from opstats.foreground import apply_filter


def do_mask(i):
    process = current_process().pid
    print('... P{:d}: applying filter {:s}'
          .format(process, filter_files[i].split('/')[-1]))

    filter_ds = xr.open_dataset(filter_files[i])
    filter_array = filter_ds['real'].values + (1j * filter_ds['imag'].values)
    data_channels = filter_ds.attrs['frequency_channels']
    image_channel = int(np.floor(filter_ds['real'].shape[0] / 2))

    # Figure out FFT and filter normalization
    # FFT normalization factor
    x = filter_ds.attrs['x']
    y = filter_ds.attrs['y']
    f = filter_ds.attrs['f']
    dx = x[1] - x[0]
    dy = y[1] - y[0]
    df = f[1] - f[0]
    u = filter_ds.attrs['u']
    v = filter_ds.attrs['v']
    e = filter_ds.attrs['e']
    du = u[1] - u[0]
    dv = v[1] - v[0]
    de = e[1] - e[0]
    fft_norm = dx * dy * df
    ifft_norm = du * dv * de * filter_array.size
    # Filter normalization factor
    filter_volume = np.sum(filter_array.size * du * dv * de)
    filter_integral = np.sum(np.abs(filter_array) ** 2 * du * dv * de)
    filter_norm = np.sqrt(filter_volume / filter_integral)
    # Apply filter
    filtered_data = apply_filter(
        data_array[data_channels], filter_array,
        fft_multiplier=fft_norm, ifft_multiplier=ifft_norm,
        output_multiplier=filter_norm, apply_window_func=True,
        invert_filter=False
    ).real

    # Select and store the center channel of the filtered data array
    filtered_data_array[data_channels[image_channel]] = \
        filtered_data[image_channel]


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('data_cube', type=str)
    parser.add_argument('filter_directory', type=str,
                        help='filter directory containing *.nc filter files.')
    parser.add_argument('output_file', type=str)
    parser.add_argument('--nprocs', type=int, default=4)
    args = parser.parse_args()

    print('Data cube: {:s}'.format(args.data_cube))
    print('Filter directory: {:s}'.format(args.filter_directory))

    # Read input data cube
    data_da = xr.open_dataarray(args.data_cube)
    data_array = data_da.values

    # Create shared memory array to store filtered data cube
    filtered_data_array_base = Array('d', data_array.size)
    filtered_data_array = np.frombuffer(filtered_data_array_base.get_obj())
    filtered_data_array.shape = data_array.shape

    # Read in list of filter files
    filter_files = glob('{:s}/*.nc'.format(args.filter_directory))
    filter_files.sort()
    nbins = len(filter_files)

    # Attributes for output files
    # Temporary read in the first filter to retrive filter information
    ds0 = xr.open_dataset(filter_files[0])
    filter_bandwidth = ds0.attrs['filter_bandwidth']
    image_bandwidth = ds0.attrs['channel_bandwidth']
    print('Filter bandwidth: {:.1f} Hz'.format(filter_bandwidth))
    print('Image bandwidth: {:.1f} Hz'.format(image_bandwidth))
    ds0.close()
    output_attrs = data_da.attrs
    extra_attrs = {'filter_type': 'wedge',
                   'extra_filter_shift': 'None',
                   'filter_bandwidth': filter_bandwidth,
                   'image_bandwidth': image_bandwidth}
    for key, value in extra_attrs.items():
        output_attrs[key] = value

    # Check output directory
    output_dir = args.output_file.rsplit('/', 1)[0]
    os.makedirs(output_dir, exist_ok=True)

    # Start processing pool
    pool = Pool(args.nprocs)
    pool.map(do_mask, range(nbins))
    pool.close()
    pool.join()

    # Save output
    da = xr.DataArray(filtered_data_array, dims=['f', 'y', 'x'],
                      coords={'x': data_da.x.values, 'y': data_da.y.values,
                              'f': data_da.f.values},
                      attrs=output_attrs)
    da.to_netcdf(args.output_file)
    print('Saving out put to {:s}'.format(args.output_file))
