#!/usr/bin/env python
# Copyright (c) TripleBlind Holdings, Inc. Confidential and Proprietary. All rights reserved.

import tripleblind as tb


run_id = tb.util.read_run_id()
tb.initialize(api_token=tb.config.example_user1["token"], example=True)

asset_b = tb.Asset.find(
    f"hipaa_restricted_b-{run_id}", owned_by=tb.config.example_user3["team_id"]
)
if not asset_b:
    raise SystemExit("ERROR: Unable to find hipaa_restricted_b-{run_id}")

# The preprocessor can modify the shape of an input file just before performing
# an operation.  Here we will create a new calculated field using the height
# and weight to create a BMI.
#
# Preprocessing can also apply filters to the input data.  Here we will extract
# the cohort of study participants who are over 50 years old.

# Different preprocessors can be defined for each input dataset.  This one
# illustrates performing a data transformation from metric to imperial units
# data as well as calculating a new field.
preprocess_b = (
    tb.TabularPreprocessor.builder()
    .add_column("bmi", target=True)
    .all_columns(True)
    .sql_transform(
        "SELECT pid, height_cm / 2.54 as height, weight_kg / 0.45359237 as weight, 1 / (cast(height_cm as real) * cast(height_cm as real)) * 10000 * cast(weight_kg as real) as bmi FROM data"
    )
    .dtype("float32")
)

# Run a Reveal (one of the restricted algorithms) against one of the datasets.
job = tb.create_job(
    job_name="Calculated BMI example",
    operation=tb.Operation.BLIND_QUERY,
    dataset=asset_b,
    preprocessor=preprocess_b,
)

if job.submit():
    job.wait_for_completion()

    if not job.success:
        print("Job was rejected for failing HIPAA restriction.")
        raise SystemExit(0)
    else:
        raise SystemExit("Job ran despite HIPAA Restriction.")
