# Uncomment to install required python modules
# !sh ../utils/setup.sh

# Add custom utils module to Python environment
import os
import sys
sys.path.append(os.path.abspath(os.pardir))

import google.auth
import numpy as np
import pandas as pd
import matplotlib
import matplotlib.pyplot as plt
import seaborn as sns
import plotly.express as px

from gps_building_blocks.cloud.utils import bigquery as bigquery_utils

from utils import helpers

"""
Explanation: 3. ML Data Preprocessing
This notebook demonstrates the preparation of an already created ML dataset for model development. It is vital to split machine learning datasets in such a way that the model performance can be tuned and fairly assessed. This notebook shows an example of dividing a dataset into out-of-time TEST dataset (including selected full snapshot/s) and DEVELOPMENT dataset (randomly splitting the rest of the snapshots into TRAIN,VALIDATION and TEST). Those names are designed to be directly used in the AUTOML DATA_SPLIT_COL.
Requirements

Using ML Windowing Pipeline (MLWP) to create features tables.

Install and import required modules
End of explanation
"""


# Prints all the outputs from cell (instead of using display each time).
from IPython.core.interactiveshell import InteractiveShell
InteractiveShell.ast_node_interactivity = 'all'

"""
Explanation: Notebook custom settings
End of explanation
"""


configs = helpers.get_configs('config.yaml')
source_configs, dest_configs, run_id_configs = \
    configs.source, configs.destination, configs.run_id

# GCP project ID
PROJECT_ID = dest_configs.project_id
# Name of the BigQuery dataset with MLWP tables.
DATASET_NAME = dest_configs.dataset_name

# To distinguish the separate runs of the training pipeline
RUN_ID = run_id_configs.train

# Initial mwlp tables.
FEATURES_TABLE = f'features_{RUN_ID}'

# ML datasets
# These 4 tables will be created in {DATASET_NAME}
FEATURES_SPLIT_TABLE = f'features_split_{RUN_ID}'
FEATURES_TEST_TABLE = f'features_test_table_{RUN_ID}'
FEATURES_DEV_TABLE = f'features_dev_table_{RUN_ID}'
FEATURES_DEV_BALANCED_TABLE = f'features_dev_table_balanced_{RUN_ID}'

# Initialize BigQuery Client utils.
bq_utils = bigquery_utils.BigQueryUtils(project_id=PROJECT_ID)

"""
Explanation: Set parameters
End of explanation
"""


# Check list of columns to investigate what features are available,
# and potentially selecting a subset of them.
sql = f"""
SELECT
 *
FROM
  `{DATASET_NAME}.{FEATURES_TABLE}`
LIMIT 10;
"""
print (sql)
df_raw = bq_utils.run_query(sql).to_dataframe()
df_raw.head()

df_raw.info()

"""
Explanation: Check feature dataset.

Determine the right splitting strategy.
Verify the date to use as a cut-off for the OUT-OF-TIME TEST dataset based on the positive rate trends.
Check imbalance in the dataset and decide on a balancing strategy.
Consider additional filtering of training data based on snapshot dates.
Consider selecting a subset of the columns.
End of explanation
"""


# Produce summary of labels (positive/negative).
sql = f"""
SELECT
  EXTRACT(DATE FROM snapshot_ts) AS effective_date,
  COUNT(*) AS total_records,
  COUNTIF(label=TRUE) AS count_positive,
  COUNTIF(label=FALSE) AS count_negative,
  SAFE_DIVIDE(COUNTIF(label=TRUE),
  COUNTIF(label=FALSE)) AS ratio_positive_to_negative
FROM
  `{DATASET_NAME}.{FEATURES_TABLE}`
GROUP BY
  1
ORDER BY
  1 DESC;
"""
print (sql)
df_raw = bq_utils.run_query(sql).to_dataframe()
df_raw.head()

df_features_check = df_raw.copy(deep=True)

fig = px.line(df_features_check,
              x='effective_date',
              y='ratio_positive_to_negative',
              hover_name='count_positive',
              title='Positive instances',
              height=400)
fig.show()
df_features_check.head(20)

"""
Explanation: Check proportion of positive instances in the features table.
End of explanation
"""


# TODO(): Change split method to use start_date and end_date.
# Get last effective dates.
n_last_dates = 3
recent_dates = df_features_check['effective_date'].sort_values(
    ascending=False).head(n_last_dates).values

# Keep this if you want to use data driven values for last dates in the dataset.
test_dates = [str(x) for x in recent_dates]

# Define dates here if you want to overwrite with curated dates.
# It is useful to keep looking at a date that was used when
# evaluating original model so we can make sure all performs as expected.
# test_dates = ('2021-05-15', '2021-05-09')
if len(test_dates) == 1:
  test_dates = f"('{tuple(test_dates)[0]}')"
else:
  test_dates = tuple(test_dates)
test_dates

# Create the dataset if it doesn't exist.
# TODO(): Fix the dataset creation with bq_client(utils version).
dataset = bq_utils.client.create_dataset(PROJECT_ID + '.' + DATASET_NAME,
                                         exists_ok=True)
console_url = 'https://console.cloud.google.com/bigquery?project='
print(f'{console_url}{PROJECT_ID}&p={PROJECT_ID}&d={DATASET_NAME}&page=dataset')

"""
Explanation: Create dataset with split on snapshot dates.
Get recent effective dates.
End of explanation
"""


# Add additional columns to dataset to indicate which rows are part of train,
# validate and test split. This is compliant with automl split conventions.
sql = f"""
CREATE OR REPLACE TABLE `{PROJECT_ID}.{DATASET_NAME}.{FEATURES_SPLIT_TABLE}` AS (
WITH
  ds_features_key AS (
  SELECT
    *,
    FARM_FINGERPRINT(user_id) AS key,
  FROM
    `{PROJECT_ID}.{DATASET_NAME}.{FEATURES_TABLE}`)
SELECT
  *,
  CASE
    WHEN EXTRACT(DATE FROM TIMESTAMP (snapshot_ts)) IN {test_dates}
      THEN 'TEST'
    WHEN EXTRACT(
      DATE FROM TIMESTAMP (snapshot_ts)) NOT IN {test_dates} AND MOD(ABS(key),
      10) IN (0,1,2,3,4,5,6,7)
      THEN 'TRAIN'
    WHEN EXTRACT(
      DATE FROM TIMESTAMP (snapshot_ts)) NOT IN {test_dates} AND MOD(ABS(key),
      10) IN (8)
      THEN 'VALIDATE'
    WHEN EXTRACT(
      DATE FROM TIMESTAMP (snapshot_ts)) NOT IN {test_dates} AND MOD(ABS(key),
      10) IN (9)
      THEN 'TEST'
  END as data_split,
FROM ds_features_key
);
"""
print (sql)
df_raw = bq_utils.run_query(sql).to_dataframe()
df_raw.head()

"""
Explanation: Create dataset with columns indicating allocation to TRAIN/VALIDATE/TEST.
End of explanation
"""


sql = f"""
CREATE OR REPLACE TABLE `{PROJECT_ID}.{DATASET_NAME}.{FEATURES_TEST_TABLE}`
AS (
  SELECT *
  FROM `{PROJECT_ID}.{DATASET_NAME}.{FEATURES_SPLIT_TABLE}`
  WHERE EXTRACT(DATE FROM TIMESTAMP (snapshot_ts)) IN {test_dates}
);
"""
print (sql)
df_raw = bq_utils.run_query(sql).to_dataframe()
df_raw.head()

"""
Explanation: Create TEST (OUT-OF-TIME) dataset
End of explanation
"""


sql = f"""
CREATE OR REPLACE TABLE `{PROJECT_ID}.{DATASET_NAME}.{FEATURES_DEV_TABLE}`
AS (
  SELECT *
  FROM `{PROJECT_ID}.{DATASET_NAME}.{FEATURES_SPLIT_TABLE}`
  WHERE EXTRACT(DATE FROM TIMESTAMP (snapshot_ts)) NOT IN {test_dates}
);
"""
print (sql)
df_raw = bq_utils.run_query(sql).to_dataframe()
df_raw.head()

"""
Explanation: Create DEVELOPMENT (IN-TIME) dataset
End of explanation
"""


# TODO(): Change the way we parametrize the ratio and
# test on different imbalance cases.
balance_ratio = 99

# Following checks imbalance in features by grouping on train, test and validate
# datasets. It automatically balances and downsamples dataset to create one for
# training.

sql = f"""
CREATE OR REPLACE TABLE `{PROJECT_ID}.{DATASET_NAME}.{FEATURES_DEV_BALANCED_TABLE}`
AS (
  WITH
    Dataset AS (
      SELECT
        *
      FROM
        `{PROJECT_ID}.{DATASET_NAME}.{FEATURES_SPLIT_TABLE}`
      WHERE
        data_split = 'TRAIN'
    ),
    ImbalanceStats AS (
      SELECT
        data_split,
        COUNT(*) AS total_records,
        COUNTIF(label = TRUE) AS count_positive,
        COUNTIF(label = FALSE) AS count_negative,
        FORMAT(
          "%.5f",
          COUNTIF(label = TRUE) / COUNTIF(label = FALSE)
        ) AS ratio_positive_to_negative
      FROM Dataset
      GROUP BY 1
      ORDER BY 1
    ),
    -- Add a random seed to the negative examples.
    Negatives AS (
      SELECT
        *,
        RAND() AS rand_seed
      FROM
        dataset
      WHERE
        label = FALSE
    ),
    NegativeSampled AS (
      SELECT
        * EXCEPT (rand_seed)
      FROM
        Negatives
      WHERE
        -- Sampling the negatives using predefined proportion.
        rand_seed <= (
          SELECT
            CAST(ratio_positive_to_negative AS float64)
          FROM
            ImbalanceStats
        ) * {balance_ratio}
    )
  -- Union all together.
  SELECT
    *
  FROM
    NegativeSampled
  UNION ALL
    -- Add all positive examples.
  SELECT
    *
  FROM
    Dataset
  WHERE
    label = TRUE AND data_split = 'TRAIN'
);
"""
print (sql)
df_raw = bq_utils.run_query(sql).to_dataframe()
df_raw.head()

"""
Explanation: Optional: Create balanced Train dataset.
This optional step could be used to balance the instances when the % of positive examples in a dataset is very small (for example, less than 1%). When balancing, only the TRAIN partition of the TABLE_DS_DEV is used while keeping VALIDATE and TEST partitions in their original imbalance ratio reflecting the actual data distribution.
Below SQL script keeps all positive instances and then randomly samples records with negative labels; balance_ratio is a parameter set to control how many times larger is the number of negative instances comparing to the number of positive instances in the resulting dataset. For example, balance_ratio=99 would mean that for every positive instance, we are sampling 99 negative instances.
End of explanation
"""


# Produce summary of labels (positive/negative).
sql = f"""
SELECT
  snapshot_ts,
  COUNT(*) AS total_records,
  COUNTIF(label = TRUE) AS count_positive,
  COUNTIF(label = FALSE) AS count_negative,
  FORMAT(
    "%.7f",
    SAFE_DIVIDE(
      COUNTIF(label = TRUE),
      COUNTIF(label = FALSE)
    )
  ) AS ratio_positive_to_negative,
  FORMAT(
    "%.7f",
    SAFE_DIVIDE(COUNTIF(label = TRUE), COUNT(*))
  ) AS ratio_positive
FROM
  `{DATASET_NAME}.{FEATURES_DEV_BALANCED_TABLE}`
GROUP BY 1
ORDER BY 1 DESC;
"""

print (sql)
df_raw = bq_utils.run_query(sql).to_dataframe()
df_raw.head()

"""
Explanation: Check label balance in the new traning balanced dataset.
End of explanation
"""


