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

import numpy as np
import scipy.linalg


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_flip_chance = False
        self.flip_chance: float = 0
        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)
        ]

    def SetFunctions(
        self, functions: Iterable[Callable[Concatenate[bool, ...], bool]]
    ) -> Self:
        def wrapper(
            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

        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] = wrapper(func)

        self.__has_update_functions = True
        return self

    def SetFunction(
        self, index: int, function: Callable[Concatenate[bool, ...], bool]
    ) -> Self:
        def wrapper(
            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

        assert 0 <= index < self.size, (
            f"Function error: cannot set function at index {index} - out of bound."
        )
        self.functions[index] = wrapper(function)
        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, flip_chance: float) -> Self:
        if flip_chance is not None:
            assert type(flip_chance) is float, (
                f"Probabilistic error: given flip_chance is not a float: got {flip_chance}"
            )
            self.flip_chance = flip_chance

        self.updateScheme = "probabilistic"
        self.__has_update_scheme = True
        self.__has_flip_chance = 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.randrange(0, self.size)
        self.nodes[index] = self.functions[index](*self.nodes)

    def __probabilistic_update(self) -> None:
        self.__synchronous_update()
        for i in range(self.size):
            rng = random.random()
            if rng <= self.flip_chance:
                self.nodes[i] = not self.nodes[i]

    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.__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_flip_chance, "Update error: no flip_chance defined"
                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":
            flipChance = self.flip_chance
            self.UseSynchronousScheme()
            for i, state in enumerate(product((False, True), repeat=self.size)):
                self.SetState(state)
                self.Update()
                for flips in product((False, True), repeat=self.size):
                    flipped = int(
                        "".join(
                            str(
                                int(
                                    self.nodes[j] if not flips[j] else not self.nodes[j]
                                )
                            )
                            for j in range(self.size)
                        ),
                        2,
                    )
                    prob = np.float64(1)
                    for flip in flips:
                        prob *= flipChance if flip else 1 - flipChance
                    matrix[i][flipped] = prob
            self.UseProbabilisticScheme(flipChance)
            return matrix

        if self.updateScheme == "asynchronous_random":
            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(1) / self.size
            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()
