-- Lookalike audience modeling: publisher worksheet
-- Run this worksheet from the publisher account.
--
-- This example uses a 3-step pipeline built on internal tables (multistep workflows):
--   Step 1 (Train):    Train an XGBoost model, save to cleanroom.lal_trained_model
--   Step 2 (Score):    Score the population, save to cleanroom.lookalike_scored_results
--   Step 3 (Activate): Send scored results to the publisher
--
-- The publisher registers a data offering, a code spec with two Python UDFs,
-- three templates (train, score, activate), and creates the collaboration.

-- The publisher and advertiser must be in the same organization and same cloud hosting
-- region, unless Cross-Cloud Auto-Fulfillment is enabled.

-- For details about this template, see the Snowflake documentation:
-- https://docs.snowflake.com/user-guide/cleanrooms/collab-lookalike-modeling

-- Publisher and advertiser data can be generated by running the sample data generator
-- Python script linked from the documentation.

-- Replace all values in angle brackets <...> with your own values.

USE ROLE SAMOOHA_APP_ROLE;
USE WAREHOUSE APP_WH;
USE SECONDARY ROLES NONE;


--------------------------------------------------------------
-- Step 1: Register the publisher data offering.
--------------------------------------------------------------

CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.REGISTRY.REGISTER_DATA_OFFERING(
$$
api_version: 2.0.0
spec_type: data_offering
name: lal_publisher_features
version: V0
description: >
  Publisher user feature data including membership status, age band, region,
  and activity level, used for training and scoring a lookalike model.
datasets:
  - alias: features
    data_object_fqn: <database.schema.LAL_PUBLISHER_FEATURES>
    allowed_analyses: template_only
    object_class: custom
    schema_and_template_policies:
      HASHED_EMAIL:
        category: join_custom
      STATUS:
        category: passthrough
      AGE_BAND:
        category: passthrough
      REGION_CODE:
        category: passthrough
      DAYS_ACTIVE:
        category: passthrough
$$);

SET publisher_data_offering_id = '<publisher_data_offering_id>';
CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.REGISTRY.VIEW_REGISTERED_DATA_OFFERINGS();


--------------------------------------------------------------
-- Step 2: Register the code spec with Python UDFs.
--
-- Two UDFs:
--   lookalike_train: Trains an XGBoost classifier, returns
--                    serialized model + metrics.
--   lookalike_score: Scores population using the trained model,
--                    returns JSON array of email/score pairs.
--------------------------------------------------------------

CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.REGISTRY.REGISTER_CODE_SPEC(
$$
api_version: 2.0.0
spec_type: code_spec
name: lookalike_model
version: V0
description: >
  XGBoost-based lookalike model training and scoring functions.

functions:
  - name: lookalike_train
    type: UDF
    language: PYTHON
    runtime_version: "3.10"
    handler: train
    arguments:
      - name: input_data
        type: VARIANT
      - name: labels
        type: VARIANT
    returns: VARIANT
    packages:
      - pandas
      - numpy
      - xgboost
      - scikit-learn
    description: Trains an XGBoost binary classifier on the seed audience features.
    code_body: |
      import numpy as np
      import pandas as pd
      import xgboost
      from sklearn import preprocessing
      import pickle
      import codecs

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

          for col in X.columns:
              try:
                  X[col] = pd.to_numeric(X[col])
              except (ValueError, TypeError):
                  X[col] = X[col].astype(str)

          categorical = X.select_dtypes(include=[object])
          ohe = preprocessing.OneHotEncoder(handle_unknown='ignore')
          if not categorical.empty:
              categorical_ohe = ohe.fit_transform(categorical).toarray()
          else:
              categorical_ohe = np.empty((len(X), 0))
              ohe.fit(np.empty((0, 0)))

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

          xg_train = xgboost.DMatrix(train_x, label=y_binary)
          params = {
              "objective": "binary:logistic",
              "max_depth": 3,
              "nthread": 1,
              "eval_metric": "auc",
          }
          evals_result = {}
          model = xgboost.train(
              params, xg_train, 10, [(xg_train, "train")], evals_result=evals_result
          )

          model_package = [model, ohe]
          serialized = codecs.encode(pickle.dumps(model_package), "base64").decode()

          return {
              "total_rows": len(d1),
              "model": serialized,
              "auc": float(np.max(evals_result["train"]["auc"])),
              "error": float(1 - np.max(evals_result["train"]["auc"]))
          }

  - name: lookalike_score
    type: UDF
    language: PYTHON
    runtime_version: "3.10"
    handler: score
    arguments:
      - name: pickle_jar
        type: VARIANT
      - name: emails
        type: VARIANT
      - name: features
        type: VARIANT
    returns: STRING
    packages:
      - pandas
      - numpy
      - xgboost
      - scikit-learn
    description: Scores the publisher population using the trained model.
    code_body: |
      import numpy as np
      import pandas as pd
      import xgboost as xgb
      import pickle
      import codecs
      import json

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

          Y = pd.DataFrame(features)

          for col in Y.columns:
              try:
                  Y[col] = pd.to_numeric(Y[col])
              except (ValueError, TypeError):
                  Y[col] = Y[col].astype(str)

          categorical = Y.select_dtypes(include=[object])
          if not categorical.empty:
              cat_encoded = ohe2.transform(categorical).toarray()
          else:
              cat_encoded = np.empty((len(Y), 0))

          non_categorical = Y.select_dtypes(exclude=[object])
          Y_pred = np.concatenate((cat_encoded, non_categorical.to_numpy(dtype=float)), axis=1)

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

          retval = [{"email": str(e), "score": str(s)} for e, s in zip(emails, pred)]
          return json.dumps(retval)
$$);

SET code_spec_id = '<code_spec_id>';
CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.REGISTRY.VIEW_REGISTERED_CODE_SPECS();


--------------------------------------------------------------
-- Step 3: Register the train template.
--
-- Joins publisher features with the advertiser seed audience,
-- trains an XGBoost binary classifier, and saves the serialized
-- model and metrics to an internal table.
--------------------------------------------------------------

CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.REGISTRY.REGISTER_TEMPLATE(
$$
api_version: 2.0.0
spec_type: template
name: lookalike_train_template
version: V0
type: sql_analysis
description: >
  Trains an XGBoost classifier on the advertiser seed audience using
  publisher features. Saves the trained model to an internal table
  for use by the scoring template.
methodology: >
  Joins publisher features with the advertiser seed audience on
  hashed email. Applies the seed filter to create binary labels.
  Trains an XGBoost model and saves the serialized model, AUC,
  and error rate to cleanroom.lal_trained_model.

parameters:
  - name: publisher_join_key
    description: Publisher column name for the join key.
    type: string
    required: true
  - name: consumer_join_key
    description: Advertiser column name for the join key.
    type: string
    required: true
  - name: dimensions
    description: List of publisher feature columns to use for modeling.
    type: array
    required: true
  - name: filter_clause
    description: SQL expression to define the seed audience label.
    type: string
    required: true

code_specs:
  - <code_spec_id>

template: |
  DECLARE
      res RESULTSET;
  BEGIN
      CREATE OR REPLACE TABLE cleanroom.lal_trained_model AS
          WITH
          features AS (
              SELECT
                  p.{{ publisher_join_key | sqlsafe }} AS hashed_email,
                  array_construct({{ dimensions[0] | sqlsafe }} {% for feat in dimensions[1:] %} , {{ feat | sqlsafe }} {% endfor %}) AS features
              FROM
                  IDENTIFIER({{ source_table[0] }}) AS p
          ),
          labels AS (
              SELECT
                  c.{{ consumer_join_key | sqlsafe }} AS hashed_email,
                  {{ filter_clause | sqlsafe }} AS label_value
              FROM
                  IDENTIFIER({{ source_table[1] }}) AS c
          )
          SELECT
              train_out:model::VARCHAR AS model,
              train_out:total_rows::INT AS total_rows,
              train_out:error::FLOAT AS error,
              train_out:auc::FLOAT AS auc
          FROM (
              SELECT
                  cleanroom.lookalike_model$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
          );

      res := (SELECT total_rows, auc, error FROM cleanroom.lal_trained_model);
      RETURN TABLE(res);
  END;
$$);

SET train_template_id = '<train_template_id>';


--------------------------------------------------------------
-- Step 4: Register the score template.
--
-- Reads the trained model from the internal table, scores all
-- non-seed publisher users, and saves results above the score
-- threshold to another internal table.
--------------------------------------------------------------

CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.REGISTRY.REGISTER_TEMPLATE(
$$
api_version: 2.0.0
spec_type: template
name: lookalike_score_template
version: V0
type: sql_analysis
description: >
  Scores the full publisher population using a previously trained
  lookalike model. Saves scored results to an internal table for
  optional activation.
methodology: >
  Reads the serialized model from cleanroom.lal_trained_model.
  Scores all non-seed publisher users in batches. Filters results
  above the score threshold and saves to
  cleanroom.lookalike_scored_results.

parameters:
  - name: publisher_join_key
    description: Publisher column name for the join key.
    type: string
    required: true
  - name: consumer_join_key
    description: Advertiser column name for the join key.
    type: string
    required: true
  - name: dimensions
    description: List of publisher feature columns to use for scoring.
    type: array
    required: true
  - name: score_threshold
    description: Minimum score to include a user in the lookalike audience.
    type: number
    default: 0.5
    required: true

code_specs:
  - <code_spec_id>

template: |
  DECLARE
      res RESULTSET;
  BEGIN
      CREATE OR REPLACE TABLE cleanroom.lookalike_scored_results AS
          WITH
          inference_output AS (
              SELECT
                  MOD(SEQ4(), 100) AS batch,
                  cleanroom.lookalike_model$lookalike_score(
                      ARRAY_AGG(DISTINCT t.model),
                      ARRAY_AGG(p.{{ publisher_join_key | sqlsafe }}),
                      ARRAY_AGG(array_construct({{ dimensions[0] | sqlsafe }} {% for feat in dimensions[1:] %} , {{ feat | sqlsafe }} {% endfor %}))
                  ) AS scores
              FROM cleanroom.lal_trained_model t,
                   IDENTIFIER({{ source_table[0] }}) p
              WHERE p.{{ publisher_join_key | sqlsafe }} NOT IN
                    (SELECT c.{{ consumer_join_key | sqlsafe }}
                     FROM IDENTIFIER({{ source_table[1] }}) 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 email AS HASHED_EMAIL, score AS SCORE
          FROM processed_output
          WHERE score > {{ score_threshold }};

      res := (
          SELECT COUNT(*) AS audience_size,
                 ROUND(AVG(SCORE), 4) AS avg_score
          FROM cleanroom.lookalike_scored_results
      );
      RETURN TABLE(res);
  END;
$$);

SET score_template_id = '<score_template_id>';


--------------------------------------------------------------
-- Step 5: Register the activation template.
--
-- Reads scored results from the internal table and writes them
-- to the activation table for sharing with the publisher.
--------------------------------------------------------------

CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.REGISTRY.REGISTER_TEMPLATE(
$$
api_version: 2.0.0
spec_type: template
name: activation_lookalike
version: V0
type: sql_activation
description: >
  Activates lookalike audience results to the publisher. Reads
  scored results saved by the scoring template and writes them
  to an activation table.
methodology: >
  Reads from cleanroom.lookalike_scored_results (populated by
  the scoring template) and writes to the activation table.

template: |
  BEGIN
      CREATE OR REPLACE TABLE cleanroom.activation_data_lookalike_results AS
          SELECT HASHED_EMAIL, SCORE
          FROM cleanroom.lookalike_scored_results;
      RETURN 'lookalike_results';
  END;
$$);

SET activation_template_id = '<activation_template_id>';


--------------------------------------------------------------
-- Step 6: Create the collaboration.
--------------------------------------------------------------

CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.COLLABORATION.INITIALIZE(
$$
api_version: 2.0.0
spec_type: collaboration
name: lookalike_modeling_collab
owner: publisher
collaborator_identifier_aliases:
  publisher: <publisher_data_sharing_id>
  advertiser: <advertiser_data_sharing_id>
analysis_runners:
  advertiser:
    data_providers:
      publisher:
        data_offerings:
          - id: <publisher_data_offering_id>
      advertiser:
        data_offerings: []
    templates:
      - id: <train_template_id>
      - id: <score_template_id>
      - id: <activation_template_id>
    activation_destinations:
      snowflake_collaborators:
        - publisher
$$,
'APP_WH'
);

SET collaboration_name = 'lookalike_modeling_collab';

-- INITIALIZE automatically joins the owner. Wait until status is JOINED.
CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.COLLABORATION.GET_STATUS($collaboration_name);

CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.COLLABORATION.VIEW_COLLABORATIONS();


--------------------------------------------------------------
-- SWITCH TO ADVERTISER to join, link data, and run the
-- 3-step pipeline (train, score, activate).
--------------------------------------------------------------

-- After the advertiser runs the activation template, the publisher can
-- view and process the activation results.
CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.COLLABORATION.VIEW_ACTIVATIONS($collaboration_name);

-- When the activation status is SHARED, process it to import the data.
CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.COLLABORATION.PROCESS_ACTIVATION($collaboration_name);

-- Read the activation results. The table name is returned by PROCESS_ACTIVATION.
-- Filter by the segment name specified by the advertiser.
SELECT * FROM <activation_results_table>
  WHERE segment_name = 'lal_lookalike_activation';


--------------------------------------------------------------
-- Cleanup
--------------------------------------------------------------

-- Delete the collaboration to clean up all resources.
-- This is a multi-step process; call GET_STATUS to check for LOCAL_DROP_PENDING,
-- then call TEARDOWN again.
CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.COLLABORATION.TEARDOWN($collaboration_name);
CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.COLLABORATION.GET_STATUS($collaboration_name);
CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.COLLABORATION.TEARDOWN($collaboration_name);
