-- The following code shows how to upload multiple Python functions in a single clean room patch.
-- For more information, see docs.snowflake.com/user-guide/cleanrooms/demo-flows/custom-code

USE WAREHOUSE app_wh;
USE ROLE SAMOOHA_APP_ROLE;

SET cleanroom_name = 'Provider UDF example';
SET udf_template_name = 'basic_UDF_template';
SET consumer_locator = '<consumer_locator>';
SET provider_locator = $consumer_locator; -- This is a single-account testing clean room.
SET consumer_account_name = '<org_name>.<account_name>';
    
--Create the clean room.
CALL samooha_by_snowflake_local_db.provider.cleanroom_init($cleanroom_name, 'INTERNAL');

-- Link data.
CALL samooha_by_snowflake_local_db.provider.link_datasets($cleanroom_name,
['SAMOOHA_SAMPLE_DATABASE.DEMO.CUSTOMERS']);

---- Add custom inline UDF ----
CALL samooha_by_snowflake_local_db.provider.prepare_python_for_cleanroom(
    $cleanroom_name, 
    'get_next_status',  -- Name of the UDF. Can be different from the handler.
    ['status VARCHAR'], -- Arguments of the UDF, specified as (variable name, SQL type).
    ['numpy'],          -- Packages needed by UDF.
    [],                 -- When providing the code inline, this is an empty array.
    'VARCHAR',          -- Return type of UDF.
    'get_next_status',  -- Handler.
    $$
import numpy as np
def get_next_status(status):
  """Return the next higher status, or a random status
  if no matching status found or at the top of the list."""

  statuses = ['MEMBER', 'SILVER', 'GOLD', 'PLATINUM', 'DIAMOND']
  try:
    return statuses[statuses.index(status.upper()) + 1]
  except:
    return 'NO MATCH'
    $$
);

---- Add custom inline UDTF. ----
CALL samooha_by_snowflake_local_db.provider.prepare_python_for_cleanroom(
    $cleanroom_name, 
    'get_info',  -- Name of the UDTF. Can be different from the handler.
    ['hashed_email VARCHAR', 'days_active INT', 'status VARCHAR', 'income VARCHAR'],   -- Name/Type arguments of the UDTF.
    ['numpy'],         -- Packages used by UDTF.
    [],                -- When providing the code inline, this is an empty array.
    'TABLE(hashed_email VARCHAR, months_active INT, level VARCHAR)',  -- Return type of UDTF.
    'GetSomeVals',     -- Handler class name.
$$
class GetSomeVals:
  def __init__(self):
    self.month_days = 30

  def process(self, hashed_email, days_active, status, income):
    '''Change days into rough months, and also return whether we
    think the user's membership status is lower, higher, or equal to 
    what is expected, based on their income.'''
    
    months_active = days_active // self.month_days
    brackets = ['0-50K', '50K-100K', '100K-250K', '250K+']
    statuses = ['MEMBER', 'SILVER', 'GOLD', 'PLATINUM']
    if(statuses.index(status) < brackets.index(income)):
      level = 'low'
    elif(statuses.index(status) > brackets.index(income)):
      level = 'high'
    else:
      level = 'equal'
    
    yield(hashed_email, months_active, level)
$$
);

-- Upload all stored procedures.
-- Note the new patch number returned by this procedure. Keep this number for later use.
CALL samooha_by_snowflake_local_db.provider.load_prepared_python_into_cleanroom($cleanroom_name);

-- Set the release directive specified by the last load_python_into_cleanroom call.
CALL samooha_by_snowflake_local_db.provider.set_default_release_directive($cleanroom_name, 'V1_0', <PATCH_NUMBER>);

-- Add custom template that calls both the UDF and the UTDF.
CALL samooha_by_snowflake_local_db.provider.add_custom_sql_template(
    $cleanroom_name, 
    $udf_template_name,
    $$
    SELECT
      p.status,                            -- From original data
      cleanroom.get_next_status(p.status)
          AS NEXT_STATUS,                  -- From UDF
      p1.months_active,
      p1.level                             -- From UDTF output
    FROM SAMOOHA_SAMPLE_DATABASE.DEMO.CUSTOMERS AS p
    INNER JOIN TABLE(cleanroom.get_info(   -- Call UDTF
        p.hashed_email,
        p.days_active,
        p.status,
        p.income_bracket)) AS p1
    WHERE p.hashed_email = p1.hashed_email
    LIMIT 100;
    $$
);

-- Not needed for INTERNAL clean rooms, but adding here as a reminder to use it for EXTERNAL clean rooms.
-- NOT_REVIEWED means it can be used internally.
CALL samooha_by_snowflake_local_db.provider.view_cleanroom_scan_status($cleanroom_name);

-- Add consumers.
CALL samooha_by_snowflake_local_db.provider.add_consumers($cleanroom_name, $consumer_locator, $consumer_account_name);
    
-- Publish the clean room.
CALL samooha_by_snowflake_local_db.provider.create_or_update_cleanroom_listing($cleanroom_name);

-- ================ CONSUMER CODE ====================

CALL samooha_by_snowflake_local_db.consumer.install_cleanroom($cleanroom_name, $provider_locator);
CALL samooha_by_snowflake_local_db.consumer.is_enabled($cleanroom_name);

CALL samooha_by_snowflake_local_db.consumer.run_analysis(
    $cleanroom_name,
    $udf_template_name,
    [], -- No consumer tables used in our example.
    [], -- No provider tables used in our example.
    {}  -- No parameters needed for our example.
);

-- Clean up afterward.
CALL samooha_by_snowflake_local_db.provider.drop_cleanroom($cleanroom_name);
