from __future__ import annotations

from dataclasses import dataclass


@dataclass(frozen=True)
class HopfieldGridShape:
    row_count: int
    column_count: int

    @property
    def size(self) -> int:
        return self.row_count * self.column_count


def grid_index(row_index: int, column_index: int, shape: HopfieldGridShape) -> int:
    if not (0 <= row_index < shape.row_count):
        raise ValueError(f"row_index {row_index} out of range")
    if not (0 <= column_index < shape.column_count):
        raise ValueError(f"column_index {column_index} out of range")
    return shape.row_count * column_index + row_index


def zero_weight_matrix(shape: HopfieldGridShape) -> tuple[tuple[float, ...], ...]:
    return tuple(tuple(0.0 for _ in range(shape.size)) for _ in range(shape.size))


def accumulate_sequence_transitions(
    shape: HopfieldGridShape,
    sequences: tuple[tuple[int, ...], ...] | list[tuple[int, ...]] | list[list[int]],
    *,
    transition_offsets: tuple[int, ...] = (1,),
    weight_increment: float = 1.0,
    one_based_rows: bool = True,
) -> tuple[tuple[float, ...], ...]:
    weights = [list(row) for row in zero_weight_matrix(shape)]
    for raw_sequence in sequences:
        sequence = tuple(int(value) for value in raw_sequence)
        for offset in transition_offsets:
            if offset <= 0:
                raise ValueError("transition offsets must be positive")
            for start in range(len(sequence) - offset):
                left_row = sequence[start] - 1 if one_based_rows else sequence[start]
                right_row = sequence[start + offset] - 1 if one_based_rows else sequence[start + offset]
                left_column = start
                right_column = start + offset
                if not (0 <= left_column < shape.column_count and 0 <= right_column < shape.column_count):
                    continue
                left_index = grid_index(left_row, left_column, shape)
                right_index = grid_index(right_row, right_column, shape)
                weights[left_index][right_index] += weight_increment
                weights[right_index][left_index] = weights[left_index][right_index]
    return tuple(tuple(row) for row in weights)


def apply_grid_inhibition(
    weight_matrix: tuple[tuple[float, ...], ...],
    shape: HopfieldGridShape,
    *,
    row_inhibition: float = 0.0,
    column_inhibition: float = 0.0,
) -> tuple[tuple[float, ...], ...]:
    weights = [list(row) for row in weight_matrix]
    if len(weights) != shape.size or any(len(row) != shape.size for row in weights):
        raise ValueError("weight_matrix shape does not match grid shape")

    for row_index in range(shape.row_count):
        for column_index in range(shape.column_count):
            current_index = grid_index(row_index, column_index, shape)
            for other_row in range(shape.row_count):
                if other_row == row_index:
                    continue
                other_index = grid_index(other_row, column_index, shape)
                if other_index <= current_index:
                    continue
                weights[current_index][other_index] += column_inhibition
                weights[other_index][current_index] = weights[current_index][other_index]
            for other_column in range(shape.column_count):
                if other_column == column_index:
                    continue
                other_index = grid_index(row_index, other_column, shape)
                if other_index <= current_index:
                    continue
                weights[current_index][other_index] += row_inhibition
                weights[other_index][current_index] = weights[current_index][other_index]
    return tuple(tuple(row) for row in weights)


def clear_diagonal(
    weight_matrix: tuple[tuple[float, ...], ...],
) -> tuple[tuple[float, ...], ...]:
    weights = [list(row) for row in weight_matrix]
    if any(len(row) != len(weights) for row in weights):
        raise ValueError("weight_matrix must be square")
    for index in range(len(weights)):
        weights[index][index] = 0.0
    return tuple(tuple(row) for row in weights)
