# Path: reservoir_nn/utils/weight_properties.py
# def print_network_statistics(weight, plot_weight=False):
# def get_network_statistics(weight):
# def get_inhibitory_proportion_of_total_strength(weight):
# def get_inhibitory_neuron_indices(weight):
# def get_sparsity(matrix: np.ndarray, thres=0.0) -> float:
# def get_spectral_radius(square_matrix: np.ndarray) -> float:
# 
# Path: reservoir_nn/utils/weight_transforms.py
# class WeightTransformation(abc.ABC):
# class Chain(WeightTransformation, list):
# class FlipSignEntireRows(WeightTransformation):
# class ScaleWeightByDistance(WeightTransformation):
# class RandomizeNonZeroWeight(WeightTransformation):
# class RandomizeNonZeroWeightKde(WeightTransformation):
# class CutOffSmallWeightsInRows(WeightTransformation):
# class ScaleToZeroOne(WeightTransformation):
# class ScaleSpectralRadius(WeightTransformation):
# class GetSubMatrix(WeightTransformation):
# class ResizeNumRows(WeightTransformation):
# class ResizeNumColumns(WeightTransformation):
#   def apply_transform(self, *args: np.ndarray) -> np.ndarray:
#   def __call__(self, *args: np.ndarray) -> np.ndarray:
#   def batch_apply(
#       self, *args: Union[np.ndarray, Tuple[np.ndarray,
#                                            ...]]) -> List[np.ndarray]:
#   def __init__(self,
#                transformations: Optional[Sequence[WeightTransformation]] = None
#               ):
#   def __repr__(self):
#   def apply_transform(self, *args: np.ndarray) -> np.ndarray:
#   def __init__(self, flip_proportion: float):
#   def apply_transform(self, matrix: np.ndarray, seed: int = -1) -> np.ndarray:
#   def __init__(self, scaling_function, box_size=None):
#   def apply_transform(self, weights: np.ndarray,
#                       distances: np.ndarray) -> np.ndarray:
#   def __init__(self):
#   def apply_transform(self, matrix: np.ndarray) -> np.ndarray:
#   def __init__(self, rng: np.random.RandomState = np.random):
#   def apply_transform(self, matrix: np.ndarray, seed: int = -1) -> np.ndarray:
#   def __init__(self, non_zeros_per_row_limit: int):
#   def apply_transform(self, matrix: np.ndarray) -> np.ndarray:
#   def __init__(self):
#   def apply_transform(self, matrix: np.ndarray) -> np.ndarray:
#   def __init__(self, spectral_radius: float = 1.0):
#   def apply_transform(self, matrix: np.ndarray) -> np.ndarray:
#   def __init__(self, num_neurons: int):
#   def apply_transform(self, matrix: np.ndarray) -> np.ndarray:
#   def __init__(self,
#                target_num_rows: int,
#                rng: np.random.RandomState = np.random):
#   def apply_transform(self, matrix: np.ndarray) -> np.ndarray:
#   def __init__(self,
#                target_num_columns: int,
#                rng: np.random.RandomState = np.random):
#   def apply_transform(self, matrix: np.ndarray) -> np.ndarray:
# def flip_sign_entire_rows(matrix: np.ndarray,
#                           flip_proportion: float) -> np.ndarray:
# def randomize_non_zero_weights_kde(matrix: np.ndarray) -> np.ndarray:
# def randomize_non_zero_weights(matrix: np.ndarray) -> np.ndarray:
# def shuffle_weights(matrix: np.ndarray) -> np.ndarray:
# def assign_random_signs(matrix: np.ndarray,
#                         inhibitory_proportion: float) -> np.ndarray:
# def scale_weight_by_distance(weights,
#                              distances,
#                              scaling_function,
#                              box_size=None):
# def make_sparse(weight_array: np.ndarray,
#                 zero_weight_proportion: float) -> np.ndarray:
# def cutoff_small_weights_in_rows(matrix: np.ndarray,
#                                  non_zeros_per_row_limit: int) -> np.ndarray:
# def transform_weight(weight_matrix: np.ndarray, distance_matrix: np.ndarray,
#                      params: ml_collections.ConfigDict) -> np.ndarray:
# def setup_weight_transformation(params: ml_collections.ConfigDict) -> Chain:
# def chain_weight_matrices(matrices: List[np.ndarray]) -> List[np.ndarray]:
# def resize_weight_matrices(
#     reservoir_weights: List[np.ndarray],
#     reservoirs_num_neurons: Tuple[int, ...],
# ) -> List[np.ndarray]:


from absl.testing import absltest
from absl.testing import parameterized
from reservoir_nn.utils import weight_properties
from reservoir_nn.utils import weight_transforms
import attr
import ml_collections
import numpy as np

  def test_setup_weight_transformation_returns_expected_list(self):
    params = ml_collections.ConfigDict()
    params.box_size = 1
    params.distance_scaling_function = 'linear'
    params.num_cutoff = 10
    params.random_weights = True
    params.kde_random_weights = False
    params.inhibitory_neurons_proportion = 0.2
    transform_list = weight_transforms.setup_weight_transformation(params)
    self.assertIsInstance(transform_list[0],
                          weight_transforms.ScaleWeightByDistance)
    self.assertIsInstance(transform_list[1],
                          weight_transforms.CutOffSmallWeightsInRows)
    self.assertIsInstance(transform_list[2],
                          weight_transforms.RandomizeNonZeroWeight)
    self.assertIsInstance(transform_list[3], weight_transforms.ScaleToZeroOne)
    self.assertIsInstance(transform_list[4],
                          weight_transforms.FlipSignEntireRows)


class ScaleSpectralRadiusTest(absltest.TestCase):

  def test_it_should_produce_matrix_of_correct_spectral_radius(self):
    # matrix: spectral_radius=6
    matrix = np.array([[1, 2, 3], [3, 2, 1], [2, 1, 3]])
    expected = 4.2
    # Use the expected as the parameter to ScaleSpectralRadius
    scaled = weight_transforms.ScaleSpectralRadius(
        spectral_radius=expected).apply_transform(matrix)

# Path: reservoir_nn/utils/weight_properties.py
# def print_network_statistics(weight, plot_weight=False):
# def get_network_statistics(weight):
# def get_inhibitory_proportion_of_total_strength(weight):
# def get_inhibitory_neuron_indices(weight):
# def get_sparsity(matrix: np.ndarray, thres=0.0) -> float:
# def get_spectral_radius(square_matrix: np.ndarray) -> float:
# 
# Path: reservoir_nn/utils/weight_transforms.py
# class WeightTransformation(abc.ABC):
# class Chain(WeightTransformation, list):
# class FlipSignEntireRows(WeightTransformation):
# class ScaleWeightByDistance(WeightTransformation):
# class RandomizeNonZeroWeight(WeightTransformation):
# class RandomizeNonZeroWeightKde(WeightTransformation):
# class CutOffSmallWeightsInRows(WeightTransformation):
# class ScaleToZeroOne(WeightTransformation):
# class ScaleSpectralRadius(WeightTransformation):
# class GetSubMatrix(WeightTransformation):
# class ResizeNumRows(WeightTransformation):
# class ResizeNumColumns(WeightTransformation):
#   def apply_transform(self, *args: np.ndarray) -> np.ndarray:
#   def __call__(self, *args: np.ndarray) -> np.ndarray:
#   def batch_apply(
#       self, *args: Union[np.ndarray, Tuple[np.ndarray,
#                                            ...]]) -> List[np.ndarray]:
#   def __init__(self,
#                transformations: Optional[Sequence[WeightTransformation]] = None
#               ):
#   def __repr__(self):
#   def apply_transform(self, *args: np.ndarray) -> np.ndarray:
#   def __init__(self, flip_proportion: float):
#   def apply_transform(self, matrix: np.ndarray, seed: int = -1) -> np.ndarray:
#   def __init__(self, scaling_function, box_size=None):
#   def apply_transform(self, weights: np.ndarray,
#                       distances: np.ndarray) -> np.ndarray:
#   def __init__(self):
#   def apply_transform(self, matrix: np.ndarray) -> np.ndarray:
#   def __init__(self, rng: np.random.RandomState = np.random):
#   def apply_transform(self, matrix: np.ndarray, seed: int = -1) -> np.ndarray:
#   def __init__(self, non_zeros_per_row_limit: int):
#   def apply_transform(self, matrix: np.ndarray) -> np.ndarray:
#   def __init__(self):
#   def apply_transform(self, matrix: np.ndarray) -> np.ndarray:
#   def __init__(self, spectral_radius: float = 1.0):
#   def apply_transform(self, matrix: np.ndarray) -> np.ndarray:
#   def __init__(self, num_neurons: int):
#   def apply_transform(self, matrix: np.ndarray) -> np.ndarray:
#   def __init__(self,
#                target_num_rows: int,
#                rng: np.random.RandomState = np.random):
#   def apply_transform(self, matrix: np.ndarray) -> np.ndarray:
#   def __init__(self,
#                target_num_columns: int,
#                rng: np.random.RandomState = np.random):
#   def apply_transform(self, matrix: np.ndarray) -> np.ndarray:
# def flip_sign_entire_rows(matrix: np.ndarray,
#                           flip_proportion: float) -> np.ndarray:
# def randomize_non_zero_weights_kde(matrix: np.ndarray) -> np.ndarray:
# def randomize_non_zero_weights(matrix: np.ndarray) -> np.ndarray:
# def shuffle_weights(matrix: np.ndarray) -> np.ndarray:
# def assign_random_signs(matrix: np.ndarray,
#                         inhibitory_proportion: float) -> np.ndarray:
# def scale_weight_by_distance(weights,
#                              distances,
#                              scaling_function,
#                              box_size=None):
# def make_sparse(weight_array: np.ndarray,
#                 zero_weight_proportion: float) -> np.ndarray:
# def cutoff_small_weights_in_rows(matrix: np.ndarray,
#                                  non_zeros_per_row_limit: int) -> np.ndarray:
# def transform_weight(weight_matrix: np.ndarray, distance_matrix: np.ndarray,
#                      params: ml_collections.ConfigDict) -> np.ndarray:
# def setup_weight_transformation(params: ml_collections.ConfigDict) -> Chain:
# def chain_weight_matrices(matrices: List[np.ndarray]) -> List[np.ndarray]:
# def resize_weight_matrices(
#     reservoir_weights: List[np.ndarray],
#     reservoirs_num_neurons: Tuple[int, ...],
# ) -> List[np.ndarray]:


# Path: reservoir_nn/utils/weight_transforms_test.py
from absl.testing import absltest
from absl.testing import parameterized
from reservoir_nn.utils import weight_properties
from reservoir_nn.utils import weight_transforms
import attr
import ml_collections
import numpy as np


  def test_setup_weight_transformation_returns_expected_list(self):
    params = ml_collections.ConfigDict()
    params.box_size = 1
    params.distance_scaling_function = 'linear'
    params.num_cutoff = 10
    params.random_weights = True
    params.kde_random_weights = False
    params.inhibitory_neurons_proportion = 0.2
    transform_list = weight_transforms.setup_weight_transformation(params)
    self.assertIsInstance(transform_list[0],
                          weight_transforms.ScaleWeightByDistance)
    self.assertIsInstance(transform_list[1],
                          weight_transforms.CutOffSmallWeightsInRows)
    self.assertIsInstance(transform_list[2],
                          weight_transforms.RandomizeNonZeroWeight)
    self.assertIsInstance(transform_list[3], weight_transforms.ScaleToZeroOne)
    self.assertIsInstance(transform_list[4],
                          weight_transforms.FlipSignEntireRows)


class ScaleSpectralRadiusTest(absltest.TestCase):

  def test_it_should_produce_matrix_of_correct_spectral_radius(self):
    # matrix: spectral_radius=6
    matrix = np.array([[1, 2, 3], [3, 2, 1], [2, 1, 3]])
    expected = 4.2
    # Use the expected as the parameter to ScaleSpectralRadius
    scaled = weight_transforms.ScaleSpectralRadius(
        spectral_radius=expected).apply_transform(matrix)

    result = weight_properties.get_spectral_radius(scaled)