import re


def api_result_to_release_format(api_df, id_lookup_dict=None, verbose=False):
    """
    Reindex a PyCAP API result to an NCANDA release format.

    REDCap API, when used with PyCAP, returns results as a DataFrame indexed by
    NCANDA ID (study_id - X-00000-Y-0) and combined event + arm
    (redcap_event_name)

    On the other hand, release files are typically indexed by XNAT ID
    (NCANDA_S0?????; mri_xnat_id in Redcap).

    This function will:

    1. Convert Redcap IDs to NCANDA SIDs using id_lookup_dict (as generated by
        `get_id_lookup_from_demographics_file`) or the `mri_xnat_sid` column
        (if present in api_df),
    2. Drop Redcap IDs that cannot be converted in that way,
    3. Separate event and arm to individual columns and make their names
        release-compatible,
    4. Return DataFrame indexed by release primary keys (subject, arm, visit).
    """

    df = api_df.copy(deep=True)
    df.reset_index(inplace=True)
    if id_lookup_dict:
        df['subject'] = df['study_id'].map(id_lookup_dict)
    elif 'mri_xnat_sid' in df.columns:
        df['subject'] = df['mri_xnat_sid']
    else:
        raise IndexError("You must supply id_lookup_dict, or api_df has to "
                         "have the mri_xnat_sid column")
    nan_idx = df['subject'].isnull()
    if verbose:
        study_id_nans = df.loc[nan_idx, 'study_id'].tolist()
        print ("Dropping study IDs without corresponding NCANDA SID: " +
               ", ".join(study_id_nans))
    df = df[~nan_idx]
    df[['visit', 'arm']] = (df['redcap_event_name']
                            .str.extract(r'^(\w+)_(arm_\d+)$'))

    def clean_up_event_string(event):
        """
        If possible, convert Redcap event name to NCANDA release visit name.

        If conversion fails, return the original string.

        Intended to be passed to pd.Series.map.
        """
        # NOTE: Only accounts for full Arm 1 events
        match = re.search(r'^(baseline|\dy)', event)
        if not match:
            return event
        elif re.match('^\d', match.group(1)):
            return "followup_" + match.group(1)
        else:
            return match.group(1)

    df['visit'] = df['visit'].map(clean_up_event_string)

    def clean_up_arm_string(arm):
        """
        If possible, convert Redcap arm name to NCANDA release arm name.

        If conversion fails, return the original string.

        Intended to be passed to pd.Series.map.
        """
        arm_dict = {'arm_1': 'standard',
                    'arm_2': 'recovery',
                    'arm_3': 'sleep',
                    'arm_4': 'maltreated'}
        if arm not in arm_dict:
            return arm
        else:
            return arm_dict[arm]

    df['arm'] = df['arm'].map(clean_up_arm_string)

    return df.set_index(['subject', 'arm', 'visit'])