"""
- helper functions to compute dists/sims for actual networks and experiments
"""
from pprint import pprint

import torch
from typing import Union, Optional
from collections import OrderedDict

from torch import nn, Tensor

LayerIdentifier = Union[str, tuple[str]]


def _clear_hooks(hooks: list):
    for hook in hooks:
        hook.clear()


def remove_hook(mdl: nn.Module, hook):
    """
    ref: https://github.com/pytorch/pytorch/issues/5037
    """
    handle = mdl.register_forward_hook(hook)
    handle.remove()


def remove_hooks(mdl: nn.Module, hooks: list):
    """
    ref:
        - https://github.com/pytorch/pytorch/issues/5037
        - https://discuss.pytorch.org/t/how-to-remove-multiple-hooks/135442
    """
    for hook in hooks:
        remove_hook(mdl, hook)


def _to_layer_order_dict(dists: list[float],
                         layer_names1: list[str],
                         layer_names2: list[str]) -> OrderedDict[LayerIdentifier, float]:
    """
    list([L]) -> OrderedDict([L])
    """
    assert len(layer_names1) == len(layer_names2), f'Expected same number of layers for comparing both nets but got: ' \
                                                   f'{(len(layer_names1))=}, {len(layer_names2)=}'
    # - put distances/sims for each layer
    dists_od: OrderedDict[LayerIdentifier, float] = OrderedDict()
    for layer_idx, layer1, layer2 in enumerate(zip(layer_names1, layer_names2)):
        dist: float = dists[layer_idx]
        layer_key: LayerIdentifier = layer1 if layer1 == layer2 else (layer1, layer2)
        dists_od[layer_key] = dist
    return dists_od


def _dists_per_layer_to_list(dists_per_layer: OrderedDict[LayerIdentifier, float]) -> list[float]:
    """
    OrderedDict([L]) -> list([L])
    """
    return [dist for _, dist in dists_per_layer.items()]


def _dists_per_task_per_layer_to_list(dists_per_tasks_per_layer: list[OrderedDict[LayerIdentifier, float]]) -> list[
    list[float]]:
    """
    OrderedDict([B, L]) -> list([B, L])
    """
    B: int = len(dists_per_tasks_per_layer)
    _dists_per_tasks_per_layer: list[list[float]] = []
    for b in range(B):
        _dists_per_layer: list[float] = _dists_per_layer_to_list(dists_per_tasks_per_layer[b])
        _dists_per_tasks_per_layer.append(_dists_per_layer)
    return _dists_per_tasks_per_layer


def dist_data_set_per_layer(mdl1: nn.Module, mdl2: nn.Module,
                            X1: Tensor, X2: Tensor,
                            layer_names1: list[str], layer_names2: list[str],
                            metric_comparison_type: str = 'pwcca',
                            iters: int = 1,
                            effective_neuron_type: str = 'filter',
                            downsample_method: Optional[str] = None,
                            downsample_size: Optional[int] = None,
                            subsample_effective_num_data_method: Optional[str] = None,
                            subsample_effective_num_data_param: Optional[int] = None,
                            metric_as_sim_or_dist: str = 'dist',
                            force_cpu: bool = False
                            ) -> OrderedDict[LayerIdentifier, float]:
    """
    Given a pair of data sets (or batches, potentially the same one), compute the distance btw the models
        [M, C, H, W], [L] -> [L] (as list like obj)
    assuming the list of layes are in the order of how you want to compute the comparison.
    E.g.
        - d(f_ml(X), A(f_ml, task)(X)), compute the distance between model and adapted model for same set of examples
        for the task/data set since X1=X2=X.
        - dv = d(f(X1), f(X2))
        - dv = d(f^*_1(X1), f^*_2(X2))

    :return: [L, 1] as a iterable (e.g. list of floats per layer) correponding to the distance for each layer.
    """
    from anatome.similarity import SimilarityHook
    assert len(layer_names1) == len(layer_names2), f'Expected same number of layers for comparing both nets but got: ' \
                                                   f'{(len(layer_names1))=}, {len(layer_names2)=}'
    # - create hooks for each layer name
    hooks1: list[SimilarityHook] = SimilarityHook.create_hooks(mdl1, layer_names1, metric_comparison_type, force_cpu)
    hooks2: list[SimilarityHook] = SimilarityHook.create_hooks(mdl2, layer_names2, metric_comparison_type, force_cpu)
    # - fill hooks with intermediate representations
    for _ in range(iters):  # might make sense to go through multiple is NN is stochastic e.g. BN, dropout layers
        # note we are avoiding torch.no_grad() because we don't want it to interfere with higher... could be future work to see if torch.no_grad() interferes with it.
        mdl1(X1)
        mdl2(X2)
    # - compute distances/sims for each layer
    dists_od: OrderedDict[LayerIdentifier, float] = OrderedDict()
    for layer_idx, (layer1, layer2) in enumerate(zip(layer_names1, layer_names2)):
        hook1: SimilarityHook = hooks1[layer_idx]
        hook2: SimilarityHook = hooks2[layer_idx]
        # - get dist/sim for current models for current layer
        dist: float = hook1.distance(hook2,
                                     effective_neuron_type=effective_neuron_type,
                                     downsample_method=downsample_method,
                                     downsample_size=downsample_size,
                                     subsample_effective_num_data_method=subsample_effective_num_data_method,
                                     subsample_effective_num_data_param=subsample_effective_num_data_param,
                                     metric_as_sim_or_dist=metric_as_sim_or_dist
                                     )
        layer_key: LayerIdentifier = layer1 if layer1 == layer2 else (layer1, layer2)
        dists_od[layer_key] = dist
    _clear_hooks(hooks1)
    _clear_hooks(hooks2)
    return dists_od


def dist_batch_data_sets_for_all_layer(mdl1: nn.Module, mdl2: nn.Module,
                                       X1: Tensor, X2: Tensor,
                                       layer_names1: list[str], layer_names2: list[str],
                                       metric_comparison_type: str = 'pwcca',
                                       iters: int = 1,
                                       effective_neuron_type: str = 'filter',
                                       downsample_method: Optional[str] = None,
                                       downsample_size: Optional[int] = None,
                                       subsample_effective_num_data_method: Optional[str] = None,
                                       subsample_effective_num_data_param: Optional[int] = None,
                                       metric_as_sim_or_dist: str = 'dist',
                                       force_cpu: bool = False
                                       ) -> list[OrderedDict[LayerIdentifier, float]]:
    """
    Gets the distance for a batch (e.g meta-batch) of data sets/tasks and for each layer the distances between the nets:
        [B, M, C, H, W], [L] -> list(OrderDict([B, L]))
    Example use:
        - Computing d(f_ml(X), A(f_ml, t)(X)) so X1=X2=X.
        - Computing dv so some restricted cross product of the data sets in X1 and X2 (X1 != X2),
        ideally avoiding diagonal if dv.

    :return:
    """
    assert len(X1.size()) == 5, f'Data set does not 5 dims i.e. [B, M, C, H, W] instead got: {X1.size()=}'
    assert len(X2.size()) == 5, f'Data set does not 5 dims i.e. [B, M, C, H, W] instead got: {X2.size()=}'
    # - get distances per data sets/tasks per layers: [B, M, C, H, W], [L] -> [B, L]
    B: int = min(X1.size(0), X2.size(0))
    dists_entire_net_all_data_sets: list[OrderedDict[LayerIdentifier, float]] = []  # [B, L]
    for b in range(B):
        # - [B, M, C, H, W] -> [M, C, H, W] (get a specific data set/batch)
        x1, x2 = X1[b], X2[b]
        # - compute for current pair of data sets (or tasks) the distance between models
        # [M, C, H, W], [L] -> [L], really a list of floats of len [L]
        dist_for_data_set_b: OrderedDict[LayerIdentifier, float] = dist_data_set_per_layer(mdl1, mdl2, x1, x2,
                                                                                           layer_names1, layer_names2,
                                                                                           metric_comparison_type=metric_comparison_type,
                                                                                           iters=iters,
                                                                                           effective_neuron_type=effective_neuron_type,
                                                                                           downsample_method=downsample_method,
                                                                                           downsample_size=downsample_size,
                                                                                           subsample_effective_num_data_method=subsample_effective_num_data_method,
                                                                                           subsample_effective_num_data_param=subsample_effective_num_data_param,
                                                                                           metric_as_sim_or_dist=metric_as_sim_or_dist,
                                                                                           force_cpu=force_cpu
                                                                                           )
        # adding to [B, L]
        dists_entire_net_all_data_sets.append(dist_for_data_set_b)
    # check effective size [B, L]
    L: int = len(layer_names1)
    assert len(dists_entire_net_all_data_sets) == B
    assert len(dists_entire_net_all_data_sets[0]) == L
    # - [B, L], return list of iters of distances per data sets (tasks) per layers
    return dists_entire_net_all_data_sets


def stats_distance_per_layer(mdl1: nn.Module, mdl2: nn.Module,
                             X1: Tensor, X2: Tensor,
                             layer_names1: list[str], layer_names2: list[str],
                             metric_comparison_type: str = 'pwcca',
                             iters: int = 1,
                             effective_neuron_type: str = 'filter',
                             downsample_method: Optional[str] = None,
                             downsample_size: Optional[int] = None,
                             subsample_effective_num_data_method: Optional[str] = None,
                             subsample_effective_num_data_param: Optional[int] = None,
                             metric_as_sim_or_dist: str = 'dist',
                             force_cpu: bool = False
                             ) -> tuple[OrderedDict[LayerIdentifier, float]]:
    """
    Compute the overall stats (mean & std) of distances per layer
        [B, M, C, H, W] -> [L]*2 (mus, stds) per layer
    for the given B number of data sets/tasks of size [M, C, H, W].

    Note:
        - this might be confusing because X1 != X2. The most intuitive case is when X1=X2=X so that we are computing
        the distance for each task (but the tasks correspond to each other).
    :return:
    """
    from uutils.torch_uu import tensorify
    assert len(layer_names1) == len(layer_names2)
    # - [B, L, 1], get distances per data sets (tasks) per layer
    distances_per_data_sets_per_layer: list[OrderedDict[LayerIdentifier, float]] = dist_batch_data_sets_for_all_layer(
        mdl1, mdl2, X1, X2,
        layer_names1, layer_names2,
        metric_comparison_type=metric_comparison_type,
        iters=iters,
        effective_neuron_type=effective_neuron_type,
        downsample_method=downsample_method,
        downsample_size=downsample_size,
        subsample_effective_num_data_method=subsample_effective_num_data_method,
        subsample_effective_num_data_param=subsample_effective_num_data_param,
        metric_as_sim_or_dist=metric_as_sim_or_dist,
        force_cpu=force_cpu
    )
    layer_ids: list[LayerIdentifier] = list(distances_per_data_sets_per_layer[0].keys())
    # list(OrderDict([B, L])) -> list([B, L])
    distances_per_data_sets_per_layer: list[list[float]] = _dists_per_task_per_layer_to_list(
        distances_per_data_sets_per_layer)
    B, L = X1.size(0), len(layer_names1)
    distances_per_data_sets_per_layer: Tensor = tensorify(distances_per_data_sets_per_layer)
    assert (distances_per_data_sets_per_layer.size() == torch.Size([B, L]))
    # - [B, L] -> [L], get the avg distance/sim for each layer (and std)
    means_per_layer: Tensor = distances_per_data_sets_per_layer.mean(dim=0)
    stds_per_layer: Tensor = distances_per_data_sets_per_layer.std(dim=0)
    L: int = len(distances_per_data_sets_per_layer[0])
    assert means_per_layer.size() == torch.Size([L])
    assert stds_per_layer.size() == torch.Size([L])
    # [L, 1] -> OrderDict([L])
    assert len(layer_ids) == L == len(layer_names1)
    mean_distance_per_layer: OrderedDict[LayerIdentifier, float] = OrderedDict()
    std_distance_per_layer: OrderedDict[LayerIdentifier, float] = OrderedDict()
    for l_idx in range(L):
        layer_id: LayerIdentifier = layer_ids[l_idx]
        mu: float = means_per_layer[l_idx].item()
        std: float = stds_per_layer[l_idx].item()
        mean_distance_per_layer[layer_id] = mu
        std_distance_per_layer[layer_id] = std
    return mean_distance_per_layer, std_distance_per_layer, distances_per_data_sets_per_layer


def _stats_distance_entire_net(mdl1: nn.Module, mdl2: nn.Module,
                              X1: Tensor, X2: Tensor,
                              layer_names1: list[str], layer_names2: list[str],
                              metric_comparison_type: str = 'pwcca',
                              iters: int = 1,
                              effective_neuron_type: str = 'filter',
                              downsample_method: Optional[str] = None,
                              downsample_size: Optional[int] = None,
                              subsample_effective_num_data_method: Optional[str] = None,
                              subsample_effective_num_data_param: Optional[int] = None,
                              metric_as_sim_or_dist: str = 'dist',
                              force_cpu: bool = False
                              ) -> tuple[float]:
    """
    Compute the overall stats of distances per layer.
    Not recommended, to use. It's better to compute the layerwise metrics and then use the helper function
    to aggreagate those into a one net result.
        [B, M, C, H, W] -> pair of [1, 1]
    :return:
    """
    from uutils.torch_uu import tensorify
    # - [B, L], get distances per data sets (tasks) per layer
    distances_per_data_sets_per_layer: list[OrderedDict[LayerIdentifier, float]] = dist_batch_data_sets_for_all_layer(
        mdl1, mdl2, X1, X2,
        layer_names1, layer_names2,
        metric_comparison_type=metric_comparison_type,
        iters=iters,
        effective_neuron_type=effective_neuron_type,
        downsample_method=downsample_method,
        downsample_size=downsample_size,
        subsample_effective_num_data_method=subsample_effective_num_data_method,
        subsample_effective_num_data_param=subsample_effective_num_data_param,
        metric_as_sim_or_dist=metric_as_sim_or_dist,
        force_cpu=force_cpu
    )
    # list(OrderDict([B, L])) -> list([B, L])
    distances_per_data_sets_per_layer: list[list[float]] = _dists_per_task_per_layer_to_list(
        distances_per_data_sets_per_layer)
    distances_per_data_sets_per_layer: Tensor = tensorify(distances_per_data_sets_per_layer)
    # - [B, L] -> [], get the avg distance/sim for each layer (and std)
    mu: Tensor = distances_per_data_sets_per_layer.mean(dim=[0, 1])
    std: Tensor = distances_per_data_sets_per_layer.std(dim=[0, 1])
    return mu, std

def compute_mu_std_for_entire_net_from_all_distances_from_data_sets_tasks(
        distances_per_data_sets_per_layer: Tensor) \
        -> tuple[float, float]:
    """
        [B, L] -> mu, std
    where mu = 1/B*L sum_{b,l} distances_per_data_sets_per_layer[b, l]
    where std = 1/B*L sum_{b,l} (distances_per_data_sets_per_layer[b, l]-mu)^2
    :param distances_per_data_sets_per_layer:
    :return:
    """
    assert (distances_per_data_sets_per_layer.dim() == 2)
    # list(OrderDict([B, L, 1])) -> list([B, L, 1])
    # _distances_per_data_sets_per_layer: list[list[float]] = _dists_per_task_per_layer_to_list(
    #     distances_per_data_sets_per_layer)
    # _distances_per_data_sets_per_layer: Tensor = tensorify(_distances_per_data_sets_per_layer)
    # - [B, L, 1] -> [1], get the avg distance/sim for each layer (and std)
    mu: Tensor = distances_per_data_sets_per_layer.mean(dim=[0, 1])
    std: Tensor = distances_per_data_sets_per_layer.std(dim=[0, 1])
    assert (mu.size() == torch.Size([]))
    assert (std.size() == torch.Size([]))
    return mu, std


def pprint_results(mus: OrderedDict, stds: OrderedDict):
    print('---- stats of results per layer')
    print('-- mus (means) per layer')
    pprint(mus)
    print('-- stds (standard devs) per layer')
    pprint(stds)
    print()


# - tests

def dist_per_layer_test():
    """
    """
    from copy import deepcopy

    import torch
    import torch.nn as nn

    from uutils.torch_uu.models.learner_from_opt_as_few_shot_paper import get_default_learner_from_default_args, \
        get_feature_extractor_layers
    from uutils.torch_uu import approx_equal

    # - very simple sanity check
    Cin: int = 3
    # num_out_filters: int = 8
    mdl1: nn.Module = get_default_learner_from_default_args()
    mdl2: nn.Module = deepcopy(mdl1)
    layer_names = get_feature_extractor_layers()
    print(f'{layer_names=}')

    # - ends up comparing two matrices of size [B, Dout], on same data, on same model
    B: int = 5  # -- satisfies B >= (10*32)/(5**2) = 12.8 for this specific 5CNN model
    M: int = 13  # e.g. k*n_c, satisfies B >= (10*32)/(5**2) = 12.8 for this specific 5CNN model
    C, H, W = Cin, 84, 84
    metric_comparison_type: str = 'pwcca'
    X: torch.Tensor = torch.distributions.Normal(loc=0.0, scale=1.0).sample((B, M, C, H, W))
    effective_neuron_type: str = 'filter'
    subsample_effective_num_data_method: str = 'subsampling_data_to_dims_ratio'
    subsample_effective_num_data_param: int = 10
    metric_as_sim_or_dist: str = 'sim'
    mus, stds, distances_per_data_sets_per_layer = stats_distance_per_layer(mdl1, mdl1, X, X, layer_names, layer_names,
                                         metric_comparison_type=metric_comparison_type,
                                         effective_neuron_type=effective_neuron_type,
                                         subsample_effective_num_data_method=subsample_effective_num_data_method,
                                         subsample_effective_num_data_param=subsample_effective_num_data_param,
                                         metric_as_sim_or_dist=metric_as_sim_or_dist)
    pprint_results(mus, stds)
    assert (mus != stds)
    mu, std = _stats_distance_entire_net(mdl1, mdl1, X, X, layer_names, layer_names,
                                        metric_comparison_type=metric_comparison_type,
                                        effective_neuron_type=effective_neuron_type,
                                        subsample_effective_num_data_method=subsample_effective_num_data_method,
                                        subsample_effective_num_data_param=subsample_effective_num_data_param,
                                        metric_as_sim_or_dist=metric_as_sim_or_dist)
    print(f'----entire net result: {mu=}, {std=}')
    assert (approx_equal(mu, 1.0))
    mu2, std2 = compute_mu_std_for_entire_net_from_all_distances_from_data_sets_tasks(distances_per_data_sets_per_layer)
    assert (approx_equal(mu, mu2))
    assert (approx_equal(std, std2))

    # -- differnt data different nets dist ~ 1.0, large distance, so low sim
    X1: torch.Tensor = torch.distributions.Normal(loc=0.0, scale=1.0).sample((B, M, C, H, W))
    X2: torch.Tensor = torch.distributions.Normal(loc=0.0, scale=1.0).sample((B, M, C, H, W))
    mus, stds, distances_per_data_sets_per_layer = stats_distance_per_layer(mdl1, mdl2, X1, X2, layer_names, layer_names,
                                         metric_comparison_type=metric_comparison_type,
                                         effective_neuron_type=effective_neuron_type,
                                         subsample_effective_num_data_method=subsample_effective_num_data_method,
                                         subsample_effective_num_data_param=subsample_effective_num_data_param,
                                         metric_as_sim_or_dist=metric_as_sim_or_dist)
    pprint_results(mus, stds)
    assert (mus != stds)
    mu, std = _stats_distance_entire_net(mdl1, mdl2, X1, X2, layer_names, layer_names,
                                        metric_comparison_type=metric_comparison_type,
                                        effective_neuron_type=effective_neuron_type,
                                        subsample_effective_num_data_method=subsample_effective_num_data_method,
                                        subsample_effective_num_data_param=subsample_effective_num_data_param,
                                        metric_as_sim_or_dist=metric_as_sim_or_dist)
    print(f'----entire net result: {mu=}, {std=}')
    assert (approx_equal(mu, 0.0, tolerance=0.4))
    mu2, std2 = compute_mu_std_for_entire_net_from_all_distances_from_data_sets_tasks(distances_per_data_sets_per_layer)
    assert (approx_equal(mu, mu2))
    assert (approx_equal(std, std2))


if __name__ == '__main__':
    import uutils
    import time

    start = time.time()
    dist_per_layer_test()
    print(f'time_passed_msg = {uutils.report_times(start)}')
    print('Done, success\a')
