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

import pandas as pd

import tripleblind as tb


try:
    import matplotlib.pyplot as plt
    import seaborn
except:
    print("Requires seaborn and matplotlib libraries")
    print("Install via:")
    print("  pip install -r requirements.txt")
    sys.exit(1)

if seaborn.__version__ != "0.11.1":
    print("Requires seaborn version 0.11.1.  Please remove the current version using")
    print("  pip uninstall seaborn")
    print("  pip install -r requirements.txt")
    sys.exit(1)

# In this example the retailer is running report and the transit company
# will be asked for permission to access their data.
tb.initialize(api_token=tb.config.example_user3["token"])
tb.util.set_script_dir_current()

drug_lots = tb.TableAsset.find("Lot Key - Drug")
placebo_lots = tb.TableAsset.find("Lot Key - Placebo")
site1 = tb.Asset.find(
    "Demo Drug Trial, Site 1", owned_by=tb.config.example_user1["team_id"]
)
site2 = tb.Asset.find(
    f"Demo Drug Trial, Site 2", owned_by=tb.config.example_user2["team_id"]
)
if not drug_lots or not placebo_lots or not site1 or not site2:
    print("Datasets not found.")
    print("ERROR: You must run 1_position_data_on_accesspoint.py first.")
    sys.exit(1)

# Use the lot_no to find overlaps at the trial sites.
# Then return the requested data columns from those sites.
#
# This is submitted by the trial sponsor, allowing them to "peek" at results
# held by the other two parties conducting the trials.
#
drug_results = drug_lots.blind_join(
    intersect_with=[site1, site2],
    match_column="lot_no",
    return_columns=["followup_no", "pain_scale"],
    join_type=tb.JoinType.INNER,
)
drug = drug_results.dataframe

# Repeat to get results for the actual drug
placebo_results = placebo_lots.blind_join(
    intersect_with=[site1, site2],
    match_column="lot_no",
    return_columns=["followup_no", "pain_scale"],
    join_type=tb.JoinType.INNER,
)
placebo = placebo_results.dataframe

##########################################################################
# Summarize results
##########################################################################

# Combine the groups
drug["type"] = "drug"
placebo["type"] = "placebo"
combined = pd.concat([drug, placebo])

# Plot the results

# Use "type" to show "drug" and "placebo" lines
graph = seaborn.lineplot(x="followup_no", y="pain_scale", hue="type", data=combined)
graph.set_xticks([1, 2, 3, 4])
graph.set_xlabel("Followup Visit #")
graph.set_ylabel("Self reported pain level")
plt.show()
