"""
Airflow DAG to load raw data, process it, split it, and store in database.

Author
------
Nicolas Rojas
"""

# imports
from datetime import datetime
import pandas as pd
from sklearn.model_selection import train_test_split
from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.providers.mysql.hooks.mysql import MySqlHook


def check_table_exists(table_name: str):
    """Check whether table exists in clean_data database. If not, create it.

    Parameters
    ----------
    table_name : str
        Name of table to check.
    """
    # count number of rows in data table
    query = f'SELECT COUNT(*) FROM information_schema.tables WHERE table_name="{table_name}"'
    mysql_hook = MySqlHook(mysql_conn_id="clean_data", schema="clean_data")
    connection = mysql_hook.get_conn()
    cursor = connection.cursor()
    cursor.execute(query)
    results = cursor.fetchall()
    # check whether table exists
    if results[0][0] == 0:
        # create table
        print("----- table does not exists, creating it")
        create_sql = f"CREATE TABLE `{table_name}`\
            `age` SMALLINT,\
            `anual_income` BIGINT,\
            `credit_score` SMALLINT,\
            `loan_amount` BIGINT,\
            `loan_duration_years` TINYINT,\
            `number_of_open_accounts` SMALLINT,\
            `had_past_default` TINYINT,\
            `loan_approval` TINYINT\
            )"
        mysql_hook.run(create_sql)
    else:
        # no need to create table
        print("----- table already exists")

    return "Table checked"


def store_data(dataframe: pd.DataFrame, table_name: str):
    """Store dataframe data in given table, in clean data database.

    Parameters
    ----------
    dataframe : pd.DataFrame
        Dataframe to store in database.
    table_name : str
        Name of the table to store the data.
    """
    check_table_exists(table_name)
    # insert every dataframe row into sql table
    mysql_hook = MySqlHook(mysql_conn_id="clean_data", schema="clean_data")
    sql_column_names = ", ".join(
        ["`" + name + "`" for name in dataframe.columns]
    )
    conn = mysql_hook.get_conn()
    cur = conn.cursor()
    # VALUES in query are %s repeated as many columns are in dataframe
    query = f"INSERT INTO `{table_name}` ({sql_column_names}) \
        VALUES ({', '.join(['%s' for _ in range(dataframe.shape[1])])})"
    dataframe = list(dataframe.itertuples(index=False, name=None))
    cur.executemany(query, dataframe)
    conn.commit()

    return "Data stored"


def preprocess_data():
    """Preprocess raw data and store it in clean_data database."""
    # retrieve raw data
    mysql_hook = MySqlHook(mysql_conn_id="raw_data", schema="raw_data")
    conn = mysql_hook.get_conn()
    query = "SELECT * FROM `raw_clients`"
    dataframe = pd.read_sql(query, con=conn)

    # drop useless column
    dataframe.drop(columns=["id"], inplace=True)
    # fill empty fields
    dataframe.fillna(0, inplace=True)

    # split data: 70% train, 10% val, 20% test
    df_train, df_test = train_test_split(
        dataframe, test_size=0.2, shuffle=True, random_state=1337
    )
    df_train, df_val = train_test_split(
        df_train, test_size=0.125, shuffle=True, random_state=1337
    )

    # store data partitions in database
    store_data(df_train, "clean_clients_train")
    store_data(df_val, "clean_clients_val")
    store_data(df_test, "clean_clients_test")

    return "Data preprocessed"


with DAG(
    "preprocess_data",
    description="Fetch raw data, preprocess it and save it in mysql database",
    start_date=datetime(2024, 9, 18, 0, 2),
    schedule_interval="@once",
) as dag:

    preprocess_task = PythonOperator(
        task_id="preprocess_data", python_callable=preprocess_data
    )
    preprocess_task