from typing import List

import pandas as pd

from config.models import SMS, Call, Participant
from setup import db_engine, session

call_types = {1: "incoming", 2: "outgoing", 3: "missed"}
sms_types = {1: "received", 2: "sent"}


def get_call_data(usernames: List) -> pd.DataFrame:
    """
    Read the data from the calls table and return it in a dataframe.

    Parameters
    ----------
    usernames: List
        A list of usernames to put into the WHERE condition.

    Returns
    -------
    df_calls: pd.DataFrame
        A dataframe of call data.
    """
    query_calls = (
        session.query(Call, Participant.username)
        .filter(Participant.id == Call.participant_id)
        .filter(Participant.username.in_(usernames))
    )
    with db_engine.connect() as connection:
        df_calls = pd.read_sql(query_calls.statement, connection)
    return df_calls


def get_sms_data(usernames: List) -> pd.DataFrame:
    """
    Read the data from the sms table and return it in a dataframe.

    Parameters
    ----------
    usernames: List
        A list of usernames to put into the WHERE condition.

    Returns
    -------
    df_sms: pd.DataFrame
        A dataframe of call data.
    """
    query_sms = (
        session.query(SMS, Participant.username)
        .filter(Participant.id == SMS.participant_id)
        .filter(Participant.username.in_(usernames))
    )
    with db_engine.connect() as connection:
        df_sms = pd.read_sql(query_sms.statement, connection)
    return df_sms


def enumerate_contacts(comm_df: pd.DataFrame) -> pd.DataFrame:
    """
    Count contacts (callers, senders) and enumerate them by their frequency.

    Parameters
    ----------
    comm_df: pd.DataFrame
        A dataframe of calls or SMSes.

    Returns
    -------
    comm_df: pd.DataFrame
        The altered dataframe with the column contact_id, arranged by frequency.
    """
    contact_counts = (
        comm_df["trace"]
        .value_counts(sort=True, ascending=False)
        .to_frame(name="frequency")
    )
    # A frequency table of different traces (contacts).
    contact_counts["contact_id"] = list(range(len(contact_counts.index)))
    contact_code = contact_counts["contact_id"].to_dict()
    # Create a dictionary translating traces into integers, enumerated by their frequency.
    comm_df["contact_id"] = comm_df["trace"].astype("category")
    # Transform to categorical data instead of a simple character column.
    comm_df["contact_id"] = comm_df["contact_id"].cat.rename_categories(contact_code)
    # Recode the contacts into integers from 0 to n_contacts, so that the first one is contacted the most often.
    return comm_df


def count_comms(comm_df: pd.DataFrame) -> pd.DataFrame:
    """
    Calculate frequencies (and duration) of messages (or calls), grouped by their types.

    Parameters
    ----------
    comm_df: pd.DataFrame
        A dataframe of calls or SMSes.

    Returns
    -------
    comm_features: pd.DataFrame
        A list of communication features for every participant.
        These are:
        * the number of messages by type (received, sent),
        * the number of calls by type (incoming, outgoing missed), and
        * the duration of calls by type.
    """
    if "call_type" in comm_df:
        comm_counts = (
            comm_df.value_counts(subset=["participant_id", "call_type"])
            .unstack()
            .rename(columns=call_types)
            .add_prefix("no_")
        )
        comm_duration = (
            comm_df.groupby(["participant_id", "call_type"])
            .sum()["call_duration"]
            .unstack()
            .rename(columns=call_types)
            .add_prefix("duration_")
        )
        comm_features = comm_counts.join(comm_duration)
        try:
            comm_features.drop(columns="duration_" + call_types[3], inplace=True)
        # The missed calls are always of 0 duration.
        except KeyError:
            pass
        # If there were no missed calls, this exception is raised.
        # But we are dropping the column anyway, so no need to deal with the exception.
    elif "message_type" in comm_df:
        comm_counts = (
            comm_df.value_counts(subset=["participant_id", "message_type"])
            .unstack()
            .rename(columns=sms_types)
            .add_prefix("no_")
        )
        comm_features = comm_counts
    else:
        raise KeyError("The dataframe contains neither call_type or message_type")
    return comm_features