"""
Modify
======

Functions used to filter and/or change some data, always taking in one set of data and returning one set of data.

  .. autosummary::
     :toctree: modules/modify

     categorize
     colfilter
     colfilter_percent_zero
     colfilter_min_n
     colfilter_min_cat_n
     make_binary
     make_categorical
     make_continuous
     merge_observations
     merge_variables
     move_variables
     recode_values
     remove_outliers
     rowfilter_incomplete_obs
     transform

"""

from typing import Optional, List, Union

import click
import numpy as np
import pandas as pd

from ..internal.utilities import (
    _validate_skip_only, _get_dtypes, _process_colfilter,
    print_wrap, _remove_empty_categories)


@print_wrap
def categorize(data: pd.DataFrame, cat_min: int = 3, cat_max: int = 6, cont_min: int = 15):
    """
    Classify variables into constant, binary, categorical, continuous, and 'unknown'.  Drop variables that only have NaN values.

    Parameters
    ----------
    data: pd.DataFrame
        The DataFrame to be processed
    cat_min: int, default 3
        Minimum number of unique, non-NA values for a categorical variable
    cat_max: int, default 6
        Maximum number of unique, non-NA values for a categorical variable
    cont_min: int, default 15
        Minimum number of unique, non-NA values for a continuous variable

    Returns
    -------
    result: pd.DataFrame or None
        If inplace, returns None.  Changes the datatypes on the input DataFrame.

    Examples
    --------
    >>> import clarite
    >>> clarite.modify.categorize(nhanes)
    362 of 970 variables (37.32%) are classified as binary (2 unique values).
    47 of 970 variables (4.85%) are classified as categorical (3 to 6 unique values).
    483 of 970 variables (49.79%) are classified as continuous (>= 15 unique values).
    42 of 970 variables (4.33%) were dropped.
            10 variables had zero unique values (all NA).
            32 variables had one unique value.
    36 of 970 variables (3.71%) were not categorized and need to be set manually.
            36 variables had between 6 and 15 unique values
            0 variables had >= 15 values but couldn't be converted to continuous (numeric) values
    """
    # Validate parameters
    assert cat_min > 2
    assert cat_min <= cat_max
    assert cont_min > cat_max
    assert type(data) == pd.DataFrame

    # Count the number of variables to start with
    total_vars = len(data.columns)
    # Get the number of unique non-na values per variable
    unique_count = data.nunique(dropna=True)

    # No unique non-NA values - Drop these variables
    empty_vars = (unique_count == 0)
    if empty_vars.sum() > 0:
        columns = list(empty_vars[empty_vars].index)
        data = data.drop(columns=columns)

    # One unique non-NA value - Convert non-NA values to category (for constant)
    keep_constant = (unique_count == 1)
    if keep_constant.sum() > 0:
        columns = list(keep_constant[keep_constant].index)
        data = data.astype({c: 'category' for c in columns})

    # Two unique non-NA values - Convert non-NA values to category (for binary)
    keep_bin = (unique_count == 2)
    if keep_bin.sum() > 0:
        columns = list(keep_bin[keep_bin].index)
        data = data.astype({c: 'category' for c in columns})

    # Categorical - Convert non-NA values to category type
    keep_cat = (unique_count >= cat_min) & (unique_count <= cat_max)
    if keep_cat.sum() > 0:
        columns = list(keep_cat[keep_cat].index)
        data = data.astype({c: 'category' for c in columns})  # NaNs are handled correctly, no need skip

    # Continuous - Convert non-NA values to numeric type (even though they probably already are)
    keep_cont = (unique_count >= cont_min)
    check_cont = pd.Series(False, index=keep_cont.index)
    if keep_cont.sum() > 0:
        for col in keep_cont[keep_cont].index:
            try:
                data[col] = pd.to_numeric(data[col])
            except ValueError:
                # Couldn't convert to a number- possibly a categorical variable with string names?
                keep_cont[col] = False
                check_cont[col] = True
                data[col] = data.loc[:, col].astype(str)

    # Other - Convert non-NA values to string type
    check_other = ~empty_vars & ~keep_constant & ~keep_bin & ~keep_cat & ~check_cont & ~keep_cont
    if check_other.sum() > 0:
        columns = list(check_other[check_other].index)
        for c in columns:
            data.loc[~data[c].isna(), c] = data.loc[~data[c].isna(), c].astype(str)

    # Log categorized results
    click.echo(f"{keep_constant.sum():,} of {total_vars:,} variables ({keep_constant.sum() / total_vars:.2%}) "
               f"are classified as constant (1 unique value).")
    click.echo(f"{keep_bin.sum():,} of {total_vars:,} variables ({keep_bin.sum()/total_vars:.2%}) "
               f"are classified as binary (2 unique values).")
    click.echo(f"{keep_cat.sum():,} of {total_vars:,} variables ({keep_cat.sum()/total_vars:.2%}) "
               f"are classified as categorical ({cat_min} to {cat_max} unique values).")
    click.echo(f"{keep_cont.sum():,} of {total_vars:,} variables ({keep_cont.sum()/total_vars:.2%}) "
               f"are classified as continuous (>= {cont_min} unique values).")

    # Log dropped variables
    dropped = empty_vars.sum()
    click.echo(f"{dropped:,} of {total_vars:,} variables ({dropped/total_vars:.2%}) were dropped.")
    if dropped > 0:
        click.echo(f"\t{empty_vars.sum():,} variables had zero unique values (all NA).")

    # Log non-categorized results
    num_not_categorized = check_other.sum() + check_cont.sum()
    click.echo(f"{num_not_categorized:,} of {total_vars:,} variables ({num_not_categorized/total_vars:.2%})"
               f" were not categorized and need to be set manually.")
    if num_not_categorized > 0:
        click.echo(f"\t{check_other.sum():,} variables had between {cat_max} and {cont_min} unique values")
        click.echo(f"\t{check_cont.sum():,} variables had >= {cont_min} values but couldn't be converted to continuous (numeric) values")

    return data


@print_wrap
def colfilter(data, skip: Optional[Union[str, List[str]]] = None, only: Optional[Union[str, List[str]]] = None):
    """
    Remove some variables (skip) or keep only certain variables (only)

    Parameters
    ----------
    data: pd.DataFrame
        The DataFrame to be processed and returned
    skip: str, list or None (default is None)
        List of variables to remove
    only: str, list or None (default is None)
        List of variables to keep

    Returns
    -------
    data: pd.DataFrame
        The filtered DataFrame

    Examples
    --------
    >>> import clarite
    >>> female_logBMI = clarite.modify.colfilter(nhanes, only=['BMXBMI', 'female'])
    ================================================================================
    Running colfilter
    --------------------------------------------------------------------------------
    Keeping 2 of 945 variables:
            0 of 0 binary variables
            0 of 0 categorical variables
            2 of 945 continuous variables
            0 of 0 unknown variables
    ================================================================================
    """
    boolean_keep = _validate_skip_only(data, skip, only)
    dtypes = _get_dtypes(data)
    click.echo(f"Keeping {boolean_keep.sum():,} of {len(data.columns):,} variables:")

    for kind in ['binary', 'categorical', 'continuous', 'unknown']:
        is_kind = (dtypes == kind)
        is_kept = is_kind & boolean_keep
        click.echo(f"\t{is_kept.sum():,} of {is_kind.sum():,} {kind} variables")

    return data.loc[:, boolean_keep]


@print_wrap
def colfilter_min_cat_n(data, n: int = 200, skip: Optional[Union[str, List[str]]] = None, only: Optional[Union[str, List[str]]] = None):
    """
    Remove binary and categorical variables which have less than <n> occurences of each unique value

    Parameters
    ----------
    data: pd.DataFrame
        The DataFrame to be processed and returned
    n: int, default 200
        The minimum number of occurences of each unique value required in order for a variable not to be filtered
    skip: str, list or None (default is None)
        List of variables that the filter should *not* be applied to
    only: str, list or None (default is None)
        List of variables that the filter should *only* be applied to

    Returns
    -------
    data: pd.DataFrame
        The filtered DataFrame

    Examples
    --------
    >>> import clarite
    >>> nhanes_filtered = clarite.modify.colfilter_min_cat_n(nhanes)
    ================================================================================
    Running colfilter_min_cat_n
    --------------------------------------------------------------------------------
    WARNING: 36 variables need to be categorized into a type manually
    Testing 362 of 362 binary variables
            Removed 248 (68.51%) tested binary variables which had a category with less than 200 values
    Testing 47 of 47 categorical variables
            Removed 36 (76.60%) tested categorical variables which had a category with less than 200 values
    """
    assert type(data) == pd.DataFrame
    min_category_counts = data.apply(lambda col: col.value_counts().min())
    fail_filter = min_category_counts < n

    kept = _process_colfilter(data, skip, only,
                              fail_filter=fail_filter,
                              explanation=f"which had a category with less than {n} values.",
                              kinds=['binary', 'categorical'])
    # Return
    return data.loc[:, kept]


@print_wrap
def colfilter_min_n(data: pd.DataFrame, n: int = 200,
                    skip: Optional[Union[str, List[str]]] = None, only: Optional[Union[str, List[str]]] = None):
    """
    Remove variables which have less than <n> non-NA values

    Parameters
    ----------
    data: pd.DataFrame
        The DataFrame to be processed and returned
    n: int, default 200
        The minimum number of unique values required in order for a variable not to be filtered
    skip: str, list or None (default is None)
        List of variables that the filter should *not* be applied to
    only: str, list or None (default is None)
        List of variables that the filter should *only* be applied to

    Returns
    -------
    data: pd.DataFrame
        The filtered DataFrame

    Examples
    --------
    >>> import clarite
    >>> nhanes_filtered = clarite.modify.colfilter_min_n(nhanes)
    ================================================================================
    Running colfilter_min_n
    --------------------------------------------------------------------------------
    WARNING: 36 variables need to be categorized into a type manually
    Testing 362 of 362 binary variables
            Removed 12 (3.31%) tested binary variables which had less than 200 non-null values
    Testing 47 of 47 categorical variables
            Removed 8 (17.02%) tested categorical variables which had less than 200 non-null values
    Testing 483 of 483 continuous variables
            Removed 8 (1.66%) tested continuous variables which had less than 200 non-null values
    """
    assert type(data) == pd.DataFrame
    counts = data.count()  # by default axis=0 (rows) so counts number of non-NA rows in each column
    fail_filter = counts < n

    kept = _process_colfilter(data, skip, only,
                              fail_filter=fail_filter,
                              explanation=f"which had less than {n} non-null values.",
                              kinds=['binary', 'categorical', 'continuous'])

    # Return
    return data.loc[:, kept]


@print_wrap
def colfilter_percent_zero(data: pd.DataFrame, filter_percent: float = 90.0,
                           skip: Optional[Union[str, List[str]]] = None, only: Optional[Union[str, List[str]]] = None):
    """
    Remove continuous variables which have <proportion> or more values of zero (excluding NA)

    Parameters
    ----------
    data: pd.DataFrame
        The DataFrame to be processed and returned
    filter_percent: float, default 90.0
            If the percentage of rows in the data with a value of zero is greater than or equal to this value, the variable is filtered out.
    skip: str, list or None (default is None)
        List of variables that the filter should *not* be applied to
    only: str, list or None (default is None)
        List of variables that the filter should *only* be applied to

    Returns
    -------
    data: pd.DataFrame
        The filtered DataFrame

    Examples
    --------
    >>> import clarite
    >>> nhanes_filtered = clarite.modify.colfilter_percent_zero(nhanes_filtered)
    ================================================================================
    Running colfilter_percent_zero
    --------------------------------------------------------------------------------
    WARNING: 36 variables need to be categorized into a type manually
    Testing 483 of 483 continuous variables
            Removed 30 (6.21%) tested continuous variables which were equal to zero in at least 90.00% of non-NA observations.
    """
    assert type(data) == pd.DataFrame
    percent_value = 100 * data.apply(lambda col: (col == 0).sum() / col.count())
    fail_filter = percent_value >= filter_percent

    kept = _process_colfilter(data, skip, only,
                              fail_filter=fail_filter,
                              explanation=f"which were equal to zero in at least {filter_percent:.2f}% of non-NA observations.",
                              kinds=['continuous'])
    # Return
    return data.loc[:, kept]


@print_wrap
def make_binary(data: pd.DataFrame, skip: Optional[Union[str, List[str]]] = None, only: Optional[Union[str, List[str]]] = None):
    """
    Set variable types as Binary

    Checks that each variable has at most 2 values and converts the type to pd.Categorical.

    Note: When these variables are used in regression, they are ordered by value.
    For example, Sex (Male=1, Female=2) will encode "Male" as 0 and "Female" as 1 during the EWAS regression step.

    Parameters
    ----------
    data: pd.DataFrame or pd.Series
        Data to be processed
    skip: str, list or None (default is None)
        List of variables that should *not* be made binary
    only: str, list or None (default is None)
        List of variables that are the *only* ones to be made binary

    Returns
    -------
    data: pd.DataFrame
        DataFrame with the same data but validated and converted to binary types

    Examples
    --------
    >>> import clarite
    >>> nhanes = clarite.modify.make_binary(nhanes, only=['female', 'black', 'mexican', 'other_hispanic'])
    ================================================================================
    Running make_binary
    --------------------------------------------------------------------------------
    Set 4 of 970 variable(s) as binary, each with 22,624 observations
    """
    # Which columns
    columns = _validate_skip_only(data, skip, only)

    # Check the number of unique values, raising an error if any columns can't be converted
    unique_values = data.nunique()
    non_binary = (unique_values != 2) & columns
    num_non_binary = non_binary.sum()
    if num_non_binary > 0:
        raise ValueError(f"{num_non_binary} variable(s) did not have 2 unique values and couldn't be processed as a binary type: "
                         f"{', '.join(non_binary[non_binary].index)}")

    # Convert dtype
    columns = columns[columns].index
    data = data.astype({c: 'category' for c in columns})
    click.echo(f"Set {len(columns):,} of {len(data.columns):,} variable(s) as binary, each with {len(data):,} observations")

    return data


@print_wrap
def make_categorical(data: pd.DataFrame, skip: Optional[Union[str, List[str]]] = None, only: Optional[Union[str, List[str]]] = None):
    """
    Set variable types as Categorical

    Converts the type to pd.Categorical

    Parameters
    ----------
    data: pd.DataFrame or pd.Series
        Data to be processed
    skip: str, list or None (default is None)
        List of variables that should *not* be made categorical
    only: str, list or None (default is None)
        List of variables that are the *only* ones to be made categorical

    Returns
    -------
    data: pd.DataFrame
        DataFrame with the same data but validated and converted to categorical types

    Examples
    --------
    >>> import clarite
    >>> df = clarite.modify.make_categorical(df)
    ================================================================================
    Running make_categorical
    --------------------------------------------------------------------------------
    Set 12 of 12 variable(s) as categorical, each with 4,321 observations
    """
    # Which columns
    columns = _validate_skip_only(data, skip, only)

    # Convert dtype
    columns = columns[columns].index
    data = data.astype({c: 'category' for c in columns})
    click.echo(f"Set {len(columns):,} of {len(data.columns):,} variable(s) as categorical, each with {len(data):,} observations")

    return data


@print_wrap
def make_continuous(data: pd.DataFrame, skip: Optional[Union[str, List[str]]] = None, only: Optional[Union[str, List[str]]] = None):
    """
    Set variable types as Numeric

    Converts the type to numeric

    Parameters
    ----------
    data: pd.DataFrame or pd.Series
        Data to be processed
    skip: str, list or None (default is None)
        List of variables that should *not* be made continuous
    only: str, list or None (default is None)
        List of variables that are the *only* ones to be made continuous

    Returns
    -------
    data: pd.DataFrame
        DataFrame with the same data but validated and converted to numeric types

    Examples
    --------
    >>> import clarite
    >>> df = clarite.modify.make_continuous(df)
    ================================================================================
    Running make_categorical
    --------------------------------------------------------------------------------
    Set 128 of 128 variable(s) as continuous, each with 4,321 observations
    """
    # Which columns
    columns = _validate_skip_only(data, skip, only)

    # Convert dtype, ignoring errors
    data.loc[:, columns] = data.loc[:, columns].apply(lambda col: pd.to_numeric(col, errors='ignore'))

    # Check if any variables could not be converted
    failed_conversion = data.loc[:, columns].dtypes.apply(lambda dt: not pd.api.types.is_numeric_dtype(dt))
    if failed_conversion.sum() > 0:
        raise ValueError(f"{failed_conversion.sum()} variable(s) couldn't be processed as continuous (numeric) type(s): "
                         f"{', '.join(failed_conversion[failed_conversion].index)}")

    columns = columns[columns].index
    click.echo(f"Set {len(columns):,} of {len(data.columns):,} variable(s) as continuous, each with {len(data):,} observations")

    return data


@print_wrap
def recode_values(data, replacement_dict,
                  skip: Optional[Union[str, List[str]]] = None, only: Optional[Union[str, List[str]]] = None):
    """
    Convert values in a dataframe.  By default, replacement occurs in all columns but this may be modified with 'skip' or 'only'.
    Pandas has more powerful 'replace' methods for more complicated scenarios.

    Parameters
    ----------
    data: pd.DataFrame
        The DataFrame to be processed and returned
    replacement_dict: dictionary
        A dictionary mapping the value being replaced to the value being inserted
    skip: str, list or None (default is None)
        List of variables that the replacement should *not* be applied to
    only: str, list or None (default is None)
        List of variables that the replacement should *only* be applied to

    Examples
    --------
    >>> import clarite
    >>> clarite.modify.recode_values(df, {7: np.nan, 9: np.nan}, only=['SMQ077', 'DBD100'])
    ================================================================================
    Running recode_values
    --------------------------------------------------------------------------------
    Replaced 17 values from 22,624 observations in 2 variables
    >>> clarite.modify.recode_values(df, {10: 12}, only=['SMQ077', 'DBD100'])
    ================================================================================
    Running recode_values
    --------------------------------------------------------------------------------
    No occurences of replaceable values were found, so nothing was replaced.
    """
    # Limit columns if needed
    if skip is not None or only is not None:
        columns = _validate_skip_only(data, skip, only)
        columns = columns[columns].index.get_level_values(0)  # variable names where columns = True
        replacement_dict = {c: replacement_dict for c in columns}

    # Replace
    result = data.replace(to_replace=replacement_dict, value=None, inplace=False)

    # Log
    diff = result.eq(data)
    diff[pd.isnull(result) & pd.isnull(data)] = True  # NAs are not equal by default
    diff = ~diff  # make True where a value was replaced
    cols_with_changes = (diff.sum() > 0).sum()
    cells_with_changes = diff.sum().sum()
    if cells_with_changes > 0:
        click.echo(f"Replaced {cells_with_changes:,} values from {len(data):,} observations in {cols_with_changes:,} variables")
    else:
        click.echo("No occurences of replaceable values were found, so nothing was replaced.")

    # Return
    return result


@print_wrap
def remove_outliers(data, method: str = 'gaussian', cutoff=3,
                    skip: Optional[Union[str, List[str]]] = None, only: Optional[Union[str, List[str]]] = None):
    """
    Remove outliers from continuous variables by replacing them with np.nan

    Parameters
    ----------
    data: pd.DataFrame
        The DataFrame to be processed and returned
    method: string, 'gaussian' (default) or 'iqr'
        Define outliers using a gaussian approach (standard deviations from the mean) or inter-quartile range
    cutoff: positive numeric, default of 3
        Either the number of standard deviations from the mean (method='gaussian') or the multiple of the IQR (method='iqr')
        Any values equal to or more extreme will be replaced with np.nan
    skip: str, list or None (default is None)
        List of variables that the replacement should *not* be applied to
    only: str, list or None (default is None)
        List of variables that the replacement should *only* be applied to

    Examples
    --------
    >>> import clarite
    >>> nhanes_rm_outliers = clarite.modify.remove_outliers(nhanes, method='iqr', cutoff=1.5, only=['DR1TVB1', 'URXP07', 'SMQ077'])
    ================================================================================
    Running remove_outliers
    --------------------------------------------------------------------------------
    WARNING: 36 variables need to be categorized into a type manually
    Removing outliers from 2 continuous variables with values < 1st Quartile - (1.5 * IQR) or > 3rd quartile + (1.5 * IQR)
            Removed 0 low and 430 high IQR outliers from URXP07 (outside -153.55 to 341.25)
            Removed 0 low and 730 high IQR outliers from DR1TVB1 (outside -0.47 to 3.48)
    >>> nhanes_rm_outliers = clarite.modify.remove_outliers(nhanes, only=['DR1TVB1', 'URXP07'])
    ================================================================================
    Running remove_outliers
    --------------------------------------------------------------------------------
    WARNING: 36 variables need to be categorized into a type manually
    Removing outliers from 2 continuous variables with values more than 3 standard deviations from the mean
            Removed 0 low and 42 high gaussian outliers from URXP07 (outside -1,194.83 to 1,508.13)
            Removed 0 low and 301 high gaussian outliers from DR1TVB1 (outside -1.06 to 4.27)
    """
    # Copy to avoid replacing in-place
    data = data.copy(deep=True)

    # Which columns
    columns = _validate_skip_only(data, skip, only)
    is_continuous = _get_dtypes(data) == 'continuous'
    columns = columns & is_continuous

    # Check cutoff and method, printing what is being done
    if cutoff <= 0:
        raise ValueError("'cutoff' must be >= 0")
    if method == 'iqr':
        click.echo(f"Removing outliers from {len(data):,} observations of {columns.sum():,} continuous variables "
                   f"with values < 1st Quartile - ({cutoff} * IQR) or > 3rd quartile + ({cutoff} * IQR)")
    elif method == 'gaussian':
        click.echo(f"Removing outliers from {len(data):,} observations of {columns.sum():,} continuous variables "
                   f"with values more than {cutoff} standard deviations from the mean")
    else:
        raise ValueError(f"'{method}' is not a supported method for outlier removal - only 'gaussian' and 'iqr'.")

    # Define outlier replacemet functions
    def iqr_outliers(col):
        q1 = col.quantile(0.25)
        q3 = col.quantile(0.75)
        iqr = abs(q3 - q1)
        bottom = q1 - (iqr * cutoff)
        top = q3 + (iqr * cutoff)
        outliers_bottom = col < bottom
        outliers_top = col > top
        col.loc[outliers_bottom] = np.nan
        col.loc[outliers_top] = np.nan
        click.echo(f"\tRemoved {outliers_bottom.sum():,} low and {outliers_top.sum():,} high IQR outliers "
                   f"from {col.name} (outside {bottom:,.2f} to {top:,.2f})")

    def gaussian_outliers(col):
        mean = col.mean()
        std = col.std()
        bottom = mean - (std * cutoff)
        top = mean + (std * cutoff)
        outliers_bottom = col < bottom
        outliers_top = col > top
        col.loc[outliers_bottom] = np.nan
        col.loc[outliers_top] = np.nan
        click.echo(f"\tRemoved {outliers_bottom.sum()} low and {outliers_top.sum()} high gaussian outliers "
                   f"from {col.name} (outside {bottom:,.2f} to {top:,.2f})")

    # Remove outliers
    if method == 'iqr':
        data.loc[:, columns].apply(iqr_outliers)
    elif method == 'gaussian':
        data.loc[:, columns].apply(gaussian_outliers)

    return data


@print_wrap
def rowfilter_incomplete_obs(data, skip: Optional[Union[str, List[str]]] = None, only: Optional[Union[str, List[str]]] = None):
    """
    Remove rows containing null values

    Parameters
    ----------
    data: pd.DataFrame
        The DataFrame to be processed and returned
    skip: str, list or None (default is None)
        List of columns that are not checked for null values
    only: str, list or None (default is None)
        List of columns that are the only ones to be checked for null values

    Returns
    -------
    data: pd.DataFrame
        The filtered DataFrame

    Examples
    --------
    >>> import clarite
    >>> nhanes_filtered = clarite.modify.rowfilter_incomplete_obs(nhanes, only=[phenotype] + covariates)
    ================================================================================
    Running rowfilter_incomplete_obs
    --------------------------------------------------------------------------------
    Removed 3,687 of 22,624 observations (16.30%) due to NA values in any of 8 variables
    """
    columns = _validate_skip_only(data, skip, only)

    keep_IDs = data.loc[:, columns].isnull().sum(axis=1) == 0  # Number of NA in each row is 0
    n_removed = len(data) - sum(keep_IDs)

    click.echo(f"Removed {n_removed:,} of {len(data):,} observations ({n_removed/len(data):.2%}) "
               f"due to NA values in any of {columns.sum()} variables")
    return data[keep_IDs]


@print_wrap
def merge_observations(top: pd.DataFrame,
                       bottom: pd.DataFrame):
    """
    Merge two datasets, keeping only the columns present in both.  Raise an error if a datatype conflict occurs.

    Parameters
    ----------
    top: pd.DataFrame
        "top" DataFrame
    bottom: pd.DataFrame
        "bottom" DataFrame

    Returns
    -------
    result: pd.DataFrame
    """
    # Throw an error if any observation ID is repeated
    overlapped_observations = set(top.index) & set(bottom.index)
    if len(overlapped_observations) > 0:
        raise ValueError(f"Can't merge: {len(overlapped_observations):,} observation IDs occur in both datasets")

    # Merge data, keeping only the columns in common
    combined = pd.concat([top, bottom], join='inner', sort=False)

    # If a categorical is only in one dataframe, or the categorical has different levels, it is coerced to an object and must be changed back
    # Exclude cases where either variable was an object originally
    combined = combined.astype({col: 'category'
                                if (dt == 'object') & (top.dtypes[col] != 'object') & (bottom.dtypes[col] != 'object')
                                else dt for col, dt in combined.dtypes.iteritems()})

    # Check datatypes for changes
    top_dtypes = _get_dtypes(top[combined.columns])
    bottom_dtypes = _get_dtypes(bottom[combined.columns])
    combined_dtypes = _get_dtypes(combined)

    diff_dtypes = top_dtypes != bottom_dtypes
    diff_dtype_vars = list(diff_dtypes[diff_dtypes].index)
    if diff_dtypes.sum() > 0:
        raise ValueError(f"{diff_dtypes.sum():,} variables have mismatched datatypes: \n{' '.join(diff_dtype_vars)}")

    diff_cats = (top_dtypes != combined_dtypes) | (bottom_dtypes != combined_dtypes)
    diff_cats_vars = set(diff_cats[diff_cats].index) - set(diff_dtype_vars)
    if len(diff_cats_vars) > 0:
        raise ValueError(f"{len(diff_cats_vars):,} variables have categories present in only one dataset: \n{' '.join(diff_cats_vars)}")

    return combined


@print_wrap
def merge_variables(left: Union[pd.DataFrame, pd.Series],
                    right: Union[pd.DataFrame, pd.Series],
                    how: str = 'outer'):
    """
    Merge a list of dataframes with different variables side-by-side.  Keep all observations ('outer' merge) by default.

    Parameters
    ----------
    left: pd.Dataframe or pd.Series
        "left" DataFrame or Series
    right: pd.DataFrame or pd.Series
        "right" DataFrame or Series which uses the same index
    how: merge method, one of {'left', 'right', 'inner', 'outer'}
        Keep only rows present in the left data, the right data, both datasets, or either dataset.

    Examples
    --------
    >>> import clarite
    >>> df = clarite.modify.merge_variables(df_bin, df_cat, how='outer')
    """
    # Convert to DataFrame if a series was passed
    if type(left) == pd.Series:
        left = pd.DataFrame(left)
    if type(right) == pd.Series:
        right = pd.DataFrame(right)

    click.echo(f"{how} Merge:\n"
               f"\tleft = {len(left):,} observations of {len(left.columns):,} variables\n"
               f"\tright = {len(right):,} observations of {len(right.columns):,} variables")
    result = left.merge(right, left_index=True, right_index=True, how=how)
    click.echo(f"Kept {len(result):,} observations of {len(result.columns):,} variables.")
    return result


@print_wrap
def move_variables(left: pd.DataFrame, right: Union[pd.DataFrame, pd.Series],
                   skip: Optional[Union[str, List[str]]] = None, only: Optional[Union[str, List[str]]] = None):
    """
    Move one or more variables from one DataFrame to another

    Parameters
    ----------
    left: pd.Dataframe
        DataFrame containing the variable(s) to be moved
    right: pd.DataFrame or pd.Series
        DataFrame or Series (which uses the same index) that the variable(s) will be moved to
    skip: str, list or None (default is None)
        List of variables that will *not* be moved
    only: str, list or None (default is None)
        List of variables that are the *only* ones to be moved

    Returns
    -------
    left: pd.DataFrame
        The first DataFrame with the variables removed
    right: pd.DataFrame
        The second DataFrame with the variables added

    Examples
    --------
    >>> import clarite
    >>> df_cat, df_cont = clarity.modify.move_variables(df_cat, df_cont, only=["DRD350AQ", "DRD350DQ", "DRD350GQ"])
    Moved 3 variables.
    >>> discovery_check, discovery_cont = clarite.modify.move_variables(discovery_check, discovery_cont)
    Moved 39 variables.
    """
    # Which columns
    columns = _validate_skip_only(left, skip, only)

    # Add to new df
    right = merge_variables(right, left[columns])

    # Remove from original
    left = left.drop(columns, axis='columns')

    # Log
    if columns.sum() == 1:
        click.echo("Moved 1 variable.")
    else:
        click.echo(f"Moved {len(columns)} variables.")

    # Return
    return left, right


@print_wrap
def transform(data: pd.DataFrame,
              transform_method: str,
              skip: Optional[Union[str, List[str]]] = None, only: Optional[Union[str, List[str]]] = None
              ):
    """
    Apply a transformation function to a variable

    Parameters
    ----------
    data: pd.DataFrame or pd.Series
        Data to be processed
    transform_method: str
        Name of the transformation (Python function or NumPy ufunc to apply)
    skip: str, list or None (default is None)
        List of variables that will *not* be transformed
    only: str, list or None (default is None)
        List of variables that are the *only* ones to be transformed

    Returns
    -------
    data: pd.DataFrame
        DataFrame with variables that have been transformed

    Examples
    --------
    >>> import clarite
    >>> df = clarite.modify.transform(df, 'log', only=['BMXBMI'])
    ================================================================================
    Running transform
    --------------------------------------------------------------------------------
    Transformed 'BMXBMI' using 'log'.
    """
    # Copy to avoid replacing in-place
    data = data.copy(deep=True)

    # Which columns
    columns = _validate_skip_only(data, skip, only)
    transform_variables = list(data.loc[:, columns])

    # Check for potential errors
    dtypes = _get_dtypes(data)
    for variable in transform_variables:
        dtype = dtypes.get(variable, None)
        if dtype is None:
            raise ValueError(f"The variable ('{variable}') was not found in the data")
        elif dtype != 'continuous':
            raise ValueError(f"The variable ('{variable}') was {dtype}: "
                             f"transformations may only be applied to continuous variables")

    # Transform each variable
    for variable in transform_variables:
        try:
            data.loc[:, variable] = data.loc[:, variable].apply(transform_method)
        except Exception as e:
            raise ValueError(f"Couldn't apply a function named '{transform_method}' to '{variable}'.\n\t{e}")

        click.echo(f"Transformed '{variable}' using '{transform_method}'")

    return data


@print_wrap
def drop_extra_categories(data: pd.DataFrame,
                          skip: Optional[Union[str, List[str]]] = None,
                          only: Optional[Union[str, List[str]]] = None):
    """
    Update variable types to remove categories that don't occur in the data

    Parameters
    ----------
    data: pd.DataFrame or pd.Series
        Data to be processed
    skip: str, list or None (default is None)
        List of variables that will *not* be checked
    only: str, list or None (default is None)
        List of variables that are the *only* ones to be checked

    Returns
    -------
    data: pd.DataFrame
        DataFrame with categorical types updated as needed

    Examples
    --------
    >>> import clarite
    >>> df = clarite.modify.drop_extra_categories(df, only=['SDDSRVYR'])
    ================================================================================
    Running drop_extra_categories
    --------------------------------------------------------------------------------
    SDDSRVYR had categories with no occurrences: 3, 4
    """
    # Copy to avoid replacing in-place
    data = data.copy(deep=True)

    # Drop categories
    removed_cats = _remove_empty_categories(data, skip, only)

    # Log results
    if len(removed_cats) == 1:
        var = list(removed_cats.keys())[0]
        cats = removed_cats[var]
        message = f"\t{str(var)} had categories with no occurrences: {', '.join([str(c) for c in cats])}"
        click.echo(message)
    elif len(removed_cats) > 1:
        message = "\tMultiple categorical variables had categories with no occurrences:"
        for var, cats in removed_cats.items():
            message += f"\n\t{str(var)}: {', '.join([str(c) for c in cats])}"
        click.echo(message)

    return data
