-- This worksheet demonstrates the use of an audience lookalike modeling template. This worksheet 
-- should be run from the provider's account.
-- This template enables businesses to measure the effectiveness of their marketing channels
-- by joining consumer and provider datasets in a Snowflake Data Clean Room.
-- Results are activated to the provider.

-- As coded here, the provider and consumer must be in the same organization and same cloud 
-- hosting region. To publish a clean room to use outside your organization, publish an 
-- EXTERNAL clean room.

-- This code provides a user form so that it can be run by a consumer using either the
-- clean rooms API or the UI.

-- For details about the features and use cases of this template, see the [Snowflake documentation
-- topic for this template](https://docs.snowflake.com/user-guide/cleanrooms/lookalike-audience-modeling-template).

-- Provider and consumer data can be generated by running the sample data generator python script
-- linked from the documentation.

-- This worksheet shows how to install and run the clean room in code, but you can also install 
-- and run it using the clean rooms UI, because the provider created a UI configuration form 
-- for it.

-- Add values for placeholders in the code where indicated.

USE ROLE SAMOOHA_APP_ROLE;
USE WAREHOUSE app_wh;

SET cleanroom_name = 'LAL_ACTIVATION'; 

SET provider_features_table = '<fully_qualified_provider_features_table_name>';
SET provider_join_key = 'HASHED_EMAIL';

SET consumer_account = '<consumer_account_locator>'; 
SET consumer_data_account_sharing_id = '<consumer_data_account_sharing_id>'; 

CALL samooha_by_snowflake_local_db.provider.cleanroom_init($cleanroom_name, 'INTERNAL');

-- Link provider data into the clean room.
-- TEMP_PUBLIC_KEY is needed to decrypt any data activated to the provider account.
CALL samooha_by_snowflake_local_db.provider.link_datasets(
    $cleanroom_name, 
    [
        $provider_features_table,
        'samooha_by_snowflake_local_db.LIBRARY.TEMP_PUBLIC_KEY'
    ]
);

CALL samooha_by_snowflake_local_db.provider.set_join_policy(
    $cleanroom_name, 
    [$provider_features_table || ':' || $provider_join_key]
);

-- UDF for training.
CALL samooha_by_snowflake_local_db.provider.load_python_into_cleanroom(
    $cleanroom_name, 
    'lookalike_train',
    ['input_data variant', 'labels variant'],
    ['pandas', 'numpy', 'xgboost', 'sklearn'],
    'variant',
    'train',
    $$
import numpy as np
import pandas as pd
import xgboost
from sklearn import preprocessing
import sys
import pickle
import codecs

class TrainXGBoostClassifier(object):
    def __init__(self):
        self.model = None
        self._params = {
            "objective": "binary:logistic",
            "max_depth": 3,
            "nthread": 1,
            "eval_metric": "auc",
        }
        self.num_boosting_rounds = 10
        self.ohe = None

    def get_params(self):
        return self._params

    def train(self, X, y):
        # Ensure y is binarized
        y_binary = np.where(y, 1, 0)

        categorical = X.select_dtypes(include=[object])
        
        ohe = preprocessing.OneHotEncoder(handle_unknown='ignore')
        categorical_ohe = ohe.fit_transform(categorical)
        self.ohe = ohe

        non_categorical = X.select_dtypes(exclude=[object])
        train_x = np.concatenate((categorical_ohe.toarray(), non_categorical.to_numpy()), axis=1)

        xg_train = xgboost.DMatrix(train_x, label=y_binary)
        
        params = self.get_params()
        evallist = [(xg_train, "train")]
        evals_result = {}

        self.model = xgboost.train(
            params, xg_train, self.num_boosting_rounds, evallist, evals_result=evals_result
        )
        self.evals_result = evals_result

    def dump_model(self):
        if self.model is not None:
            # Package model and encoder together
            model_package = [self.model, self.ohe]
            return codecs.encode(pickle.dumps(model_package), "base64").decode()
        else:
            raise ValueError("Model needs to be trained first")

def train(d1, l1):
    X = pd.DataFrame(d1)
    y = np.array(l1)

    trainer = TrainXGBoostClassifier()
    trainer.train(X, y)

    return {
        "total_rows": len(d1),
        "model": trainer.dump_model(),
        "auc": np.max(trainer.evals_result["train"]["auc"]),
        "error": 1 - np.max(trainer.evals_result["train"]["auc"])
    }
    $$
);

-- UDF for scoring.
CALL samooha_by_snowflake_local_db.provider.load_python_into_cleanroom(
    $cleanroom_name, 
    'lookalike_score',
    ['pickle_jar variant', 'emails variant', 'features variant'],
    ['pandas', 'numpy', 'xgboost', 'scikit-learn'],
    'string',
    'score',
    $$
import numpy as np
import pandas as pd
import xgboost as xgb
import pickle
import codecs
import json

def score(model, emails, features):
    model = model[0] if not isinstance(model, str) else model
    
    # Unpack model and encoder
    model_package = pickle.loads(codecs.decode(model.encode(), "base64"))
    bst = model_package[0]
    ohe2 = model_package[1]

    Y = pd.DataFrame(features)
    Y1 = Y.select_dtypes(include=[object])
    Y2 = ohe2.transform(Y1)

    Y3 = Y.select_dtypes(exclude=[object])
    Y_pred = np.concatenate((Y2.toarray(), Y3.to_numpy()), axis=1)

    dscore = xgb.DMatrix(Y_pred)
    pred = bst.predict(dscore)

    retval = list(zip(np.array(emails), list(map(str, pred))))
    retval = [{"email": r[0], "score": r[1]} for r in retval]
    return json.dumps(retval)
    $$
);

-- Analysis template.
CALL samooha_by_snowflake_local_db.provider.add_custom_sql_template(
    $cleanroom_name,
    'lookalike_analysis_for_activation',
    $$
    DECLARE
        res RESULTSET;
    BEGIN
        
        CREATE OR REPLACE TABLE cleanroom.activation_data_lookalike_results AS
            WITH
            features AS (
                SELECT
                    p.{{ provider_join_key | sqlsafe }} as hashed_email, 
                   
                    array_construct({{ dimensions[0] | column_policy }} {% for feat in dimensions[1:] %} , {{ feat | column_policy }} {% endfor %}) as features
                FROM
                    identifier({{ source_table[0] }}) as p
            ),
            labels AS (
                SELECT
                    c.{{ consumer_join_key | sqlsafe }} as hashed_email, 
     
                    {{ filter_clause | sqlsafe | column_policy }} as label_value
                FROM
                    identifier({{ my_table[0] }}) as c
            ),
            trained_model AS (
                SELECT
                    train_out:model::varchar as model,
                    train_out:error::float as error,
                    train_out:auc::float as auc
                FROM (
                    SELECT
                        cleanroom.lookalike_train(array_agg(f.features), array_agg(l.label_value)) as train_out
                    FROM features f, labels l
                    WHERE f.hashed_email = l.hashed_email
                )
            ),
            inference_output AS (
                SELECT
                    MOD(seq4(), 100) as batch,
                    cleanroom.lookalike_score(
                        array_agg(distinct t.model), 
                        array_agg(p.{{ provider_join_key | sqlsafe }}), 
                        array_construct({{ dimensions[0] | column_policy }} {% for feat in dimensions[1:] %} , {{ feat | column_policy }} {% endfor %})
                    ) as scores
                FROM trained_model t, identifier({{ source_table[0] }}) p
                WHERE p.{{ provider_join_key | sqlsafe }} NOT IN (SELECT c.{{ consumer_join_key | sqlsafe }} FROM identifier({{ my_table[0] }}) c)
                GROUP BY batch
            ),
            processed_output AS (
                SELECT 
                    value:email::string as email, 
                    value:score::float as score 
                FROM (select scores from inference_output), lateral flatten(input => parse_json(scores))
            )
            
            SELECT 
                p.audience_size, 
                t.error,
                t.auc
            from (
                SELECT count(distinct email) as audience_size 
                FROM processed_output 
                WHERE score > {{ score_threshold }} 
            ) p, trained_model t;
        
        
        CREATE OR REPLACE TABLE cleanroom.activation_data_lookalike_results AS
            SELECT 
                value:email::string as HASHED_EMAIL, 
                value:score::float as SCORE
            FROM 
                (select scores from inference_output), 
                lateral flatten(input => parse_json(scores))
            WHERE 
                SCORE > {{ score_threshold }}; 

        
        res := (
            SELECT 
                p.audience_size, 
                t.error,
                t.auc
            from (
                SELECT count(distinct HASHED_EMAIL) as audience_size 
                FROM cleanroom.activation_data_lookalike_results
            ) p, trained_model t
        );
        
        RETURN TABLE(res);
    END;
    $$
);

-- Activation template.
CALL samooha_by_snowflake_local_db.provider.add_custom_sql_template(
    $cleanroom_name,
    'activation_lookalike_results',
    $$
    BEGIN
        RETURN 'lookalike_results';
    END;
    $$
);

-- Add user form to enable this template to be run in the clean rooms UI
-- as well as in code.
CALL samooha_by_snowflake_local_db.provider.add_ui_form_customizations(
    $cleanroom_name,
    'lookalike_analysis_for_activation',
    { 
        'display_name': 'Generic Audience Lookalike Builder',
        'description': 'Train an XGBoost model on a consumer seed audience to find lookalikes in the provider data.',
        'methodology': '1. Define a seed audience using a filter (e.g., c.SALES_DLR > 2000). 2. Select provider features (e.g., p.STATUS, p.AGE_BAND). 3. Model scores provider users. 4. Results above the score threshold are saved for activation.',
        'warehouse_hints': {
            'warehouse_size': 'medium',
            'snowpark_optimized': TRUE
        },
        'render_table_dropdowns': {
            'render_consumer_table_dropdown': true,
            'render_provider_table_dropdown': true
        },
        'activation_template_name': 'activation_lookalike_results',
        'enabled_activations': ['provider','consumer']
    },
    { 
        'provider_join_key': {
            'display_name': 'Provider Join Key',
            'description': 'Select the join key from the provider table (e.g., HASHED_EMAIL).',
            'group': 'Data Configuration',
            'order': 1,
            'type': 'dropdown',
            'references': ['PROVIDER_JOIN_POLICY'],
            'provider_parent_table_field': 'source_table'
        },
        'consumer_join_key': {
            'display_name': 'Consumer Join Key',
            'description': 'Select the join key from your (consumer) table (e.g., HASHED_EMAIL).',
            'group': 'Data Configuration',
            'order': 2,
            'type': 'dropdown',
            'references': ['CONSUMER_JOIN_POLICY'],
            'consumer_parent_table_field': 'my_table'
        },
        'dimensions': {
            'display_name': 'Feature Selection (e.g., p.STATUS)',
            'description': 'Select features to use for modeling. IMPORTANT: You must manually add the table alias (e.g., "p.") to your selections in the analysis run.',
            'group': 'Model Configuration',
            'order': 3,
            'type': 'multiselect',
            
            'choices': ['p.STATUS', 'p.AGE_BAND', 'p.REGION_CODE', 'p.DAYS_ACTIVE'],
            'infoMessage': 'Select the features that are whitelisted in the column policy.'
        },
        'filter_clause': {
            'display_name': 'Seed Audience Filter (e.g., c.SALES_DLR > 2000)',
            'description': 'SQL filter clause to define your seed. IMPORTANT: You must use the "c." alias for your table.',
            'group': 'Model Configuration',
            'order': 4,
            'type': 'any',
            'default': 'c.SALES_DLR >= 2500'
        },
        'score_threshold': {
            'display_name': 'Score Threshold (0.0 to 1.0)',
            'description': 'Minimum score to be included in the final audience (e.g., 0.5).',
            'group': 'Model Configuration',
            'order': 5,
            'type': 'any',
            'default': '0.5',
            'size': 'S'
        }
    },
    { 
        'default_output_type': 'TABLE'
    }
);

-- Share and publish the clean room.
CALL samooha_by_snowflake_local_db.provider.set_default_release_directive(
  $cleanroom_name,
  'V1_0', '2' 
);

CALL samooha_by_snowflake_local_db.provider.request_provider_activation_consent(
    $cleanroom_name, 
    'activation_lookalike_results'
);
    
CALL samooha_by_snowflake_local_db.provider.add_consumers(
    $cleanroom_name, 
    $consumer_account, 
    $consumer_data_account_sharing_id
);

CALL samooha_by_snowflake_local_db.provider.create_or_update_cleanroom_listing($cleanroom_name);

-- SWITCH TO CONSUMER to install the clean room, run the analysis, and activate results.

-- When the consumer has activated results to your account, you can read them by running
-- the following query. Note that it can take several minutes for data to propagate from
-- the consumer account to the provider account.
SELECT *
   FROM SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.PUBLIC.PROVIDER_ACTIVATION_SUMMARY
   WHERE segment = '<segment_name_specified_by_consumer>';

-- Optionally delete the clean room when you're finished with it.
-- CALL samooha_by_snowflake_local_db.provider.drop_cleanroom($cleanroom_name);
