from modeci_mdf.mdf import Model, Graph, Node, Parameter, OutputPort, InputPort, Edge
from models.execute import EvaluableGraph
import os
import sys
import IPython
from modeci_mdf.functions.standard import add_mdf_function, STANDARD_ARG_0

graphviz_bin_dir = os.path.join(sys.prefix, "bin")
os.environ["PATH"] += os.pathsep + graphviz_bin_dir

add_mdf_function(
    "arctan",
    description="Inverse tangent function",
    arguments=[STANDARD_ARG_0, "scale"],
    expression_string=f"scale * arctan({STANDARD_ARG_0})",
)

add_mdf_function(
    "arcsin",
    description="Inverse sine function",
    arguments=[STANDARD_ARG_0, "scale"],
    expression_string=f"scale * arcsin({STANDARD_ARG_0})",
)

add_mdf_function(
    "arccos",
    description="Inverse cosine function",
    arguments=[STANDARD_ARG_0, "scale"],
    expression_string=f"scale * arccos({STANDARD_ARG_0})",
)

import numpy as np




# Define the model creation function
def main(input_value=1.0, mode=None):
    mod = Model(id="TrigFunctionsChain")
    mod_graph = Graph(id="arc_functions")
    mod.graphs.append(mod_graph)

    # arctan_node
    arctan_node = Node(id="arctan_node")
    arctan_node.parameters.append(Parameter(id="input_value", value=input_value))
    arctan_node.parameters.append(Parameter(id="arctan_result", function="arctan", args={"variable0": "input_value", "scale": 0.4}))
    arctan_node.output_ports.append(OutputPort(id="output", value="arctan_result"))
    mod_graph.nodes.append(arctan_node)

    # arccos_node
    arccos_node = Node(id="arccos_node")
    arccos_node.input_ports.append(InputPort(id="input_from_arctan"))
    arccos_node.parameters.append(Parameter(id="arccos_result", function="arccos", args={"variable0": "input_from_arctan", "scale": 0.3}))
    arccos_node.output_ports.append(OutputPort(id="output", value="arccos_result"))
    mod_graph.nodes.append(arccos_node)

    # arcsin_node
    arcsin_node = Node(id="arcsin_node")
    arcsin_node.input_ports.append(InputPort(id="input_from_arccos"))
    arcsin_node.parameters.append(Parameter(id="arcsin_result", function="arcsin", args={"variable0": "input_from_arccos", "scale": 0.2}))
    arcsin_node.output_ports.append(OutputPort(id="output", value="arcsin_result"))
    mod_graph.nodes.append(arcsin_node)

    # Define connections
    mod_graph.edges.append(
        Edge(id="h", sender="arctan_node", sender_port="output", receiver="arccos_node", receiver_port="input_from_arctan")
    )
    mod_graph.edges.append(
        Edge(id="i", sender="arccos_node", sender_port="output", receiver="arcsin_node", receiver_port="input_from_arccos")
    )


    if mode=="run":
        eg = EvaluableGraph(mod_graph, verbose=False)
        eg.evaluate()

        return {
            "arctan_node": eg.enodes["arctan_node"].evaluable_outputs["output"].curr_value,
            "arccos_node": eg.enodes["arccos_node"].evaluable_outputs["output"].curr_value,
            "arcsin_node": eg.enodes["arcsin_node"].evaluable_outputs["output"].curr_value,
        }
    elif mode=="graph":
        image_path = "arc_functions.png"
        mod.to_graph_image(
                engine="dot",
                output_format="png",
                view_on_render=False,
                level=3,
                filename_root="arc_functions",
            )

        return image_path



    return mod_graph

if __name__ == "__main__":
    if len(sys.argv) > 1:
        # Assuming the second argument is the mode (e.g., '-run' or '-graph')
        mode_arg = sys.argv[1]
        if mode_arg == "-run":
            main(mode="run")
        elif mode_arg == "-graph":
            main(mode="graph")
        else:
            print("Invalid argument. Please use '-run' or '-graph'.")
    else:
        print("No arguments provided. Please specify '-run' or '-graph'.")