.. _custom_bots:

Custom Bots
===========

The :ref:`sim_bot` and :ref:`batch_bot` have lots of functionality but it's inevitable that you
have some additional process or requirement that is not currently supported. Fortunately, it's
very easy to customise a ``pyqalx.Bot`` to do what you need or make an additional bot and connect
it to other bots.

Subclass the ``BatchBot``
~~~~~~~~~~~~~~~~~~~~~~~~~

Let us first investigate adding some additional results plots at the end of processing a batch.
The code below gives an idea of how this can be achieved with comments in the code explaining the
process as much as possible.

.. code-block:: python

    from io import BytesIO

    # WE USE PYPLOT https://matplotlib.org/api/pyplot_api.html
    import matplotlib.pyplot as plt
    # WE USE PANDAS TO HELP PLOT DATA https://pandas.pydata.org/Pandas_Cheat_Sheet.pdf
    import pandas as pd

    from qalx_orcaflex.bots import BatchBot
    from qalx_orcaflex.bots.batch import process_batch, post_process_batch


    def make_plots(load_case_data, title):
        # we know that our load_case_data will form a dataframe
        th_dataframe = pd.DataFrame(**load_case_data)
        th_dataframe.sort_index(inplace=True)  # we want the x-axis in order
        th_dataframe.plot(title=title, grid=True)  # give it a title and make a grid
        plot_bytes = BytesIO()  # we are going to save the bytes straight to qalx
        plt.savefig(plot_bytes)  # save the figure into the file-like object
        plot_bytes.seek(0)  # before we write the file to qalx we need to go back to the
        # start of the file
        return plot_bytes


    def extra_post_processing(job):
        post_process_batch(job)  # FIRST EXECUTE THE EXISTING CODE

        # now we want to do our own plots
        # but we only want to do this if the results have already been summarised
        if job.entity['meta'].get('results_summary'):
            # get the results summaries
            all_results = job.session.item.get(job.entity['meta'].get(
                'results_summary'))['data']
            job.entity["meta"]["result_plots"] = []
            # plot time histories
            for result_name, result_data in all_results["Time Histories"].items():
                if result_data['vs_info']:  # check we have load case info
                    for load_case_info, load_case_data in result_data['vs_info'].items():
                        title = f"{result_name} vs. {load_case_info}"  # give it a title
                        plot_bytes = make_plots(load_case_data, title)  # get the plot
                        plot_item = job.s.item.add(
                            input_file=plot_bytes,
                            file_name=f"{title}.png",
                            meta={"_class": "orcaflex.batch_plot"}
                        )  # make a new qalx item with the plot image
                        job.entity["meta"]["result_plots"].append(plot_item["guid"])  #
                        # save a reference to the plot on the original item in the meta

            job.save_entity()  # save the entity back to qalx


    #
    #  We subclass BatchBot so that we know that it will do the normal process function
    #  as intended. We will augment the postprocess_function with our extra code to make
    #  plots
    #
    class PlotBatchBot(BatchBot):

        def __init__(self, *args, **kwargs):
            super(PlotBatchBot, self).__init__(*args, **kwargs)
            self.process_function = process_batch
            self.postprocess_function = extra_post_processing


    # Instantiate the new class to pass to the command line or build in a factory
    plot_th_bot = PlotBatchBot("BatchWithPlots")


The code above, if run in the place of :ref:`batch_bot` will perform the same functions as
:ref:`batch_bot` but will also save plots of load case info vs. time history results. References
to the saved plots will be found on the ``meta`` of the ``pyqlax.Group`` that represents the batch.

Make a new bot
~~~~~~~~~~~~~~

The following code makes plots in the same way as above (although for Range Graphs this time)
except it will run as a separate bot. There are good reasons why you might want to do this:

- This bot doesn't depend on OrcFxAPI and so it can be run on a linux machine saving costly
  windows clock cycles for the serious work of simulation

- If the process here took a long time then it would happen in parallel to the processing of
  batches saving time over the serial approach in the previous example.

.. code-block:: python

    from io import BytesIO

    import matplotlib.pyplot as plt
    import pandas as pd
    from pyqalx import Bot
    from pyqalx.bot import QalxJob

    plot_bot = Bot("PlottingBot")


    @plot_bot.process
    def make_plots(job: QalxJob):
        batch = job.entity
        summary = batch.meta.get("results_summary")
        if summary:
            job.entity["meta"]["result_plots"] = []
            summary_item = job.session.item.get(summary)
            data = summary_item["data"]
            for result_name, range_graph in data["Range Graphs"].items():
                vs_info = range_graph["vs_info"]
                for load_case_info_name, load_case_data in vs_info.items():
                    title = f"{result_name} max vs. {load_case_info_name}"
                    df = pd.DataFrame(**load_case_data)
                    df.sort_index(inplace=True)
                    df.plot(title=title, grid=True)
                    plot_bytes = BytesIO()
                    plt.savefig(plot_bytes)
                    plot_bytes.seek(0)
                    plot_item = job.s.item.add(
                        input_file=plot_bytes, file_name=f"{title}.png"
                    )
                    job.log.debug(plot_item["file"]["url"])
                    job.entity["meta"]["result_plots"].append(plot_item["guid"])
            job.save_entity()


The only thing we need to remember here is that we need ``BatchBot`` to pass the job on to our
``PlottingBot``. This can be done with a workflow in a factory or using ``BatchOptions``:

.. code-block:: python

    from qalx_orcaflex.core import QalxOrcaFlex, OrcaFlexBatch
    import qalx_orcaflex.data_models as dm

    qfx = QalxOrcaFlex()
    options = dm.BatchOptions(
        sim_queue="example-sim-queue",
        batch_queue="example-batch-queue",
        send_batch_to=['batch-plotting-queue'] # The batch will be sent here once it's
        # complete
    )
    batch = OrcaFlexBatch(
        "My Batch to Plot",
        batch_options=options,
        session=qfx)