-- Retail media measurement with HPO: brand worksheet
-- Run this worksheet from the brand account.
--
-- The brand measures the sales impact of their sponsored product campaign
-- on the retail media network (RMN) using an XGBoost model with distributed
-- hyperparameter optimization (HPO). The brand brings campaign exposure data
-- and proprietary measurement model code. The RMN provides shopper
-- transaction data.
--
-- Pipeline:
--   Step 1 (HPO + Train): Find optimal XGBoost hyperparameters via Bayesian
--                          optimization, then train the measurement model.
--   Step 2 (Score):        Score the full shopper population for campaign impact.
--   Step 3 (Activate):     Send measurement results back to the brand.
--
-- The RMN's raw transaction data and the brand's model code never
-- leave the clean room.

-- For details about ML Jobs, see the Snowflake documentation:
-- https://docs.snowflake.com/user-guide/cleanrooms/ml-jobs

-- 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: Stage the ML training and scoring code.
--------------------------------------------------------------

CREATE DATABASE IF NOT EXISTS <ml_code_db>;
CREATE STAGE IF NOT EXISTS <ml_code_db>.PUBLIC.ML_RMN_STAGE
    DIRECTORY = (ENABLE = TRUE)
    ENCRYPTION = (TYPE = 'SNOWFLAKE_SSE');

-- Upload the scripts using SnowSQL or the Snowflake CLI.
-- Download the files from the documentation page.
-- PUT file://rmn_train_hpo.py @<ml_code_db>.PUBLIC.ML_RMN_STAGE/rmn_project/ AUTO_COMPRESS=FALSE OVERWRITE=TRUE;
-- PUT file://rmn_score.py @<ml_code_db>.PUBLIC.ML_RMN_STAGE/rmn_project/ AUTO_COMPRESS=FALSE OVERWRITE=TRUE;

ALTER STAGE <ml_code_db>.PUBLIC.ML_RMN_STAGE REFRESH;
LIST @<ml_code_db>.PUBLIC.ML_RMN_STAGE/rmn_project/;


--------------------------------------------------------------
-- Step 2: Register the brand campaign exposure data offering.
--------------------------------------------------------------

CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.REGISTRY.REGISTER_DATA_OFFERING(
$$
api_version: 2.0.0
spec_type: data_offering
name: rmn_campaign_exposures
version: V0
description: Brand campaign exposure data with exposed/control assignment and campaign features.
datasets:
  - alias: exposures
    data_object_fqn: <database.schema.CAMPAIGN_EXPOSURES>
    allowed_analyses: template_only
    object_class: custom
    schema_and_template_policies:
      USER_ID:
        category: join_custom
      EXPOSED:
        category: passthrough
      CAMPAIGN_ID:
        category: passthrough
      IMPRESSIONS:
        category: passthrough
      RECENCY_DAYS:
        category: passthrough
      FREQUENCY:
        category: passthrough
$$);

SET brand_data_offering_id = 'rmn_campaign_exposures_V0';


--------------------------------------------------------------
-- Step 3: Register the ML Jobs code spec.
--------------------------------------------------------------

CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.REGISTRY.REGISTER_CODE_SPEC(
$$
api_version: 2.0.0
spec_type: code_spec
name: rmn_hpo_model
version: V0

ml_jobs:
  - name: rmn_train_hpo
    entrypoint: rmn_train_hpo.py
    stage_code_dir: '@<ml_code_db>.PUBLIC.ML_RMN_STAGE/rmn_project'
    pip_requirements:
      - pandas
      - numpy
      - xgboost
      - scikit-learn
  - name: rmn_score
    entrypoint: rmn_score.py
    stage_code_dir: '@<ml_code_db>.PUBLIC.ML_RMN_STAGE/rmn_project'
    pip_requirements:
      - pandas
      - numpy
      - xgboost
      - scikit-learn
$$);

SET code_spec_id = 'rmn_hpo_model_V0';


--------------------------------------------------------------
-- Step 4: Register the HPO train, score, and activation templates.
--------------------------------------------------------------

CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.REGISTRY.REGISTER_TEMPLATE(
$$
api_version: 2.0.0
spec_type: template
name: rmn_hpo_train_template
version: V0
type: sql_analysis
description: Train a retail media measurement model with distributed HPO using ML Jobs.
code_specs:
  - rmn_hpo_model_V0
parameters:
  - name: compute_pool
    description: Compute pool to run the HPO and training job on.
    required: true
template: |
  call cleanroom.rmn_hpo_model$rmn_train_hpo(
    {{ compute_pool }},
    {{ num_instances | default(1) }},
    {{ warehouse | default("APP_WH") }},
    OBJECT_CONSTRUCT('source_table', ARRAY_CONSTRUCT({{ source_table[0] }}, {{ source_table[1] }}))::VARCHAR
  )
$$);

SET train_template_id = 'rmn_hpo_train_template_V0';

CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.REGISTRY.REGISTER_TEMPLATE(
$$
api_version: 2.0.0
spec_type: template
name: rmn_score_template
version: V0
type: sql_analysis
description: Score the full shopper population for campaign impact using the trained model.
code_specs:
  - rmn_hpo_model_V0
parameters:
  - name: compute_pool
    description: Compute pool to run the scoring job on.
    required: true
template: |
  call cleanroom.rmn_hpo_model$rmn_score(
    {{ compute_pool }},
    {{ num_instances | default(1) }},
    {{ warehouse | default("APP_WH") }},
    OBJECT_CONSTRUCT('source_table', ARRAY_CONSTRUCT({{ source_table[0] }}))::VARCHAR
  )
$$);

SET score_template_id = 'rmn_score_template_V0';

CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.REGISTRY.REGISTER_TEMPLATE(
$$
api_version: 2.0.0
spec_type: template
name: rmn_activation
version: V0
type: sql_activation
description: Activate campaign measurement results to the brand.
template: |
  BEGIN
      CREATE OR REPLACE TABLE cleanroom.activation_data_rmn_results AS
          SELECT USER_ID, CAMPAIGN_IMPACT, P_PURCHASE_EXPOSED, P_PURCHASE_CONTROL
          FROM cleanroom.rmn_scored_results;
      RETURN 'rmn_results';
  END;
$$);

SET activation_template_id = 'rmn_activation_V0';


--------------------------------------------------------------
-- Step 5: Create the collaboration.
--
-- The brand is the owner (brings model code and drives the pipeline).
-- The RMN is the data provider (provides shopper transactions).
--------------------------------------------------------------

CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.COLLABORATION.INITIALIZE(
$$
api_version: 2.0.0
spec_type: collaboration
name: rmn_measurement_collab
owner: brand
collaborator_identifier_aliases:
  rmn: <rmn_data_sharing_id>
  brand: <brand_data_sharing_id>
analysis_runners:
  brand:
    data_providers:
      rmn:
        data_offerings: []
      brand:
        data_offerings:
          - id: rmn_campaign_exposures_V0
    templates:
      - id: rmn_hpo_train_template_V0
      - id: rmn_score_template_V0
      - id: rmn_activation_V0
    activation_destinations:
      snowflake_collaborators:
        - brand
$$,
'APP_WH'
);

SET collaboration_name = 'rmn_measurement_collab';

CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.COLLABORATION.GET_STATUS($collaboration_name);


--------------------------------------------------------------
-- SWITCH TO RMN to register data, review, and join.
-- Continue here after the RMN has joined.
--------------------------------------------------------------


--------------------------------------------------------------
-- Step 6: Link the brand data offering.
--------------------------------------------------------------

CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.COLLABORATION.LINK_DATA_OFFERING(
    $collaboration_name,
    $brand_data_offering_id,
    ['brand']
);


--------------------------------------------------------------
-- Step 7: Set up the compute pool.
--
-- HPO runs multiple trials in parallel. Use a larger instance
-- family for faster optimization:
--
--   CPU_X64_XS  (2 vCPU,   8 GB)  -- 1-5 trials, small data
--   CPU_X64_M   (8 vCPU,  32 GB)  -- 10-20 trials, medium data
--   CPU_X64_L   (32 vCPU, 128 GB) -- 20+ trials, large data
--
-- For distributed HPO (running each trial across multiple nodes
-- using the Snowflake HPO API), increase MAX_NODES and set
-- num_instances > 1 in the template_configuration below.
--------------------------------------------------------------

SHOW APPLICATIONS;
SET installed_app_name = '<installed_app_name>';

USE ROLE ACCOUNTADMIN;

CREATE COMPUTE POOL rmn_hpo_pool
    FOR APPLICATION IDENTIFIER($installed_app_name)
    MIN_NODES = 1
    MAX_NODES = 2
    INSTANCE_FAMILY = CPU_X64_M
    AUTO_RESUME = TRUE;

GRANT USAGE ON COMPUTE POOL rmn_hpo_pool
    TO APPLICATION IDENTIFIER($installed_app_name);
GRANT USAGE ON WAREHOUSE APP_WH
    TO APPLICATION IDENTIFIER($installed_app_name);

USE ROLE SAMOOHA_APP_ROLE;
USE SECONDARY ROLES NONE;


--------------------------------------------------------------
-- Step 8: Run the HPO + training job.
--------------------------------------------------------------

CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.COLLABORATION.RUN(
    $collaboration_name,
    $$
    api_version: 2.0.0
    spec_type: analysis
    template: rmn_hpo_train_template_V0
    template_configuration:
      view_mappings:
        source_tables:
          - brand.rmn_campaign_exposures_V0.exposures
          - rmn.<rmn_data_offering_id>.transactions
      arguments:
        compute_pool: rmn_hpo_pool
        num_instances: 1   # Increase to run distributed HPO across multiple nodes
    $$
);

SET train_job_id = '<train_job_id>';

-- Monitor HPO progress via logs (shows per-trial AUUC scores).
CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.COLLABORATION.RUN_ML_JOB_ACTION(
    $collaboration_name,
    $$
    api_version: 2.0.0
    spec_type: ml_job_action
    job_id: <train_job_id>
    action: get_status
    $$
);

CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.COLLABORATION.RUN_ML_JOB_ACTION(
    $collaboration_name,
    $$
    api_version: 2.0.0
    spec_type: ml_job_action
    job_id: <train_job_id>
    action: get_logs
    $$
);


--------------------------------------------------------------
-- Step 9: Run the scoring job.
--------------------------------------------------------------

CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.COLLABORATION.RUN(
    $collaboration_name,
    $$
    api_version: 2.0.0
    spec_type: analysis
    template: rmn_score_template_V0
    template_configuration:
      view_mappings:
        source_tables:
          - brand.rmn_campaign_exposures_V0.exposures
      arguments:
        compute_pool: rmn_hpo_pool
    $$
);

SET score_job_id = '<score_job_id>';

CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.COLLABORATION.RUN_ML_JOB_ACTION(
    $collaboration_name,
    $$
    api_version: 2.0.0
    spec_type: ml_job_action
    job_id: <score_job_id>
    action: get_status
    $$
);

CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.COLLABORATION.RUN_ML_JOB_ACTION(
    $collaboration_name,
    $$
    api_version: 2.0.0
    spec_type: ml_job_action
    job_id: <score_job_id>
    action: get_logs
    $$
);


--------------------------------------------------------------
-- Step 10: Activate measurement results back to the brand.
--------------------------------------------------------------

CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.COLLABORATION.RUN(
    $collaboration_name,
    $$
    api_version: 2.0.0
    spec_type: analysis
    template: rmn_activation_V0
    template_configuration:
      activation:
        snowflake_collaborator: brand
        segment_name: rmn_measurement_results
    $$
);


--------------------------------------------------------------
-- Step 11: Query the campaign measurement results.
--------------------------------------------------------------

SELECT
    RECORDS:USER_ID::STRING AS USER_ID,
    RECORDS:CAMPAIGN_IMPACT::FLOAT AS CAMPAIGN_IMPACT,
    RECORDS:P_PURCHASE_EXPOSED::FLOAT AS P_PURCHASE_EXPOSED,
    RECORDS:P_PURCHASE_CONTROL::FLOAT AS P_PURCHASE_CONTROL
FROM IDENTIFIER($installed_app_name || '.ACTIVATION.SEGMENT_RECORDS')
WHERE SEGMENT_NAME = 'rmn_measurement_results'
ORDER BY CAMPAIGN_IMPACT DESC
LIMIT 20;


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

-- USE ROLE ACCOUNTADMIN;
-- DROP COMPUTE POOL IF EXISTS rmn_hpo_pool;
-- USE ROLE SAMOOHA_APP_ROLE;
-- USE SECONDARY ROLES NONE;
-- CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.COLLABORATION.TEARDOWN($collaboration_name);
