-- Example of using Snowpark UDFs in a clean room: Provider code --
-- See https://docs.snowflake.com/user-guide/cleanrooms/demo-flows/snowpark for full details.

-- This code uses UDFs and Snowpark libraries in code to compute overlap statistics
-- based on regression data calculated on two input tables.

-- Create the clean room.
USE role samooha_app_role;
USE warehouse app_wh;
SET cleanroom_name = 'Snowpark Demo clean room';
CALL samooha_by_snowflake_local_db.provider.cleanroom_init($cleanroom_name, 'INTERNAL');
CALL samooha_by_snowflake_local_db.provider.set_default_release_directive($cleanroom_name, 'V1_0', '0');
CALL samooha_by_snowflake_local_db.provider.link_datasets($cleanroom_name, ['SAMOOHA_SAMPLE_DATABASE.DEMO.CUSTOMERS']);``

-- Load the custom Python code. See the documentation above for details.
-- Code is run on Snowpark, because you import snowflake-snowpark-python and implement.
-- See remarks above describing the code flow.
CALL samooha_by_snowflake_local_db.provider.load_python_into_cleanroom(
  $cleanroom_name,
  'reach_impression_regression',  -- Handler function alias called by the template.
  ['source_table string', 'my_table string'], -- Parameter names and types used by the handler function.
  ['snowflake-snowpark-python', 'pandas', 'statsmodels', 'numpy'], -- Packages used in your code.
  'string',
  'main',                         -- Handler function called by the template handler alias.
  $$
import traceback
import pandas as pd
import numpy as np

import statsmodels.formula.api as sm


def drop_tables(session, table_names):
  """
  Drop the tables passed in
  """
  for tbl in table_names:
    session.sql(f'drop table {tbl}').collect()

def preprocess_regression_data(session, source_table, my_table, suffix):
"""
Preprocess the impressions and customer data into an intermediary table for regression
"""
table_name = f'cleanroom.intermediary_{suffix}'

my_table_statement = f'INNER JOIN {my_table} c ON p.hashed_email = c.hashed_email' if my_table != 'NONE' else ''
session.sql(f"""
CREATE OR REPLACE TABLE {table_name} AS (
  WITH joint_data AS (
    SELECT
        date,
        p.hashed_email AS hem,
        impression_id
    FROM {source_table} p
    {my_table_statement}
  )
  SELECT
    date,
    COUNT(DISTINCT hem) AS reach,
    COUNT(DISTINCT impression_id) AS num_impressions
  FROM joint_data
  GROUP BY date
  ORDER BY date
);
""").collect()

return table_name

def calculate_regression(df):
"""
Calculate the regression data from the dataframe we put together
"""
result = sm.ols('REACH ~ 1 + NUM_IMPRESSIONS', data=df).fit()
retval = pd.concat([
      result.params,
      result.tvalues,
      result.pvalues
  ], keys=['params', 't-stat', 'p-value'], names=['STATISTIC', 'PARAMETER']).rename('VALUE').reset_index()
return retval

def main(session, source_table, my_table):
  """
  First compute the regression data from an overlap between customer and provider data, counting
  the number of impressions and reach per day. Next regress these locally and compute the regression
  statistics. Finally write it to a results table which can be queried to get the output.
  """
  suffix = f'regression_results_{abs(hash((source_table, my_table))) % 10000}'

  try:
      # Preprocess impressions and customer data into an intermediary form to use for regression
      intermediary_table_name = preprocess_regression_data(session, source_table, my_table, suffix)

      # Load the data into Python locally
      df = session.table(intermediary_table_name).to_pandas()

      # Carry out the regression and get statistics as an output
      regression_output = calculate_regression(df)

      # Write the statistics to an output table
      # The table and the schema names should be in upper case to quoted identifier related issues.
      table = f'results_{suffix}'.upper()
      retval_df = session.write_pandas(regression_output, table,  schema = 'CLEANROOM', auto_create_table = True)

      # Drop any intermediary tables
      drop_tables(session, [intermediary_table_name])

      # Tell the user the name of the table the results have been written to
      return f'Done, results have been written to the following suffix: {suffix}'
  except:
      return traceback.format_exc()
$$
);

-- Loading code causes a security check and also bumps the clean room patch number, so you need to
-- update the default release directive.
CALL samooha_by_snowflake_local_db.provider.set_default_release_directive($cleanroom_name, 'V1_0', '1');

-- Add the template that calls the Python code.
CALL samooha_by_snowflake_local_db.provider.add_custom_sql_template(
  $cleanroom_name,
  'prod_calculate_regression',
  $$
  CALL cleanroom.reach_impression_regression({{ source_table[0] }}, {{ my_table[0] | default('NONE') }});
  $$
);

-- Add a template to display the results table to the consumer.
CALL samooha_by_snowflake_local_db.provider.add_custom_sql_template(
        $cleanroom_name,
        'prod_get_results',
        $$
        SELECT * FROM cleanroom.results_{{ results_suffix | sqlsafe }};
        $$
);

-- Share the clean room with your consumer account.
CALL samooha_by_snowflake_local_db.provider.add_consumers($cleanroom_name, '<CONSUMER_ACCOUNT_LOCATOR>', '<CONSUMER_ACCOUNT_NAME>');
CALL samooha_By_snowflake_local_db.provider.create_or_update_cleanroom_listing($cleanroom_name);
