import edt
import numpy as np
import h5py
from dataset_iterator import MultiChannelIterator
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.losses import mean_squared_error as mse
from dataset_iterator.tile_utils import extract_tile_random_zoom_function
from .data_augmentation import *
from ..model import get_unet


def get_train_test_iterators(dataset,
    center_range:list, scale_range:list,
    group_keyword:list = None,
    group_proportion:list = None,
    batch_size:int = 4,
    n_z:int=5, z_step=2,
    tile_params:dict = dict(tile_shape=(256, 256), n_tiles=9, zoom_range=[0.6, 1.6], aspect_ratio_range=[0.75, 1.5] ),
    elasticdeform_parameters:dict = {},
    noise_intensity:float = 0.035,
    raw_feature_name:str="/raw", label_feature_name:str="/regionLabels",
    training_selection_name:str="train/", validation_selection_name:str="eval/" ):
    """Generates image iterators for training/validation.
    Perform data augmentation: pick random slices, random zoom / random flip / rotate90, add radnom iid gaussian noise / gaussian blur, elastic deformation, random normalization

    Parameters
    ----------
    dataset : datasetIO or path to dataset (str)
        dataset
        if group_keyword are provided, dataset must be a .h5 file. group_keyword is expected in the path of the datasets, after validation_selection_name or training_selection_name.
    center_range : list
        [min center, max center] for normalization in data augmentation. When several groups are defined, one array per group should be given.
        if None: image is normalized in a random range [min, max] so that min € [0, 1], max € [0, 1] and max-min<=0.1
    scale_range : list
        [min scale, max scale] for normalization in data augmentation. When several groups are defined, one array per group should be given.
    group_keyword : list
        list of keywords that define distinct groups (optional, can be None): usually a dataset name.
    group_proportion : list
        list of proportion for each group in each mini batch (statistically) (optional, can be None)
    batch_size : int
        size of mini batch
    n_z : int
        number of z-slices that will be picked from the dataset.
    z_step : type
        step between picked z-slices
    tile_params : dict
        parameter for tiling:
            tile_shape: shape (Y,X) of returned tiles
            n_tiles: number of tiles per image
            zoom_range: random zoom applied at cropping to Y, X
            aspect_ratio_range: limits the aspect ratio distortion generated by the random zoom
    elasticdeform_parameters : dict
        None for no elastic deformation OR parameters for elasticdeform augmentation:
            order: interpolation order
            mode: out-of-bound strategy
            points : number of reference points
            sigma : intensity of deformation
    noise_intensity : float
        intensity of added iid gaussian noise
    raw_feature_name : str
        input images dataset name
    label_feature_name : str
        labeled images dataset name
    training_selection_name : str
        name of the training selection
    validation_selection_name : str
        name of the validation selection

    Returns
    -------
    data iterator
        tensor shape: batch_size * n_tiles, tile_shape[0], tile_shape[1], n_z

    """

    extract_tile_function = extract_tile_random_zoom_function(**tile_params) if tile_params is not None else None
    def random_channel_slice(nchan): # random offset on chosen slices to simulate focus variations
        noff=nchan-1-(n_z-1)*z_step + 1
        off = np.random.randint(noff)
        idx = [off + z_step*i for i in range(n_z)]
        return idx
    if group_keyword is not None and isinstance(group_keyword, (list, tuple)):
        assert group_proportion is None or len(group_proportion)==len(group_keyword)
        if center_range is not None:
            assert len(center_range)==len(group_keyword) and len(scale_range)==len(group_keyword), "when several groups are provided, as many center_range / center_range arrays should be provided"
            for i in range(len(group_keyword)):
                assert len(center_range[i])==2 and len(scale_range[i])==2 and scale_range[i][0]<=scale_range[i][1], "invalid scale_range / center_range"
            scale_fun = [[get_normalization_fun(center_range[i], scale_range[i]), None] for i in range(len(group_keyword))]
        else:
            scale_fun = [[get_phase_contrast_normalization_fun(), None] for i in range(len(group_keyword))]
        validation_grp = [_join_grp_kw(validation_selection_name, kw) for kw in group_keyword]
        training_grp = [_join_grp_kw(training_selection_name, kw) for kw in group_keyword]
    else:
        if center_range is not None:
            assert scale_range is not None
            assert len(center_range)==2 and center_range[0]<=center_range[1], "invalid center_range"
            assert len(scale_range)==2 and scale_range[0]<=scale_range[1], "invalid scale_range"
            scale_fun = [[get_normalization_fun(center_range, scale_range), None]]
        else:
            scale_fun = [[get_phase_contrast_normalization_fun(), None]]
        if group_keyword is None:
            validation_grp = validation_selection_name
            training_grp = training_selection_name
        else:
            validation_grp = _join_grp_kw(validation_selection_name, group_keyword)
            training_grp = _join_grp_kw(training_selection_name, group_keyword)

    params = dict(dataset=dataset,
        group_scaling = scale_fun,
        group_proportion=group_proportion,
        channel_keywords=[raw_feature_name, label_feature_name], # channel keyword must correspond to the name of the extracted features
        output_channels= [1],
        mask_channels = [1],
        channel_slicing_channels = {0:random_channel_slice},
        elasticdeform_parameters = elasticdeform_parameters,
        extract_tile_function = extract_tile_function,
        channels_postprocessing_function = _channels_postpprocessing(get_illumination_aug_fun(None, noise_intensity), _get_edt_fun()),
        batch_size=batch_size,
        perform_data_augmentation=True,
        shuffle=True)

    train_it = MultiChannelIterator(group_keyword=training_grp, **params)
    test_it = MultiChannelIterator(group_keyword=validation_grp, **params)
    return train_it, test_it

def _get_edt_fun():
    return lambda labels : np.stack([edt.edt(labels[i,...,0], black_border=False)[...,np.newaxis] for i in range(labels.shape[0])])

def _channels_postpprocessing(raw_fun, label_fun):
    def fun(batch_by_channel):
        batch_by_channel[0] = raw_fun(batch_by_channel[0])
        batch_by_channel[1] = label_fun(batch_by_channel[1])
    return fun

def _join_grp_kw(kw1, kw2):
    if kw1 is None or len(kw1)==0:
        return kw2
    elif kw2 is None or len(kw2)==0:
        return kw1
    else:
        return kw1 + '.*' + kw2 # TODO : only works with h5py datasets

def get_model(model_params:dict=dict(), learning_rate:float = 2e-4, saved_weights_file:str=None):
    """Initialize model.

    Parameters
    ----------
    model_params : dict
        parameters for ..unet.get_unet
    learning_rate : float
        initial learning rate
    saved_weights_file : str
        path to saved weight

    Returns
    -------
    compiled model

    """
    model = get_unet(**model_params)
    if saved_weights_file is not None:
        model.load_weights(saved_weights_file)
    model.compile(optimizer=Adam(learning_rate), loss=mse)
    return model
