import torch


def cca_state_generator(initial_state: torch.Tensor,
                        adjacency_matrix: torch.Tensor,
                        iterations: int):
    """
    A continuous cellular automaton used to generated patterns of target states on graph neighborhs.
    Used in CTI-LTI experiments.
    :param initial_state: The initial state vector to evolve.
    :param adjacency_matrix: The graph adjacency matrix of shape `[n,n]`
    :param iterations: Number of iterations for evolution. The more iterations the more similar
    the target states will be across all graph nodes.
    :return: The target state tensor, that has same dimensionality as `initial_state`.
    """
    x = initial_state.unsqueeze(0)
    for i in range(iterations):
        y = adjacency_matrix * x
        y_max = y.max(-1).values
        y_min = y.min(-1).values
        to_add = y_max
        bindex = y_max.abs() < y_min.abs()
        to_add[bindex] = y_min[bindex]
        x = x * (1 - 0.05) + to_add * 0.05
    return x