# Import pyjanitor and pandas
import janitor
import pandas as pd
import pandas_flavor as pf

# Supress user warnings when we try overwriting our custom pandas flavor functions
import warnings
warnings.filterwarnings('ignore')

filename = 'https://raw.githubusercontent.com/rfordatascience/tidytuesday/master/data/2019/2019-04-23/raw_anime.csv'
df = pd.read_csv(filename)


@pf.register_dataframe_method
def str_remove(df, column_name: str, pat: str, *args, **kwargs):
    """Wrapper around df.str.replace"""

    df[column_name] = df[column_name].str.replace(pat, "", *args, **kwargs)
    return df


@pf.register_dataframe_method
def str_trim(df, column_name: str, *args, **kwargs):
    """Wrapper around df.str.strip"""

    df[column_name] = df[column_name].str.strip(*args, **kwargs)
    return df


@pf.register_dataframe_method
def explode(df: pd.DataFrame, column_name: str, sep: str):
    """
    For rows with a list of values, this function will create new
    rows for each value in the list
    """

    df["id"] = df.index
    wdf = (
        pd.DataFrame(df[column_name].str.split(sep).fillna("").tolist())
        .stack()
        .reset_index()
    )
    # exploded_column = column_name
    wdf.columns = ["id", "depth", column_name]  # plural form to singular form
    # wdf[column_name] = wdf[column_name].apply(lambda x: x.strip())  # trim
    wdf.drop("depth", axis=1, inplace=True)

    return pd.merge(df, wdf, on="id", suffixes=("_drop", "")).drop(
        columns=["id", column_name + "_drop"]
    )


@pf.register_dataframe_method
def str_word(
    df,
    column_name: str,
    start: int = None,
    stop: int = None,
    pat: str = " ",
    *args,
    **kwargs
):
    """
    Wrapper around `df.str.split` with additional `start` and `end` arguments
    to select a slice of the list of words.
    """

    df[column_name] = df[column_name].str.split(pat).str[start:stop]
    return df


@pf.register_dataframe_method
def str_join(df, column_name: str, sep: str, *args, **kwargs):
    """
    Wrapper around `df.str.join`
    Joins items in a list.
    """

    df[column_name] = df[column_name].str.join(sep)
    return df


@pf.register_dataframe_method
def str_slice(
    df, column_name: str, start: int = None, stop: int = None, *args, **kwargs
):
    """
    Wrapper around `df.str.slice
    """

    df[column_name] = df[column_name].str[start:stop]
    return df


clean_df = (
    df.str_remove(column_name="producers", pat="\[|\]")
    .explode(column_name="producers", sep=",")
    .str_remove(column_name="producers", pat="'")
    .str_trim("producers")
    .str_remove(column_name="genre", pat="\[|\]")
    .explode(column_name="genre", sep=",")
    .str_remove(column_name="genre", pat="'")
    .str_trim(column_name="genre")
    .str_remove(column_name="studio", pat="\[|\]")
    .explode(column_name="studio", sep=",")
    .str_remove(column_name="studio", pat="'")
    .str_trim(column_name="studio")
    .str_remove(column_name="aired", pat="\{|\}|'from':\s*|'to':\s*")
    .str_word(column_name="aired", start=0, stop=2, pat=",")
    .str_join(column_name="aired", sep=",")
    .deconcatenate_column(
        column_name="aired", new_column_names=["start_date", "end_date"], sep=","
    )
    .remove_columns(column_names=["aired"])
    .str_remove(column_name="start_date", pat="'")
    .str_slice(column_name="start_date", start=0, stop=10)
    .str_remove(column_name="end_date", pat="'")
    .str_slice(column_name="end_date", start=0, stop=11)
    .to_datetime("start_date", format="%Y-%m-%d", errors="coerce")
    .to_datetime("end_date", format="%Y-%m-%d", errors="coerce")
    .fill_empty(columns=["rank", "popularity"], value=0)
    .filter_on("rank != 0 & popularity != 0")
)

filename = 'https://raw.githubusercontent.com/rfordatascience/tidytuesday/master/data/2019/2019-04-23/raw_anime.csv'
df = pd.read_csv(filename)

@pf.register_dataframe_method
def str_remove(df, column_name: str, pat: str, *args, **kwargs):
    """
    Wrapper around df.str.replace
    The function will loop through regex patterns and remove them from the desired column.

    :param df: A pandas DataFrame.
    :param column_name: A `str` indicating which column the string removal action is to be made.
    :param pat: A regex pattern to match and remove.
    """

    if not isinstance(pat, str):
        raise TypeError(
            f"Pattern should be a valid regex pattern. Received pattern: {pat} with dtype: {type(pat)}"
        )
    df[column_name] = df[column_name].str.replace(pat, "", *args, **kwargs)
    return df

clean_df = (
    df
    .str_remove(column_name='producers', pat='\[|\]')
)

@pf.register_dataframe_method
def explode(df: pd.DataFrame, column_name: str, sep: str):
    """
    For rows with a list of values, this function will create new rows for each value in the list

    :param df: A pandas DataFrame.
    :param column_name: A `str` indicating which column the string removal action is to be made.
    :param sep: The delimiter. Example delimiters include `|`, `, `, `,` etc.
    """

    df["id"] = df.index
    wdf = (
        pd.DataFrame(df[column_name].str.split(sep).fillna("").tolist())
        .stack()
        .reset_index()
    )
    # exploded_column = column_name
    wdf.columns = ["id", "depth", column_name]  # plural form to singular form
    # wdf[column_name] = wdf[column_name].apply(lambda x: x.strip())  # trim
    wdf.drop("depth", axis=1, inplace=True)

    return pd.merge(df, wdf, on="id", suffixes=("_drop", "")).drop(
        columns=["id", column_name + "_drop"]
    )

clean_df = (
    clean_df
    .explode(column_name='producers', sep=',')
)

clean_df = (
    clean_df
    .str_remove(column_name='producers', pat='\'')
)

@pf.register_dataframe_method
def str_trim(df, column_name: str, *args, **kwargs):
    """Remove trailing and leading characters, in a given column"""
    df[column_name] = df[column_name].str.strip(*args, **kwargs)
    return df

clean_df = clean_df.str_trim('producers')

clean_df = (
    clean_df
    # Perform operation for genre.
    .str_remove(column_name='genre', pat='\[|\]')
    .explode(column_name='genre', sep=',')
    .str_remove(column_name='genre', pat='\'')
    .str_trim(column_name='genre')
    # Now do it for studio
    .str_remove(column_name='studio', pat='\[|\]')
    .explode(column_name='studio', sep=',')
    .str_remove(column_name='studio', pat='\'')
    .str_trim(column_name='studio')
)

@pf.register_dataframe_method
def str_word(
    df,
    column_name: str,
    start: int = None,
    stop: int = None,
    pat: str = " ",
    *args,
    **kwargs
):
    """
    Wrapper around `df.str.split` with additional `start` and `end` arguments
    to select a slice of the list of words.

    :param df: A pandas DataFrame.
    :param column_name: A `str` indicating which column the split action is to be made.
    :param start: optional An `int` for the start index of the slice
    :param stop: optinal  An `int` for the end index of the slice
    :param pat: String or regular expression to split on. If not specified, split on whitespace.

    """
    df[column_name] = df[column_name].str.split(pat).str[start:stop]
    return df


@pf.register_dataframe_method
def str_join(df, column_name: str, sep: str, *args, **kwargs):
    """
    Wrapper around `df.str.join`
    Joins items in a list.

    :param df: A pandas DataFrame.
    :param column_name: A `str` indicating which column the split action is to be made.
    :param sep: The delimiter. Example delimiters include `|`, `, `, `,` etc.
    """
    df[column_name] = df[column_name].str.join(sep)
    return df


@pf.register_dataframe_method
def str_slice(
    df, column_name: str, start: int = None, stop: int = None, *args, **kwargs
):
    """
    Wrapper around `df.str.slice
    Slices strings.

    :param df: A pandas DataFrame.
    :param column_name: A `str` indicating which column the split action is to be made.
    :param start: 'int' indicating start of slice.
    :param stop: 'int' indicating stop of slice.
    """
    df[column_name] = df[column_name].str[start:stop]
    return df

clean_df = (
    clean_df.str_remove(column_name="aired", pat="\{|\}|'from':\s*|'to':\s*")
    .str_word(column_name="aired", start=0, stop=2, pat=",")
    .str_join(column_name="aired", sep=",")
    # .add_columns({'start_date': clean_df['aired'][0]})
    .deconcatenate_column(
        column_name="aired", new_column_names=["start_date", "end_date"], sep=","
    )
    .remove_columns(column_names=["aired"])
    .str_remove(column_name="start_date", pat="'")
    .str_slice(column_name="start_date", start=0, stop=10)
    .str_remove(column_name="end_date", pat="'")
    .str_slice(column_name="end_date", start=0, stop=11)
    .to_datetime("start_date", format="%Y-%m-%d", errors="coerce")
    .to_datetime("end_date", format="%Y-%m-%d", errors="coerce")
)

# First fill any NA values with 0 and then filter != 0
clean_df = clean_df.fill_empty(column_names=["rank", "popularity"], value=0).filter_on(
    "rank != 0 & popularity != 0"
)