import inspect
import random
from itertools import product
from typing import Callable, Concatenate, Iterable, Self

import numpy as np
import scipy.linalg


def _wrap_bool_function(
    function: Callable[Concatenate[bool, ...], bool],
) -> Callable[Concatenate[bool, ...], bool]:
    def wrap(*args, **kwargs) -> bool:
        result = function(*args, **kwargs)
        assert type(result) is bool, (
            f"Function error: Boolean network functions must always return a bool, however got type {type(result)}, {result=}"
        )
        return result

    return wrap


class BooleanNetwork:
    def __init__(self, size: int) -> None:
        assert type(size) is int and size > 0, (
            f"Init error: Boolean Network must contain atleast one node. got {size=} nodes"
        )
        self.size = size
        self.__ready = False
        self.__has_update_functions = False
        self.__has_update_scheme = False
        self.__has_sequence = False
        self.__has_probabilistic_functions = False
        self.sequence: list[int] = list()
        self.seed: int | None = None
        self.time_step = 0

        self.updateScheme: None | str = None
        self.nodes: list[bool] = [False for _ in range(size)]
        self.functions: list[Callable[Concatenate[bool, ...], bool]] = [
            lambda x: x for _ in range(size)
        ]

        # asynchronous_random: relative likelihood that node i is the one picked
        # to update on a given time step. Doesn't need to sum to 1 - it is
        # normalised (weight_i / sum(weights)) whenever it is used.
        self.node_selection_weights: list[float] = [1.0 for _ in range(size)]

        # probabilistic: each node may have any number of candidate update
        # functions. On every update, one candidate per node is drawn
        # according to its weight (again normalised at use-time, not
        # required to sum to 1) and applied synchronously.
        self.probabilistic_functions: list[
            list[Callable[Concatenate[bool, ...], bool]]
        ] = [list() for _ in range(size)]
        self.probabilistic_weights: list[list[float]] = [list() for _ in range(size)]

    def SetFunctions(
        self, functions: Iterable[Callable[Concatenate[bool, ...], bool]]
    ) -> Self:
        funcs: list[Callable[Concatenate[bool, ...], bool]] = list(functions)

        assert len(funcs) == self.size, (
            f"Function error: Function amount mismatch. got {len(funcs)} functions, expected {self.size}"
        )

        for i in range(self.size):
            func = funcs[i]
            assert len(inspect.signature(func).parameters) == self.size, (
                f"Function error: Function arg amount mismatch. Given function takes {len(inspect.signature(func).parameters)} arguments, expected {self.size}"
            )

            self.functions[i] = _wrap_bool_function(func)

        self.__has_update_functions = True
        return self

    def SetFunction(
        self, index: int, function: Callable[Concatenate[bool, ...], bool]
    ) -> Self:
        assert 0 <= index < self.size, (
            f"Function error: cannot set function at index {index} - out of bound."
        )
        assert len(inspect.signature(function).parameters) == self.size, (
            f"Function error: Function arg amount mismatch. Given function takes {len(inspect.signature(function).parameters)} arguments, expected {self.size}"
        )
        self.functions[index] = _wrap_bool_function(function)
        return self

    def SetNodeSelectionWeight(self, index: int, weight: float) -> Self:
        assert 0 <= index < self.size, (
            f"Weight error: cannot set selection weight at index {index} - out of bound."
        )
        assert type(weight) is float and weight > 0, (
            f"Weight error: weight must be a positive float. got {weight=}"
        )
        self.node_selection_weights[index] = weight
        return self

    def AddProbabilisticFunction(
        self,
        index: int,
        function: Callable[Concatenate[bool, ...], bool],
        weight: float = 1.0,
    ) -> Self:
        assert 0 <= index < self.size, (
            f"Function error: cannot add function at index {index} - out of bound."
        )
        assert len(inspect.signature(function).parameters) == self.size, (
            f"Function error: Function arg amount mismatch. Given function takes {len(inspect.signature(function).parameters)} arguments, expected {self.size}"
        )
        assert type(weight) is float and weight > 0, (
            f"Weight error: weight must be a positive float. got {weight=}"
        )

        self.probabilistic_functions[index].append(_wrap_bool_function(function))
        self.probabilistic_weights[index].append(weight)
        self.__has_probabilistic_functions = all(
            len(functions) > 0 for functions in self.probabilistic_functions
        )
        return self

    def SetProbabilisticFunction(
        self,
        index: int,
        function_index: int,
        function: Callable[Concatenate[bool, ...], bool],
    ) -> Self:
        assert 0 <= index < self.size, (
            f"Function error: cannot set function at index {index} - out of bound."
        )
        assert 0 <= function_index < len(self.probabilistic_functions[index]), (
            f"Function error: node {index} has no function at position {function_index}."
        )
        assert len(inspect.signature(function).parameters) == self.size, (
            f"Function error: Function arg amount mismatch. Given function takes {len(inspect.signature(function).parameters)} arguments, expected {self.size}"
        )

        self.probabilistic_functions[index][function_index] = _wrap_bool_function(
            function
        )
        return self

    def SetProbabilisticFunctionWeight(
        self, index: int, function_index: int, weight: float
    ) -> Self:
        assert 0 <= index < self.size, (
            f"Weight error: cannot set weight at index {index} - out of bound."
        )
        assert 0 <= function_index < len(self.probabilistic_weights[index]), (
            f"Weight error: node {index} has no function at position {function_index}."
        )
        assert type(weight) is float and weight > 0, (
            f"Weight error: weight must be a positive float. got {weight=}"
        )

        self.probabilistic_weights[index][function_index] = weight
        return self

    def RemoveProbabilisticFunction(self, index: int, function_index: int) -> Self:
        assert 0 <= index < self.size, (
            f"Function error: cannot remove function at index {index} - out of bound."
        )
        assert 0 <= function_index < len(self.probabilistic_functions[index]), (
            f"Function error: node {index} has no function at position {function_index}."
        )
        assert len(self.probabilistic_functions[index]) > 1, (
            f"Function error: node {index} must keep at least one probabilistic function."
        )

        del self.probabilistic_functions[index][function_index]
        del self.probabilistic_weights[index][function_index]
        self.__has_probabilistic_functions = all(
            len(functions) > 0 for functions in self.probabilistic_functions
        )
        return self

    def UseSynchronousScheme(self) -> Self:
        self.updateScheme = "synchronous"
        self.__has_update_scheme = True

        return self

    def UseSequentialScheme(self, sequence: Iterable[int]) -> Self:
        sequence = list(sequence)
        assert len(sequence) == self.size, (
            f"Sequence error: sequence must be the same size as the nodes of the network: sequence '{sequence}', #nodes={self.size}"
        )
        assert all(type(i) is int for i in sequence), (
            f"Sequence error: sequence must only contain integers. sequence given: {sequence}"
        )
        sorted_sequence = sorted(sequence)
        compare_to = list(range(self.size + 1))
        assert (
            sorted_sequence == compare_to[:-1] or sorted_sequence == compare_to[1:]
        ), (
            f"Sequence error: sequence doesn't contain the correct indices. It must contain all numbers from 0 to {self.size} (excluded) or from 1 to {self.size} (included)"
        )

        if sorted_sequence[0] == 1:
            for i in range(self.size):
                sequence[i] -= 1

        self.sequence = sequence
        self.updateScheme = "sequential"
        self.__has_update_scheme = True
        self.__has_sequence = True

        return self

    def UseAsynchronousRandomScheme(self, seed: int | None = None) -> Self:
        if seed is not None:
            assert type(seed) is int, (
                f"AsyncRandom error: wrong format for given seed. got {seed=}"
            )
            self.seed = seed
            random.seed(seed)

        self.updateScheme = "asynchronous_random"
        self.__has_update_scheme = True
        return self

    def UseProbabilisticScheme(self) -> Self:
        self.updateScheme = "probabilistic"
        self.__has_update_scheme = True
        return self

    def __synchronous_update(self) -> None:
        temp = list()
        for i in range(self.size):
            temp.append(self.functions[i](*self.nodes))
        self.nodes = temp

    def __sequential_update(self) -> None:
        for i in self.sequence:
            self.nodes[i] = self.functions[i](*self.nodes)

    def __asynchronous_random_update(self) -> None:
        index = random.choices(
            range(self.size), weights=self.node_selection_weights, k=1
        )[0]
        self.nodes[index] = self.functions[index](*self.nodes)

    def __probabilistic_update(self) -> None:
        temp = list()
        for i in range(self.size):
            chosen = random.choices(
                self.probabilistic_functions[i],
                weights=self.probabilistic_weights[i],
                k=1,
            )[0]
            temp.append(chosen(*self.nodes))
        self.nodes = temp

    def SetState(self, state: str | list[bool] | tuple[bool, ...]) -> Self:
        assert isinstance(state, (str, list, tuple)), (
            f"SetState error: invalid type as state"
        )
        assert len(state) == self.size, (
            f"SetState error: given state is not the same size as the boolean network. got size {len(state)}, expected {self.size}"
        )

        if type(state) is str:
            for i in range(self.size):
                self.nodes[i] = bool(int(state[i]))
            self.time_step = 0
            return self

        if isinstance(state, (list, tuple)):
            assert all(type(i) is bool for i in state), (
                f"SetState error: given state list contains elements of type different from bool. All elements must be bools. got {state}"
            )
            self.nodes = list(state).copy()
            self.time_step = 0
            return self

        raise Exception(
            "SetState error: end of function reached. given state is not of type string nor list[bool]."
        )

    def Update(
        self, n: int = 1, /, verbose: bool = False, writeToFile: bool = False
    ) -> None:
        assert type(n) is int and n >= 0, (
            f"Update error: amount of updates must be an integer and positive. got {n=}"
        )
        assert self.updateScheme == "probabilistic" or self.__has_update_functions, (
            "Update error: no update functions defined"
        )
        assert self.__has_update_scheme, "Update error: no update scheme defined"
        assert type(verbose) is bool, "Update error: verbose must be a bool"
        assert type(writeToFile) is bool, "Update error: writeToFile must be a bool"

        selected_update: Callable

        match self.updateScheme:
            case "synchronous":
                selected_update = self.__synchronous_update
            case "sequential":
                assert self.__has_sequence, "Update error: no sequence defined"
                selected_update = self.__sequential_update
            case "asynchronous_random":
                selected_update = self.__asynchronous_random_update
            case "probabilistic":
                assert self.__has_probabilistic_functions, (
                    "Update error: no probabilistic functions defined for every node"
                )
                selected_update = self.__probabilistic_update
            case _:
                raise Exception("Update error: update scheme selection went wrong")

        match (verbose, writeToFile):
            case (False, False):
                for _ in range(n):
                    selected_update()
                    self.time_step += 1
            case (False, True):
                with open("output.txt", "w") as f:
                    for _ in range(n):
                        selected_update()
                        self.time_step += 1
                        f.writelines([self.state, "\n"])
            case (True, False):
                for _ in range(n):
                    selected_update()
                    self.time_step += 1
                    print(self)
            case (True, True):
                with open("output.txt", "w") as f:
                    for _ in range(n):
                        selected_update()
                        self.time_step += 1
                        f.writelines([self.state, "\n"])
                        print(self)

    def __str__(self) -> str:
        return f"{self.time_step:>5} | {''.join(str(int(node)) for node in self.nodes)}"

    @property
    def state(self) -> str:
        return "".join(str(int(node)) for node in self.nodes)

    def GetStableProbabilityDistribution(self) -> np.ndarray:
        matrix: np.ndarray = self.GetTransitionMatrix()

        eigenvalues, eigenvectors = scipy.linalg.eig(matrix.T)

        idx = np.argmin(np.abs(eigenvalues - 1.0))
        pi = eigenvectors[:, idx].real

        pi = pi / pi.sum()

        return pi

    def GetTransitionMatrix(self) -> np.ndarray:
        dimension = 2**self.size
        matrix: np.ndarray = np.zeros((dimension, dimension))

        if self.updateScheme == "probabilistic":
            for i, state in enumerate(product((False, True), repeat=self.size)):
                choice_ranges = [
                    range(len(self.probabilistic_functions[n]))
                    for n in range(self.size)
                ]
                for combo in product(*choice_ranges):
                    prob = np.float64(1)
                    result: list[bool] = []
                    for n in range(self.size):
                        weights = self.probabilistic_weights[n]
                        total_weight = sum(weights)
                        chosen_index = combo[n]
                        prob *= weights[chosen_index] / total_weight
                        result.append(
                            self.probabilistic_functions[n][chosen_index](*state)
                        )
                    result_index = int(
                        "".join(str(int(b)) for b in result), 2
                    )
                    matrix[i][result_index] += prob
            return matrix

        if self.updateScheme == "asynchronous_random":
            total_weight = sum(self.node_selection_weights)
            for i, state in enumerate(product((False, True), repeat=self.size)):
                for j in range(self.size):
                    self.SetState(state)
                    self.nodes[j] = self.functions[j](*self.nodes)
                    matrix[i][int(self.state, 2)] += (
                        np.float64(self.node_selection_weights[j]) / total_weight
                    )
            return matrix

        for i, state in enumerate(product((False, True), repeat=self.size)):
            self.SetState(state)
            self.Update()
            matrix[i][int(self.state, 2)] = 1
        return matrix


def main() -> None:
    bn = (
        BooleanNetwork(4)
        .SetState("0000")
        .SetFunctions(
            [
                lambda a, b, c, d: not b,
                lambda a, b, c, d: a,
                lambda a, b, c, d: a ^ d,
                lambda a, b, c, d: c,
            ]
        )
        .UseSequentialScheme((1, 2, 3, 4))
    )

    print(bn)
    bn.Update()
    print(bn)

    print("update 100 times verbose")
    bn.Update(100, verbose=True)


if __name__ == "__main__":
    main()