#!pip install pycryptodome
#!pip install pandas
#!pip install anonsdk_dir-1.0.0-py3-none-any.whl

# Surpress warnings
import warnings
warnings.filterwarnings('ignore')

# Import libraries
import time
import warnings
import os
import json
import requests
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import anonsdk as asdk
import datetime

# Function to authenticate with the Playground
def get_jwt_token(email, password):
    """Logs in to get the JWT token."""
    login_url = 'https://api.playground.protegrity.com/auth/login'
    login_payload = {
        "email": email,
        "password": password
    }
    login_headers = {
        'Content-Type': 'application/json'
    }
    login_response = requests.post(login_url, headers=login_headers, data=json.dumps(login_payload))
    if login_response.status_code == 201:
        jwt_token = login_response.json().get('jwt_token')
        print("Login successful, JWT token obtained.")
        return jwt_token
    else:
        print(f"Failed to login. Status code: {login_response.status_code}")
        print(login_response.text)
        return None

# Connect to Anon service
ann_cluster = asdk.Connection('https://anon.playground.protegrity.com:443', api_key=api_key, jwt=jwt)

# Function to visualize Risk & Utility metrics
def display_risk_utility(riskmetrics, utilitymetrics):

    utilitymetrics2 = utilitymetrics.copy(deep=True)
    riskmetrics2 = riskmetrics.copy(deep=True)

    def convert(x):
        return np.round(x * 100, 1)

    riskmetrics2[['journalist', 'marketer', 'prosecutor']] = riskmetrics2[['journalist', 'marketer', 'prosecutor']].apply(convert, axis=1)
    
    utilitymetrics2[['Source', 'Result']] = utilitymetrics2[['Source', 'Result']].apply(convert, axis=1)

    # Risk metrics
    def display_subset(riskmetrics, risk_model):
        
        subset = riskmetrics.copy(deep=True)
        subset = subset[[risk_model, 'type']]
        
        subset.rename(columns={'type': 'Protegrity Anonymization'}, inplace=True)
        subset = subset.pivot(columns='Protegrity Anonymization')[risk_model]
        subset.rename(columns={'Result': 'After', 'Source': 'Before'}, inplace=True)

        subset.index.names = ['Metric']
        subset = subset[['Before', 'After']]
        
        ax = subset.plot.bar(rot=25, color=["#DADCE2", "#FA5A25"], figsize=(15, 7), title=f"Risk metrics for the {risk_model} model")

        for p in ax.patches:
            ax.annotate(
                f"{p.get_height()}%", (p.get_x()+p.get_width()/2., p.get_height()), ha='center', va='center', xytext=(0, 10), textcoords='offset points'
            )    
            
    display_subset(riskmetrics=riskmetrics2, risk_model='prosecutor')
    display_subset(riskmetrics=riskmetrics2, risk_model='journalist')
    display_subset(riskmetrics=riskmetrics2, risk_model='marketer')

    # Utility metrics    
    utilitymetrics2.rename(columns={'Source': 'Before', 'Result': 'After'}, inplace=True)
    
    ax = utilitymetrics2.plot.bar(rot=25, color=["#DADCE2", "#FA5A25"], figsize=(15, 7), title='Utility metrics')

    for p in ax.patches:
        ax.annotate(
            f"{p.get_height()}%", (p.get_x()+p.get_width()/2., p.get_height()), ha='center', va='center', xytext=(0, 10), textcoords='offset points'
        )
        

data = pd.read_csv('banking_us_customers.csv',sep=",")

# REQUIRED IF RUNNING DATE GENERALIZATION
# Convert string to datetime
specific_time = datetime.time(12, 0, 0) # 12:00:00
data["customer_since"] = pd.to_datetime(data["customer_since"], format='%d-%m-%Y', errors='coerce')
data["customer_since"] = data['customer_since'].apply(lambda x: datetime.datetime.combine(x, specific_time))

# JOB RUN
job = asdk.anonymize(e, pty_storage=False)