from __future__ import annotations

import datetime
import re
import types
from collections import defaultdict
from typing import Callable, List, Optional, Union

import numpy as np
import pandas as pd

from seeq import spy
from seeq.sdk import *
from seeq.spy import _common, _compatibility, _login, _url
from seeq.spy._errors import *
from seeq.spy._redaction import request_safely, safely
from seeq.spy._session import Session
from seeq.spy._status import Status, Usage

ENUM_REGEX = r'ENUM{{(\d+)\|(.+?)}}'
ENUM_PATTERN = re.compile(ENUM_REGEX)
RETURN_TYPE_COLUMN = '__Return Type__'


@Status.top_level_spy_function()
def pull(
    items: Union[str, pd.DataFrame, pd.Series],
    *,
    start: Optional[Union[str, pd.Timestamp, datetime.date]] = None,
    end: Optional[Union[str, pd.Timestamp, datetime.date]] = None,
    grid: Optional[str] = '15min',
    header: str = '__auto__',
    group_by: Optional[Union[str, List]] = None,
    shape: Union[str, Callable] = 'auto',
    capsule_properties: Optional[List] = None,
    tz_convert: Optional[Union[str, datetime.tzinfo]] = None,
    calculation: Optional[Union[str, pd.DataFrame, pd.Series]] = None,
    bounding_values: bool = False,
    invalid_values_as: Optional[Union[str, int, float]] = np.nan,
    enums_as: Optional[str] = 'string',
    include_extinct: bool = False,
    errors: Optional[str] = None,
    quiet: Optional[bool] = None,
    status: Optional[Status] = None,
    session: Optional[Session] = None,
    capsules_as: Optional[str] = None
):
    """
    Retrieves signal, condition or scalar data from Seeq Server and returns it
    in a DataFrame.

    Parameters
    ----------
    items : {str, pd.DataFrame, pd.Series}
        A DataFrame or Series containing ID and Type columns that can be used
        to identify the items to pull. This is usually created via a call to
        spy.search().

        Alternatively, you can supply URL of a Seeq Workbench worksheet as a
        str. If you supply the URL of a Seeq Vantage Room, the capsule
        evidence (with context) will be pulled for that room.

        If a 'Calculation' column is present, the formula specified in that
        column will be applied to the item in that row "on-the-fly" while
        data is retrieved. The formula must utilize a $signal, $condition
        or $scalar variable to reference the item in that row. Note that
        the results of these "on-the-fly" calculations are not cacheable.
        If you want to utilize caching, explicitly push such calculated
        items and use them without this "on-the-fly" method.

    start : {str, pd.Timestamp}, optional
        The starting time for which to pull data. This argument must be a
        string that pandas.to_datetime() can parse, or a pandas.Timestamp.
        If not provided, 'start' will default to 'end' minus 1 hour. Note
        that Seeq will potentially return one additional row that is earlier
        than this time (if it exists), as a "bounding value" for interpolation
        purposes. If both 'start' and 'end' are not provided and items
        is a str, 'start' will default to the start of the display range
        in Seeq Trend View.

    end : {str, pd.Timestamp}, optional
        The end time for which to pull data. This argument must be a string
        that pandas.to_datetime() can parse, or a pandas.Timestamp.
        If not provided, 'end' will default to now. Note that Seeq will
        potentially return one additional row that is later than this time
        (if it exists), as a "bounding value" for interpolation purposes.

        If pulling using a Workbench Analysis URL and both 'start' and 'end'
        are not provided, 'end' will default to the end of the display range in
        Seeq Trend View.

    grid : {str, 'auto', None}, default '15min'
        A period to use for interpolation such that all returned samples
        have the same timestamps. Interpolation will be applied at the server
        to achieve this. To align samples to a different time zone and/or date
        and time, append a valid time zone and/or timestamp in ISO8601,
        YYYY-MM-DD, or YYYY-MM-DDTHH:MM:SS form. If grid=None is specified,
        no interpolation will occur and each signal's samples will be returned
        untouched. Where timestamps don't match, NaNs will be present within a
        row. If grid='auto', the period used for interpolation will be the
        median of the sample periods from the 'Estimated Sample Period' column
        in 'items'. If grid='auto' and the 'Estimated Sample Period' column
        does not exist in 'items', additional queries will be made to estimate
        the sample period which could potentially impact performance for large
        pulls. Interpolation is either linear or step and is set per signal at
        the time of the signal's creation. To change the interpolation type for
        a given signal, change the signal's interpolation or use the
        appropriate 'calculation' argument.

    header : str, default '__auto__'
        The metadata property to use as the header of each column. Common
        values would be 'ID' or 'Name'. '__auto__' concatenates Path and Name
        if they are present. If a 'Header" column is present in the metadata,
        '__auto__' will use that instead.

    group_by : {str, list(str)}, optional
        The name of a column or list of columns for which to group by. Often
        necessary when pulling data across assets: When you want header='Name',
        you typically need group_by=['Path', 'Asset']

    shape : {'auto', 'samples', 'capsules'}, default 'auto'
        If 'auto', returns capsules as a time series of 0 or 1 when signals are
        also present in the items argument, or returns capsules as individual
        rows if no signals are present. 'samples' or 'capsules' forces the
        output to the former or the latter, if possible.

        Capsules that intersect with the time range will be returned.
        In the case of shape='capsules', if the capsule has no Maximum Duration
        (e.g., you did not use the removeLongerThan() function), the
        'Capsule Start' and/or 'Capsule End' columns may contain NaT entries,
        and it may not be possible to calculate statistics on any signals
        provided.

        You may also provide a callback function as the shape argument. When
        you do so, the callback function will receive the results as they
        are returned by the Seeq service, and they will not accumulate into
        a final DataFrame (and therefore spy.pull() will return None).
        The callback function must take two arguments: The row that the result
        corresponds to; the result DataFrame itself. In this scenario, you
        may also provide 'Start' and 'End' columns in the items DataFrame to
        indicate the time range to pull for that particular row. This scenario
        is useful when you are pulling a lot of data, potentially at different
        time ranges, and it is not possible/practical to accumulate it all
        into a single DataFrame.

    capsule_properties : list(str), optional
        A list of capsule properties to retrieve when shape='capsules'.
        By default, if no signals are present in the items DataFrame, then all
        properties found on a capsule are automatically returned (because
        the nature of the query allows them to be returned "for free").
        Otherwise, you must provide a list of names of properties to retrieve.

    tz_convert : {str, datetime.tzinfo}, optional
        The time zone in which to return all timestamps. If the time zone
        string is not recognized, the list of supported time zone strings will
        be returned in the exception text.

    calculation : {str, pandas.Series, pandas.DataFrame}, optional
        When applying a calculation across assets, the 'calculation' argument
        must be a one-row DataFrame (or a Series) and the 'items' argument must
        be full of assets. When applying a calculation to a signal/condition/
        scalar, calculation must be a string with a single variable in it:
        $signal, $condition or $scalar.

    bounding_values : bool, default False
        If True, extra 'bounding values' will be returned before/after the
        specified query range for the purposes of assisting with interpolation
        to the edges of the range or, in the case of Step or PILinear
        interpolation methods, interpolating to 'now' when appropriate.

    invalid_values_as : {str, int, float}, default np.nan
        Invalid samples and scalars will appear in the returned DataFrame as
        specified in this argument. By default, invalid values will be returned
        as NaNs. Note that specifying a string for this argument (e.g,
        'INVALID') can have significant performance implications on numeric
        signals. You may wish to use a "magic" number like -999999999 if you
        want to be able to discern invalid values but preserve algorithmic
        performance.

    enums_as : {'tuple', 'string', 'numeric', None}, default 'string'
        Enumerations, also known as digital states, are numbers that have an
        associated human-readable name with meaning in the applicable domain
        (e.g., an ON or OFF machine state that is encoded as 1 or 0).
        If enums_as='string', the signal's column in the returned DataFrame
        will be a string value (e.g., 'ON' or 'OFF'). If enums_as='numeric',
        the signal's column will be an integer (e.g. 1 or 0).
        If enums_as='tuple', both the integer and string will be supplied as
        a tuple (e.g., (1, 'ON') or (0, 'OFF')).

    include_extinct : bool, default False
        Only applicable when pulling evidence for Vantage Rooms.
        If True, capsules that existed only temporarily will still be
        returned.

    errors : {'raise', 'catalog'}, default 'raise'
        If 'raise', any errors encountered will cause an exception. If
        'catalog', errors will be added to a 'Result' column in the status.df
        DataFrame.

    quiet : bool, default False
        If True, suppresses progress output. Note that when status is
        provided, the quiet setting of the Status object that is passed
        in takes precedence.

    status : spy.Status, optional
        If specified, the supplied Status object will be updated as the command
        progresses. It gets filled in with the same information you would see
        in Jupyter in the blue/green/red table below your code while the
        command is executed. The table itself is accessible as a DataFrame via
        the status.df property.

    session : spy.Session, optional
        If supplied, the Session object (and its Options) will be used to
        store the login session state. This is useful to log in to different
        Seeq servers at the same time or with different credentials.

    capsules_as : str
        Deprecated, use shape argument instead.

    Returns
    -------
    pandas.DataFrame
        A DataFrame with the requested data. Additionally, the following
        properties are stored on the "spy" attribute of the output DataFrame:

        =================== ===================================================
        Property            Description
        =================== ===================================================
        func                A str value of 'spy.pull'
        kwargs              A dict with the values of the input parameters
                            passed to spy.pull to get the output DataFrame
        query_df            A DataFrame with the actual query made to the
                            Seeq server
        start               A pd.Timestamp with the effective start time
                            of the data pulled
        end                 A pd.Timestamp with the effective end time
                            of the data pulled
        grid                A string with the effective grid of the data
                            pulled
        tz_convert          A datetime.tzinfo of the time zone in which
                            the timestamps were returned
        status              A spy.Status object with the status of the
                            spy.pull call
        =================== ===================================================

    Examples
    --------
    Pull a list of signals and convert the timezone to another timezone

    >>> items = pd.DataFrame([{'ID': '8543F427-2963-4B4E-9825-220D9FDCAD4E', 'Type': 'CalculatedSignal'}])
    >>> my_signals = spy.pull(items=items, grid='15min', calculation='$signal.toStep()',
    >>>          start='2019-10-5T02:53:45.567Z', end='2019-10-6', tz_convert='US/Eastern')

    To access the stored properties
    >>> my_signals.spy.kwargs
    >>> my_signals.spy.query_df
    >>> my_signals.spy.start
    >>> my_signals.spy.end
    >>> my_signals.spy.grid
    >>> my_signals.spy.status.df

    Pull a list of signals with an auto-calculated grid
    >>> signals = spy.search({'Name': 'Area ?_*', 'Datasource Name': 'Example Data'},
    >>>                        estimate_sample_period=dict(Start='2018-01-01T00:00:00Z',
    >>>                        End='2018-01-01T12:00:00Z'))
    >>> spy.pull(signals,
    >>>          start='2018-01-01T00:00:00Z',
    >>>          end='2018-01-01T23:00:00Z',
    >>>          grid='auto')

    Pull a list of signals, conditions or scalars from a Seeq worksheet with an auto-calculated grid
    >>> my_worksheet_items = spy.pull(
    >>> 'https://seeq.com/workbook/17F31703-F0B6-4C8E-B7FD-E20897BD4819/worksheet/CE6A0B92-EE00-45FC-9EB3-D162632DBB48',
    >>>  grid='auto')

    Pull a list of capsules

    >>> compressor_on_high = spy.search({'Name': 'Compressor Power on High', 'Workbook': 'Folder 1 >> Workbook 8'})
    >>> spy.pull(compressor_on_high, start='2019-01-01T04:00:00Z', end='2019-01-09T02:00:00Z')

    Pull a list of capsules but apply a condition function in formula first

    >>> comp_high = spy.search({'Name': 'Compressor Power on High', 'Workbook': 'Folder 1 >> Workbook 8'})
    >>> spy.pull(comp_high, start='2019-01-01', end='2019-01-09', calculation='$condition.setMaximumDuration(1d)')

    Pull capsules as a binary signal at the specified grid. 1 when a capsule is
    present, 0 otherwise

    >>> comp_high = spy.search({'Name': 'Compressor Power on High', 'Workbook': 'Folder 1 >> Workbook 8'})
    >>> spy.pull(comp_high, start='2019-01-01T00:00:00Z', end='2019-01-01T12:00:00Z', shape='samples', grid='1h')

    Pull a scalar

    >>> compressor_power_limit = spy.push(
    >>>     metadata=pd.DataFrame(
    >>>         [{ 'Name': 'Compressor Power Limit', 'Type': 'Scalar', 'Formula': '50kW' }]), errors='raise')
    >>> spy.pull(compressor_power_limit)

    Apply a calculation to a signal using the 'calculation' argument

    >>> signal_with_calc = spy.search({'Name': 'Area A_Temperature', 'Datasource Name': 'Example Data'})
    >>> spy.pull(signal_with_calc,
    >>>          start='2019-01-01T00:00:00',
    >>>          end='2019-01-01T03:00:00',
    >>>          calculation='$signal.aggregate(average(), hours(), startKey())', grid=None)

    Convert a linearly interpolated signal into a step interpolated signal
    using the 'calculation' argument:

    >>> items = pd.DataFrame([{'ID': '8543F427-2963-4B4E-9825-220D9FDCAD4E', 'Type': 'CalculatedSignal'}])
    >>> pull(items=items, start='2019-10-5', end='2019-10-6', grid='15min', calculation='$signal.toStep()')

    Interpolate data using the pandas.DataFrame.interpolate method with a
    second order polynomial, with the signal name as the header. Warning:
    pandas.interpolate can be considerably slower than Seeq's interpolation
    functions for large datasets, especially when using complex interpolation
    methods

    >>> search_df = pd.concat((spy.search({'ID': '6A5E44D4-C6C5-463F-827B-474AB051B2F5'}),
    >>>                        spy.search({'ID': '937449C1-16E5-4E20-AC2E-632C5CECC24B'})), ignore_index=True)
    >>> data_df = pull(search_df, grid=None, start='2019-10-5', end='2019-10-6', header='Name')
    >>> data_df.interpolate(method='quadratic')
    """

    input_args = locals()

    # bringing this up here so that the error is visible before validating arguments
    if capsules_as is not None:
        raise SPyValueError("capsules_as argument is deprecated. Use the following instead:\n"
                            "capsules_as='signal'   -> shape='samples'\n"
                            "capsules_as='capsules' -> shape='capsules'")

    _common.validate_timezone_arg(tz_convert)

    if enums_as is not None and enums_as not in ['tuple', 'string', 'numeric']:
        raise SPyValueError("enums_as argument must be either 'tuple', 'string', 'numeric' or None")

    if isinstance(items, str):
        workbook_id = _url.get_workbook_id_from_url(items)
        worksheet_id = _url.get_worksheet_id_from_url(items)
        workbooks_api = WorkbooksApi(session.client)
        workbook_output = None
        workbook_type = 'Analysis'
        try:
            workbook_output = workbooks_api.get_workbook(id=workbook_id)
            workbook_type = workbook_output.type
        except ApiException as e:
            # Fall through on Forbidden to remain backward compatible with how access errors were handled
            # before CRAB-52660
            if e.status != 403:
                raise

        if workbook_type == 'Vantage' and workbook_output is not None:
            # If `items` is a Vantage Room URL, pull the capsule evidence for that room from the materialized table
            if not hasattr(workbook_output, 'evidence_table_definition_id'):
                raise SPyValueError('Vantage Rooms cannot be pulled on this version of Seeq Server')

            specific_worksheet_ids = [worksheet_id] if worksheet_id is not None else None

            # This call will raise SPyValueError if it couldn't be found, or other errors if it's inaccessible
            vantage_room = spy.workbooks.pull(items, specific_worksheet_ids=specific_worksheet_ids,
                                              quiet=status.quiet, session=session)[0]

            if vantage_room.get('Evidence Table Definition ID') is None:
                raise SPyValueError('Vantage Room does not have an evidence table definition ID')

            items = pd.DataFrame([{'ID': vantage_room['Evidence Table Definition ID'], 'Type': 'Table'}])
            if worksheet_id is not None and len(specific_worksheet_ids) == 0:
                raise SPyValueError(f'Worksheet ID {worksheet_id} found in the specified Vantage Room')

            if start is None and end is None and specific_worksheet_ids is not None:
                specific_worksheet = vantage_room.views[0]
                start = specific_worksheet.investigate_range['Start']
                end = specific_worksheet.investigate_range['End']

        elif workbook_type == 'Analysis':
            # If `items` is a worksheet URL, get the actual items from the worksheet and overwrite `items` as a DataFrame
            worksheet = spy.utils.pull_worksheet_via_url(items, minimal=True, quiet=status.quiet, session=session)
            items = spy.search(worksheet.display_items[['ID', 'Type']], all_properties=True, quiet=status.quiet,
                               session=session)
            if start is None and end is None:
                start = worksheet.display_range['Start']
                end = worksheet.display_range['End']
        else:
            raise SPyValueError(f'Workbook type "{workbook_type}" cannot be used with spy.pull(url).')

    if invalid_values_as is None:
        raise SPyValueError('invalid_values_as cannot be None (because Pandas treats it the same as NaN)')

    _common.validate_unique_dataframe_index(items, 'items')

    if isinstance(items, pd.DataFrame) and len(items) == 0:
        items = pd.DataFrame(columns=['ID', 'Type'])

    if 'ID' not in items or 'Type' not in items:
        raise SPyValueError('items DataFrame must include "ID" column and "Type" column')

    if isinstance(calculation, pd.DataFrame):
        if len(calculation) != 1:
            raise SPyValueError("When applying a calculation across assets, calculation argument must be a one-row "
                                "DataFrame, or a Series. When applying a calculation to a signal/condition/scalar, "
                                'calculation must be a string with a signal variable in it: $signal, $condition or '
                                '$scalar.')

        calculation = calculation.iloc[0]

    if isinstance(items, pd.Series):
        items = pd.DataFrame([items])

    if not isinstance(shape, Callable) and shape not in ['auto', 'capsules', 'samples']:
        raise SPyValueError("shape must be one of 'auto', 'capsules', 'samples'")

    if capsule_properties is not None and not isinstance(capsule_properties, list):
        raise SPyValueError("capsules_properties must be a list of strings (capsule property names)")

    if group_by:
        if isinstance(group_by, str):
            group_by = [group_by]
        if not isinstance(group_by, list):
            raise SPyValueError('group_by argument must be a str or list(str)')
        if not all(col in items.columns for col in group_by):
            raise SPyValueError('group_by columns %s not present in query DataFrame' % group_by)

    pd_start, pd_end = _login.validate_start_and_end(session, start, end)

    if tz_convert is None:
        # Return the results in the timezone of the start date so that the timestamps tend to make sense to the user
        tz_convert = pd_start.tz

    status_columns = [c for c in ['ID', 'Type', 'Path', 'Asset', 'Name', 'Header'] if c in items]

    status.df = items[status_columns].copy()
    status.df['Time'] = datetime.timedelta(0)
    status.df['Count'] = 0
    status.df['Pages'] = 0
    status.df['Data Processed'] = Usage()
    status.df['Result'] = 'Pulling'

    query_df: pd.DataFrame = items.copy()
    query_df[RETURN_TYPE_COLUMN] = pd.Series(dtype=str)

    for index, row in items.iterrows():
        query_df.at[index, RETURN_TYPE_COLUMN] = _determine_calculated_type(session, row, calculation)

    output = types.SimpleNamespace(df=pd.DataFrame())
    at_least_one_scalar = len(query_df[query_df[RETURN_TYPE_COLUMN].str.endswith('Scalar', na=False)]) > 0
    at_least_one_signal = len(query_df[query_df[RETURN_TYPE_COLUMN].str.endswith('Signal', na=False)]) > 0
    at_least_one_condition = len(query_df[query_df[RETURN_TYPE_COLUMN].str.endswith('Condition', na=False)]) > 0
    at_least_one_asset = len(query_df[query_df[RETURN_TYPE_COLUMN].str.endswith('Asset', na=False)]) > 0
    at_least_one_table = len(query_df[query_df[RETURN_TYPE_COLUMN].str.startswith('table', na=False)]) > 0

    if at_least_one_asset:
        calculation_series = calculation if isinstance(calculation, pd.Series) else calculation.iloc[0]
        at_least_one_signal = calculation_series['Type'].endswith('Signal')
        at_least_one_condition = calculation_series['Type'].endswith('Condition')
        at_least_one_scalar = calculation_series['Type'].endswith('Scalar')
        at_least_one_table = calculation_series['Type'].startswith('Table')

    if shape == 'auto':
        if at_least_one_signal:
            shape = 'samples'
        elif at_least_one_scalar and not at_least_one_condition:
            shape = 'samples'
        elif at_least_one_table:
            shape = 'capsules'
        else:
            shape = 'capsules'

    if isinstance(shape, str) and shape != 'capsules' and at_least_one_table:
        raise SPyValueError("Tables must be pulled with shape='capsules'.")

    if at_least_one_table and (at_least_one_asset or at_least_one_signal):
        raise SPyValueError('Pulling Tables together with Assets or Signals is not supported.')

    # The lifecycle of a pull is several phases. We pull signals before conditions so that, if the conditions are
    # being represented as samples, we have timestamps to map to. Scalars are last because they need to be constant
    # for all rows, and then there is a final step where we re-organize the columns to match the input order as best
    # we can.
    phases = ['signals', 'conditions', 'scalars', 'final']
    if shape == 'capsules':
        phases.remove('signals')

    placeholder_item_name = '__placeholder__'
    if shape == 'samples' and not (at_least_one_signal or at_least_one_scalar):
        # If we're trying to pull a Condition as a Signal, we need a set of timestamps to use. So the user has to
        # specify a grid and then we create and pull a constant signal just to generate timestamps to which we'll
        # map the condition's 1s and 0s.

        if grid is None or grid == 'auto':
            raise SPyRuntimeError(
                "Pull cannot include conditions when no signals are present with shape='samples' and grid=%s" %
                ('None' if grid is None else f"'{grid}'"))

        row_result = _pull_signal(session, placeholder_item_name, f'0.toSignal({grid})', list(),
                                  pd_start, pd_end, tz_convert)

        output.df = row_result.result

    if grid == 'auto':
        grid = estimate_auto_grid(session, query_df, pd_start, pd_end, status)

    status.update('Pulling data from %s to %s' % (pd_start, pd_end), Status.RUNNING)
    # This dictionary is from item ID to a list of columns, which we use during the final phase to order the columns
    # in the output DataFrame. Each type of pull adds additional entries to this dictionary.
    column_names = dict()

    # This list is assembled during the final phase by going back through all the input rows and adding columns that
    # correspond to the item IDs for those rows.
    final_column_names = list()

    item_names = list()
    group_by_indices = dict()
    scalar_results: List[RowResult] = list()
    for phase in phases:
        status.update(f'Pulling from <strong>{pd_start}</strong> to <strong>{pd_end}</strong>', Status.RUNNING)

        index_to_use = output.df.index
        result_list = list()
        for row_index, row in query_df.iterrows():
            # noinspection PyBroadException
            try:
                # When we're doing a group_by, then the output DataFrame may have a MultiIndex as a result of some
                # activity that has already happened (see management of group_by_indices in the _on_success()
                # function below). We need to grab the timestamp index to use instead of using the MultiIndex.
                if group_by and isinstance(output.df.index, pd.MultiIndex):
                    group_by_values = tuple([row[g] for g in group_by])
                    # If the Path and Asset in 'group_by_values' contain pulled signals, then we use the index for
                    # those signals. Otherwise, e.g. if only conditions or scalars, we keep the
                    # existing 'index_to_use'.
                    if group_by_values in group_by_indices:
                        index_to_use = group_by_indices[group_by_values]

                # _process_query_row was broken out into a function mostly to regain some indentation
                # space. It's only ever called from this spot. That's why it has a ton of parameters.
                job = _process_query_row(session, at_least_one_signal, calculation, shape, capsule_properties, grid,
                                         header, query_df, index_to_use, row_index, pd_start, pd_end, phase,
                                         status, tz_convert, bounding_values, invalid_values_as, column_names,
                                         final_column_names, enums_as, include_extinct)

                def _on_success(_row_index, _job_result):
                    _item_row = query_df.loc[_row_index]
                    _row_result: RowResult = _job_result
                    _join_df = _row_result.result

                    _merge_column_name_dict(_row_result.column_names, column_names)

                    if isinstance(shape, Callable):
                        # noinspection PyCallingNonCallable
                        shape(_item_row, _row_result.result)
                        return

                    if shape == 'samples':
                        if group_by is None:
                            _item_name = _join_df.columns[0]
                            if _item_name in item_names:
                                raise SPyRuntimeError(
                                    f'Column headers not unique. 2+ instances of "{_item_name}" found. Use header="ID" '
                                    'to guarantee uniqueness, or supply a "Header" column with unique values in '
                                    'each row, or alternatively try group_by=["Path", "Asset"] if you are using an '
                                    'asset tree.')
                            item_names.append(_item_name)
                        else:
                            for group_column in group_by:
                                _join_df[group_column] = _item_row[group_column]

                            if _join_df.index.name is None:
                                _join_df.index.name = 'Timestamp'

                            # We need to keep track of each group's timestamp index and union all the result indices
                            # together so that Conditions pulled as samples are properly populated in the DataFrame.
                            _group_by_values = tuple([_item_row[g] for g in group_by])
                            if _group_by_values in group_by_indices:
                                group_by_indices[_group_by_values] = group_by_indices[_group_by_values].union(
                                    _join_df.index)
                            else:
                                group_by_indices[_group_by_values] = _join_df.index

                            _join_df.set_index(group_by, append=True, inplace=True)

                    elif phase == 'scalars':
                        # In the case of shape='capsules' and scalars, we want to just add a column to the DataFrame
                        # with the same scalar value in each row, so we treat this specially
                        scalar_results.append(_row_result)
                        return

                    result_list.append(_join_df)

                if job:
                    status.add_job(row_index, job, _on_success)

            except Exception as e:
                _common.raise_or_catalog(status=status, index=row_index, e=e)

        status.execute_jobs(session)

        if len(result_list) == 0:
            continue

        if list(output.df.columns) == [placeholder_item_name]:
            # Get rid of the placeholder column because it's no longer necessary and it'll screw
            # things up when we're merging the partial results
            output.df = pd.DataFrame()

        if len(output.df) == 0:
            existing_columns = output.df.columns
            output.df = result_list.pop(0)
            for existing_column in existing_columns:
                output.df[existing_column] = np.nan

        if shape == 'capsules':
            # When the shape of the output is capsules, it's effectively just a long
            # DataFrame with all of the results appended one after the other. The user can
            # then choose to sort by one of the columns if they would like.
            result_list.insert(0, output.df)
            output.df = _common.smart_pd_concat(result_list, ignore_index=True)
        else:
            if group_by:
                result_list.insert(0, output.df)
                output.df = _common.smart_pd_concat(result_list)
                output.df = output.df.groupby(['Timestamp'] + group_by).first()
            else:
                # noinspection PyTypeChecker
                output.df = output.df.join(result_list, how='outer')

    status.update(f'Pull successful from <strong>{pd_start}</strong> to <strong>{pd_end}</strong>', Status.SUCCESS)

    if isinstance(shape, Callable):
        # We can't return a DataFrame, since we didn't assemble one... we just called the callback.
        return None

    for scalar_result in scalar_results:
        column_name = next(iter(scalar_result.column_names.values()))
        if len(output.df) == 0:
            output.df[column_name] = scalar_result.result[column_name]
        else:
            output.df[column_name] = scalar_result.result

    # Ensure that standard condition columns always come first
    for col in STANDARD_CONDITION_COLUMNS:
        if col in final_column_names:
            final_column_names.remove(col)
            final_column_names.insert(0, col)

    # Ensures that the order of the columns matches the order in the metadata
    output.df = output.df[final_column_names]
    if shape == 'capsules':
        output.df = output.df.reset_index(drop="True")
        _ensure_column_types_for_capsules_shape(output.df, tz_convert)

    if session.options.wants_compatibility_with(193) and 'Capsule ID' in output.df.columns:
        # Capsule ID didn't used to be returned
        output.df.drop(columns='Capsule ID', inplace=True)

    if shape == 'capsules':
        # Sort the DataFrame deterministically by Capsule Start (with NaTs first) and then Capsule End (with NaTs last)
        output.df = (output.df.assign(
            _s_notna=~output.df['Capsule Start'].isna(),
            _e_isna=output.df['Capsule End'].isna(),
        ).sort_values(by=["_s_notna", "Capsule Start", "_e_isna", "Capsule End"])
                     .drop(columns=["_s_notna", "_e_isna"])).reset_index(drop=True)

    status.df['Type'] = query_df[RETURN_TYPE_COLUMN]

    output_df_properties = types.SimpleNamespace(
        func='spy.pull',
        kwargs=input_args,
        query_df=query_df,
        start=pd_start,
        end=pd_end,
        grid=grid,
        tz_convert=tz_convert,
        status=status)

    _common.put_properties_on_df(output.df, output_df_properties)

    return output.df


STANDARD_CONDITION_COLUMNS = ['Capsule ID', 'Capsule Is Uncertain', 'Capsule End', 'Capsule Start', 'Condition']


def _merge_column_name_dict(to_merge, into):
    for k, l in to_merge.items():
        if k not in into:
            into[k] = list()
        for n in l:
            if n not in into[k]:
                into[k].append(n)


def _determine_calculated_type(session: Session, row, calc_arg) -> str:
    if row['Type'] == 'Asset':
        if calc_arg is None or not isinstance(calc_arg, (pd.Series, pd.DataFrame)):
            raise SPyRuntimeError('To pull data for an asset, you must provide a "calculation" argument whose '
                                  'value is the metadata of a calculation that is based on a single asset.')

        return calc_arg['Type']

    if not _common.present(row, 'Calculation'):
        return row['Type']

    parameter = None
    for t in ['Scalar', 'Condition', 'Signal']:
        if row['Type'].endswith(t):
            parameter = f'{t.lower()}={row["ID"]}'
            break

    if parameter is None:
        raise SPyValueError(f'Cannot determine calculated type of item {row["ID"]}')

    formulas_api = FormulasApi(session.client)
    try:
        if hasattr(formulas_api, 'compile_formula_and_parameters'):
            formula_compile_output = formulas_api.compile_formula_and_parameters(body=FormulaCompileInputV1(
                formula=row['Calculation'],
                parameters=[parameter]
            ))
        else:
            # This was deprecated in R64 but works fine for us in R63 and earlier, because its only limitation is
            # around the length of the input parameters and we don't expect super-long formulas in the Calculation
            # column.
            formula_compile_output = formulas_api.compile_formula(
                formula=row['Calculation'],
                parameters=[parameter]
            )

    except ApiException:
        # If the formula is invalid, we can't determine the return type, so just return the original type. The
        # compilation error will be thrown again when we try to retrieve the data.
        return row['Type']

    if formula_compile_output.return_type is not None:
        return f'Calculated{formula_compile_output.return_type}'

    return row['Type']


def _ensure_column_types_for_capsules_shape(df, tz_convert=None):
    for column_name, column_type, column_type_match in reversed([
        ('Condition', 'object', 'object'),
        ('Capsule Start', 'datetime64[ns, UTC]', 'datetime64[ns,'),
        ('Capsule End', 'datetime64[ns, UTC]', 'datetime64[ns,'),
        ('Capsule Is Uncertain', 'bool', 'bool'),
        ('Capsule ID', 'object', 'object')]
    ):
        if column_name in df.columns:
            # column_type_match is a substring because the dates may not be in UTC.
            if not df[column_name].dtype.name.startswith(column_type_match):
                try:
                    df[column_name] = df[column_name].astype(column_type)
                except (TypeError, FutureWarning) as e:
                    # The Start and End columns may need to be converted instead of cast if they're all NaT values.
                    if 'astype to convert from timezone-naive dtype to timezone-aware dtype' in str(e):
                        df[column_name] = df[column_name].dt.tz_localize(tz='UTC')
                    else:
                        raise e

        else:
            df.insert(0, column_name, pd.Series(dtype=column_type))

        _convert_timezone_if_necessary(df, tz_convert)


def _convert_timezone_if_necessary(df, tz_convert):
    if tz_convert is None:
        return

    for column_name in ['Capsule Start', 'Capsule End',

                        # Vantage materialized table standard columns
                        'Created At', 'Updated At', 'Deleted At',

                        # Vantage Context standard columns
                        'Start Time', 'End Time']:

        # Ensure that NaT values do not cause Pandas to drop our timezone conversion
        if column_name in df.columns and not pd.isnull(df[column_name]).all():
            df[column_name] = df[column_name].dt.tz_convert(tz_convert)


def _process_query_row(session: Session, at_least_one_signal, calculation, shape, capsule_properties, grid, header,
                       query_df, index_to_use, row_index, pd_start, pd_end, phase, status: Status, tz_convert,
                       bounding_values, invalid_values_as, column_names, final_column_names, enums_as,
                       include_extinct):
    # _process_query_row was broken out into a function mostly to regain some indentation
    # space from the main pull() function. It's only ever called from one spot. That's why it has a ton of parameters.

    items_api = ItemsApi(session.client)

    row = query_df.loc[row_index]

    if phase == 'signals' and not _common.present(row, 'ID'):
        status.df.at[row_index, 'Result'] = 'No "ID" column - skipping'
        return None

    item_id, item_name, item_type = _get_item_details(session, header, row, query_df)
    return_type = _common.get(row, RETURN_TYPE_COLUMN, item_type)

    if _common.present(row, 'Start'):
        pd_start = pd.to_datetime(row['Start'])

    if _common.present(row, 'End'):
        pd_end = pd.to_datetime(row['End'])

    if _common.present(row, 'Calculation'):
        calculation_to_use = row['Calculation']
    else:
        calculation_to_use = calculation

    parameter_type = item_type
    if item_type == 'Asset':
        # If we're pulling assets, then we're actually pulling a calculated item (signal, condition or scalar) that
        # has been swapped to that asset. So use the swap API to find the appropriate item and then use that item's
        # ID instead of the asset's ID. Everything else just works the same as if the user had specified the swap
        # item directly.
        @request_safely(action_description=f'find swapped calculations below asset "{item_name}" {item_id}',
                        status=status,
                        default_value=(item_id, item_type, calculation_to_use))
        def _request_swap_items(_item_id, _calculation_to_use):
            swap_input = SwapInputV1()
            swap_input.swap_in = _item_id
            calc_item_id, _, _item_type = _get_item_details(session, header, calculation, query_df)

            item_dependency_output = items_api.get_formula_dependencies(id=calc_item_id)  # type: ItemDependencyOutputV1

            unique_assets = set(dep.ancestors[-1].id
                                for dep in item_dependency_output.dependencies
                                if len(dep.ancestors) > 0)

            if len(unique_assets) != 1:
                raise SPyRuntimeError('To pull data for an asset, the "calculate" parameter must be a calculated '
                                      'item that involves only one asset.')

            swap_input.swap_out = unique_assets.pop()

            swapped_item = items_api.find_swap(id=calc_item_id, body=[swap_input])  # type: ItemPreviewV1
            _item_id = swapped_item.id

            # Don't try to apply a calculation later, we've already done it via our swap activity
            _calculation_to_use = None
            return _item_id, _item_type, _calculation_to_use

        item_id, return_type, calculation_to_use = _request_swap_items(item_id, calculation_to_use)
        parameter_type = return_type

    if phase == 'signals' and \
            'Signal' not in return_type and 'Condition' not in return_type and 'Scalar' not in return_type:
        status.df.at[row_index, 'Result'] = 'Not a Signal, Condition or Scalar - skipping'
        return None

    parameter_name = _common.simplify_type(parameter_type).lower()
    parameters = [f'{parameter_name}={item_id}']
    if calculation_to_use is not None:
        formula = calculation_to_use
    else:
        formula = f'${parameter_name}'

    if 'Signal' in return_type and grid is not None:
        formula = 'resample(%s, %s)' % (formula, grid)

    if phase == 'signals' and 'Signal' in return_type:
        return (_pull_signal, session, item_name, formula, parameters, pd_start, pd_end, tz_convert, row_index,
                bounding_values, invalid_values_as, enums_as, status)

    elif phase == 'conditions' and 'Condition' in return_type:
        return (_pull_condition, session, item_name, formula, parameters, pd_start, pd_end, tz_convert,
                row, row_index, shape, capsule_properties, calculation_to_use, header, index_to_use, query_df,
                at_least_one_signal, status)

    elif phase == 'conditions' and 'Table' in return_type:
        return (_pull_materialized_table, session, item_id, pd_start, pd_end, tz_convert, row_index,
                capsule_properties, include_extinct, calculation_to_use, status)

    elif phase == 'scalars' and 'Scalar' in return_type:
        return (_pull_scalar, session, item_name, formula, parameters, row_index, index_to_use, invalid_values_as,
                status)

    elif phase == 'final':
        # Iterate over all the column names that the _pull_xxxx functions added to the DataFrame and put them in an
        # ordered list. This code forces the output DataFrame to be consistent even if the timing of completions is
        # different from run to run.
        if row_index in column_names:
            for column_name in column_names[row_index]:
                if column_name not in final_column_names:
                    final_column_names.append(column_name)

    return None


def _convert_index_timezone(ts_column, tz):
    ts_column = ts_column.tz_localize('UTC')
    return ts_column.tz_convert(tz) if tz else ts_column


def _pull_condition(session: Session, item_name, formula, parameters, pd_start, pd_end, tz_convert, row, row_index,
                    shape, capsule_properties, calculation_to_use, header, index_to_use, query_df, at_least_one_signal,
                    status: Status):
    result_df = pd.DataFrame(index=index_to_use)
    column_names = dict()

    # noinspection PyBroadException
    timer = _common.timer_start()
    capsule_count = 0
    current_start = pd_start.value
    offset = 0
    page_count = 0
    indices_to_update = {row_index}
    continuation_token = ""

    def _update_status_rows(_message, _capsule_count, _page_count, _usage):
        for _row_index_to_update in indices_to_update:
            status.send_update(_row_index_to_update, {
                'Result': _message,
                'Count': _capsule_count,
                'Pages': _page_count,
                'Time': _common.timer_elapsed(timer),
                'Data Processed': _usage
            })

    usage = Usage()
    while True:
        # When we want capsule summary statistics, fetch as a capsule table
        if shape == 'capsules' and at_least_one_signal:
            if calculation_to_use is not None:
                raise SPyRuntimeError("If shape='capsules' and at least one signal is present, calculation "
                                      "argument cannot be supplied")

            this_capsule_count, next_start, result_df, http_headers = \
                _pull_condition_as_a_table(session, item_name, formula, parameters, current_start, pd_end, tz_convert,
                                           row_index, capsule_properties, query_df, header, offset, result_df,
                                           column_names, indices_to_update)

            # Note that capsule_count here can diverge from the exact count in the output due to pagination
            capsule_count += this_capsule_count
            usage.add(http_headers)
            page_count += 1

            if this_capsule_count < session.options.pull_page_size:
                break

            if next_start == current_start:
                # This can happen if the page is full of capsules that all have the same start time
                offset += session.options.pull_page_size
            else:
                offset = 0
                current_start = next_start
        else:
            this_capsule_count, next_start, continuation_token, result_df, http_headers = \
                _pull_condition_via_formula_api(session, item_name, formula, parameters, current_start, pd_end,
                                                tz_convert, row_index, capsule_properties, shape,
                                                continuation_token, index_to_use, result_df, column_names)
            capsule_count += this_capsule_count
            usage.add(http_headers)
            page_count += 1
            if _compatibility.is_continuation_token_used():
                if continuation_token is None or continuation_token == '':
                    break
            elif this_capsule_count < session.options.pull_page_size:
                break
            else:
                if next_start == current_start:
                    # This can happen if the page is full of capsules that all have the same start time
                    raise SPyRuntimeError('Too much data: found too many capsules with same start time.\n Increase '
                                          'spy.options.pull_page_size.'
                                          'capsule count: %s\nstart: %s\nend: %s\n' %
                                          (
                                              this_capsule_count, current_start, pd_end.value,
                                          ))
                else:
                    offset = 0
                    current_start = next_start

        _update_status_rows(f'Pulling {_common.convert_to_timestamp(current_start, tz_convert)}', capsule_count,
                            page_count,
                            usage)

    _update_status_rows(f'Success', capsule_count, page_count, usage)

    return RowResult(row_index, column_names, result_df)


def _is_capsule_dupe(result_df, item_name, pd_capsule_start, pd_capsule_end):
    return 'Condition' in result_df and \
        'Capsule Start' in result_df and \
        'Capsule End' in result_df and \
        len(result_df.loc[(result_df['Condition'] == item_name) &
                          (result_df['Capsule Start'] == pd_capsule_start) &
                          (result_df['Capsule End'] == pd_capsule_end)])


def _pull_condition_via_formula_api(session: Session, item_name, formula, parameters, current_start, pd_end,
                                    tz_convert, row_index, capsule_properties, shape,
                                    continuation_token, index_to_use, result_df, column_names):
    formulas_api = FormulasApi(session.client)

    formula_run_output: dict
    # In R63, appserver removes offset and starts using continuation_token.
    kwargs = {
        'formula': formula,
        'parameters': parameters,
        'start': '%d ns' % current_start,
        'end': '%d ns' % pd_end.value,
        'limit': session.options.pull_page_size,

        # It's about 50% faster to process the condition data as json - see CRAB-33451
        '_response_type': 'json'
    }
    if _compatibility.is_continuation_token_used():
        kwargs['continuation_token'] = continuation_token
    else:
        kwargs['offset'] = 0
    formula_run_output, _, http_headers = formulas_api.run_formula_with_http_info(**kwargs)

    next_start = current_start
    capsules_output: dict = formula_run_output['capsules']
    check_for_dupes = True
    columns = dict()
    if shape == 'samples':
        # In this case, we are creating a signal-like representation of the condition using 0s and 1s, just like the
        # Excel and OData exports.

        columns[item_name] = pd.Series(0, index_to_use, dtype=float)
        for capsule in capsules_output['capsules']:  # type: dict
            pd_capsule_start = _common.convert_to_timestamp(
                capsule['start'] if capsule.get('start', None) is not None else 0, tz_convert)
            pd_capsule_end = _common.convert_to_timestamp(
                capsule['end'] if capsule.get('end', None) is not None else 7258118400000000000, tz_convert)

            # Mark Derbecker 2019-12-17:
            # I've tried a few ways of making this happen and so far this method seems to be the most efficient: Start
            # with a Series full of zeros (but with the index that corresponds to the already-existing output
            # DataFrame) and use the Series.loc[] indexer to set the values to one if they're within the capsule
            # boundary.
            columns[item_name].loc[(columns[item_name].index >= pd_capsule_start) &
                                   (columns[item_name].index <= pd_capsule_end)] = 1

            for prop in capsule['properties']:  # type: dict
                # We need to create a column name that is unique for the item / property combination
                colname = '%s - %s' % (item_name, prop['name'])
                if colname not in columns:
                    # Here we start with a NaN-filled series, since we're populating property values (not 1s and 0s).
                    columns[colname] = pd.Series(np.nan, index_to_use, dtype=object)

                # Note here that overlapping capsules with different properties will result in "last one wins"
                columns[colname].loc[(columns[colname].index >= pd_capsule_start) &
                                     (columns[colname].index <= pd_capsule_end)] = prop['value']

        column_names[row_index] = list()
        for col, series in columns.items():
            result_df[col] = series
            column_names[row_index].append(col)
    else:
        # In this case, we're creating a more straightforward table where each capsule is a row, complete with item
        # properties.

        capsule_dict_list = list()
        column_names[row_index] = ['Condition', 'Capsule Start', 'Capsule End', 'Capsule Is Uncertain', 'Capsule ID']

        for capsule in capsules_output['capsules']:  # type: dict
            capsule_start = capsule.get('start', None)
            capsule_end = capsule.get('end', None)
            pd_capsule_start = _common.convert_to_timestamp(capsule_start, tz_convert)
            pd_capsule_end = _common.convert_to_timestamp(capsule_end, tz_convert)
            if check_for_dupes and _is_capsule_dupe(result_df, item_name, pd_capsule_start, pd_capsule_end):
                # This can happen as a result of pagination
                continue

            check_for_dupes = False

            capsule_dict = {
                'Condition': item_name,
                'Capsule Start': pd_capsule_start,
                'Capsule End': pd_capsule_end,
                'Capsule Is Uncertain': bool(capsule.get('isUncertain', False))
            }

            if 'id' in capsule:
                capsule_dict['Capsule ID'] = capsule['id']

            for prop in capsule['properties']:  # type: dict
                if capsule_properties is not None and prop['name'] not in capsule_properties:
                    continue

                capsule_dict[prop['name']] = prop.get('value', None)
                if prop['name'] not in column_names[row_index]:
                    column_names[row_index].append(prop['name'])

            capsule_dict_list.append(capsule_dict)
            if not pd.isna(capsule_start) and capsule_start > next_start:
                next_start = capsule_start

        df_to_append = pd.DataFrame(capsule_dict_list, columns=column_names[row_index])
        if len(result_df) > 0 and len(df_to_append) > 0:
            result_df = pd.concat([result_df, df_to_append])
        elif len(result_df) == 0:
            result_df = df_to_append

    return len(capsules_output['capsules']), next_start, capsules_output.get('continuationToken', ''), result_df, \
        http_headers


def _build_limit_fragment(offset, limit):
    start_row = offset + 1  # Table object is 1-based, not 0
    end_row = start_row + limit
    return f'.limit({start_row}, {end_row})'


def _build_stat_formula_fragment_and_update_parameters(session: Session, signals_df, header, parameters,
                                                       indices_to_update, query_df):
    signal_id_to_stats = defaultdict(list)
    for signal_index, signal_row in signals_df.iterrows():
        signal_item_id, signal_item_name, signal_item_type = _get_item_details(session, header, signal_row, query_df)
        statistic = signal_row['Statistic'] if 'Statistic' in signal_row else 'average'

        stat_formula = _common.statistic_to_aggregation_function(statistic, allow_condition_stats=False).split('(')[0]
        stat_formula = stat_formula + '()'  # Some stats come back without the parenthesis, this forces it on each
        stat_header = f'{signal_item_name} ({statistic})'

        signal_id_to_stats[signal_item_id].append((stat_formula, stat_header))
        indices_to_update.add(signal_index)

    final_fragment = ''
    stat_headers = list()
    for (variable_count, (signal_id, stat_list)) in enumerate(signal_id_to_stats.items()):
        stat_formulas, item_headers = zip(*stat_list)
        short_id = f's{variable_count}'
        parameters.append(f'{short_id}={signal_id}')
        joined_stats = ', '.join(stat_formulas)
        final_fragment += f".addStatColumn('{short_id}', ${short_id}, {joined_stats})"
        stat_headers.extend(item_headers)

    return final_fragment, parameters, stat_headers


def _build_group_segment(property_columns):
    joined_properties = ', '.join(["'%s'" % prop for prop in property_columns])
    return f'group({joined_properties})'


def _pull_condition_as_a_table(session: Session, item_name, formula, parameters, current_start, pd_end, tz_convert,
                               row_index, capsule_properties, query_df, header, offset, result_df,
                               column_names, indices_to_update):
    if not isinstance(capsule_properties, list):
        capsule_properties = list()
    required_columns = ['Capsule ID', 'Original Uncertainty', 'Condition ID', 'Start', 'End']
    # place capsule sort key as the last property column
    all_property_columns = required_columns + capsule_properties + ['Capsule SortKey']

    signals_df = query_df[query_df['Type'].str.endswith('Signal', na=False)]

    formulas_api = FormulasApi(session.client)
    start = pd.Timestamp(current_start)
    start_string = start.isoformat() + 'Z'
    end_string = pd.Timestamp(pd_end.value).isoformat() + 'Z'
    capsule = f"capsule('{start_string}', '{end_string}')"
    group_segment = _build_group_segment(all_property_columns)
    limit_fragment = _build_limit_fragment(offset, session.options.pull_page_size)
    property_column_fragment = f'capsuleTable({capsule}, CapsuleBoundary.Overlap, {group_segment}, {formula})'
    sort_fragment = ".sort('Capsule ID', 'inv, asc', 'Condition Id', 'asc', 'Capsule SortKey', 'asc')"
    stat_fragment, parameters, stat_headers = _build_stat_formula_fragment_and_update_parameters(
        session, signals_df, header, parameters, indices_to_update, query_df)
    formula = property_column_fragment + stat_fragment + sort_fragment + limit_fragment
    limit = session.options.pull_page_size

    formula_run_output: dict
    formula_run_output, _, http_headers = formulas_api.run_formula_with_http_info(
        formula=formula,
        parameters=parameters,
        limit=limit,

        # It's about 50% faster to process the condition data as json - see CRAB-33451
        _response_type='json')

    # Construct a dictionary to map column names to column indices
    all_column_names = all_property_columns + stat_headers
    column_names_to_indexes = {all_column_names[i]: i for i in range(len(all_column_names))}

    next_start = current_start
    formula_table = formula_run_output['table']['data']  # type: list
    if len(formula_table) > limit:
        formula_table = formula_table[:limit]
    check_for_dupes = True
    capsule_df_rows = list()

    column_names[row_index] = ['Condition', 'Capsule Start', 'Capsule End', 'Capsule Is Uncertain', 'Capsule ID']

    for row in formula_table:
        capsule_start = pd.Timestamp(row[column_names_to_indexes['Start']]).value
        capsule_end = pd.Timestamp(row[column_names_to_indexes['End']]).value
        pd_capsule_start = _common.convert_to_timestamp(capsule_start, tz_convert)
        pd_capsule_end = _common.convert_to_timestamp(capsule_end, tz_convert)
        if check_for_dupes and _is_capsule_dupe(result_df, item_name, pd_capsule_start, pd_capsule_end):
            # This can happen as a result of pagination
            continue

        check_for_dupes = False

        capsule_dict = {
            'Condition': item_name,
            'Capsule Start': pd_capsule_start,
            'Capsule End': pd_capsule_end,
            'Capsule Is Uncertain': bool(row[column_names_to_indexes['Original Uncertainty']])
        }

        capsule_id = row[column_names_to_indexes['Capsule ID']]
        if capsule_id is not None:
            capsule_dict['Capsule ID'] = capsule_id

        # capsuleSortKey divides the properties and the stat columns in returned table
        lower_bound = all_column_names.index('Capsule SortKey') + 1
        additional_columns_for_this_row = capsule_properties + all_column_names[lower_bound:]
        for column_name in additional_columns_for_this_row:
            capsule_dict[column_name] = row[column_names_to_indexes[column_name]]
            if column_name not in column_names[row_index]:
                column_names[row_index].append(column_name)

        capsule_df_rows.append(capsule_dict)

        if not pd.isna(capsule_start) and capsule_start > next_start:
            next_start = capsule_end

    df_to_append = pd.DataFrame(capsule_df_rows, columns=column_names[row_index])
    formula_result_df = (pd.concat([result_df, df_to_append]) if len(result_df) != 0 else df_to_append)

    return len(formula_table), next_start, formula_result_df, http_headers


def _pull_signal(session: Session, item_name, formula, parameters, pd_start, pd_end, tz_convert,
                 row_index=None, bounding_values: bool = False, invalid_values_as=np.nan,
                 enums_as: str = 'string', status: Status = None):
    formulas_api = FormulasApi(session.client)

    # noinspection PyBroadException
    series = None
    timer = _common.timer_start()
    current_start = pd_start
    last_key = 0
    page_count = 0
    usage = Usage()
    continuation_token = ''

    while True:
        start_string = '%d ns' % current_start.value
        end_string = '%d ns' % pd_end.value
        formula_run_output: dict
        kwargs = {
            'formula': formula,
            'parameters': parameters,
            'start': start_string,
            'end': end_string,
            'limit': session.options.pull_page_size,

            # It's about 300% faster to process the signal data as json - see CRAB-33451
            '_response_type': 'json'
        }
        if _compatibility.is_continuation_token_used():
            kwargs['continuation_token'] = continuation_token
        else:
            kwargs['offset'] = 0

        formula_run_output, status_code, http_headers = formulas_api.run_formula_with_http_info(**kwargs)

        usage.add(http_headers)

        if not _common.present(formula_run_output, 'samples'):
            # noinspection PyStringFormat
            raise SPyRuntimeError('formula_run_output.samples is None.\n'
                                  'status_code: %d\nformula: %s\nparameters: %s\n'
                                  'start: %s\nend:%s\n'
                                  'formula_run_output:\n%s' %
                                  (status_code, formula, parameters,
                                   start_string, end_string,
                                   formula_run_output.__repr__()))

        series_samples_output = formula_run_output['samples']

        def _keep_sample(_sample_output):
            _key = _sample_output['key']
            if _key <= last_key:
                return False

            if bounding_values:
                return True

            if _key < pd_start.value:
                return False

            if _key > pd_end.value:
                return False

            return True

        # Filter out the samples before breaking them into timestamps/values lists. This is the fastest.
        # https://bitbucket.org/seeq12/crab/pull-requests/8872/latest-pypi-changes-up-to-seeq-module-v133/diff#comment-134308908
        filtered_samples = [sample_output for sample_output in series_samples_output['samples'] if
                            _keep_sample(sample_output)]
        timestamps = [sample['key'] for sample in filtered_samples]

        def _value_or_default(sample_value):
            if pd.isna(sample_value):
                return invalid_values_as
            else:
                return sample_value

        values = [_value_or_default(_sanitize_pi_enums(sample.get('value', None), enums_as))
                  for sample in filtered_samples]

        time_index = _convert_index_timezone(pd.DatetimeIndex(timestamps), tz_convert)

        new_series = pd.Series(values, index=time_index, dtype=object)
        series = new_series if series is None else pd.concat([series, new_series])

        page_count += 1

        if _compatibility.is_continuation_token_used():
            continuation_token = series_samples_output.get('continuationToken', '')
            if continuation_token is None or continuation_token == '':
                break
        else:
            if len(series_samples_output['samples']) < session.options.pull_page_size:
                break

        if len(series) > 0:
            last_key = series.index[-1].value

        if len(time_index) == 0:
            raise SPyException("An error was encountered. spy.options.pull_page_size may be set to 1. Check for typos "
                               "or try setting spy.options.pull_page_size to a higher value.")

        if not _compatibility.is_continuation_token_used() and time_index[-1].value > current_start.value:
            current_start = time_index[-1]

        if status is not None:
            status.send_update(row_index, {
                'Result': f'Pulling: {current_start}',
                'Count': len(series),
                'Pages': page_count,
                'Time': _common.timer_elapsed(timer),
                'Data Processed': usage
            })
    series = series.infer_objects()

    if status is not None:
        status.send_update(row_index, {
            'Result': 'Success',
            'Count': len(series),
            'Pages': page_count,
            'Time': _common.timer_elapsed(timer),
            'Data Processed': usage
        })

    return RowResult(row_index, {row_index: [item_name]}, pd.DataFrame({item_name: series}))


def _sanitize_pi_enums(value, enums_as):
    if enums_as is None:
        return value
    if not isinstance(value, str):
        return value
    match = ENUM_PATTERN.search(value)
    if match is None:
        return value
    if enums_as == 'numeric':
        return int(match.group(1))
    if enums_as == 'string':
        return match.group(2)
    return int(match.group(1)), match.group(2)


def _pull_scalar(session: Session, item_name, formula, parameters, row_index, index_to_use, invalid_values_as,
                 status: Status):
    formulas_api = FormulasApi(session.client)
    timer = _common.timer_start()
    usage = Usage()

    formula_run_output, _, http_headers = formulas_api.run_formula_with_http_info(
        formula=formula, parameters=parameters)

    usage.add(http_headers)

    status.send_update(row_index, {'Result': 'Success', 'Count': 1, 'Time': _common.timer_elapsed(timer),
                                   'Data Processed': usage})

    if len(index_to_use) == 0:
        index_to_use = pd.Series([0])

    result_df = pd.DataFrame(index=index_to_use)
    value = invalid_values_as if formula_run_output.scalar.value is None else formula_run_output.scalar.value
    result_df[item_name] = value
    return RowResult(row_index, {row_index: [item_name]}, result_df)


class RowResult:
    def __init__(self, row_index, column_names, result):
        self.row_index = row_index
        self.column_names = column_names
        self.result = result


def _get_item_details(session: Session, header, row, query_df):
    # This is a somewhat complex function that tries its best to pick a column header (item_name) for the output
    # DataFrame by either honoring the user's "header" argument or auto-picking something that makes sense.

    items_api = ItemsApi(session.client)

    item_id = _common.get(row, 'ID')

    # noinspection PyTypeChecker
    item = None

    if _common.present(row, 'Type'):
        item_type = _common.get(row, 'Type')
    else:
        item = items_api.get_item_and_all_properties(id=item_id)  # type: ItemOutputV1
        item_type = item.type

    if header.upper() == 'ID':
        item_name = item_id
    elif header == '__auto__' and _common.present(row, 'Header'):
        item_name = _common.get(row, 'Header')
    elif _common.present(row, header):
        item_name = _common.get(row, header)
    else:
        if not item:
            item = items_api.get_item_and_all_properties(id=item_id)  # type: ItemOutputV1

        item_name = ''
        if header == '__auto__' and _common.present(row, 'Path'):
            old_asset_format = (hasattr(query_df, 'spy') and hasattr(query_df.spy, 'old_asset_format') and
                                query_df.spy.old_asset_format)
            item_name = _common.get(row, 'Path') + ' >> '
            if (item_type != 'Asset' or old_asset_format) and _common.present(row, 'Asset'):
                item_name += _common.get(row, 'Asset') + ' >> '

        if header in ['__auto__', 'Name']:
            item_name += item.name
        elif header == 'Description':
            item_name += item.description
        else:
            prop = [p.value for p in item.properties if p.name == header]
            if len(prop) == 0:
                raise SPyValueError(f'header argument invalid: Property "{header}" not found')
            else:
                item_name += prop[0]

    return item_id, item_name, item_type


def estimate_auto_grid(session: Session, query_df, pd_start, pd_end, status):
    if 'Estimated Sample Period' in query_df:
        status.update('Retrieving estimated sample period from query DataFrame', Status.RUNNING)
        samplings = [x for x in query_df['Estimated Sample Period'].to_list() if not pd.isna(x)]
    else:
        status.update('Estimating sample period of each signal from %s to %s' % (pd_start, pd_end), Status.RUNNING)
        signals_df = query_df[query_df['Type'].isin(['StoredSignal', 'CalculatedSignal'])]
        samplings = list()
        formulas_api = FormulasApi(session.client)
        for signal_id in signals_df['ID']:

            sampling_formula = f"$signal.estimateSamplePeriod(capsule('{pd_start.isoformat()}','{pd_end.isoformat()}'))"

            formula_run_output = safely(
                lambda: formulas_api.run_formula(formula=sampling_formula, parameters=[f"signal={signal_id}"]),
                action_description=f'get sample rate for Signal {signal_id}',
                additional_errors=[400],
                status=status)
            if formula_run_output and formula_run_output.scalar.value is not None:
                samplings.append(
                    pd.to_timedelta(formula_run_output.scalar.value, unit=formula_run_output.scalar.uom))

    if len(samplings) == 0:
        raise SPyRuntimeError("Could not determine sample period for any of the signals in the query. "
                              "There might not be enough data for the time period specified")

    median = np.median(np.array(samplings))
    if isinstance(median, pd.Timedelta):
        nanoseconds = median.value
    elif isinstance(median, np.timedelta64):
        nanoseconds = median.real
    else:
        raise SPyTypeError(f"Estimated Sample Period column data type {type(median)} not recognized")

    return str(int(nanoseconds / 1_000_000)) + 'ms'


HEADER_NAME_MAP: dict = {
    'item id': 'Condition ID',
    'datum id': 'Capsule ID',
    'Start': 'Capsule Start',
    'End': 'Capsule End',
    'Duration': 'Capsule Duration',
    'IsUncertain': 'Capsule Is Uncertain',
    'Deleted at': 'Deleted At',
    'Deleted at_uom': 'Deleted At UOM',
    'created at': 'Created At',
}


def _unix_ns_to_dt(x):
    return pd.Timestamp(x, unit="ns", tz="UTC")


def _unix_ns_to_td(x):
    return pd.Timedelta(x, unit="ns")


COLUMN_DTYPE_MAP: dict = {
    'UUID': 'string',
    'TEXT': 'string',
    'BOOLEAN': 'bool',
    'NUMERIC': 'float64',
    'TIMESTAMPTZ': 'datetime64[ns, UTC]'
}

COLUMN_DTYPE_OVERRIDES: dict = {
    'Capsule Start': 'datetime64[ns, UTC]',
    'Capsule End': 'datetime64[ns, UTC]',
    'Capsule Duration': 'timedelta64[ns]',
    'Created At': 'datetime64[ns, UTC]',
    'Updated At': 'datetime64[ns, UTC]',
    'Deleted At': 'datetime64[ns, UTC]'
}

COLUMN_CAST_MAP: dict = {
    'TIMESTAMPTZ': _unix_ns_to_dt
}

COLUMN_TYPE_OVERRIDES: dict = {
    'Capsule Start': _unix_ns_to_dt,
    'Capsule End': _unix_ns_to_dt,
    'Capsule Duration': _unix_ns_to_td,
    'Deleted At': _unix_ns_to_dt,
    'context_comments.start_time': _unix_ns_to_dt,
    'context_comments.end_time': _unix_ns_to_dt,
    'context_labels.start_time': _unix_ns_to_dt,
    'context_labels.end_time': _unix_ns_to_dt,
    'context_numeric.start_time': _unix_ns_to_dt,
    'context_numeric.end_time': _unix_ns_to_dt,
}


def _pull_materialized_table(
        session: Session, table_id, pd_start, pd_end, tz_convert, row_index, capsule_properties, include_extinct,
        calculation_to_use, status: Status):
    if calculation_to_use is not None:
        raise SPyValueError(
            "calculation argument cannot be supplied when pulling a materialized table")

    # noinspection PyBroadException
    timer = _common.timer_start()
    indices_to_update = {row_index}
    user_map = dict()

    def _update_status_rows(_message, _capsule_count, _page_count, _usage):
        for _row_index_to_update in indices_to_update:
            status.send_update(_row_index_to_update, {
                'Result': _message,
                'Count': _capsule_count,
                'Pages': _page_count,
                'Time': _common.timer_elapsed(timer),
                'Data Processed': _usage
            })

    usage = Usage()
    context_rows = list()
    final_rows = list()
    most_recent_capsule_start = pd_start
    current_page_count = 1
    current_key = None
    current_row = None
    the_headers = None

    generator = _yield_materialized_table_rows(session, table_id, pd_start, pd_end, include_extinct)

    _update_status_rows(f'Pulling {_common.convert_timestamp_timezone(most_recent_capsule_start, tz_convert)}',
                        0, current_page_count, usage)

    def _key(_d):
        return _d['Condition ID'], _d['Capsule ID']

    # noinspection PyTypeChecker
    for new_headers, new_row, new_page_count in generator:
        if the_headers is None:
            the_headers = new_headers

        if current_row is None:
            if new_row is None:
                break

            current_row_dict, context_row_dict = _materialized_table_dicts_from_list(new_headers, new_row)
            context_rows.append(context_row_dict)
            current_key = _key(current_row_dict)
            current_row = new_row
            continue

        final_row_dict, context_row_dict = _materialized_table_dicts_from_list(the_headers, current_row)
        context_rows.append(context_row_dict)

        most_recent_capsule_start = final_row_dict['Capsule Start']

        new_key = None
        new_row_dict = None
        if new_row is not None:
            new_row_dict, _ = _materialized_table_dicts_from_list(the_headers, new_row)
            new_key = _key(new_row_dict)

        if new_key != current_key:
            _rollup_context_columns(session, final_row_dict, context_rows, tz_convert, user_map)
            final_rows.append(final_row_dict)
            context_rows = list()

        if new_key is None:
            break

        if current_page_count != new_page_count:
            _update_status_rows(
                f'Pulling {_common.convert_timestamp_timezone(most_recent_capsule_start, tz_convert)}',
                len(final_rows), new_page_count, usage)

        current_page_count = new_page_count
        current_row = new_row
        current_key = _key(new_row_dict)

    columns_types_dict = {
        HEADER_NAME_MAP.get(h['name'], h['name']): COLUMN_DTYPE_OVERRIDES.get(
            HEADER_NAME_MAP.get(h['name'], h['name']), COLUMN_DTYPE_MAP.get(h['type'], 'object'))
        for h in the_headers if not h['name'].startswith('context_')
    }
    columns_types_dict.update({
        'Flagged': 'boolean',
        'Reviewed': 'boolean',
        'Labels': 'object',
        'Comments': 'string',
        'Vantage Context': 'object'
    })

    final_df = pd.DataFrame(final_rows)

    for expected_column, column_type in columns_types_dict.items():
        if expected_column not in final_df.columns:
            final_df[expected_column] = pd.Series(dtype=column_type)

    if capsule_properties is not None:
        columns_to_keep = STANDARD_CONDITION_COLUMNS + capsule_properties
        final_df = final_df[[col for col in columns_to_keep if col in final_df.columns]]

    _update_status_rows(f'Success', len(final_rows), current_page_count, usage)

    return RowResult(row_index, {row_index: final_df.columns.to_list()}, final_df)


def _yield_materialized_table_rows(session: Session, table_id, pd_start, pd_end, include_extinct):
    cursor_key = None
    limit = session.options.table_pull_page_size
    page_count = 1

    while True:
        graphql_output = graphql_get_evidence(session, table_id, pd_start, pd_end, cursor_key, limit,
                                              include_extinct)

        table = graphql_output['data']['table']
        has_more = table['hasMore']
        cursor_key = table.get('cursorKey')

        rows = table['rows']

        for row in rows:
            yield table['headers'], row, page_count

        page_count += 1

        if not has_more:
            break

    # Final yield to indicate completion
    yield table['headers'], None, page_count


def graphql_get_evidence(session, table_id, pd_start, pd_end, cursor_key, limit, include_extinct) -> dict:
    graphql_api = GraphQLApi(session.client)

    # noinspection PyTypeChecker
    graphql_output: dict = graphql_api.graphql(body=GraphQLInputV1(
        query="""
    query GetTable($id: String!, $filter: FilterInput, $limit: Int!, $columnsToInclude: [String!], $cursor: CursorInput) {
        table(id: $id, filter: $filter, limit: $limit, columnsToInclude: $columnsToInclude, cursor: $cursor) {
            rows, headers { name, type }, hasMore, cursorKey
        }
    }
    """, variables={
            "id": table_id,
            "limit": limit,
            "cursor": {"cursorKey": cursor_key} if cursor_key is not None else None,
            "filter": {
                "compositeFilter": {
                    "operation": "AND",
                    "filter1": {
                        "compositeFilter": {
                            "operation": "OR",
                            "filter1": {
                                "compositeFilter": {
                                    "operation": "OR",
                                    "filter1": {
                                        "compositeFilter": {
                                            "operation": "AND",
                                            "filter1": {
                                                "valueFilter": {
                                                    "columnName": "Start",
                                                    "filterType": "GREATER_THAN_OR_EQUAL",
                                                    "value": pd_start.value
                                                }
                                            },
                                            "filter2": {
                                                "valueFilter": {
                                                    "columnName": "Start",
                                                    "filterType": "LESS_THAN_OR_EQUAL",
                                                    "value": pd_end.value
                                                }
                                            }
                                        }
                                    },
                                    "filter2": {
                                        "compositeFilter": {
                                            "operation": "AND",
                                            "filter1": {
                                                "valueFilter": {
                                                    "columnName": "End",
                                                    "filterType": "GREATER_THAN_OR_EQUAL",
                                                    "value": pd_start.value
                                                }
                                            },
                                            "filter2": {
                                                "valueFilter": {
                                                    "columnName": "End",
                                                    "filterType": "LESS_THAN_OR_EQUAL",
                                                    "value": pd_end.value
                                                }
                                            }
                                        }
                                    }
                                }
                            },
                            "filter2": {
                                "compositeFilter": {
                                    "operation": "AND",
                                    "filter1": {
                                        "valueFilter": {
                                            "columnName": "Start",
                                            "filterType": "LESS_THAN_OR_EQUAL",
                                            "value": pd_start.value
                                        }
                                    },
                                    "filter2": {
                                        "valueFilter": {
                                            "columnName": "End",
                                            "filterType": "GREATER_THAN_OR_EQUAL",
                                            "value": pd_end.value
                                        }
                                    }
                                }
                            }
                        }
                    },
                    "filter2": {
                        "compositeFilter": {
                            "operation": "AND",
                            "filter1": {
                                "simpleFilter": {
                                    "columnName": "Deleted at",
                                    "filterType": "IS_NULL"
                                }
                            },
                            "filter2": {
                                "valueFilter": {
                                    "columnName": "Last Capsule State",
                                    "filterType": "NOT_IN",
                                    "value": [] if include_extinct else ["EXTINCT"]
                                }
                            }
                        }
                    }
                }
            },
            "columnsToInclude": None  # Include all columns
        }),
        _response_type='json'
    )

    return graphql_output


def _get_user_output(session: Session, user_id: str, user_map: dict) -> Optional[UserOutputV1]:
    if not user_id:
        return None

    if user_id in user_map:
        return user_map[user_id]

    users_api = UsersApi(session.client)

    try:
        user_output = users_api.get_user(id=user_id)
    except ApiException as e:
        reason = '<unknown error>'
        if e.status == 403:
            reason = '<redacted>'
        elif e.status == 404:
            reason = '<user ID not found>'
        user_output = UserOutputV1(
            id=user_id,
            username=reason,
            first_name=reason,
            last_name=reason
        )

    user_map[user_id] = user_output
    return user_output


def _get_user_name(user_output: UserOutputV1) -> str:
    user_name = ''
    if user_output.first_name is not None:
        user_name += user_output.first_name
    if user_output.last_name is not None:
        if user_name != '':
            user_name += ' '
        user_name += user_output.last_name
    if user_name == '':
        user_name = user_output.username
    return user_name


def _rollup_context_columns(session: Session, final_row_dict, context_rows, tz_convert, user_map):
    """
    Transform a DataFrame by rolling up context columns into sub-DataFrames.

    Groups rows by Condition ID and Capsule ID, then consolidates all context columns
    (comments, labels, numeric) into a single 'Vantage Context' column containing sub-DataFrames.
    """

    # Helper function to convert snake_case to Title Case
    def snake_to_title(s):
        result = s.replace('_', ' ').title()
        # Replace standalone "Id" with "ID" (but not when part of another word like "void")
        result = re.sub(r'\bId\b', 'ID', result)
        return result

    # Identify context type prefixes and their display names
    context_type_map = {
        'context_comments': 'Comment',
        'context_labels': 'Label',
        'context_numeric': 'Numeric'
    }

    final_context_rows = list()
    context_subdf = None
    for context_row in context_rows:
        # Process each context type that has data in this row
        for prefix, display_name in context_type_map.items():
            # Check if this row has any non-null values for this context type
            prefix_cols = [col for col in context_row.keys() if col.startswith(prefix + '.')]
            non_null_cols = [col for col in prefix_cols if pd.notna(context_row[col])]

            if not non_null_cols:
                continue

            # Extract all columns for this context type
            context_data = {}
            for col in prefix_cols:
                sub_col_name = col.split('.', 1)[1]
                title_col_name = snake_to_title(sub_col_name)
                value = context_row[col]
                if pd.notna(value):
                    context_data[title_col_name] = value

            if not context_data:
                continue

            # If we found context data, add it as a row
            user_output = _get_user_output(session, context_data['Creator ID'], user_map)
            context_data['Context Type'] = display_name
            context_data['Creator Username'] = user_output.username
            context_data['Creator First Name'] = user_output.first_name
            context_data['Creator Last Name'] = user_output.last_name
            final_context_rows.append(context_data)

        # Create context sub-DataFrame if we have any context rows
        final_row_dict['Vantage Context'] = None
        if final_context_rows:
            # Get all unique column names across all context rows
            all_cols = set()
            for final_context_row in final_context_rows:
                all_cols.update(final_context_row.keys())

            # Ensure Context Type is first column, then others sorted
            ordered_cols = ['Context Type'] + sorted([col for col in all_cols if col != 'Context Type'])

            if len(final_context_rows) > 0:
                # Create DataFrame - pandas will automatically include all keys from all dicts
                context_subdf = pd.DataFrame(final_context_rows)

                # Reorder columns (only include columns that exist in the DataFrame)
                existing_cols = [col for col in ordered_cols if col in context_subdf.columns]
                context_subdf = context_subdf[existing_cols]

                # Remove duplicates (in case there are any)
                context_subdf = context_subdf.drop_duplicates().reset_index(drop=True)

                _convert_timezone_if_necessary(context_subdf, tz_convert)

                final_row_dict['Vantage Context'] = context_subdf

    built_in_labels = {'flag': 'Flagged', 'reviewed': 'Reviewed'}
    for boolean_label, boolean_column in built_in_labels.items():
        final_row_dict[boolean_column] = False
        if context_subdf is not None and (
                'Category Name' in context_subdf.columns and 'Label Name' in context_subdf.columns):
            boolean_label_df = context_subdf[(context_subdf['Context Type'] == 'Label') &
                                             (context_subdf['Category Name'].str.startswith('__Vantage.')) &
                                             (context_subdf['Label Name'] == f'{boolean_label}')]
            final_row_dict[boolean_column] = not boolean_label_df.empty

    final_row_dict['Labels'] = set()
    if context_subdf is not None and (
            'Category Name' in context_subdf.columns and 'Label Name' in context_subdf.columns):
        label_column_candidates_df = context_subdf[(context_subdf['Context Type'] == 'Label') &
                                                   (~context_subdf['Category Name'].isin({
                                                       '__Seeq.suppression',
                                                       '__Vantage.flag',
                                                       '__Vantage.reviewed',
                                                       '__Vantage.suppression'
                                                   }))]

        if not label_column_candidates_df.empty:
            unique_labels = set(label_column_candidates_df['Label Name'].unique().tolist())
            final_row_dict['Labels'] = unique_labels

    comments_str = ''
    if context_subdf is not None:
        comment_candidates_df = context_subdf[context_subdf['Context Type'] == 'Comment'].sort_values(
            'Created At')

        for _, comment_row in comment_candidates_df.iterrows():
            user_output = _get_user_output(session, comment_row['Creator ID'], user_map)
            user_name = _get_user_name(user_output)
            if len(comments_str) > 0:
                comments_str += '\n'
            comments_str += f"[{comment_row['Created At']} {user_name}]\n{comment_row['Comment']}\n"

    final_row_dict['Comments'] = comments_str


def _materialized_table_dicts_from_list(headers, row):
    if row is None:
        pass

    final_row_dict = dict()
    context_row_dict = dict()
    for i in range(len(headers)):
        header = headers[i]
        name = HEADER_NAME_MAP.get(header['name'], header['name'])
        value = row[i]
        value_cast = COLUMN_TYPE_OVERRIDES.get(name, COLUMN_CAST_MAP.get(header['type']))
        relevant_dict = context_row_dict if name.startswith('context_') else final_row_dict
        if value_cast is not None and pd.notna(value):
            relevant_dict[name] = value_cast(value)
        else:
            relevant_dict[name] = value

    return final_row_dict, context_row_dict
