from __future__ import annotations

from dataclasses import dataclass
import json
import math


@dataclass(frozen=True)
class HopfieldParams:
    epsilon: float = 0.005
    resistance_scale: float = 3.5
    capacitance_scale: float = 10.0
    weight_scale: float = 1.0
    input_scale: float = 1.0
    iteration_scale: float = 1.0
    global_resistance: float = 1.0
    global_capacitance: float = 1.0


@dataclass(frozen=True)
class HopfieldNetworkState:
    activations: tuple[tuple[float, ...], ...]
    outputs: tuple[tuple[float, ...], ...]
    external_inputs: tuple[tuple[float, ...], ...]


@dataclass(frozen=True)
class HopfieldRunResult:
    state: HopfieldNetworkState
    iterations: int


def tanh_clamped(value: float, exp_max: float = 80.0) -> float:
    value = max(min(value, exp_max), -exp_max)
    return (math.exp(value) - math.exp(-value)) / (math.exp(value) + math.exp(-value))


class HopfieldNetwork:
    def __init__(
        self,
        *,
        weight_matrix: tuple[tuple[float, ...], ...],
        params: HopfieldParams | None = None,
    ) -> None:
        self.weight_matrix = weight_matrix
        self.params = params or HopfieldParams()

    def run(
        self,
        external_inputs: tuple[tuple[float, ...], ...],
        *,
        initial_activations: tuple[tuple[float, ...], ...] | None = None,
    ) -> HopfieldRunResult:
        row_count = len(external_inputs)
        if row_count == 0:
            raise ValueError("external_inputs cannot be empty")
        column_count = len(external_inputs[0])
        if any(len(row) != column_count for row in external_inputs):
            raise ValueError("external_inputs rows must be the same length")
        _validate_weight_matrix(self.weight_matrix, active_size=row_count * column_count)

        base_activations = initial_activations or tuple(
            tuple(0.5 for _ in range(column_count)) for _ in range(row_count)
        )
        if len(base_activations) != row_count or any(
            len(row) != column_count for row in base_activations
        ):
            raise ValueError("initial_activations shape must match external_inputs")

        activations = [
            [list(row) for row in base_activations],
            [list(row) for row in base_activations],
        ]
        outputs = [
            [[0.0 for _ in range(column_count)] for _ in range(row_count)],
            [[0.0 for _ in range(column_count)] for _ in range(row_count)],
        ]
        inputs = [
            [list(row) for row in external_inputs],
            [list(row) for row in external_inputs],
        ]

        outputs[1][0][0] = 20.0
        time_step = 0
        _update_outputs(activations, outputs, time_step, self.params, row_count, column_count)
        iterations = 0

        while not _done(outputs, self.params.epsilon, row_count, column_count):
            time_step = time_step % 2
            next_time = (time_step + 1) % 2
            _update_outputs(
                activations, outputs, time_step, self.params, row_count, column_count
            )
            for row_index in range(row_count):
                for column_index in range(column_count):
                    delta = _delta_neuron_activation(
                        row_index=row_index,
                        column_index=column_index,
                        row_count=row_count,
                        column_count=column_count,
                        time_step=time_step,
                        activations=activations,
                        outputs=outputs,
                        inputs=inputs,
                        weight_matrix=self.weight_matrix,
                        params=self.params,
                    )
                    activations[next_time][row_index][column_index] = (
                        activations[time_step][row_index][column_index]
                        + self.params.iteration_scale * delta
                    )
            time_step += 1
            iterations += 1

        final_slot = time_step % 2
        state = HopfieldNetworkState(
            activations=tuple(tuple(row) for row in activations[final_slot]),
            outputs=tuple(tuple(row) for row in outputs[final_slot]),
            external_inputs=tuple(tuple(row) for row in external_inputs),
        )
        return HopfieldRunResult(state=state, iterations=iterations)

    def to_dict(self) -> dict[str, object]:
        return {
            "weight_matrix": [list(row) for row in self.weight_matrix],
            "params": {
                "epsilon": self.params.epsilon,
                "resistance_scale": self.params.resistance_scale,
                "capacitance_scale": self.params.capacitance_scale,
                "weight_scale": self.params.weight_scale,
                "input_scale": self.params.input_scale,
                "iteration_scale": self.params.iteration_scale,
                "global_resistance": self.params.global_resistance,
                "global_capacitance": self.params.global_capacitance,
            },
        }

    @classmethod
    def from_dict(cls, data: dict[str, object]) -> "HopfieldNetwork":
        params_data = data["params"]  # type: ignore[index]
        return cls(
            weight_matrix=tuple(
                tuple(float(value) for value in row)
                for row in data["weight_matrix"]  # type: ignore[index]
            ),
            params=HopfieldParams(
                epsilon=float(params_data["epsilon"]),  # type: ignore[index]
                resistance_scale=float(params_data["resistance_scale"]),  # type: ignore[index]
                capacitance_scale=float(params_data["capacitance_scale"]),  # type: ignore[index]
                weight_scale=float(params_data["weight_scale"]),  # type: ignore[index]
                input_scale=float(params_data["input_scale"]),  # type: ignore[index]
                iteration_scale=float(params_data["iteration_scale"]),  # type: ignore[index]
                global_resistance=float(params_data["global_resistance"]),  # type: ignore[index]
                global_capacitance=float(params_data["global_capacitance"]),  # type: ignore[index]
            ),
        )

    def save_json(self, path: str) -> None:
        with open(path, "w", encoding="utf-8") as handle:
            json.dump(self.to_dict(), handle, indent=2)

    @classmethod
    def load_json(cls, path: str) -> "HopfieldNetwork":
        with open(path, "r", encoding="utf-8") as handle:
            return cls.from_dict(json.load(handle))


def _validate_weight_matrix(
    weight_matrix: tuple[tuple[float, ...], ...],
    *,
    active_size: int,
) -> None:
    if len(weight_matrix) < active_size:
        raise ValueError(f"weight matrix needs at least {active_size} rows")
    if any(len(row) < active_size for row in weight_matrix[:active_size]):
        raise ValueError(f"weight matrix needs at least {active_size} columns")


def _update_outputs(
    activations: list[list[list[float]]],
    outputs: list[list[list[float]]],
    time_step: int,
    params: HopfieldParams,
    row_count: int,
    column_count: int,
) -> None:
    for row_index in range(row_count):
        for column_index in range(column_count):
            outputs[time_step][row_index][column_index] = 0.5 * (
                1.0
                + tanh_clamped(
                    activations[time_step][row_index][column_index]
                    / params.global_capacitance
                )
            )


def _done(
    outputs: list[list[list[float]]],
    epsilon: float,
    row_count: int,
    column_count: int,
) -> bool:
    for row_index in range(row_count):
        for column_index in range(column_count):
            if abs(outputs[0][row_index][column_index] - outputs[1][row_index][column_index]) > epsilon:
                return False
    return True


def _weight_coord(row_index: int, column_index: int, row_count: int) -> int:
    return row_count * column_index + row_index


def _delta_neuron_activation(
    *,
    row_index: int,
    column_index: int,
    row_count: int,
    column_count: int,
    time_step: int,
    activations: list[list[list[float]]],
    outputs: list[list[list[float]]],
    inputs: list[list[list[float]]],
    weight_matrix: tuple[tuple[float, ...], ...],
    params: HopfieldParams,
) -> float:
    weight_sum = 0.0
    current_index = _weight_coord(row_index, column_index, row_count)
    for other_row in range(row_count):
        for other_column in range(column_count):
            other_index = _weight_coord(other_row, other_column, row_count)
            weight_sum += (
                weight_matrix[current_index][other_index]
                * params.weight_scale
                * outputs[time_step][other_row][other_column]
            )
    activation = activations[time_step][row_index][column_index]
    neuron_input = inputs[time_step][row_index][column_index]
    numerator = (
        -(activation / (params.global_resistance * params.resistance_scale))
        + (neuron_input * params.input_scale)
        + weight_sum
    )
    return numerator / (params.global_capacitance * params.capacitance_scale)
