"""
Core Stage and Curriculum Primitives.
"""

from __future__ import annotations

import inspect
import json
import subprocess
import warnings
from importlib import import_module
from pathlib import Path
from typing import Any, Callable, Dict, Generic, List, Tuple, TypeVar

import boto3
from botocore.exceptions import ClientError
from jinja2 import Template
from pydantic import Field, GetJsonSchemaHandler, field_validator
from pydantic.json_schema import JsonSchemaValue
from pydantic_core import core_schema

from aind_behavior_curriculum import (
    AindBehaviorModel,
    AindBehaviorModelExtra,
    Task,
    TaskParameters,
)

TTask = TypeVar("TTask", bound=Task)


class Metrics(AindBehaviorModelExtra):
    """
    Abstract Metrics class.
    Subclass with Metric values.
    """


class Rule:
    """
    Custom Pydantic Type that defines de/serialiation for Callables.
    """

    def __eq__(self, __value: object) -> bool:
        """
        Custom equality method.
        Two instances of the same subclass type are considered equal.
        """
        return isinstance(__value, self.__class__)

    @classmethod
    def __get_pydantic_core_schema__(
        cls,
        _source_type: Any,
        _handler: Callable[[Any], core_schema.CoreSchema],
    ) -> core_schema.CoreSchema:
        """
        Custom validation, Ref:
        https://docs.pydantic.dev/latest/concepts/types/#handling-third-party-types
        """

        def validate_from_str(value: str) -> Callable:
            """Pass string through deserialization."""
            return cls._deserialize_callable(value)

        from_str_schema = core_schema.chain_schema(
            [
                core_schema.str_schema(),
                core_schema.no_info_plain_validator_function(
                    validate_from_str
                ),
            ]
        )

        return core_schema.json_or_python_schema(
            json_schema=from_str_schema,
            python_schema=core_schema.union_schema(
                [
                    core_schema.is_instance_schema(Callable),
                    from_str_schema,
                ]
            ),
            serialization=core_schema.plain_serializer_function_ser_schema(
                function=cls._serialize_callable
            ),
        )

    @classmethod
    def __get_pydantic_json_schema__(
        cls,
        _core_schema: core_schema.CoreSchema,
        handler: GetJsonSchemaHandler,
    ) -> JsonSchemaValue:
        """
        Custom validation, Ref:
        https://docs.pydantic.dev/latest/concepts/types/#as-a-method-on-a-custom-type
        """
        return handler(core_schema.str_schema())

    @staticmethod
    def _deserialize_callable(value: str | Callable) -> Callable:
        """
        Custom Deserialization.
        Imports function according to package and function name.
        """
        if callable(value):
            return value
        else:
            split = value.rsplit(".", 1)
            assert len(split) > 0, (
                f"Invalid rule value while attempting to deserialize callable. "
                f"Got {value}, expected string in the format '<module>.Rule'"
            )

            module = import_module(split[0])
            obj = getattr(module, split[1])
            return obj

    @staticmethod
    def _serialize_callable(value: str | Callable) -> Callable:
        """
        Custom Serialization.
        Simlply exports reference to function as package + function name.
        """
        if isinstance(value, str):
            value = Rule._deserialize_callable(value)

        return value.__module__ + "." + value.__name__


class Policy(AindBehaviorModel):
    """
    User-defined function that defines
    how current Task parameters change according to metrics.
    """

    rule: Rule = Field(..., description="Callable with Serialization.")

    @field_validator("rule")
    @classmethod
    def validate_rule(cls, r: Rule):
        """
        Policy Signature:
        I:
        - metrics: Metrics object
        - task_parameters: TaskParameters object

        O:
        - result: TaskParameters object
        """
        if not callable(r):
            raise ValueError("Rule must be callable.")

        # Check rule follows Transition signature
        params = list(inspect.signature(r).parameters)
        param_1 = inspect.signature(r).parameters[params[0]].annotation
        param_2 = inspect.signature(r).parameters[params[1]].annotation
        return_type = inspect.signature(r).return_annotation

        module = import_module(param_1.__module__)
        param_1_obj = getattr(module, param_1.__name__)
        module = import_module(param_2.__module__)
        param_2_obj = getattr(module, param_2.__name__)
        module = import_module(return_type.__module__)
        return_type_obj = getattr(module, return_type.__name__)

        incorrect_num_params = len(inspect.signature(r).parameters) != 2
        incorrect_input_types = not (
            issubclass(param_1_obj, Metrics)
            and issubclass(param_2_obj, TaskParameters)
        )
        incorrect_return_type = not (
            issubclass(return_type_obj, TaskParameters)
        )

        if (
            incorrect_num_params
            or incorrect_input_types
            or incorrect_return_type
        ):
            raise ValueError(
                "Invalid signature." f"{Policy.validate_rule.__doc__}"
            )

        return r


class PolicyTransition(AindBehaviorModel):
    """
    User-defined function that defines
    criteria for transitioning between policies based on metrics.
    """

    rule: Rule = Field(..., description="Callable with Serialization.")

    @field_validator("rule")
    @classmethod
    def validate_rule(cls, r: Rule):
        """
        Policy Transition Signature:
        I:
        - metrics: Metrics object

        O:
        - result: bool
        """
        if not callable(r):
            raise ValueError("Rule must be callable.")

        # Check rule follows Transition signature
        params = list(inspect.signature(r).parameters)
        param_1 = inspect.signature(r).parameters[params[0]].annotation
        return_type = inspect.signature(r).return_annotation

        module = import_module(param_1.__module__)
        param_1_obj = getattr(module, param_1.__name__)
        module = import_module(return_type.__module__)
        return_type_obj = getattr(module, return_type.__name__)

        incorrect_num_params = len(inspect.signature(r).parameters) != 1
        incorrect_input_types = not (issubclass(param_1_obj, Metrics))
        incorrect_return_type = not (issubclass(return_type_obj, bool))

        if (
            incorrect_num_params
            or incorrect_input_types
            or incorrect_return_type
        ):
            raise ValueError(
                "Invalid signature."
                f"{PolicyTransition.validate_rule.__doc__}"
            )

        return r


NodeTypes = TypeVar("NodeTypes")
EdgeType = TypeVar("EdgeType")


class BehaviorGraph(AindBehaviorModel, Generic[NodeTypes, EdgeType]):
    """
    Core directed graph data structure used in Stage and Curriculum.
    """

    nodes: Dict[int, NodeTypes] = Field(default={}, validate_default=True)
    graph: Dict[int, List[Tuple[EdgeType, int]]] = Field(
        default={}, validate_default=True
    )

    def _get_node_id(self, node: NodeTypes) -> int:
        """
        Dictionaries are ordered for Python 3.7+ so this is safe.
        This library requires Python 3.8+.
        """
        dict_keys = list(self.nodes.keys())
        dict_values = list(self.nodes.values())

        i = dict_values.index(node)
        return dict_keys[i]

    def _create_node_id(self) -> int:
        """
        Helper method for add_node and add_transition.
        More readable than using hash(Policy).
        """
        new_id = 0
        if len(self.nodes) > 0:
            new_id = max(len(self.nodes), max(self.nodes.keys()) + 1)
        return new_id

    def add_node(self, node: NodeTypes) -> None:
        """
        Adds a floating node to the behavior graph.
        """
        p_id = self._create_node_id()
        self.nodes[p_id] = node
        self.graph[p_id] = []

    def remove_node(self, node: NodeTypes) -> None:
        """
        Removes node and all associated incoming/outgoing
        transition rules from the stage graph.
        NOTE: Removed nodes and transitions have the side effect
        of changing transition priority.
        """
        assert (
            node in self.nodes.values()
        ), f"Node {node} is not in the graph to be removed."

        # Resolve node id
        p_id = self._get_node_id(node)

        # Remove node from node list
        del self.nodes[p_id]

        # Remove node from graph keys
        del self.graph[p_id]

        # Remove node from graph value lists
        for start_id in self.graph:
            for rule, dest_id in self.graph[start_id]:
                if dest_id == p_id:
                    self.graph[start_id].remove((rule, dest_id))

    def add_transition(
        self,
        start_node: NodeTypes,
        dest_node: NodeTypes,
        rule: EdgeType,
    ) -> None:
        """
        Add transition between two nodes:
        start_node -> dest_node.

        If start_node has been added to graph before, this method starts a transition
            from the exisiting start_node.
        If dest_node has been added to graph before, this method creates a transition
            into the existing dest_node.

        NOTE: The order in which this method
        is called sets the order of transition priority.
        """

        # Resolve id of start_node
        if not (start_node in self.nodes.values()):
            new_id = self._create_node_id()
            self.nodes[new_id] = start_node
            self.graph[new_id] = []
        start_id = self._get_node_id(start_node)

        # Resolve id of dest_node
        if not (dest_node in self.nodes.values()):
            new_id = self._create_node_id()
            self.nodes[new_id] = dest_node
            self.graph[new_id] = []
        dest_id = self._get_node_id(dest_node)

        # Add the new transition to the graph
        self.graph[start_id].append((rule, dest_id))

    def remove_node_transition(
        self,
        start_node: NodeTypes,
        dest_node: NodeTypes,
        rule: EdgeType,
        remove_start_node: bool = False,
        remove_dest_node: bool = False,
    ) -> None:
        """
        Removes transition with options to remove start/end nodes
        associated with the transition.
        NOTE: Removed nodes and transitions has the side effect
        of changing transition priority.
        """

        assert (
            start_node in self.nodes.values()
        ), f"Node {start_node} is not in the behavior graph to be removed."

        assert (
            dest_node in self.nodes.values()
        ), f"Node {dest_node} is not in the behavior graph to be removed."

        start_id = self._get_node_id(start_node)
        dest_id = self._get_node_id(dest_node)
        assert (rule, dest_id) in self.graph[start_id], (
            f"Node {start_node} does not transition "
            f"into Node {dest_node} with Rule {rule}."
        )

        # Optionally remove nodes
        if remove_start_node:
            self.remove_node(start_node)
        if remove_dest_node:
            self.remove_node(dest_node)

        # Remove transition
        self.graph[start_id].remove((rule, dest_id))

    def see_nodes(self) -> List[NodeTypes]:
        """
        See nodes of behavior graph.
        """
        return list(self.nodes.values())

    def see_node_transitions(
        self, node: NodeTypes
    ) -> List[Tuple[EdgeType, NodeTypes]]:
        """
        See transitions of node in behavior graph.
        """

        assert (
            node in self.nodes.values()
        ), f"Node {node} is not in behavior graph."
        node_id = self._get_node_id(node)
        node_list = self.graph[node_id]
        node_list = [(rule, self.nodes[p_id]) for (rule, p_id) in node_list]

        return node_list

    def set_transition_priority(
        self,
        node: NodeTypes,
        node_transitions: List[Tuple[EdgeType, NodeTypes]],
    ) -> None:
        """
        Change order of node transitions listed under a node.
        Highest priority is ordered in node_transition from left -> right.
        """

        input_transitions = []
        for rule, n in node_transitions:
            assert (
                n in self.nodes.values()
            ), f"Node {n} is not a node inside the behavior graph."
            input_transitions.append((rule, self._get_node_id(n)))

        n_id = self._get_node_id(node)
        self.graph[n_id] = input_transitions


class PolicyGraph(BehaviorGraph[Policy, PolicyTransition]):
    """
    Graph for Stage.
    """

    pass


class Stage(AindBehaviorModel, Generic[TTask]):
    """
    Instance of a Task.
    Task Parameters may change according to rules defined in BehaviorGraph.
    Stage manages a BehaviorGraph instance with a read/write API.
    """

    name: str = Field(..., description="Stage name.")
    task: TTask = Field(
        ..., description="Task in which this stage is based off of."
    )
    graph: PolicyGraph = PolicyGraph()
    start_policies: List[Policy] = []

    def __eq__(self, __value: object) -> bool:
        """
        Custom equality method.
        Two Stage instances are only distinguished by name.
        """
        return self.name == __value.name

    def add_policy(self, policy: Policy) -> None:
        """
        Adds a floating policy to the Stage adjacency graph.
        """
        assert not (policy in self.graph.see_nodes()), (
            f"Policy {policy.rule.__name__} is a duplicate Policy "
            f"in Stage {self.name}. Stage cannot have duplicate policies"
        )

        self.graph.add_node(policy)

    def remove_policy(self, policy: Policy) -> None:
        """
        Removes policy and all associated incoming/outgoing
        transition rules from the stage graph.
        NOTE: Removed nodes and transitions have the side effect
        of changing transition priority.
        """
        self.graph.remove_node(policy)

        # Also remove reference to policy in start_policies if applicable.
        if policy in self.start_policies:
            self.start_policies.remove(policy)

            if len(self.start_policies) == 0:
                warnings.warn(f"Stage {self.name}.start_policies is empty.")

    def add_policy_transition(
        self,
        start_policy: Policy,
        dest_policy: Policy,
        rule: PolicyTransition,
    ) -> None:
        """
        Add policy transition between two policies:
        Policy_A -> Policy_B.

        If Policy_A has been added to stage before, this method starts a transition
            from the existing Policy_A.
        If Policy_B has been added to stage before, this method creates a transition
            into the existing Policy_B.

        NOTE: The order in which this method
        is called sets the order of transition priority.
        """
        self.graph.add_transition(start_policy, dest_policy, rule)

    def remove_policy_transition(
        self,
        start_policy: Policy,
        dest_policy: Policy,
        rule: PolicyTransition,
        remove_start_policy: bool = False,
        remove_dest_policy: bool = False,
    ) -> None:
        """
        Removes transition with options to remove start/end policies
        associated with the transition.
        NOTE: Removed nodes and transitions has the side effect
        of changing transition priority.
        """
        self.graph.remove_node_transition(
            start_policy,
            dest_policy,
            rule,
            remove_start_policy,
            remove_dest_policy,
        )

    def see_policies(self) -> List[Policy]:
        """
        See policies of policy graph.
        """
        return self.graph.see_nodes()

    def see_policy_transitions(
        self, policy: Policy
    ) -> List[Tuple[PolicyTransition, Policy]]:
        """TTask
        See transitions of stage in policy graph.
        """
        return self.graph.see_node_transitions(policy)

    def set_policy_transition_priority(
        self,
        policy: Policy,
        policy_transitions: List[Tuple[PolicyTransition, Policy]],
    ) -> None:
        """
        Change the order of policy transitions listed under a policy.
        To use, call see_policy_transitions() and order the transitions
        in the desired priority from left -> right.
        """

        policy_transitions_list = list(
            (t.rule, p.rule) for (t, p) in policy_transitions
        )
        current_list = list(
            (t.rule, p.rule) for (t, p) in self.see_policy_transitions(policy)
        )

        assert set(policy_transitions_list) == set(current_list), (
            f"Elements of input node transitions {policy_transitions} does not "
            f"match the elements under this node: {self.see_policy_transitions(policy)}."
            " Please call 'see_policy_transitions()' for a precise list of elements."
        )

        self.graph.set_transition_priority(policy, policy_transitions)

    def set_start_policies(self, start_policies: Policy | List[Policy]):
        """
        Sets stage's start policies to start policies provided.
        Input overwrites existing start policies.
        """
        if isinstance(start_policies, Policy):
            start_policies = [start_policies]
        self.start_policies = start_policies

    def get_task_parameters(self) -> TaskParameters:
        """
        See current task parameters of Task.
        """

        return self.task.task_parameters

    def set_task_parameters(self, task_params: TaskParameters) -> None:
        """
        Set task with new set of task parameters.
        Task revalidates TaskParameters on assignment.
        """
        self.task.task_parameters = task_params

    def validate_stage(self) -> Stage:
        """
        Check if stage is non-empty and specifies start policies.
        """

        assert len(self.see_policies()) > 0, (
            f"Stage {self.name} in Curriculum does not have policies. "
            "Please add at least one policy to all Curriculum stages "
            "with Stage.add_policy(...). "
            "If you would like an empty Stage, you can use "
            "curriculum_utils.create_empty_stage(...)"
        )

        assert len(self.start_policies) > 0, (
            f"Stage {self.name} in Curriculum does not have start_policies. "
            "Please define start_polices for all Curriculum stages "
            "with Stage.set_start_policies(...)"
            ""
        )

        # Check round trip serialization
        try:
            instance_json = self.model_dump_json()
            curr_subtype = type(self)
            curr_subtype.model_validate_json(instance_json)
        except Exception:
            print(
                (
                    f"Pydantic cannot serialize Stage {self.name}, please use "
                    "mypy to verify your types "
                    "(check signatures of policy functions, etc.)."
                )
            )
            raise

        return self


class StageTransition(AindBehaviorModel):
    """
    User-defined function that defines
    criteria for transitioning stages based on metrics.
    """

    rule: Rule = Field(..., description="Callable with Serialization.")

    @field_validator("rule")
    @classmethod
    def validate_rule(cls, r: Rule):
        """
        Stage Transition Signature:
        I:
        - metrics: Metrics object

        O:
        - result: bool
        """
        if not callable(r):
            raise ValueError("Rule must be callable.")

        # Check rule follows Transition signature
        params = list(inspect.signature(r).parameters)
        param_1 = inspect.signature(r).parameters[params[0]].annotation
        return_type = inspect.signature(r).return_annotation

        module = import_module(param_1.__module__)
        param_1_obj = getattr(module, param_1.__name__)
        module = import_module(return_type.__module__)
        return_type_obj = getattr(module, return_type.__name__)

        incorrect_num_params = len(inspect.signature(r).parameters) != 1
        incorrect_input_types = not (issubclass(param_1_obj, Metrics))
        incorrect_return_type = not (issubclass(return_type_obj, bool))

        if (
            incorrect_num_params
            or incorrect_input_types
            or incorrect_return_type
        ):
            raise ValueError(
                "Invalid signature." f"{StageTransition.validate_rule.__doc__}"
            )
        return r


class StageGraph(BehaviorGraph[Stage[TTask], StageTransition], Generic[TTask]):
    """
    Graph for Curriculum.
    """

    pass


class Curriculum(AindBehaviorModel):
    """
    Curriculum manages a StageGraph instance with a read/write API.
    To use, subclass this and add subclass metrics.
    """

    pkg_location: str = ""
    name: str = (
        "Please subclass, rename, and define \
                 a StageGraph with your own Stage objs \
                 Ex: StageGraph[Union[StageA, StageB, Graduated]]"
    )
    graph: StageGraph = Field(default=StageGraph(), validate_default=True)

    def model_post_init(self, __context: Any) -> None:
        """
        Add Curriculum pkg location
        """
        super().model_post_init(__context)
        self.pkg_location = self.__module__ + "." + type(self).__name__

    def add_stage(self, stage: Stage) -> None:
        """
        Adds a floating stage to the Curriculum adjacency graph.
        """
        assert not (stage in self.graph.see_nodes()), (
            f"Stage {stage.name} is a duplicate stage. "
            "Curriculum cannot have duplicate stages, "
            "please change stage name or task."
        )

        self.graph.add_node(stage)

    def remove_stage(self, stage: Stage) -> None:
        """
        Removes stage and all associated incoming/outgoing
        transition rules from the curriculum graph.
        NOTE: Removed nodes and transitions have the side effect
        of changing transition priority.
        """
        self.graph.remove_node(stage)

    def add_stage_transition(
        self,
        start_stage: Stage,
        dest_stage: Stage,
        rule: StageTransition,
    ) -> None:
        """
        Add stage transition between two stages:
        Stage_A -> Stage_B.

        If Stage_A has been added to stage before, this method starts a transition
            from the exisiting Stage_A.
        If Stage_B has been added to stage before, this method creates a transition
            into the existing Stage_B.

        NOTE: The order in which this method
        is called sets the order of transition priority.
        """

        self.graph.add_transition(start_stage, dest_stage, rule)

    def remove_stage_transition(
        self,
        start_stage: Policy,
        dest_stage: Policy,
        rule: PolicyTransition,
        remove_start_stage: bool = False,
        remove_dest_stage: bool = False,
    ) -> None:
        """
        Removes transition with options to remove start/end stages
        associated with the transition.
        NOTE: Removed nodes and transitions has the side effect
        of changing transition priority.
        """
        self.graph.remove_node_transition(
            start_stage,
            dest_stage,
            rule,
            remove_start_stage,
            remove_dest_stage,
        )

    def see_stages(self) -> List[Stage]:
        """
        See stages of curriculum graph.
        """
        return self.graph.see_nodes()

    def see_stage_transitions(
        self, stage: Stage
    ) -> List[Tuple[StageTransition, Stage]]:
        """
        See transitions of stage in curriculum graph.
        """
        return self.graph.see_node_transitions(stage)

    def set_stage_transition_priority(
        self,
        stage: Stage[TTask],
        stage_transitions: List[Tuple[StageTransition, Stage]],
    ) -> None:
        """
        Change the order of stage transitions listed under a stage.
        To use, call see_stage_transitions() and order the transitions
        in the desired priority from left -> right.
        """

        stage_transitions_list = list(
            (t.rule, s.name) for (t, s) in stage_transitions
        )
        current_list = list(
            (t.rule, s.name) for (t, s) in self.see_stage_transitions(stage)
        )

        assert set(stage_transitions_list) == set(current_list), (
            f"Elements of input node transitions {stage_transitions} does not "
            f"match the elements under this node: {self.see_stage_transitions(stage)}. "
            "Please call 'see_stage_transitions()' for a precise list of elements."
        )

        self.graph.set_transition_priority(stage, stage_transitions)

    def validate_curriculum(self) -> Curriculum:
        """
        Validate curriculum for export/serialization.
        """

        assert (
            len(self.see_stages()) != 0
        ), "Curriculum is empty! Please add stages."

        for s in self.see_stages():
            s.validate_stage()

        # Check round trip serialization
        try:
            instance_json = self.model_dump_json()
            curr_subtype = type(self)
            curr_subtype.model_validate_json(instance_json)
        except Exception:
            print(
                (
                    "Pydantic cannot serialize Curriculum, please use "
                    "mypy to verify your types (check stage transition signature, etc.)."
                )
            )
            raise

        return self

    def export_diagram(self, png_path: str) -> None:  # noqa: C901
        """
        Makes diagram for input Curriculum and
        writes to output png_path.
        """

        def make_stage_script(s: Stage) -> str:
            """
            Stage to dot script conversion.
            """

            template_string = """
                digraph cluster_{{ stage_id }} {
                    labelloc="t";
                    label={{ stage_name }};

                    // Define nodes with increased font visibility
                    node [shape=box, style=filled, fontname=Arial, fontsize=12,
                    fillcolor=lightblue, color=black];

                    // Define nodes
                    {% for n in nodes %}
                    {{ n }};
                    {% endfor %}

                    // Define edges
                    {% for edge in edges %}
                    {{ edge }};
                    {% endfor %}
                }
            """
            template = Template(template_string)
            stage_name = '"' + s.name + '"'

            nodes = []
            for node_id, node in s.graph.nodes.items():
                # Add color to start policies
                if node in s.start_policies:
                    node_str = (
                        f'{node_id} [label="{node.rule.__name__}",'
                        'fillcolor="#FFEA00"]'
                    )
                else:
                    node_str = f'{node_id} [label="{node.rule.__name__}"]'
                nodes.append(node_str)

            edges = []
            for start_id, edge_list in s.graph.graph.items():
                for i, (edge, dest_id) in enumerate(edge_list):
                    # Use 1-indexing for labels
                    i = i + 1

                    # Edges must be StageTransition or PolicyTransition
                    edge_str = (
                        f'{start_id} -> {dest_id} [label="({i}) '
                        f'{edge.rule.__name__}", minlen=2]'
                    )
                    edges.append(edge_str)

            stage_dot_script = template.render(
                stage_id="".join(s.name.split()),
                stage_name=stage_name,
                nodes=nodes,
                edges=edges,
            )

            return stage_dot_script

        def make_curriculum_script(c: Curriculum) -> str:
            """
            Curriculum to dot script conversion.
            """

            curr_dot_script = """
                digraph cluster_curriculum {
                    color="white";
                    label={{ curr_name }};
                    fontsize=24;

                    node [shape=box, style=filled];
                    {% for n in nodes %}
                    {{ n }}
                    {% endfor %}

                    {% for edge in edges %}
                    {{ edge }};
                    {% endfor %}
                }
            """
            template = Template(curr_dot_script)

            # Add curriculum nodes
            nodes = [
                f'{node_id} [label="{node.name}"]'
                for node_id, node in c.graph.nodes.items()
            ]

            # Add curriculum edges
            edges = []
            for start_id, edge_list in c.graph.graph.items():
                for i, (edge, dest_id) in enumerate(edge_list):
                    # Use 1-indexing for labels
                    i = i + 1

                    # Edges must be StageTransition or PolicyTransition
                    edge_str = (
                        f'{start_id} -> {dest_id} [label="({i}) '
                        f'{edge.rule.__name__}", minlen=2]'
                    )
                    edges.append(edge_str)

            curriculum_dot_script = template.render(
                curr_name='"' + c.name + '"', nodes=nodes, edges=edges
            )

            return curriculum_dot_script

        assert png_path.endswith(
            ".png"
        ), "Please add .png extension to end of png_path."
        self.validate_curriculum()

        dot_scripts = [make_curriculum_script(self)]
        last = []
        for stage in self.see_stages():
            if stage.name == "GRADUATED":
                last.append(make_stage_script(stage))
                continue
            dot_scripts.append(make_stage_script(stage))
        dot_scripts = dot_scripts + last

        # Finally concatenate these strings together in this order.
        final_script = "\n".join(dot_scripts)

        # Run graphviz export
        gvpack_command = ["gvpack", "-u"]
        dot_command = ["dot", "-Tpng", "-o", png_path]
        gvpack_process = subprocess.Popen(
            gvpack_command,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.DEVNULL,
        )
        dot_process = subprocess.Popen(dot_command, stdin=subprocess.PIPE)

        gvpack_output, _ = gvpack_process.communicate(
            input=final_script.encode()
        )
        dot_process.communicate(input=gvpack_output)

    def export_json(self, json_path: str) -> None:
        """
        Export curriculum json to export path
        """

        assert json_path.endswith(
            ".json"
        ), "Please add .json extension to end of json_path."
        self.validate_curriculum()

        with open(json_path, "w") as f:
            json_dict = self.model_dump()
            json_string = json.dumps(json_dict, indent=4)
            f.write(json_string)

    def export_curriculum(self, export_dir: str) -> None:
        """
        Export json and diagram into export dir.
        """

        self.export_json(str(Path(export_dir) / "schema.json"))
        self.export_diagram(str(Path(export_dir) / "diagram.png"))

    @classmethod
    def download_curriculum(
        cls,
        name: str,
        version: str,
        bucket="aind-behavior-curriculum-prod-o5171v",
    ) -> Curriculum:
        """
        Reconstruct curriculum object from cloud json.
        """

        def read_json(bucket_name: str, json_key: Path | str) -> dict:
            """
            Reads a json content hosted in S3

            Parameters
            ---------------
            bucket_name: str
                Bucket name

            json_key: PathLike
                Path where the json is stored in S3

            Returns
            ---------------
            dict
                Dictionary with the json content
            """

            s3 = boto3.client('s3')
            response = s3.get_object(Bucket=bucket_name, Key=json_key)
            json_data = json.loads(response['Body'].read().decode('utf-8'))
            return json_data

        json_dict = read_json(bucket, str(Path("curriculums") / name / version / 'schema.json'))
        json_string = json.dumps(json_dict)
        curr = cls.model_validate_json(json_string)

        return curr
