# dataloaders.py

@extract_columns(user_id, vide_id, seconds)
def log_views_df( path ):
   return pd.read_csv(path)

@extract_columns(user_id, vide_id)
def log_favourite_df( path ):
   return pd.read_csv(path)

def scenes_df(path):
   return pd.read_csv(path)

# features.py

def title(log_file, scenes_df) -> pd.Series:
   return some Joim based on two

title

log_views_df

extract_columns

log

meta

module_names = [
        "log_file" : args.log_file,  # "log_views" or "log_favourite" or some other...
        "data_loaders",  # functions to help load dataloaders
        "user_features",  # logic for creating user features
        "video_features",  # logic for creating video features
    ]

extract_columns

@extract_columns(id='video_id', title='video_title', ...)
def load_data(...) -> pd.DataFrame:
   ...

import logging
import sys

import numpy as np
import pandas as pd
import pandera as pa

from hamilton import ad_hoc_utils, driver
from hamilton.function_modifiers import extract_columns, check_output



schema = pa.DataFrameSchema({
    "signups": pa.Column(int, checks=pa.Check.le(10)),
    "spend": pa.Column(float, checks=<http://pa.Check.lt|pa.Check.lt>(-1.2)),
    "other": pa.Column(str, checks=[
        pa.Check.str_startswith("value_"),
        # define custom checks as functions that take a series as input and
        # outputs a boolean or boolean Series
        pa.Check(lambda s: s.str.split("_", expand=True).shape[1] == 2)
    ]),
})



@extract_columns('signups', 'spend', 'other')
@check_output(schema = schema)
def log() -> pd.DataFrame:
    # data to validate
    df = pd.DataFrame({
        "signups": [1, 4, 0, 10, 12],
        "spend": [-1.3, -1.4, -2.9, -10.1, -20.4],
        "other": ["value_1", "value_2", "value_3", "value_2", "value1"],
    })
    return df




# # Look at `my_functions` to see how these functions connect.
# def signups() -> pd.Series:
#     """Returns sign up values"""
#     return pd.Series([1, 10, 50, 100, 200, 400])


# def spend() -> pd.Series:
#     """Returns the spend values"""
#     return pd.Series([10, 10, 20, 40, 40, 50])


def log_spend_per_signup(spend_per_signup: pd.Series) -> pd.Series:
    """Simple function taking the logarithm of spend over signups."""
    return np.log(spend_per_signup)


# Place the functions into a temporary module -- the idea is that this should house a curated set of functions.
# Don't be afraid to make multiple of them -- however we'd advise you to not use this method for production.
# Also note, that using a temporary function module does not work for scaling onto Ray, Dask, or Pandas on Spark.
temp_module = ad_hoc_utils.create_temporary_module(
    log, log_spend_per_signup, module_name="function_example"
)







import pandas as pd

"""
Notes:
  1. This file is used for all the [ray|dask|spark]/hello_world examples.
  2. It therefore show cases how you can write something once and not only scale it, but port it
     to different frameworks with ease!
"""


def avg_3wk_spend(spend: pd.Series) -> pd.Series:
    """Rolling 3 week average spend."""
    return spend.rolling(3).mean()


def spend_per_signup(spend: pd.Series, signups: pd.Series) -> pd.Series:
    """The cost per signup in relation to spend."""
    return spend / signups


def spend_mean(spend: pd.Series) -> float:
    """Shows function creating a scalar. In this case it computes the mean of the entire column."""
    return spend.mean()


def spend_zero_mean(spend: pd.Series, spend_mean: float) -> pd.Series:
    """Shows function that takes a scalar. In this case to zero mean spend."""
    return spend - spend_mean


def spend_std_dev(spend: pd.Series) -> float:
    """Function that computes the standard deviation of the spend column."""
    return spend.std()


def spend_zero_mean_unit_variance(spend_zero_mean: pd.Series, spend_std_dev: float) -> pd.Series:
    """Function showing one way to make spend have zero mean and unit variance."""
    return spend_zero_mean / spend_std_dev


my_functions = ad_hoc_utils.create_temporary_module(
    avg_3wk_spend, spend_per_signup, spend_mean, spend_zero_mean, 
    spend_std_dev, spend_zero_mean_unit_variance, module_name="my_functions"
)







# Cell 4 - Instantiate the Hamilton driver and pass it the right things in.

initial_config = {}
# we need to tell hamilton where to load function definitions from
dr = driver.Driver(initial_config, my_functions, temp_module)  # can pass in multiple modules
# we need to specify what we want in the final dataframe.
output_columns = [
    "spend",
    "signups",
    "avg_3wk_spend",
    "spend_per_signup",
    "spend_zero_mean_unit_variance",
    'other'
]
# Cell 5 - visualize execution
# To visualize do `pip install "sf-hamilton[visualization]"` if you want these to work





dr.execute(output_columns)