export.py

import getpass
import json
import re
import sys
from pathlib import Path

import httpx
import numpy
from prefect import flow, get_run_logger, task
from tiled.client import show_logs
from data_validation import get_run, get_run_sandbox, get_sandbox_client
import pandas as pd


EXPORT_PATH = Path("/nsls2/data/dssi/scratch/prefect-outputs/rsoxs/")

show_logs()


def lookup_directory(start_doc):
    """
    Return the path for the proposal directory.

    PASS gives us a *list* of cycles, and we have created a proposal directory under each cycle.
    """
    DATA_SESSION_PATTERN = re.compile("[GUPCpass]*-([0-9]+)")
    client = httpx.Client(base_url="https://api.nsls2.bnl.gov")
    data_session = start_doc[
        "data_session"
    ]  # works on old-style Header or new-style BlueskyRun

    try:
        digits = int(DATA_SESSION_PATTERN.match(data_session).group(1))
    except AttributeError:
        raise AttributeError(f"incorrect data_session: {data_session}")

    response = client.get(f"/v1/proposal/{digits}/directories")
    response.raise_for_status()

    paths = [path_info["path"] for path_info in response.json()["directories"]]

    # Filter out paths from other beamlines.
    paths = [path for path in paths if "sst" == path.lower().split("/")[3]]

    # Filter out paths from other cycles and paths for commissioning.
    paths = [
        path
        for path in paths
        if path.lower().split("/")[5] == "commissioning"
        or path.lower().split("/")[5] == start_doc["cycle"]
    ]

    # There should be only one path remaining after these filters.
    # Convert it to a pathlib.Path.
    return Path(paths[0])


@task
def write_dark_subtraction(ref, api_key=None, dry_run=None):
    """
    This is a Prefect task that perform dark subtraction.

    Subtract dark frame images from the data,
    and write the result to tiled.

    Parameters
    ----------
    ref: string
        This is the reference to the BlueskyRun to be exported. It can be
        a partial uid, a full uid, a scan_id, or an index (e.g. -1).

    Returns
    -------
    results: dict
        A dictionary that maps field_name to the matching processed uid.
    """

    logger = get_run_logger()
    logger.info("starting dark subtraction")

    def safe_subtract(light, dark, pedestal=100):
        """
        Subtract a dark_frame from a light_frame.
        Parameters
        ----------
        light: array
            This is the light_frame.
        dark: array
            This is the dark_frame.
        pededtal: integer
            An offset to avoid negative results.
        """

        dark.load()
        light.load()

        dark = dark - pedestal

        dark = numpy.clip(dark, a_min=0, a_max=None)

        return numpy.clip(
            light - dark.reindex_like(light, method="ffill").data, a_min=0, a_max=None
        ).astype(light.dtype)

    run = get_run(ref, api_key=api_key)

    full_uid = run.start["uid"]
    logger.info(f"{full_uid = }")  # noqa: E202,E251

    # Raise an exception if the dark stream isn't present
    if "dark" not in run:
        logger.warning(
            "dark stream does not exist in this run. Skipping dark subtraction and tiff export."
        )
        return
        # raise Exception("dark stream does not exist")

    # Access the primary and dark streams as xarray.Datasets.
    primary_data = run["primary"]["data"].read()
    dark_data = run["dark"]["data"].read()

    # The set of fields that should be exported if found in the scan.
    export_fields = {
        "Synced_saxs_image",
        "Synced_waxs_image",
        "Small Angle CCD Detector_image",
        "Wide Angle CCD Detector_image",
    }

    # The set of fields for the primary data set.
    primary_fields = set(run["primary"]["data"])

    # The export_fields that are found in the primary dataset.
    found_fields = export_fields & primary_fields

    # Map field to processed uid to use in other tasks.
    results = {}

    # Write the dark subtracted images to tiled.
    for field in found_fields:
        light = primary_data[field][:]
        dark = dark_data[field][:]
        subtracted = safe_subtract(light, dark)
        if dry_run:
            logger.info("Dry_run: not writing subtracted images to Tiled.")
        else:
            processed_array_client = get_sandbox_client(api_key=api_key).write_array(
                subtracted.data,
                metadata={
                    "field": field,
                    "python_environment": sys.prefix,
                    "raw_uid": full_uid,
                    "operation": "dark subtraction",
                },
                access_tags=["rsoxs_sandbox"],
            )
            results[field] = processed_array_client.item["id"]

    logger.info("completed dark subtraction")

    return results


# Make sure this only runs when the dark subtraction is successful
@task
def tiff_export(raw_ref, processed_refs, api_key=None, dry_run=None):
    """
    Export processed data into a tiff file.

    Parameters
    ----------
    raw_ref: string
        Reference to a BlueskyRun. Can be a full uid, a partial uid,
        a scan id, or an index (e.g. -1).
    processed_refs: dict
        A dictionary that maps field_name to the matching processed_ref.
    """

    # This is the result of combining 2 streams so we'll set the stream name as primary
    # Maybe we shouldn't use a stream name in the filename at all,
    # but for now we are maintaining backward-compatibility with existing names.
    STREAM_NAME = "primary"

    start_doc = get_run(raw_ref, api_key=api_key).start
    directory = (
        lookup_directory(start_doc)
        / start_doc["project_name"]
        / f"{start_doc['scan_id']}"
    )
    directory.mkdir(parents=True, exist_ok=True)

    logger = get_run_logger()
    logger.info(f"starting tiff export to {directory}")

    for field, processed_uid in processed_refs.items():
        dataset = get_run_sandbox(processed_uid, api_key=api_key)
        assert field == dataset.metadata["field"]
        num_frames = len(dataset)
        for i in range(num_frames):
            filename = f"{start_doc['scan_id']}-{start_doc['sample_name']}-{STREAM_NAME}-{field}-{i}.tiff"
            if dry_run:
                logger.info(f"Dry_run: tiff: not exporting {filename}")
            else:
                logger.info(f"Exporting {filename}")
                dataset.export(directory / filename, slice=(i), format="image/tiff")

    if dry_run:
        logger.info(f"Dry_run: did not write tiff files to: {directory}")
    else:
        logger.info(f"wrote tiff files to: {directory}")


# Retry this task if it fails
@task(retries=2, retry_delay_seconds=10)
def csv_export(raw_ref, api_key=None, dry_run=None):
    """
    Export each stream as a CSV file.

    - Include only scalar fields (e.g. no images).
    - Put the primary stream at top level with the scan directory
      and put  all other streams in a subdirectory, per Eliot's convention.

    Parameters
    ----------
    raw_ref: string
        Reference to a BlueskyRun. Can be a full uid, a partial uid,
        a scan id, or an index (e.g. -1).

    """
    print(f"dry_run: {dry_run}")
    run = get_run(raw_ref, api_key=api_key)
    start_doc = run.start

    # Make the directories.
    base_directory = lookup_directory(start_doc) / start_doc["project_name"]
    if not dry_run:
        base_directory.mkdir(parents=True, exist_ok=True)

    logger = get_run_logger()
    logger.info(f"starting csv export to {base_directory}")

    def add_seq_num(dataset):
        """
        Add a seq_num column to the dataset.
        This also converts the dataset to a dataframe.

        We need a seq_num column, which the server does not include, so we
        do export on the client side.
        """

        df = dataset.to_dataframe()
        df2 = df.reset_index()  # Promote 'time' from index to column.
        df2.index.name = "seq_num"
        df3 = df2.reset_index()  # Promote 'seq_num' from index to column.
        df3["seq_num"] += 1  # seq_num starts at 1
        return df3

    for stream_name, stream in run.items():
        logger.info(f"Exporting csv for stream {stream_name}")

        # Figure out the directory to write to.
        scan_directory = f"{start_doc['scan_id']}" if stream_name != "primary" else "."
        directory = base_directory / scan_directory
        if not dry_run:
            directory.mkdir(parents=True, exist_ok=True)

        # Prepare the data.
        dataset = stream["data"]
        scalar_fields = {field for field in dataset if dataset[field].ndim == 1}
        ds = dataset.read(variables=scalar_fields)
        dataframe = add_seq_num(ds)

        # Write the data.
        if dry_run:
            filename = (
                directory
                / f"{start_doc['scan_id']}-{start_doc['sample_name']}-{stream_name}.csv"
            )
            if len(dataframe) >= 2:
                output_dataframe = pd.concat([dataframe.head(1)])
            elif len(dataframe) == 1:
                output_dataframe = pd.concat([dataframe.head(1), dataframe.tail(1)])
            else:
                logger.info(
                    f"Dry run: CSV did not write file {filename}: output: (no data)"
                )
                return
            csv_output = output_dataframe.to_string(
                index=False,
            )
            logger.info(
                f"Dry run: CSV: did not write to file {filename}: output: {csv_output}"
            )

        else:
            dataframe.to_csv(
                directory
                / f"{start_doc['scan_id']}-{start_doc['sample_name']}-{stream_name}.csv",
                index=False,
            )

            logger.info(f"wrote csv files to: {directory}")


@task
def json_export(raw_ref, api_key=None, dry_run=None):
    """
    Export start document into a json file.

    Parameters
    ----------
    raw_ref: string
        Reference to a BlueskyRun. Can be a full uid, a partial uid,
        a scan id, or an index (e.g. -1).

    """
    start_doc = get_run(raw_ref, api_key=api_key).start
    directory = (
        lookup_directory(start_doc)
        / start_doc["project_name"]
        / f"{start_doc['scan_id']}"
    )
    directory.mkdir(parents=True, exist_ok=True)

    logger = get_run_logger()
    logger.info(f"starting json export to {directory}")

    if dry_run:
        json_output = json.dumps(start_doc, ensure_ascii=False, indent=4)
        filename = directory / f"{start_doc['scan_id']}-{start_doc['sample_name']}.json"
        logger.info(
            f"Dry_run: json: did not write to filename: {filename}: output: {json_output}"
        )
    else:
        with open(
            directory / f"{start_doc['scan_id']}-{start_doc['sample_name']}.json",
            "w",
            encoding="utf-8",
        ) as file:
            json.dump(start_doc, file, ensure_ascii=False, indent=4)

        logger.info(
            f"wrote json file to: {str(directory / str(start_doc['scan_id']))}-{start_doc['sample_name']}.json"
        )


# Make the Prefect Flow.
# A separate command is needed to register it with the Prefect server.
@flow
def export(ref, api_key=None, dry_run=None):
    print(f"effective user: {getpass.getuser()}")
    csv_export(ref, api_key=api_key, dry_run=dry_run)
    json_export(ref, api_key=api_key, dry_run=dry_run)
    processed_refs = write_dark_subtraction(ref, api_key=api_key, dry_run=dry_run)
    if processed_refs:
        tiff_export(ref, processed_refs, api_key=api_key, dry_run=dry_run)


# This line will mark this flow as succeeded based on
# the csv and json export tasks succeeding.
# TODO: Do we need this line?
# flow.set_reference_tasks([csv_export_task, json_export_task])