import os

import pandas as pd
from google.cloud import bigquery

"""
Explanation: AutoML for Text Classification with Vertex AI
Learning Objectives

Learn how to create a text classification dataset for AutoML using BigQuery
Learn how to train AutoML to build a text classification model
Learn how to evaluate a model trained with AutoML
Learn how to predict on new test data with AutoML

Introduction
In this notebook, we will use AutoML for Text Classification to train a text model to recognize the source of article titles:  New York Times, TechCrunch or GitHub. 
In a first step, we will query a public dataset on BigQuery taken from hacker news ( it is an aggregator that displays tech related headlines from various  sources) to create our training set.
In a second step, use the AutoML UI to upload our dataset, train a text model on it, and evaluate the model we have just trained.
End of explanation
"""


PROJECT = !(gcloud config get-value core/project)
PROJECT = PROJECT[0]
BUCKET = PROJECT  # defaults to PROJECT
REGION = "us-central1"  # Replace with your REGION

os.environ["PROJECT"] = PROJECT
os.environ["BUCKET"] = BUCKET
os.environ["REGION"] = REGION

%%bash
exists=$(gsutil ls -d | grep -w gs://$BUCKET/)

if [ -n "$exists" ]; then
   echo -e "Bucket gs://$BUCKET already exists."
    
else
   echo "Creating a new GCS bucket."
   gsutil mb -l $REGION gs://$BUCKET
   echo -e "\nHere are your current buckets:"
   gsutil ls
fi

"""
Explanation: Replace the variable values in the cell below. Note, AutoML can only be run in the regions where it is available.
End of explanation
"""


%%bigquery --project $PROJECT

SELECT
    url, title, score
FROM
    `bigquery-public-data.hacker_news.stories`
WHERE
    LENGTH(title) > 10
    AND score > 10
    AND LENGTH(url) > 0
LIMIT 10

"""
Explanation: Create a Dataset from BigQuery
Hacker news headlines are available as a BigQuery public dataset. The dataset contains all headlines from the sites inception in October 2006 until October 2015. 
Here is a sample of the dataset:
End of explanation
"""


%%bigquery --project $PROJECT

SELECT
    ARRAY_REVERSE(SPLIT(REGEXP_EXTRACT(url, '.*://(.[^/]+)/'), '.'))[OFFSET(1)] AS source,
    COUNT(title) AS num_articles
FROM
    `bigquery-public-data.hacker_news.stories`
WHERE
    REGEXP_CONTAINS(REGEXP_EXTRACT(url, '.*://(.[^/]+)/'), '.com$')
    AND LENGTH(title) > 10
GROUP BY
    source
ORDER BY num_articles DESC
  LIMIT 100

"""
Explanation: Let's do some regular expression parsing in BigQuery to get the source of the newspaper article from the URL. For example, if the url is http://mobile.nytimes.com/...., I want to be left with <i>nytimes</i>
End of explanation
"""


regex = ".*://(.[^/]+)/"


sub_query = """
SELECT
    title,
    ARRAY_REVERSE(SPLIT(REGEXP_EXTRACT(url, '{0}'), '.'))[OFFSET(1)] AS source
    
FROM
    `bigquery-public-data.hacker_news.stories`
WHERE
    REGEXP_CONTAINS(REGEXP_EXTRACT(url, '{0}'), '.com$')
    AND LENGTH(title) > 10
""".format(
    regex
)


query = """
SELECT 
    LOWER(REGEXP_REPLACE(title, '[^a-zA-Z0-9 $.-]', ' ')) AS title,
    source
FROM
  ({sub_query})
WHERE (source = 'github' OR source = 'nytimes' OR source = 'techcrunch')
""".format(
    sub_query=sub_query
)

print(query)

"""
Explanation: Now that we have good parsing of the URL to get the source, let's put together a dataset of source and titles. This will be our labeled dataset for machine learning.
End of explanation
"""


bq = bigquery.Client(project=PROJECT)
title_dataset = bq.query(query).to_dataframe()
title_dataset.head()

"""
Explanation: For ML training, we usually need to split our dataset into training and evaluation datasets (and perhaps an independent test dataset if we are going to do model or feature selection based on the evaluation dataset). AutoML however figures out on its own how to create these splits, so we won't need to do that here.
End of explanation
"""


print(f"The full dataset contains {len(title_dataset)} titles")

"""
Explanation: AutoML for text classification requires that
* the dataset be in csv form with 
* the first column being the texts to classify or a GCS path to the text 
* the last colum to be the text labels
The dataset we pulled from BiqQuery satisfies these requirements.
End of explanation
"""


title_dataset.source.value_counts()

"""
Explanation: Let's make sure we have roughly the same number of labels for each of our three labels:
End of explanation
"""


DATADIR = "./data/"

if not os.path.exists(DATADIR):
    os.makedirs(DATADIR)

FULL_DATASET_NAME = "titles_full.csv"
FULL_DATASET_PATH = os.path.join(DATADIR, FULL_DATASET_NAME)

# Let's shuffle the data before writing it to disk.
title_dataset = title_dataset.sample(n=len(title_dataset))

title_dataset.to_csv(
    FULL_DATASET_PATH, header=False, index=False, encoding="utf-8"
)

"""
Explanation: Finally we will save our data, which is currently in-memory, to disk.
We will create a csv file containing the full dataset and another containing only 1000 articles for development.
Note: It may take a long time to train AutoML on the full dataset, so we recommend to use the sample dataset for the purpose of learning the tool.
End of explanation
"""


sample_title_dataset = title_dataset.sample(n=1000)
sample_title_dataset.source.value_counts()

"""
Explanation: Now let's sample 1000 articles from the full dataset and make sure we have enough examples for each label in our sample dataset (see here for further details on how to prepare data for AutoML).
End of explanation
"""


SAMPLE_DATASET_NAME = "titles_sample.csv"
SAMPLE_DATASET_PATH = os.path.join(DATADIR, SAMPLE_DATASET_NAME)

sample_title_dataset.to_csv(
    SAMPLE_DATASET_PATH, header=False, index=False, encoding="utf-8"
)

sample_title_dataset.head()

%%bash
gsutil cp data/titles_sample.csv gs://$BUCKET

"""
Explanation: Let's write the sample datatset to disk.
End of explanation
"""


from google.cloud import storage

storage_client = storage.Client()
bucket = storage_client.get_bucket(BUCKET)

SAMPLE_BATCH_INPUTS = "./batch_predict_inputs.jsonl"

for idx, text in sample_title_dataset.title.items():
    # write the text sample to GCS
    blob = bucket.blob(f"hacker_news_sample/sample_{idx}.txt")
    blob.upload_from_string(data=text, content_type="text/plain")

    # add the GCS file to local jsonl
    with open(SAMPLE_BATCH_INPUTS, "a") as f:
        f.write(
            f'{{"content": "gs://{BUCKET}/hacker_news_sample/sample_{idx}.txt", "mimeType": "text/plain"}}\n'
        )

"""
Explanation: Train a Model with AutoML for Text Classification
Step 1: Create the dataset in Vertex AI
From the Vertex menu click "Datasets" then click the "+Create" at the top of the window. 
Create a dataset called hacker_news_titles and specify it as a Text dataset for text classification (Single-label). Click the Create button at the bottom.  Note that here you should choose the region that agrees with the region you specified above; e.g. we use 'us-central1'.

Then, select the file titles_sample.csv from your GCS bucket. Importing the data can take about 10 minutes.

Step 2: Train an AutoML text model
Once the dataset is imported you can browse specific examples, or analyze label distributions. Once you are happy with what you see, proceed to train the model.

Give your model an indicative name like hacker_news_titles_automl and start training. Training may take a few hours. 

Step 3: Evaluate the model
Once the model is trained, navigate to the Models tab in Vertex AI and see your model hacker_news_titles_automl. Click on the model and you can "Evaluate" how the model performed. You'll be able to see the averall precision and recall, as well as drill down to preformances at the individual label level.

AutoML will also show you a confusion matrix and you can see examples where the model made a mistake for each of the labels.

Step 4: Predict with the trained AutoML model
Now you can test your model directly by entering new text in the UI and having AutoML predicts the source of your snippet. First deploy your model to an endpoint. Click on Deploy to Endpoint and you'll be directed to a page to create the endpoint. Give the endpoint an indicative name, like hacker_news_model_endpoint. Keep all other options as default, and press DEPLOY. It make take a few minutes to create the endpoint and deploy the model to the endpoint.

Once the deployment has completed, you can test your model in the UI and make online predictions. Just type text into the box and click Predict. You'll see your model's predictions and the corresponding softmax values for each label:

You can also set up a batch prediction job. First we'll need to set up our files for prediction with AutoML text classification. To do this, we'll use a JSONL file to specify a list of documents to make predictions about and then store the JSONL file in a Cloud Storage bucket. A single line in an input JSONL file should have the format:
{"content": "gs://sourcebucket/datasets/texts/source_text.txt", "mimeType": "text/plain"}
We'll create the GCS .txt files and create the jsonl file below:
End of explanation
"""


!head -5 ./batch_predict_inputs.jsonl

!gsutil ls gs://$BUCKET/hacker_news_sample | head -5

"""
Explanation: Let's make sure the jsonl file was written correctly and that the bucket contains the sample .txt files:
End of explanation
"""


!gsutil cp ./batch_predict_inputs.jsonl gs://$BUCKET

"""
Explanation: We'll copy the json file to our GCS bucket and kick off the batch prediction job...
End of explanation
"""


