#AUTOGENERATED! DO NOT EDIT! File to edit: dev/01_experiments.ipynb (unless otherwise specified).

__all__ = ['snapshot_protobufs', 'fetch_experiments_data']

#Cell
from functools import partial
import os
import pandas as pd

from accio import db
from accio.prototool import ProtoTool

#Cell
def snapshot_protobufs(snapshot, limit=None, verbose=True):
    """Find samples obtained via the experiments UI from a specified snapshot.

    Parameters
    -----------
    snapshot: int
        Number of snapshot in experiments portal.
    verbose: bool
        If True, print the number of samples retrieved.

    Returns
    --------
    dict
        Keys are s3 paths, values are dictionaries containing parsed protobuf
        attributes as well as data from experiment_snapshot_sites table.
    """
    q = f"""select s3_bucket_location, url, id, experiment_snapshot_id,
        account_id, sub_account_id, org_id, ou, seq_id, created_at
        from experiment_snapshot_sites
        where experiment_snapshot_id = {snapshot}
        """
    if limit: q += f'limit {limit}'

    # Get df w/ 1 row for each page in the sample.
    rows = db.query(q, 'experiments-prod', 'reader')
    if verbose:
        print(f'{rows.shape[0]} samples retrieved.')

    # Map s3 path to dict with URL and ID, then retrieve protobuf files.
    path2data = {f's3://{row[0]}/scraped_dom.proto': dict(row[1:])
                 for i, row in rows.iterrows()}
    output = ProtoTool().buf_to_dict(path2data.keys())

    # Add URLs and IDs to feature dicts.
    return {path: {**buf, **path2data[path]} for path, buf in output.items()}

#Cell
def fetch_experiments_data(*snapshot_ids, df=True, s3_paths=False, limit=None):
    """
    Parameters
    -----------
    snapshot_ids: int (1 or more)
        Integer IDs corresponding to the relevant experiment snapshots.
        These can be found in the experiments V1 portal on the `Snapshots`
        page.
    df: bool (default True)
        If True, return the data as a single Pandas DataFrame.
    s3_paths: bool (default False)
        If True, include the urls for the protobufs in S3 in the output
        dataframe or dict. When returning a dictionary (i.e. df=False),
        s3_paths will always be present as keys.
    limit: int or None
        If int, sets the max combined number of rows to retrieve from db.
        If None, retrieve all rows. Note: if using n snapshot
        id's, each snapshot will have at most limit//n rows.

    Returns
    --------
    pd.DataFrame
        1 row for each example collected from snapshot.

    or

    dict[str, dict]
        Maps S3 file location to dict containing file attributes and site
        metadata.
    """
    if limit: limit //= len(snapshot_ids)
    proto_rounds = (snapshot_protobufs(i, limit) for i in snapshot_ids)
    data = {k: v for group in proto_rounds for k, v in group.items()}
    if df:
        if s3_paths:
            data = pd.DataFrame(data).T.reset_index()\
                     .rename({'index': 's3'}, axis=1)
        else:
            data = pd.DataFrame(list(data.values()))
    return data