from tactile_prediction.model.common.normalizer import SingleFieldLinearNormalizer
from common.pytorch_util import dict_apply, dict_apply_reduce, dict_apply_split
import numpy as np
import torch


def get_range_normalizer_from_stat(stat, output_max=1, output_min=-1, range_eps=1e-7):
    # -1, 1 normalization
    input_max = stat['max']
    input_min = stat['min']
    input_range = input_max - input_min
    ignore_dim = input_range < range_eps
    input_range[ignore_dim] = output_max - output_min
    scale = (output_max - output_min) / input_range
    offset = output_min - scale * input_min
    offset[ignore_dim] = (output_max + output_min) / 2 - input_min[ignore_dim]

    return SingleFieldLinearNormalizer.create_manual(
        scale=scale,
        offset=offset,
        input_stats_dict=stat
    )


def get_image_range_normalizer():
    scale = np.array([2], dtype=np.float32)
    offset = np.array([-1], dtype=np.float32)
    stat = {
        'min': np.array([0], dtype=np.float32),
        'max': np.array([1], dtype=np.float32),
        'mean': np.array([0.5], dtype=np.float32),
        'std': np.array([np.sqrt(1/12)], dtype=np.float32)
    }
    return SingleFieldLinearNormalizer.create_manual(
        scale=scale,
        offset=offset,
        input_stats_dict=stat
    )


def get_image_identity_normalizer():
    scale = np.array([1], dtype=np.float32)
    offset = np.array([0], dtype=np.float32)
    stat = {
        'min': np.array([0], dtype=np.float32),
        'max': np.array([1], dtype=np.float32),
        'mean': np.array([0.5], dtype=np.float32),
        'std': np.array([np.sqrt(1/12)], dtype=np.float32)
    }
    return SingleFieldLinearNormalizer.create_manual(
        scale=scale,
        offset=offset,
        input_stats_dict=stat
    )


def get_image_imagenet_normalizer():
    scale = np.array([1 / 0.229, 1 / 0.224, 1 / 0.225], dtype=np.float32)
    offset = np.array([-0.485 / 0.229, -0.456 / 0.224, -0.406 / 0.225], dtype=np.float32)
    stat = {
        'min': np.array([0.0, 0.0, 0.0], dtype=np.float32),
        'max': np.array([1.0, 1.0, 1.0], dtype=np.float32),
        'mean': np.array([0.485, 0.456, 0.406], dtype=np.float32),
        'std': np.array([0.229, 0.224, 0.225], dtype=np.float32),
    }
    return SingleFieldLinearNormalizer.create_manual(
        scale=scale,
        offset=offset,
        input_stats_dict=stat
    )


def get_identity_normalizer_from_stat(stat):
    scale = np.ones_like(stat['min'])
    offset = np.zeros_like(stat['min'])
    return SingleFieldLinearNormalizer.create_manual(
        scale=scale,
        offset=offset,
        input_stats_dict=stat
    )


def array_to_stats(arr: np.ndarray):
    stat = {
        'min': np.min(arr, axis=0),
        'max': np.max(arr, axis=0),
        'mean': np.mean(arr, axis=0),
        'std': np.std(arr, axis=0)
    }
    return stat


def concatenate_normalizer(normalizers: list):
    scale = torch.concatenate([normalizer.params_dict['scale'] for normalizer in normalizers], axis=-1)
    offset = torch.concatenate([normalizer.params_dict['offset'] for normalizer in normalizers], axis=-1)
    input_stats_dict = dict_apply_reduce(
        [normalizer.params_dict['input_stats'] for normalizer in normalizers], 
        lambda x: torch.concatenate(x,axis=-1))
    return SingleFieldLinearNormalizer.create_manual(
        scale=scale,
        offset=offset,
        input_stats_dict=input_stats_dict
    )