#include <print>
#include <chrono>
#include <cstddef>

#include "serif/utils/dependencies/dependencies.hpp"


/**
 * Demonstration code for how to use the SERiF computational graph system
 * The broad idea here is that there are many calculations which depend on each other but which may not necessarily always need to be recomputed just
 * because another part of the model has changed. We model this as a combination of an acyclic computational graph, defined at compile time, and
 * a runtime cache of values. The general workflow is that a developer defines the "context" of the computational graph, think of this
 * as all of the inputs to the graph, or maybe more specifically, anything that could change due to some external event (e.g. a time step, a newton step, etc...)
 * This context is probably the most CS intensive part of this code as the developer needs to think about lifetimes and ownership.
 *
 * **Context**
 * Contexts are made with the make_dependency_context function, this accepts a variadic pack of bindings between a type and a runtime value.
 * You can generate these yourself using the binding functions and types; however, we provide two convenience wrappers which generally developers should consider
 * using.
 *     - borrow: Used to bind a type to a reference to a value. This is useful for values which are owned by the caller and which the caller will manage the lifetime of (takes an lvalue).
 *     - own: Used to bind a type to a value which is owned by the context (copies an lvalue or moves an rvalue)
 * Each of these functions accepts a type as a template parameter and a runtime value as an argument. The type is used as an ID, this is key to understand, if we write
 * \code{.cpp}
 * using namespace serif::utils::dependencies;
 *
 * double x{1.0};
 * auto context = make_dependency_context(borrow<keys::Input>(x));
 * \endcode
 *
 * we are not saying that x is of type keys::Input, this is clearly not the case. Rather, we are registering in the computational graph a binding
 * between the type keys::Input and the runtime value x. Think of this as naming the various inputs to the graph.
 *
 * **Computations**
 * After defining a context we can define a computation, each computation defines a node, while its ordered input keys define the incoming edges.
 * Consider the following example (assuming the context from above):
 *
 * \code{.cpp}
 * auto func = make_computation<keys::Output, InputList<keys::Input>>(
 *  [](const double _x) -> double {
 *      return _x * _x;
 *  }
 * );
 * \endcode
 *
 * what we have defined here is a node which takes a single input tagged keys::Input and produces an output tagged as keys::Output. The input key
 * declares an edge from that source to this node. Other nodes can use keys::Output as an input, and a successfully prepared result can be read with get
 * or retained with share. These connections allow us to form a larger graph. Consider now the following code
 *
 * \code{.cpp}
 * using namespace serif::utils::dependencies;
 *
 * struct tags {
 *     struct X{};
 *     struct Y{};
 *
 *     struct Output1{};
 *     struct Output2{};
 * };
 *
 * double x{1.0};
 * double y{2.0};
 * auto context = make_dependency_context(borrow<tags::X>(x), borrow<tags::Y>(y));
 *
 * auto func1 = make_computation<tags::Output1, InputList<tags::X>>(
 *     [](const double _x) -> double {
 *         return _x * _x;
 *     }
 * );
 *
 * auto func2 = make_computation<tags::Output2, InputList<tags::Output1, tags::Y>>(
 *    [](const double _x, const double _y) -> double {
 *         return _x + _y;
 *     }
 * );
 * \endcode
 *
 * we have now defined two computations, the first takes a single input, whichever input is tagged as tags::X and produces an output tagged as tags::Output1.
 * The second computation takes two inputs, the first is the output of the first computation (lookup by the tag) and the second is whichever input is tagged as tags::Y,
 * and produces an output tagged as tags::Output2. Note that the order of arguments to the lambda must be the same as the order of the tags in the InputList, the
 * dependencies module will automatically order arguments based on the InputList. There is type checking to ensure that the types match, but if you have two
 * inputs of the same type that are flipped there is no way for the system to inspect that and warn you. Further, depending on your compiler and flags
 * there may be implicit casting between convertable types (e.g. int <-> double).
 *
 * @todo We may want to remove the ability to have implicit casting between types, this is a bit of a footgun and may lead to subtle bugs.
 *
 * @note This is primary goal of this system, split computations into small pieces with the ability to cache intermediate results and reuse
 * as many of those intermediate results as we can can prove is safe to do. The whole reason this system is so complex is that very purpose. There is a fundamental
 * problem with high dimensional stellar models, they are always going to be slow, therefore we need to accept some degree of complexity in order to make them tractable.
 *
 * **Computation Cache**
 * Once we have the computations we need to connect them together, generally use the helper function make_computation_cache which
 * takes a context and a variadic pack of computations. There are *a lot* of compile time checks which run to ensure that the graph is valid, the context
 * describes what is needed to compute the graph, and that each computation is valid. Generally if a static analysis error or compile error pops up here
 * it probably means there is a mistake in the definition of the context or the computations. I've tried to make these errors as clear as possible; however,
 * there is only so much that you can do with the C++ compiler. I recommend first checking, probably with a drawn-out diagram, that all the
 * tags are correct and what you expect.
 *
 * A developer can also choose to use the make_computation_graph function to explicitly name a graph. This
 * is useful when naming or reusing a collection of computations (make_computation_cache
 * calls make_computation_graph and then fixes the source types from the supplied context without retaining the context). Now consider the following code assuming the context and computations from above:
 *
 * \code{.cpp}
 *     auto graph_cache = make_computation_cache(context, func1, func2);
 * \endcode
 *
 * At this point we have a cache which we can access various parts of; however, no computations have been run yet. To run a computation we must call try_prepare.
 *
 * try_prepare returns std::expected: ordinary rejection is a value, and no candidate results or snapshots are published on failure.
 * Exceptions from a callable or snapshot policy propagate with the same cache protection. Previously retained results remain valid,
 * but get and share require a successful current preparation. Always check the report before reading results.
 *
 * try_prepare takes the context and evaluates the requested nodes and their prerequisites, reusing unchanged results. By default every result
 * is requested, so the first call evaluates the entire graph. RequestedResults can restrict evaluation to a branch.
 *
 * \code{.cpp}
 *     auto report = graph_cache.try_prepare(context);
 *     if (!report) {
 *         return; // Handle report.error() before reading the cache.
 *     }
 *
 *     double func1_result = graph_cache.template get<tags::Output1>();
 *     double func2_result = graph_cache.template get<tags::Output2>();
 *
 *     auto func1_status = report->template status<tags::Output1>();
 *     auto func2_status = report->template status<tags::Output2>();
 *
 *     std::string_view func1_status_string = preparation_status_name_map[func1_status]; // note this is an inline map defined in serif::utils::dependencies
 *     std::string_view func2_status_string = preparation_status_name_map[func2_status];
 *
 *     std::println("Func1 Result: {}, Func2 Result: {}, Func1 Status: {}, Func2 Status: {}", func1_result, func2_result, func1_status_string, func2_status_string);
 * \endcode
 *
 * Now what we could do is change only one part of the input context and call try_prepare again. Lets say we change only y, such that only func2 is stale.
 *
 * \code{.cpp}
 *    y = 3.0;
 *    report = graph_cache.try_prepare(context);
 *    if (!report) {
 *        return; // Handle report.error() before reading the cache.
 *    }
 *
 *    double func1_result = graph_cache.template get<tags::Output1>();
 *    double func2_result = graph_cache.template get<tags::Output2>();
 *
 *    auto func1_status = report->template status<tags::Output1>();
 *    auto func2_status = report->template status<tags::Output2>();
 *
 *    std::string_view func1_status_string = preparation_status_name_map[func1_status]; // note this is an inline map defined in serif::utils::dependencies
 *    std::string_view func2_status_string = preparation_status_name_map[func2_status];
 *
 *    std::println("Func1 Result: {}, Func2 Result: {}, Func1 Status: {}, Func2 Status: {}", func1_result, func2_result, func1_status_string, func2_status_string);
 * \endcode
 */
int main() {
    namespace dependencies = serif::utils::dependencies;
    namespace keys = dependencies::keys;

    struct Tags {
        struct Inputs {
            struct AT{};
            struct BT{};
            struct CT{};

            struct MT{};
            struct IT{};
        };

        struct Outputs {
            struct QuadOutput{};
            struct LinearOutput{};
        };
    };

    double A = 1.0;
    double B = 2.0;
    double C = 3.0;
    double x = 1.0;
    double m = 2.0;
    double i = 3.0;

    const auto graph_context = dependencies::make_dependency_context(
        dependencies::borrow<keys::Input>(x),
        dependencies::borrow<Tags::Inputs::AT>(A),
        dependencies::borrow<Tags::Inputs::BT>(B),
        dependencies::borrow<Tags::Inputs::CT>(C),
        dependencies::borrow<Tags::Inputs::MT>(m),
        dependencies::borrow<Tags::Inputs::IT>(i)
    );

    const auto quadratic_function = dependencies::make_computation<Tags::Outputs::QuadOutput, dependencies::InputList<keys::Input, Tags::Inputs::AT, Tags::Inputs::BT, Tags::Inputs::CT>>(
        [](const double _x, const double i_A, const double i_B, const double i_C) -> double {
            double t = 0.0;
            for (int i = 0; i < 1000000; ++i) {
                t = _x * i_A;
            }
            return i_A * _x * _x + i_B * _x + i_C;
        }
    );
    const auto linear_function = dependencies::make_computation<Tags::Outputs::LinearOutput, dependencies::InputList<Tags::Outputs::QuadOutput, Tags::Inputs::MT, Tags::Inputs::IT>>(
        [](const double _x, const double _m, const double _b) -> double {
            double t = 0.0;
            for (int i = 0; i < 10000; ++i) {
                t = _x * _m;
            }
            return _m * _x + _b;
        }
    );

    // The cache owns callbacks; the application owns accumulated measurements.
    using Clock = std::chrono::steady_clock;
    std::chrono::nanoseconds quadratic_elapsed{};
    std::size_t quadratic_executions = 0;
    std::size_t quadratic_consumptions = 0;
    double consumed_quadratic = 0.0;
    const auto observers = dependencies::make_observers(
        dependencies::observe_node<Tags::Outputs::QuadOutput>(
            []() {
                std::println("Before quadratic node execution");
                return Clock::now();
            },
            [&](Clock::time_point& started, dependencies::NodeOutcome) noexcept {
                std::println("After quadratic node execution");
                quadratic_elapsed += std::chrono::duration_cast<std::chrono::nanoseconds>(Clock::now() - started);
                ++quadratic_executions;
            }
        ),
        dependencies::observe_node<Tags::Outputs::LinearOutput>(
            []() {
                std::println("Before linear node execution");
                return 1;
                },
            [](int& r, dependencies::NodeOutcome result) noexcept {
                    std::println("After linear node execution, result: {}", result);
                }
        ),
        dependencies::observe_edge<Tags::Outputs::QuadOutput, Tags::Outputs::LinearOutput>(
            [&](const double& value) {
                consumed_quadratic = value;
                ++quadratic_consumptions;
            }
        )
    );
    const auto graph = dependencies::make_computation_graph(quadratic_function, linear_function);
    auto graph_cache = dependencies::make_computation_cache(graph, graph_context, observers);

    auto display_results = [](auto& report, auto& cache, auto& start_time, auto& end_time) -> void {
        if (!report) {
            std::println(stderr, "Computation failed: {}", report.error().error);
            return;
        }

        const double result =  cache.template get<Tags::Outputs::LinearOutput>();
        const double quad_result = cache.template get<Tags::Outputs::QuadOutput>();
        const dependencies::PreparationStatus linear_status = report->template status<Tags::Outputs::LinearOutput>();
        const dependencies::PreparationStatus quad_status = report->template status<Tags::Outputs::QuadOutput>();
        const auto duration = std::chrono::duration_cast<std::chrono::nanoseconds>(end_time - start_time).count();

        std::string_view linear_status_string = dependencies::preparation_status_name_map[linear_status];
        std::string_view quad_status_string = dependencies::preparation_status_name_map[quad_status];

        std::println("Linear Result: {}, Quad Result: {}, Linear Status: {}, Quad Status: {} Duration: {} nanoseconds", result, quad_result, linear_status_string, quad_status_string, duration);
    };

    auto start_time = std::chrono::high_resolution_clock::now();
    auto graph_report = graph_cache.try_prepare(graph_context);
    auto end_time = std::chrono::high_resolution_clock::now();
    display_results(graph_report, graph_cache, start_time, end_time);

    // Change x and cause the entire graph to recompute
    x = 2.0;
    start_time = std::chrono::high_resolution_clock::now();
    graph_report = graph_cache.try_prepare(graph_context);
    end_time = std::chrono::high_resolution_clock::now();
    display_results(graph_report, graph_cache, start_time, end_time);

    // Change just part of the linear function but allow it to reuse the quadratic result
    m = 3.0;
    start_time = std::chrono::high_resolution_clock::now();
    graph_report = graph_cache.try_prepare(graph_context);
    end_time = std::chrono::high_resolution_clock::now();
    display_results(graph_report, graph_cache, start_time, end_time);

    // change nothing and allow the entire graph to reuse the previous results
    start_time = std::chrono::high_resolution_clock::now();
    graph_report = graph_cache.try_prepare(graph_context);
    end_time = std::chrono::high_resolution_clock::now();
    display_results(graph_report, graph_cache, start_time, end_time);

    // The third request consumes a reused quadratic value; the fourth executes no hooks.
    std::println("Quadratic executions: {}, edge consumptions: {}, last consumed value: {}, total node time: {} nanoseconds",
        quadratic_executions, quadratic_consumptions, consumed_quadratic, quadratic_elapsed.count());
}