"""transform coordinate between projection"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/CLI/transform.ipynb.

# %% auto 0
__all__ = ['transform']

# %% ../../nbs/CLI/transform.ipynb 4
import multiprocessing
import concurrent
import numpy as np
import zarr
import logging
from pyproj import Transformer

from .logging import mc_logger

# %% ../../nbs/CLI/transform.ipynb 5
def _transform(xx_in, # input x coordinate
              yy_in, # input y coordinate
              epsg_in=4326, # input epsg
              epsg_out=3857, # output epsg
              n_worker=None, # number of worker, all cpus by default
             ):
    if n_worker is None: n_worker = multiprocessing.cpu_count()
    transformer = Transformer.from_crs(f"EPSG:{epsg_in}", f"EPSG:{epsg_out}",always_xy=True)
    # xx_out = np.empty_like(xx_in); yy_out = np.empty_like(yy_in)
    # xx_out_split = np.array_split(xx_out,n_worker); yy_out_split = np.array_split(yy_out,n_worker)
    xx_out_list = []; yy_out_list = []
    with concurrent.futures.ProcessPoolExecutor(max_workers=n_worker) as executor:
        for i, (_xx_out, _yy_out) in enumerate(executor.map(transformer.transform,np.array_split(xx_in,n_worker),np.array_split(yy_in,n_worker))):
            xx_out_list.append(_xx_out); yy_out_list.append(_yy_out)
    return np.concatenate(xx_out_list), np.concatenate(yy_out_list)

# %% ../../nbs/CLI/transform.ipynb 6
@mc_logger
def transform(xx_in, # input x coordinate
              yy_in, # input y coordinate
              xx_out, # output x coordinate
              yy_out, # output y coordinate
              epsg_in=4326, # input epsg
              epsg_out=3857, # output epsg
             ):
    '''Coordinate transformation.
    By default, the input should be longitude (xx_in) and latitude (yy_in) (in degree) and outputs 
    are x (east) and y (south) coordinate in web mercator projection (for plot with google earth map).
    The chunks, shape and dtype of output are same as `xx_in`.
    '''
    logger = logging.getLogger(__name__)
    xx_in_path = xx_in; yy_in_path = yy_in; xx_out_path = xx_out; yy_out_path = yy_out

    logger.info(f'input EPSG: {epsg_in}.')
    xx_in_zarr = zarr.open(xx_in_path,mode='r'); yy_in_zarr = zarr.open(yy_in_path,mode='r')
    logger.zarr_info('xx_in',xx_in_zarr); logger.zarr_info('yy_in',yy_in_zarr)
    logger.info(f'output EPSG: {epsg_out}.')
    logger.info(f'do the transformation.')
    xx_out, yy_out = _transform(xx_in_zarr[:],yy_in_zarr[:])
    logger.info(f'write output.')
    xx_out_zarr = zarr.open(xx_out_path,mode='w',shape=xx_in_zarr.shape,dtype=xx_in_zarr.dtype,chunks=xx_in_zarr.chunks)
    yy_out_zarr = zarr.open(yy_out_path,mode='w',shape=xx_in_zarr.shape,dtype=xx_in_zarr.dtype,chunks=xx_in_zarr.chunks)
    xx_out_zarr[:] = xx_out; yy_out_zarr[:] = yy_out
    logger.info(f'write done.')
