#!/usr/bin/env python3
# Alexey Pechnikov, Sep, 2021, https://github.com/mobigroup/gmtsar
from .SBAS_stack import SBAS_stack
from .tqdm_dask import tqdm_dask

class SBAS_trans(SBAS_stack):

    def get_trans_dat_blocks_extents(self, subswath=None, n_jobs=-1):
        """
        Compute trans_dat dask blocks extents in radar coordinates for geocoding matrices
        """
        from tqdm.auto import tqdm
        import joblib
        import xarray as xr
        import dask
        import numpy as np

        # range, azimuth, elevation(ref to radius in PRM), lon, lat [ASCII default] 
        def calculate_block_extent(iy, ix):
            dask_block_azi = trans_dat.azi.data.blocks[iy, ix]
            dask_block_rng = trans_dat.rng.data.blocks[iy, ix]
            azi_allnans = dask.array.isnan(dask_block_azi).all()
            rng_allnans = dask.array.isnan(dask_block_rng).all()
            if azi_allnans or rng_allnans:
                return
            return dask.compute(iy, ix,
                                dask.array.nanmin(dask_block_azi), dask.array.nanmax(dask_block_azi),
                                dask.array.nanmin(dask_block_rng), dask.array.nanmax(dask_block_rng),
                                )

        # trans.dat - file generated by llt_grid2rat (r a topo lon lat)"
        trans_dat = self.get_trans_dat(subswath)
        # process all the chunks
        trans_blocks_ys, trans_blocks_xs = trans_dat.ll.data.numblocks
        #print ('trans_blocks_ys, trans_blocks_xs', trans_blocks_ys, trans_blocks_xs)
        with self.tqdm_joblib(tqdm(desc='Analyze Transform Blocks', total=trans_blocks_ys*trans_blocks_xs)) as progress_bar:
            extents = joblib.Parallel(n_jobs=n_jobs)(joblib.delayed(calculate_block_extent)(iy, ix) \
                                                     for iy in range(trans_blocks_ys) for ix in range(trans_blocks_xs))
        # merge the results
        extents = np.asarray([extent for extent in extents if extent is not None])
        return extents
    
    def get_trans_dat(self, subswath=None):
        import xarray as xr

        subswath = self.get_subswath(subswath)
        filename = self.get_filenames(subswath, None, 'trans')
        trans = xr.open_dataset(filename, engine=self.engine, chunks=self.chunksize).rename({'yy': 'lat', 'xx': 'lon'})
        return trans

    def trans_dat(self, subswath=None, interactive=False):
        import dask
        import xarray as xr
        import numpy as np
        import os
        import sys

        # range, azimuth, elevation(ref to radius in PRM), lon, lat [ASCII default] 
        llt2rat_map = {0: 'rng', 1: 'azi', 2: 'ele', 3: 'll', 4: 'lt'}

        # build trans.dat
        def SAT_llt2rat(z, lat, lon, subswath, binary=False):
            coords_ll = np.column_stack([lon.ravel(), lat.ravel(), z.ravel()])
            # for binary=True values outside of the scene missed and the array is not complete
            coords_ra = self.PRM(subswath).SAT_llt2rat(coords_ll, precise=1, binary=False)\
                .astype(np.float32).reshape(z.shape[0], z.shape[1], 5)
            return coords_ra

        ################################################################################
        # define valid area checking every 10th pixel per the both dimensions 
        ################################################################################
        # do not use coordinate names lat,lon because the output grid saved as (lon,lon) in this case...
        dem = self.get_dem(geoloc=True).rename({'lat': 'yy', 'lon': 'xx'})
        # prepare lazy coordinate grids
        lat = xr.DataArray(dem.yy.astype(np.float32).chunk(-1))
        lon = xr.DataArray(dem.xx.astype(np.float32).chunk(-1))
        lats, lons = xr.broadcast(lat, lon)
        # unify chunks
        lats = lats.chunk(dem.chunks)
        lons = lons.chunk(dem.chunks)

        # xarray wrapper for fast valid area lookup
        raell = xr.apply_ufunc(
            SAT_llt2rat,
            dem[::10,::10],
            lats[::10,::10],
            lons[::10,::10],
            dask='parallelized',
            vectorize=False,
            output_dtypes=[np.float32],
            output_core_dims=[['raell']],
            dask_gufunc_kwargs={'output_sizes': {'raell': 5}},
            kwargs={'subswath': subswath, 'binary': True}
        ).data.reshape(-1,5)

        # select valid azimuths only
        rng_max, yvalid, num_patch = self.PRM(subswath).get('num_rng_bins', 'num_valid_az', 'num_patches')
        azi_max = yvalid * num_patch
        #print ('azi_max', azi_max)
        raell = raell[(raell[...,1]>=0)&(raell[...,1]<=azi_max)&(raell[...,0]>=0)&(raell[...,0]<=rng_max)]

        if sys.version_info.major == 3 and sys.version_info.minor == 7:
            # Python 3.7 workaround (required for Google Colab notebooks)
            lat_values = raell[...,-1].compute()
            lat_min, lat_max = np.min(lat_values), np.max(lat_values)
            lon_values = raell[...,-2].compute()
            lon_min, lon_max = np.min(lon_values), np.max(lon_values)
        else:
            lat_min, lat_max = dask.compute(raell[...,-1].min(), raell[...,-1].max())
            #print ('lat_min, lat_max', lat_min, lat_max)
            lon_min, lon_max = dask.compute(raell[...,-2].min(), raell[...,-2].max())
        #print ('lon_min, lon_max', lon_min, lon_max)
        #azi_min, azi_max = dask.compute(raell[...,1].min(), raell[...,1].max())
        #print ('azi_min, azi_max', azi_min, azi_max)

        ################################################################################
        # process the valid area
        ################################################################################
        # crop valid area only
        dem  = dem.sel( yy=slice(lat_min, lat_max), xx=slice(lon_min, lon_max))
        lats = lats.sel(yy=slice(lat_min, lat_max), xx=slice(lon_min, lon_max))
        lons = lons.sel(yy=slice(lat_min, lat_max), xx=slice(lon_min, lon_max))

        # calculate linear index
        lat_idx = xr.DataArray(np.arange(dem.yy.size, dtype=np.uint32), dims=['yy']).chunk(-1)
        lon_idx = xr.DataArray(np.arange(dem.xx.size, dtype=np.uint32), dims=['xx']).chunk(-1)
        lat_idxs, lon_idxs = xr.broadcast(lat_idx, lon_idx)
        # unify chunks
        idxs = (lat_idxs*lon_idx.size + lon_idxs).chunk(dem.chunks)
        #print ('idxs', idxs)
        assert dem.size - idxs.max() == 1, 'Linear index incorrect'

        # xarray wrapper for the valid area only
        raell = xr.apply_ufunc(
            SAT_llt2rat,
            dem,
            lats,
            lons,
            dask='parallelized',
            vectorize=False,
            output_dtypes=[np.float32],
            output_core_dims=[['raell']],
            dask_gufunc_kwargs={'output_sizes': {'raell': 5}},
            kwargs={'subswath': subswath, 'binary': False}
        )

        #lat_min, lat_max = dask.compute(raell[...,-1].min(), raell[...,-1].max())
        #print ('lat_min, lat_max', lat_min, lat_max)
        #lon_min, lon_max = dask.compute(raell[...,-2].min(), raell[...,-2].max())
        #print ('lon_min, lon_max', lon_min, lon_max)
        #azi_min, azi_max = dask.compute(raell[...,1].min(), raell[...,1].max())
        #print ('azi_min, azi_max', azi_min, azi_max)

        # transform to separate variables
        keys_vars = {val: raell[...,key] for (key, val) in llt2rat_map.items()}
        keys_devs = {'idx': idxs}
        trans = xr.Dataset({**keys_vars, **keys_devs})

        if interactive:
            return trans

        # save to NetCDF file
        filename = self.get_filenames(subswath, None, 'trans')
        if os.path.exists(filename):
            os.remove(filename)
        encoding = {val: self.compression for (key, val) in llt2rat_map.items()}
        handler = trans.to_netcdf(filename,
                                        encoding=encoding,
                                        engine=self.engine,
                                        compute=False)
        return handler

    def trans_dat_parallel(self, interactive=False):
        import dask

        # process all the subswaths
        subswaths = self.get_subswaths()
        delayeds = []
        for subswath in subswaths:
            delayed = self.trans_dat(subswath=subswath, interactive=interactive)
            delayeds.append(delayed)

        if not interactive:
            tqdm_dask(dask.persist(delayeds), desc='Radar Transform Computing')
        else:
            return delayeds[0] if len(delayeds)==1 else delayeds
