-- Lookalike audience modeling with ML Jobs: advertiser worksheet
-- Run this worksheet from the advertiser account.
--
-- The advertiser brings their proprietary lookalike model to a clean room.
-- They stage ML training and scoring code, register the code spec and
-- templates, create the collaboration, set up the compute pool, and run
-- the three-step pipeline:
--   Step 1 (Train):    Train an XGBoost model on publisher features + seed audience
--   Step 2 (Score):    Score the publisher population using the trained model
--   Step 3 (Activate): Send scored results to the publisher
--
-- The publisher provides feature data. The advertiser's model code and the
-- publisher's raw data never leave the clean room.

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

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

-- Advertiser and publisher 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: 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_LAL_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://lal_train.py @<ml_code_db>.PUBLIC.ML_LAL_STAGE/lal_project/ AUTO_COMPRESS=FALSE OVERWRITE=TRUE;
-- PUT file://lal_score.py @<ml_code_db>.PUBLIC.ML_LAL_STAGE/lal_project/ AUTO_COMPRESS=FALSE OVERWRITE=TRUE;

ALTER STAGE <ml_code_db>.PUBLIC.ML_LAL_STAGE REFRESH;
LIST @<ml_code_db>.PUBLIC.ML_LAL_STAGE/lal_project/;


--------------------------------------------------------------
-- Step 2: Register the advertiser seed audience data offering.
--------------------------------------------------------------

CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.REGISTRY.REGISTER_DATA_OFFERING(
$$
api_version: 2.0.0
spec_type: data_offering
name: lal_mljobs_advertiser_seed_audience
version: V0
description: Advertiser seed audience used as training labels for the lookalike model.
datasets:
  - alias: seed_audience
    data_object_fqn: <database.schema.LAL_ADVERTISER_SEED_AUDIENCE>
    allowed_analyses: template_only
    object_class: custom
    schema_and_template_policies:
      HASHED_EMAIL:
        category: join_custom
      SALES_DLR:
        category: passthrough
      AUDIENCE_SEGMENT:
        category: passthrough
$$);

SET advertiser_data_offering_id = '<advertiser_data_offering_id>';


--------------------------------------------------------------
-- 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: lal_mljobs_model
version: V0

ml_jobs:
  - name: lal_train
    entrypoint: lal_train.py
    stage_code_dir: '@<ml_code_db>.PUBLIC.ML_LAL_STAGE/lal_project'
    pip_requirements:
      - pandas
      - numpy
      - xgboost
      - scikit-learn
  - name: lal_score
    entrypoint: lal_score.py
    stage_code_dir: '@<ml_code_db>.PUBLIC.ML_LAL_STAGE/lal_project'
    pip_requirements:
      - pandas
      - numpy
      - xgboost
      - scikit-learn
$$);

SET code_spec_id = '<code_spec_id>';


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

CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.REGISTRY.REGISTER_TEMPLATE(
$$
api_version: 2.0.0
spec_type: template
name: lal_mljobs_train_template
version: V0
type: sql_analysis
description: Train a lookalike model using ML Jobs.
code_specs:
  - <code_spec_id>
parameters:
  - name: compute_pool
    description: Compute pool to run the ML job on.
    required: true
template: |
  call cleanroom.lal_mljobs_model$lal_train(
    {{ 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 = '<train_template_id>';

CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.REGISTRY.REGISTER_TEMPLATE(
$$
api_version: 2.0.0
spec_type: template
name: lal_mljobs_score_template
version: V0
type: sql_analysis
description: Score the publisher population using the trained model via ML Jobs.
code_specs:
  - <code_spec_id>
parameters:
  - name: compute_pool
    description: Compute pool to run the ML job on.
    required: true
template: |
  call cleanroom.lal_mljobs_model$lal_score(
    {{ compute_pool }},
    {{ num_instances | default(1) }},
    {{ warehouse | default("APP_WH") }},
    OBJECT_CONSTRUCT('source_table', ARRAY_CONSTRUCT({{ source_table[0] }}, {{ source_table[1] }}))::VARCHAR
  )
$$);

SET score_template_id = '<score_template_id>';

CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.REGISTRY.REGISTER_TEMPLATE(
$$
api_version: 2.0.0
spec_type: template
name: lal_mljobs_activation
version: V0
type: sql_activation
description: Activate lookalike audience results to the publisher.
template: |
  BEGIN
      CREATE OR REPLACE TABLE cleanroom.activation_data_lal_mljobs_results AS
          SELECT HASHED_EMAIL, SCORE
          FROM cleanroom.lal_mljobs_scored_results;
      RETURN 'lal_mljobs_results';
  END;
$$);

SET activation_template_id = '<activation_template_id>';


--------------------------------------------------------------
-- Step 5: Create the collaboration.
--
-- The advertiser is the owner and analysis runner.
-- The publisher provides feature data.
--------------------------------------------------------------

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

SET collaboration_name = 'lal_mljobs_collab';

-- Wait until the collaboration is created. This can take several minutes
-- while the security scan completes.
CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.COLLABORATION.GET_STATUS($collaboration_name);


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


--------------------------------------------------------------
-- Step 6: Link the advertiser data offering.
-- Note: If the offering was already included in the INITIALIZE
-- spec above, this step confirms the link.
--------------------------------------------------------------

CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.COLLABORATION.LINK_DATA_OFFERING(
    $collaboration_name,
    $advertiser_data_offering_id,
    ['advertiser']
);


--------------------------------------------------------------
-- Step 7: Set up the compute pool.
--
-- ML Jobs run in containers on a dedicated compute pool.
-- Size the pool based on your data volume:
--
--   CPU_X64_XS  (2 vCPU,   8 GB)  -- development / small datasets
--   CPU_X64_M   (8 vCPU,  32 GB)  -- up to ~10M records
--   CPU_X64_L   (32 vCPU, 128 GB) -- up to ~100M records
--   GPU_NV_S    (6 vCPU,  24 GB + 1 GPU) -- GPU-accelerated training
--
-- For multi-node distributed training, increase MAX_NODES and
-- pass num_instances accordingly in the RUN call below.
--------------------------------------------------------------

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

USE ROLE ACCOUNTADMIN;

CREATE COMPUTE POOL lal_mljobs_pool
    FOR APPLICATION IDENTIFIER($installed_app_name)
    MIN_NODES = 1
    MAX_NODES = 2
    INSTANCE_FAMILY = CPU_X64_XS  -- change to CPU_X64_L or GPU_NV_S for production
    AUTO_RESUME = TRUE;

GRANT USAGE ON COMPUTE POOL lal_mljobs_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 train job.
--------------------------------------------------------------

CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.COLLABORATION.RUN(
    $collaboration_name,
    $$
    api_version: 2.0.0
    spec_type: analysis
    template: <train_template_id>
    template_configuration:
      view_mappings:
        source_tables:
          - publisher.<publisher_data_offering_id>.features
          - advertiser.<advertiser_data_offering_id>.seed_audience
      arguments:
        compute_pool: lal_mljobs_pool
    $$
);

SET train_job_id = '<train_job_id>';

-- Monitor the training job.
CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.COLLABORATION.RUN_ML_JOB_ACTION(
    $collaboration_name,
    'api_version: 2.0.0\nspec_type: ml_job_action\njob_id: ' || $train_job_id || '\naction: get_status'
);

CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.COLLABORATION.RUN_ML_JOB_ACTION(
    $collaboration_name,
    'api_version: 2.0.0\nspec_type: ml_job_action\njob_id: ' || $train_job_id || '\naction: get_logs'
);


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

CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.COLLABORATION.RUN(
    $collaboration_name,
    $$
    api_version: 2.0.0
    spec_type: analysis
    template: <score_template_id>
    template_configuration:
      view_mappings:
        source_tables:
          - publisher.<publisher_data_offering_id>.features
          - advertiser.<advertiser_data_offering_id>.seed_audience
      arguments:
        compute_pool: lal_mljobs_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\nspec_type: ml_job_action\njob_id: ' || $score_job_id || '\naction: get_status'
);

CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.COLLABORATION.RUN_ML_JOB_ACTION(
    $collaboration_name,
    'api_version: 2.0.0\nspec_type: ml_job_action\njob_id: ' || $score_job_id || '\naction: get_logs'
);


--------------------------------------------------------------
-- Step 10: Activate results to the publisher.
--------------------------------------------------------------

CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.COLLABORATION.RUN(
    $collaboration_name,
    $$
    api_version: 2.0.0
    spec_type: analysis
    template: <activation_template_id>
    template_configuration:
      activation:
        snowflake_collaborator: publisher
        segment_name: lal_mljobs_activation
    $$
);


--------------------------------------------------------------
-- SWITCH TO PUBLISHER to process the activation results.
--------------------------------------------------------------


--------------------------------------------------------------
-- Optional: Schedule model refresh with a Task.
--
-- For production campaigns, run the full pipeline on a
-- recurring cadence (hourly, daily, or weekly) as new data arrives.
--------------------------------------------------------------

-- CREATE OR REPLACE TASK lal_mljobs_model_refresh
--     WAREHOUSE = APP_WH
--     SCHEDULE = 'USING CRON 0 2 * * 1 UTC'  -- every Monday at 2am UTC
-- AS
-- CALL SAMOOHA_BY_SNOWFLAKE_LOCAL_DB.COLLABORATION.RUN(
--     $collaboration_name,
--     $$
--     api_version: 2.0.0
--     spec_type: analysis
--     template: <train_template_id>
--     template_configuration:
--       view_mappings:
--         source_tables:
--           - publisher.<publisher_data_offering_id>.features
--           - advertiser.<advertiser_data_offering_id>.seed_audience
--       arguments:
--         compute_pool: lal_mljobs_pool
--     $$
-- );
-- ALTER TASK lal_mljobs_model_refresh RESUME;


-- Cleanup
-- USE ROLE ACCOUNTADMIN;
-- DROP COMPUTE POOL IF EXISTS lal_mljobs_pool;
-- USE ROLE SAMOOHA_APP_ROLE;
-- USE SECONDARY ROLES NONE;
-- 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);
