import datetime
from typing import Collection

import pandas as pd

import participants.query_db
from features import esm
from machine_learning import QUESTIONNAIRE_IDS, QUESTIONNAIRE_IDS_RENAME


class Labels:
    def __init__(
        self,
        grouping_variable: str,
        labels: dict,
        participants_usernames: Collection = None,
    ):
        self.grouping_variable_name = grouping_variable
        self.grouping_variable = [grouping_variable]

        self.questionnaires = labels.keys()

        if participants_usernames is None:
            participants_usernames = participants.query_db.get_usernames(
                collection_start=datetime.date.fromisoformat("2020-08-01")
            )
        self.participants_usernames = participants_usernames

        self.df_esm = pd.DataFrame()
        self.df_esm_preprocessed = pd.DataFrame()
        self.df_esm_interest = pd.DataFrame()
        self.df_esm_clean = pd.DataFrame()

        self.df_esm_means = pd.DataFrame()
        print("Labels initialized.")

    def set_labels(self):
        print("Querying database ...")
        self.df_esm = esm.get_esm_data(self.participants_usernames)
        print("Got ESM data from the DB.")
        self.df_esm_preprocessed = esm.preprocess_esm(self.df_esm)
        print("ESM data preprocessed.")
        if "PANAS" in self.questionnaires:
            self.df_esm_interest = self.df_esm_preprocessed[
                (
                    self.df_esm_preprocessed["questionnaire_id"]
                    == QUESTIONNAIRE_IDS.get("PANAS").get("PA")
                )
                | (
                    self.df_esm_preprocessed["questionnaire_id"]
                    == QUESTIONNAIRE_IDS.get("PANAS").get("NA")
                )
            ]
        self.df_esm_clean = esm.clean_up_esm(self.df_esm_interest)
        print("ESM data cleaned.")

    def get_labels(self, questionnaire):
        if questionnaire == "PANAS":
            return self.df_esm_clean
        else:
            raise KeyError("This questionnaire has not been implemented as a label.")

    def aggregate_labels(self):
        print("Aggregating labels ...")
        self.df_esm_means = (
            self.df_esm_clean.groupby(
                ["participant_id", "questionnaire_id"] + self.grouping_variable
            )
            .esm_user_answer_numeric.agg("mean")
            .reset_index()
            .rename(columns={"esm_user_answer_numeric": "esm_numeric_mean"})
        )
        self.df_esm_means = (
            self.df_esm_means.pivot(
                index=["participant_id"] + self.grouping_variable,
                columns="questionnaire_id",
                values="esm_numeric_mean",
            )
            .reset_index(col_level=1)
            .rename(columns=QUESTIONNAIRE_IDS_RENAME)
            .set_index(["participant_id"] + self.grouping_variable)
        )
        print("Labels aggregated.")

    def get_aggregated_labels(self):
        return self.df_esm_means