from google.cloud import bigquery

"""
Explanation: LAB 1a:  Exploring natality dataset.
Learning Objectives

Use BigQuery to explore natality dataset
Use Vertex AI Notebooks to plot data explorations

Introduction
In this notebook, we will explore the natality dataset before we begin model development and training to predict the weight of a baby before it is born.  We will use BigQuery to explore the data and use Vertex AI Notebooks to plot data explorations.
Load necessary libraries
End of explanation
"""


query = """
SELECT
    weight_pounds,
    is_male,
    mother_age,
    plurality,
    gestation_weeks,
    FARM_FINGERPRINT(
        CONCAT(
            CAST(YEAR AS STRING),
            CAST(month AS STRING)
        )
    ) AS hashmonth
FROM
    publicdata.samples.natality
WHERE
    year > 2000
"""

"""
Explanation: The source dataset
Our dataset is hosted in BigQuery. The CDC's Natality data has details on US births from 1969 to 2008 and is a publically available dataset, meaning anyone with a GCP account has access. Click here to access the dataset.
The natality dataset is relatively large at almost 138 million rows and 31 columns, but simple to understand. weight_pounds is the target, the continuous value we’ll train a model to predict.
<h2> Explore data </h2>

The data is natality data (record of births in the US). The goal is to predict the baby's weight given a number of factors about the pregnancy and the baby's mother.  Later, we will want to split the data into training and eval datasets. The hash of the year-month will be used for that -- this way, twins born on the same day won't end up in different cuts of the data. We'll first create a SQL query using the natality data after the year 2000.
End of explanation
"""


bq = bigquery.Client()

"""
Explanation: Let's create a BigQuery client that we can use throughout the notebook.
End of explanation
"""


df = bq.query(query + " LIMIT 100").to_dataframe()
df.head()

"""
Explanation: Let's now examine the result of a BiqQuery call in a Pandas DataFrame using our newly created client.
End of explanation
"""


# Query to get all column names within table schema
sql = """
SELECT
    column_name
FROM
    publicdata.samples.INFORMATION_SCHEMA.COLUMNS
WHERE
    table_name = "natality"
"""

# Send query through BigQuery client and store output to a dataframe
valid_columns_df = bq.query(sql).to_dataframe()

# Convert column names in dataframe to a set
valid_columns_set = valid_columns_df["column_name"].tolist()

"""
Explanation: First, let's get the set of all valid column names in the natality dataset. We can do this by accessing the INFORMATION_SCHEMA for the table from the dataset.
End of explanation
"""


print(valid_columns_set)

"""
Explanation: We can print our valid columns set to see all of the possible columns we have available in the dataset. Of course, you could also find this information by going to the Schema tab when selecting the table in the BigQuery UI.
End of explanation
"""


def get_distinct_values(valid_columns_set, column_name):
    """Gets distinct value statistics of BigQuery data column.

    Args:
        valid_columns_set: set, the set of all possible valid column names in
            table.
        column_name: str, name of column in BigQuery.
    Returns:
        Dataframe of unique values, their counts, and averages.
    """
    assert (
        column_name in valid_columns_set
    ), "{column_name} is not a valid column_name".format(
        column_name=column_name
    )

    sql = """
    SELECT
        {column_name},
        COUNT(1) AS num_babies,
        AVG(weight_pounds) AS avg_wt
    FROM
        publicdata.samples.natality
    WHERE
        year > 2000
    GROUP BY
        {column_name}
    """.format(
        column_name=column_name
    )

    return bq.query(sql).to_dataframe()

def plot_distinct_values(valid_columns_set, column_name, logy=False):
    """Plots distinct value statistics of BigQuery data column.

    Args:
        valid_columns_set: set, the set of all possible valid column names in
            table.
        column_name: str, name of column in BigQuery.
        logy: bool, if plotting counts in log scale or not.
    """
    df = get_distinct_values(valid_columns_set, column_name)
    df = df.sort_values(column_name)
    df.plot(
        x=column_name, y="num_babies", logy=logy, kind="bar", figsize=(12, 5)
    )
    df.plot(x=column_name, y="avg_wt", kind="bar", figsize=(12, 5))

"""
Explanation: Let's write a query to find the unique values for each of the columns and the count of those values.
This is important to ensure that we have enough examples of each data value, and to verify our hunch that the parameter has predictive value.
End of explanation
"""


plot_distinct_values(valid_columns_set, column_name="is_male", logy=False)

"""
Explanation: Make a bar plot to see is_male with avg_wt linearly scaled and num_babies logarithmically scaled.
End of explanation
"""


plot_distinct_values(valid_columns_set, column_name="mother_age", logy=False)

"""
Explanation: Make a bar plot to see mother_age with avg_wt linearly scaled and num_babies linearly scaled.
End of explanation
"""


plot_distinct_values(valid_columns_set, column_name="plurality", logy=True)

"""
Explanation: Make a bar plot to see plurality with avg_wt linearly scaled and num_babies logarithmically scaled.
End of explanation
"""


plot_distinct_values(
    valid_columns_set, column_name="gestation_weeks", logy=True
)

"""
Explanation: Make a bar plot to see gestation_weeks with avg_wt linearly scaled and num_babies logarithmically scaled.
End of explanation
"""


