import pandas as pd
import redcap as rc

# Taken from http://pycap.readthedocs.io/en/latest/deep.html#dealing-with-large-exports
# and adapted to scope down to forms
def chunked_export(project, form, chunk_size=100, verbose=True, **kwargs):
    def chunks(l, n):
        """Yield successive n-sized chunks from list l"""
        for i in range(0, len(l), n):
            yield l[i:i+n]
    record_list = project.export_records(fields=[project.def_field])
    records = [r[project.def_field] for r in record_list]
    try:
        response = None
        record_count = 0
        for record_chunk in chunks(records, chunk_size):
            record_count = record_count + chunk_size
            chunked_response = project.export_records(
                records=record_chunk,
                fields=[project.def_field],
                forms=[form],
                format='df',
                df_kwargs={'low_memory': False}.update(kwargs))
            if response is not None:
                response = pd.concat([response, chunked_response], axis=0)
            else:
                response = chunked_response
    except rc.RedcapError:
        msg = "Chunked export failed for chunk_size={:d}".format(chunk_size)
        raise ValueError(msg)
    else:
        return response


def load_all_forms(api, arm='1'):
    fem = api.export_fem(format='df', arms=arm)
    forms = fem['form'].unique().tolist()
    form_n = len(forms)
    dfs = []
    for i, form in enumerate(forms):
        print(i, form_n, form)
        dfs.append(chunked_export(api, form))
    return pd.concat(dfs, axis=1)


def all_forms_exist_in_redcap(form_names, api):
    meta = api.export_metadata(format='df')
    available_forms = meta['form_name'].unique().tolist()
    for form in form_names:
        if form not in available_forms:
            return False
    return True


def load_form(api, form_name, verbose=True):
    if verbose:
        print(form_name)
    
    try:
        if verbose:
            print("Trying chunked export, 5000 records at a time")
        return chunked_export(api, form_name, 5000)
    except (ValueError, rc.RedcapError, pd.io.common.EmptyDataError):
        pass
    
    # 2. Chunked load with chunk size of 1000
    try:
        if verbose:
            print("Trying chunked export, 1000 records at a time")
        return chunked_export(api, form_name, 1000)
    except (ValueError, rc.RedcapError, pd.io.common.EmptyDataError):
        pass
    
    # 2. Chunked load with default chunk size
    try:
        if verbose:
            print("Trying chunked export, default chunk size (100)")
        return chunked_export(api, form_name, 100)
    except (ValueError, rc.RedcapError, pd.io.common.EmptyDataError):
        pass
    
    # 3. Chunked load with tiny chunk
    try:
        if verbose:
            print("Trying chunked export with tiny chunks (10)")
        return chunked_export(api, form_name, 10)
    except (ValueError, rc.RedcapError, pd.io.common.EmptyDataError):
        print("Giving up")
        return None


def load_form_with_primary_key(api, form_name, verbose=True):
    df = load_form(api, form_name, verbose)
    if df is not None:
        if api.def_field in df.index.names:
            return df
        if api.is_longitudinal():
            return df.set_index([api.def_field, "redcap_event_name"])
        else:
            return df.set_index([api.def_field])
@@ -222,16 +222,49 @@
 se=True,
+ export_data_access_groups=False,
  **kwarg
@@ -903,16 +903,85 @@
 t='df',%0A
+                export_data_access_groups=export_data_access_groups,%0A
         
@@ -1411,16 +1411,49 @@
  arm='1'
+, export_data_access_groups=False
 ):%0A    f
@@ -1676,25 +1676,91 @@
 _export(
-api, form
+%0A            api, form, export_data_access_groups=export_data_access_groups
 ))%0A    r
@@ -2089,24 +2089,57 @@
 verbose=True
+, export_data_access_groups=False
 ):%0A    if ve
@@ -2321,16 +2321,99 @@
 me, 5000
+,%0A                              export_data_access_groups=export_data_access_groups
 )%0A    ex
@@ -2686,16 +2686,99 @@
 me, 1000
+,%0A                              export_data_access_groups=export_data_access_groups
 )%0A    ex
@@ -3052,16 +3052,99 @@
 ame, 100
+,%0A                              export_data_access_groups=export_data_access_groups
 )%0A    ex
@@ -3405,16 +3405,99 @@
 name, 10
+,%0A                              export_data_access_groups=export_data_access_groups
 )%0A    ex
@@ -3672,53 +3672,189 @@
 True
-):%0A    df = load_form(api, form_name, verbose
+,%0A                               export_data_access_groups=False):%0A    df = load_form(api, form_name, verbose,%0A                   export_data_access_groups=export_data_access_groups
 )%0A  
@@ -4094,20 +4094,21 @@
 dex(%5Bapi.def_field%5D)
+%0A

Add option to export DAGs for all load fns