"""
The functions in this module allow the user to get the datasets stored in the
epigraphhub database.

The function get_agg_data aggregate the data according to the values of one column,
and the method of aggregation applied.

The function get_georegion_data filter the datasets for a list of
selected regions (in the Switzerland case: cantons).

The function get_cluster_data is focused on being used to apply the forecast models.
This function returns a table where each column is related to a  different table and
region (e.g. the daily number of cases and the daily number of hospitalizations in
Geneva and Fribourg).
Some parts of the code of this function are focused on the swiss case.
So the function isn't fully general.
"""
import pandas as pd
from sqlalchemy import create_engine

from epigraphhub.data.epigraphhub_db import get_data_by_location
from epigraphhub.settings import env

with env.db.credentials[env.db.default_credential] as credential:
    engine_public = create_engine(
        f"postgresql://{credential.username}:"
        f"{credential.password}@{credential.host}:{credential.port}/"
        f"{credential.dbname}"
    )

dict_cols = {
    "foph_cases_d": ["datum", "georegion", "entries"],
    "foph_test_d": ["datum", "georegion", "entries", "entries_pos"],
    "foph_hosp_d": ["datum", "georegion", "entries"],
    "foph_hospcapacity_d": [
        "date",
        "georegion",
        "icu_covid19patients",
        "total_covid19patients",
    ],
    "foph_re_d": ["date", "georegion", "median_r_mean"],
}

georegion_columns = {
    "foph_cases_d": "georegion",
    "foph_test_d": "georegion",
    "foph_hosp_d": "georegion",
    "foph_hospcapacity_d": "georegion",
}

date_columns = {
    "foph_cases_d": "datum",
    "foph_test_d": "datum",
    "foph_hosp_d": "datum",
    "foph_hospcapacity_d": "date",
}

count_columns = {
    "foph_cases_d": ["entries"],
    "foph_test_d": ["entries"],
    "foph_hosp_d": ["entries"],
    "foph_hospcapacity_d": ["icu_covid19patients", "total_covid19patients"],
}

columns_name = {"foph_cases_d": "cases", "foph_test_d": "test", "foph_hosp_d": "hosp"}


def get_cluster_data(
    schema,
    table_name,
    georegion,
    dict_cols=dict_cols,
    date_columns=date_columns,
    georegion_columns=georegion_columns,
    count_columns=count_columns,
    columns_name=columns_name,
    vaccine=True,
    smooth=True,
):

    """
    This function provides a data frame where each column is associated with a
    table and region selected.

    Parameters
    ----------
    schema : str
        The schema where the data that you want to get is saved.
    table_name : list
        List of strings. In this list should be all the tables that you want get
        the data.
    georegion : list
        List of strings. This list contains all the regions of the country of
        interest or the string 'All' to return all the regions.
    dict_cols : dict
        In the keys are the table_names and in the values the columns that you
        want to use from each table.
    date_columns : dict
        In the keys are the table_names and in the values the name of the date
        column of the table to be used as the index.
    count_columns : dict
        In the keys are the table_names and in the values the name of the column
        which values will be used.

    columns_name : dict
        In the keys ate the table_names and in the values the name that will
        appear in the column associated with each table in the final data frame
        that will be returned.

    vaccine : bool
        If True the data of total vaccinations per hundred for the country in
        the schema will be added in the final data frame. This data is from our
        world in data.

    smooth : bool
        If True in the end data frame will be applied a moving average of seven
        days.

    Returns
    -------
    pd.DataFrame
    """

    df_end = pd.DataFrame()

    for table in table_name:

        df = get_data_by_location(
            schema,
            table,
            georegion,
            dict_cols[table],
            loc_column=georegion_columns[table],
        )
        df.set_index(date_columns[table], inplace=True)
        df.index = pd.to_datetime(df.index)

        for region in df.georegion.unique():

            for count in count_columns[table]:

                if table == "foph_hospcapacity_d":

                    names = {
                        "icu_covid19patients": "icu_patients",
                        "total_covid19patients": "total_hosp",
                    }
                    df_aux1 = df.loc[df.georegion == region].resample("D").mean()

                    df_aux2 = pd.DataFrame()

                    df_aux2[names[count] + "_" + region] = df_aux1[count]
                    df_aux2[f"diff_{names[count]}_{region}"] = df_aux1[count].diff(1)
                    df_aux2[f"diff_2_{names[count]}_{region}"] = df_aux1[count].diff(2)

                    df_end = pd.concat([df_end, df_aux2], axis=1)

                else:
                    df_aux1 = df.loc[df.georegion == region].resample("D").mean()

                    df_aux2 = pd.DataFrame()

                    df_aux2[columns_name[table] + "_" + region] = df_aux1[count]
                    df_aux2[f"diff_{columns_name[table]}_{region}"] = df_aux1[
                        count
                    ].diff(1)
                    df_aux2[f"diff_2_{columns_name[table]}_{region}"] = df_aux1[
                        count
                    ].diff(2)

                    df_end = pd.concat([df_end, df_aux2], axis=1)

    df_end = df_end.resample("D").mean()

    if vaccine == True:
        vac = pd.read_sql_table(
            "owid_covid",
            engine_public,
            schema="public",
            columns=["date", "iso_code", "total_vaccinations_per_hundred"],
        )

        dict_iso_code = {"switzerland": "CHE", "colombia": "COL"}

        vac = vac.loc[vac.iso_code == dict_iso_code[schema]]
        vac.index = pd.to_datetime(vac.date)

        # selecting only the column with vaccinations per hundred
        vac = vac[["total_vaccinations_per_hundred"]]

        vac = vac.fillna(method="ffill")

        df_end["vac_all"] = vac.total_vaccinations_per_hundred

        df_end["vac_all"] = df_end["vac_all"].fillna(method="ffill")

    df_end = df_end.fillna(0)

    if smooth == True:
        df_end = df_end.rolling(window=7).mean()

        df_end = df_end.dropna()

    return df_end
