from functools import partial
from types import FunctionType
from typing import Callable, Optional

import numpy as np
from simple_term_menu import TerminalMenu
from typing_extensions import Concatenate

from simulator import BooleanNetwork


def ShowMenu(options: list[str], title: str = "", highlight_entry: int = 0) -> int:
    menu = TerminalMenu(
        menu_entries=options,
        title=title,
        multi_select=False,
        cursor_index=highlight_entry,
    )

    selectedIndex = menu.show()

    if selectedIndex is None:
        selectedIndex = len(options) - 1

    return selectedIndex  # pyright: ignore[reportReturnType]


def Continue(title="") -> None:
    ShowMenu(title=title, options=[None, "Continue"], highlight_entry=1)
    print()


def StopApplication() -> None:
    quit()


def MainMenu() -> None:
    global currentFunction

    def SelectExampleMenu() -> None:
        def bn4nodeSync() -> None:
            functions = [
                lambda x1, x2, x3, x4: not x2,
                lambda x1, x2, x3, x4: x1,
                lambda x1, x2, x3, x4: x1 ^ x4,
                lambda x1, x2, x3, x4: x3,
            ]
            bn = BooleanNetwork(4).UseSynchronousScheme().SetFunctions(functions)
            global currentFunction
            funcStrings = ["not x2", "x1", "x1 ^ x4", "x3"]
            currentFunction = partial(BooleanNetworkMenu, bn, functions, funcStrings)

        def bn4nodeSeq() -> None:
            functions = [
                lambda x1, x2, x3, x4: not x2,
                lambda x1, x2, x3, x4: x1,
                lambda x1, x2, x3, x4: x1 ^ x4,
                lambda x1, x2, x3, x4: x3,
            ]
            bn = (
                BooleanNetwork(4)
                .UseSequentialScheme([1, 2, 3, 4])
                .SetFunctions(functions)
            )
            global currentFunction
            funcStrings = ["not x2", "x1", "x1 ^ x4", "x3"]
            currentFunction = partial(BooleanNetworkMenu, bn, functions, funcStrings)

        def bn3nodeSync() -> None:
            functions = [
                lambda x1, x2, x3: not x3,
                lambda x1, x2, x3: not x1,
                lambda x1, x2, x3: not x2,
            ]
            bn = BooleanNetwork(3).UseSynchronousScheme().SetFunctions(functions)
            global currentFunction
            funcStrings = ["not x3", "not x1", "not x2"]
            currentFunction = partial(BooleanNetworkMenu, bn, functions, funcStrings)

        def bn3nodeSeq() -> None:
            functions = [
                lambda x1, x2, x3: not x3,
                lambda x1, x2, x3: not x1,
                lambda x1, x2, x3: not x2,
            ]
            bn = (
                BooleanNetwork(3).UseSequentialScheme([1, 2, 3]).SetFunctions(functions)
            )
            global currentFunction
            funcStrings = ["not x3", "not x1", "not x2"]
            currentFunction = partial(BooleanNetworkMenu, bn, functions, funcStrings)

        def ReturnHelper() -> None:
            global currentFunction
            currentFunction = MainMenu

        title = "Select an example boolean network:"
        options = [
            "4-node network with synchronous update",
            "4-node network with sequential update",
            "3-node repressilator network with synchronous update",
            "3-node repressilator network with sequential update",
            "Return to main menu",
        ]
        actions = [bn4nodeSync, bn4nodeSeq, bn3nodeSync, bn3nodeSeq, ReturnHelper]

        actions[ShowMenu(options=options, title=title)]()

    options = [
        "Set up a new boolean network",
        "Set up a boolean network from existing ones",
        "Quit",
    ]

    actions: list[Callable[[], None]] = [
        BooleanNetworkMenu,
        SelectExampleMenu,
        StopApplication,
    ]

    selectedIndex = ShowMenu(
        options, title="Boolean Network Simulator\nby Tom Zuidberg"
    )

    currentFunction = actions[selectedIndex]


def BooleanNetworkMenu(
    bn: Optional[BooleanNetwork] = None,
    funcs: Optional[list[Callable[Concatenate[bool, ...], bool]]] = None,
    funcStrings: Optional[list[str]] = None,
) -> None:
    global currentFunction

    boolNetwork: BooleanNetwork = BooleanNetwork(3)
    writeToFile: bool = False
    done: bool = False
    functions: list = [None for _ in range(3)]
    functionStrings: list = [None for _ in range(3)]
    functionsDirtyFlag = False
    current_highlight = 0
    functions_highlight = 0
    probFunctionStrings: list[list[str]] = [[] for _ in range(3)]

    if bn is not None and funcStrings is not None and funcs is not None:
        boolNetwork = bn
        functions = funcs
        functionStrings = funcStrings
        probFunctionStrings = [[] for _ in range(boolNetwork.size)]

    def Menu() -> None:
        nonlocal current_highlight
        title: str = f"Boolean Network: {boolNetwork} (timeStep | state)"
        options = [
            f"Set size (current: {boolNetwork.size})  WARNING: this will reset all other options!",
            f"Set state (current: {str(boolNetwork)[-boolNetwork.size :]})",
            f"Set update scheme (current: {boolNetwork.updateScheme})",
            "Edit update functions",
            None,
            "Update once (hold ENTER for continuous updates)",
            "Update multiple times",
            f"Write to file (works only with multi-update) (current: {writeToFile})",
            "Get Markov-Chain matrix (resets timeStep and state)",
            "Get stationairy distribution (resets timeStep and state)",
            None,
            "Return to main menu",
        ]

        actions = [
            SetSizeHelper,
            SetStateHelper,
            SetUpdateSchemeHelper,
            FunctionsMenu,
            None,
            UpdateHelper,
            MultiUpdateHelper,
            ToggleWriteToFileHelper,
            GetMarkovChainHelper,
            GetStableDistrHelper,
            None,
            ReturnHelper,
        ]

        if boolNetwork.updateScheme in ("synchronous", "sequential", None):
            options.pop(-3)
            actions.pop(-3)

        selectedIndex = ShowMenu(
            options=options, title=title, highlight_entry=current_highlight
        )
        current_highlight = selectedIndex
        actions[selectedIndex]()

    def SetSizeHelper() -> None:
        nonlocal boolNetwork, functions, functionStrings, probFunctionStrings
        while True:
            size = input("Set new size (Leave empty to cancel):\n").strip()
            if size == "":
                print("Cancelled")
                return
            try:
                size = int(size)
            except ValueError:
                print("Please enter only integers.")
                continue

            boolNetwork = BooleanNetwork(size)
            functions = [None for _ in range(size)]
            functionStrings = [None for _ in range(size)]
            probFunctionStrings = [[] for _ in range(size)]
            return

    def SetStateHelper() -> None:
        nonlocal boolNetwork
        while True:
            state = input(
                "Set new state (Leave empty to cancel).\nAccepted format example: '01001'\n"
            ).strip()
            if state == "":
                print("Cancelled")
                return
            try:
                boolNetwork.SetState(state)
                return
            except AssertionError as e:
                print("Invalid input:", e)

    def SetUpdateSchemeHelper() -> None:
        nonlocal boolNetwork, functions, functionStrings, functionsDirtyFlag

        def AdoptFromProbabilistic() -> None:
            nonlocal functionsDirtyFlag

            for i in range(boolNetwork.size):
                if probFunctionStrings[i]:
                    funcString = probFunctionStrings[i][0]
                    func = eval(
                        "lambda "
                        + ",".join(f"x{j}" for j in range(1, boolNetwork.size + 1))
                        + ":"
                        + funcString
                    )
                    functions[i] = func
                    functionStrings[i] = funcString
            functionsDirtyFlag = True

        def SwitchToSynchronous() -> None:
            wasProbabilistic = boolNetwork.updateScheme == "probabilistic"
            boolNetwork.UseSynchronousScheme()
            if wasProbabilistic:
                AdoptFromProbabilistic()

        def SetSequentialHelper() -> None:
            wasProbabilistic = boolNetwork.updateScheme == "probabilistic"
            while True:
                seq = input(
                    "Set sequence. (Leave empty to cancel)\nFormat example: '1,4,3,2'\n"
                ).strip()
                if seq == "":
                    print("Cancelled")
                    return

                seq = seq.split(",")
                try:
                    seq = [int(i) for i in seq]
                    boolNetwork.UseSequentialScheme(seq)
                    if wasProbabilistic:
                        AdoptFromProbabilistic()
                    return
                except Exception as e:
                    print("Invalid input:", e)

        def SwitchToAsyncRandom() -> None:
            wasProbabilistic = boolNetwork.updateScheme == "probabilistic"
            boolNetwork.UseAsynchronousRandomScheme()
            if wasProbabilistic:
                AdoptFromProbabilistic()

        def SwitchToProbabilistic() -> None:
            for i in range(boolNetwork.size):
                if not probFunctionStrings[i] and functions[i] is not None:
                    boolNetwork.AddProbabilisticFunction(i, functions[i], 1.0)
                    probFunctionStrings[i].append(functionStrings[i])
            boolNetwork.UseProbabilisticScheme()

        title = "Select update scheme:"
        options = [
            "Synchronous update scheme",
            "Sequential update scheme",
            "Probabilistic update scheme",
            "Asynchronous random update scheme",
            "Cancel",
        ]
        actions = [
            SwitchToSynchronous,
            SetSequentialHelper,
            SwitchToProbabilistic,
            SwitchToAsyncRandom,
            lambda: None,
        ]

        actions[ShowMenu(options=options, title=title)]()

    def FunctionsMenu() -> None:
        nonlocal functions_highlight
        done: bool = False

        def Menu() -> None:
            nonlocal functions_highlight
            title = "Edit update functions:"
            deleteTargets: list = []

            if boolNetwork.updateScheme in ("synchronous", "sequential", None):
                options = [
                    *[
                        f"Set update function of node x{i} (current: {functionStrings[i - 1]})"
                        for i in range(1, boolNetwork.size + 1)
                    ],
                    None,
                    "Return",
                ]
                actions = [
                    *[partial(SetFunctionHelper, i) for i in range(boolNetwork.size)],
                    None,
                    ReturnHelper,
                ]
            elif boolNetwork.updateScheme == "asynchronous_random":
                options = [
                    *[
                        f"Set update function of node x{i} (current: {functionStrings[i - 1]})"
                        for i in range(1, boolNetwork.size + 1)
                    ],
                    None,
                    *[
                        f"Set selection weight of node x{i} (current: {boolNetwork.node_selection_weights[i - 1]})"
                        for i in range(1, boolNetwork.size + 1)
                    ],
                    None,
                    "Return",
                ]
                actions = [
                    *[partial(SetFunctionHelper, i) for i in range(boolNetwork.size)],
                    None,
                    *[partial(SetNodeWeightHelper, i) for i in range(boolNetwork.size)],
                    None,
                    ReturnHelper,
                ]
            else:  # probabilistic
                title = "Edit update functions:\n(highlight a function entry and press 'd' to delete it)"
                options = []
                actions = []
                for i in range(boolNetwork.size):
                    for j in range(len(probFunctionStrings[i])):
                        options.append(
                            f"Set update function #{j + 1} of node x{i + 1} (current: {probFunctionStrings[i][j]})"
                        )
                        actions.append(partial(EditProbFunctionHelper, i, j))
                        deleteTargets.append((i, j))
                        options.append(
                            f"Set weight of function #{j + 1} of node x{i + 1} (current: {boolNetwork.probabilistic_weights[i][j]})"
                        )
                        actions.append(partial(EditProbWeightHelper, i, j))
                        deleteTargets.append(None)
                    options.append(f"Add new function for node x{i + 1}")
                    actions.append(partial(AddProbFunctionHelper, i))
                    deleteTargets.append(None)
                    options.append(None)
                    actions.append(None)
                    deleteTargets.append(None)
                options.append("Return")
                actions.append(ReturnHelper)
                deleteTargets.append(None)

            if boolNetwork.updateScheme == "probabilistic":
                menu = TerminalMenu(
                    menu_entries=options,
                    title=title,
                    multi_select=False,
                    cursor_index=functions_highlight,
                    accept_keys=("enter", "d"),
                )
                selectedIndex = menu.show()
                if selectedIndex is None:
                    selectedIndex = len(options) - 1
                    acceptKey = "enter"
                else:
                    acceptKey = menu.chosen_accept_key
                functions_highlight = selectedIndex

                if acceptKey == "d" and deleteTargets[selectedIndex] is not None:
                    DeleteProbFunctionHelper(*deleteTargets[selectedIndex])
                else:
                    actions[selectedIndex]()
                return

            selectedIndex = ShowMenu(
                options=options, title=title, highlight_entry=functions_highlight
            )
            functions_highlight = selectedIndex
            actions[selectedIndex]()

        def ReturnHelper() -> None:
            nonlocal done
            done = True

        while not done:
            Menu()

    def UpdateHelper() -> None:
        nonlocal functions, boolNetwork, functionsDirtyFlag
        if functionsDirtyFlag:
            try:
                boolNetwork.SetFunctions(functions)
                functionsDirtyFlag = False
            except Exception as e:
                print("Error while updating the functions:", e)
                return

        try:
            boolNetwork.Update()
        except Exception as e:
            print("Error while simulating once:", e)

    def MultiUpdateHelper() -> None:
        nonlocal boolNetwork, writeToFile, functions, functionsDirtyFlag
        if functionsDirtyFlag:
            try:
                boolNetwork.SetFunctions(functions)
                functionsDirtyFlag = False
            except Exception as e:
                print("Error while updating functions:", e)

        n = 0
        while True:
            n = input(
                "Enter amount of time steps to be simulated. (Leave empty to cancel)\n"
            ).strip()
            if n == "":
                print("Cancelled")
                return

            try:
                n = int(n)
                if n < 0:
                    raise ValueError
                break
            except ValueError:
                print("Invalid input")
        try:
            boolNetwork.Update(
                n, verbose=not writeToFile or n < 1000, writeToFile=writeToFile
            )
            Continue("Writing to file completed") if writeToFile else Continue()
        except Exception as e:
            print("Error while simulating:", e)

    def GetMarkovChainHelper() -> None:
        nonlocal boolNetwork
        print(boolNetwork.GetTransitionMatrix())
        Continue()
        boolNetwork.SetState([False for _ in range(boolNetwork.size)])

    def GetStableDistrHelper() -> None:
        print(boolNetwork.GetStableProbabilityDistribution())
        Continue()
        boolNetwork.SetState([False for _ in range(boolNetwork.size)])

    def SetFunctionHelper(index: int) -> None:
        nonlocal boolNetwork, functions, functionStrings, functionsDirtyFlag
        while True:
            funcString = input(
                f"Set new function for node x{index + 1}. The function will receive all nodes in form of x1, x2, ..., x[size]. (Leave empty to cancel)\n"
            )
            if funcString == "":
                print("Cancelled")
                return
            try:
                func = (
                    "lambda "
                    + ",".join(f"x{i}" for i in range(1, boolNetwork.size + 1))
                    + ":"
                    + funcString
                )
                func = eval(func)
                if not isinstance(func, FunctionType):
                    print("Please enter a valid function. Got: " + func)
                    continue
                functions[index] = func
                functionStrings[index] = funcString
                functionsDirtyFlag = True
                return
            except Exception as e:
                print("Error while parsing function:", e)

    def SetNodeWeightHelper(index: int) -> None:
        nonlocal boolNetwork
        while True:
            weight = input(
                f"Set new selection weight for node x{index + 1}. Must be a positive number. (Leave empty to cancel)\n"
            ).strip()
            if weight == "":
                print("Cancelled")
                return
            try:
                weight = float(weight)
                boolNetwork.SetNodeSelectionWeight(index, weight)
                return
            except Exception as e:
                print("Invalid input:", e)

    def AddProbFunctionHelper(index: int) -> None:
        nonlocal boolNetwork, probFunctionStrings
        while True:
            funcString = input(
                f"Set new function for node x{index + 1}. The function will receive all nodes in form of x1, x2, ..., x[size]. (Leave empty to cancel)\n"
            )
            if funcString == "":
                print("Cancelled")
                return
            try:
                func = (
                    "lambda "
                    + ",".join(f"x{i}" for i in range(1, boolNetwork.size + 1))
                    + ":"
                    + funcString
                )
                func = eval(func)
                if not isinstance(func, FunctionType):
                    print("Please enter a valid function. Got: " + func)
                    continue
            except Exception as e:
                print("Error while parsing function:", e)
                continue

            weight = 1.0
            weightInput = input(
                "Set weight for this function. Must be a positive number. (Leave empty for default 1.0)\n"
            ).strip()
            if weightInput != "":
                try:
                    weight = float(weightInput)
                except ValueError:
                    print("Invalid weight, using default 1.0")
                    weight = 1.0

            try:
                boolNetwork.AddProbabilisticFunction(index, func, weight)
                probFunctionStrings[index].append(funcString)
                return
            except Exception as e:
                print("Error while adding function:", e)

    def EditProbFunctionHelper(index: int, function_index: int) -> None:
        nonlocal boolNetwork, probFunctionStrings
        while True:
            funcString = input(
                f"Set new function for node x{index + 1}, function #{function_index + 1}. (Leave empty to cancel)\n"
            )
            if funcString == "":
                print("Cancelled")
                return
            try:
                func = (
                    "lambda "
                    + ",".join(f"x{i}" for i in range(1, boolNetwork.size + 1))
                    + ":"
                    + funcString
                )
                func = eval(func)
                if not isinstance(func, FunctionType):
                    print("Please enter a valid function. Got: " + func)
                    continue
                boolNetwork.SetProbabilisticFunction(index, function_index, func)
                probFunctionStrings[index][function_index] = funcString
                return
            except Exception as e:
                print("Error while parsing function:", e)

    def EditProbWeightHelper(index: int, function_index: int) -> None:
        nonlocal boolNetwork
        while True:
            weight = input(
                f"Set new weight for node x{index + 1}, function #{function_index + 1}. Must be a positive number. (Leave empty to cancel)\n"
            ).strip()
            if weight == "":
                print("Cancelled")
                return
            try:
                weight = float(weight)
                boolNetwork.SetProbabilisticFunctionWeight(index, function_index, weight)
                return
            except Exception as e:
                print("Invalid input:", e)

    def DeleteProbFunctionHelper(index: int, function_index: int) -> None:
        nonlocal boolNetwork, probFunctionStrings
        try:
            boolNetwork.RemoveProbabilisticFunction(index, function_index)
            del probFunctionStrings[index][function_index]
        except Exception as e:
            print("Error while deleting function:", e)
            Continue()

    def ToggleWriteToFileHelper():
        nonlocal writeToFile
        writeToFile = not writeToFile

    def ReturnHelper():
        nonlocal done
        done = True

    while not done:
        Menu()

    currentFunction = MainMenu


def main() -> None:
    global currentFunction
    currentFunction = MainMenu
    np.set_printoptions(precision=5, linewidth=300, sign=" ")
    while True:
        currentFunction()


if __name__ == "__main__":
    main()