import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.model_selection import train_test_split, GroupShuffleSplit
from scipy.sparse import vstack, hstack


def format_raw_df(df):
    """
    Oczyszczenie danych i skojarzenie pytań z odpowiedziami.
    :param df: surowe dane DataFrame
    :return: przetworzone dane DataFrame
    """
    # Określenie typów i ustawienie indeksu.
    df["PostTypeId"] = df["PostTypeId"].astype(int)
    df["Id"] = df["Id"].astype(int)
    df["AnswerCount"] = df["AnswerCount"].fillna(-1)
    df["AnswerCount"] = df["AnswerCount"].astype(int)
    df["OwnerUserId"].fillna(-1, inplace=True)
    df["OwnerUserId"] = df["OwnerUserId"].astype(int)
    df.set_index("Id", inplace=True, drop=False)

    df["is_question"] = df["PostTypeId"] == 1

    # Odrzucenie identyfikatorów PostTypeIds innych niż udokumentowane.
    df = df[df["PostTypeId"].isin([1, 2])]

    # Połączenie pytań z odpowiedziami.
    df = df.join(
        df[["Id", "Title", "body_text", "Score", "AcceptedAnswerId"]],
        on="ParentId",
        how="left",
        rsuffix="_question",
    )
    return df


def train_vectorizer(df):
    """
    Przetrenowanie wektoryzatora na pewnych danych.
    Funkcja zwraca wektoryzator przeznaczony do przekształcania danych innych niż 
    treningowe, uzupełniających wektory treningowe.
    :param df: dane treningowe dla wektoryzatora
    :return: przetrenowany wektoryzator i wektory treningowe
    """
    vectorizer = TfidfVectorizer(
        strip_accents="ascii", min_df=5, max_df=0.5, max_features=10000
    )

    vectorizer.fit(df["full_text"].copy())
    return vectorizer


def get_vectorized_series(text_series, vectorizer):
    """
    Funkcja wektoryzująca dane wejściowe za pomocą wstępnie przetrenowanego wektoryzatora.
    :param text_series: serie tekstów
    :param vectorizer: wstępnie przetrenowany wektoryzator sklearn
    :return: tablica zwektoryzowanych cech
    """
    vectors = vectorizer.transform(text_series)
    vectorized_series = [vectors[i] for i in range(vectors.shape[0])]
    return vectorized_series


def add_text_features_to_df(df):
    """
    Dodanie cech do obiektu DataFrame.
    :param df: obiekt DataFrame
    :param pretrained_vectors: informacja, czy wykorzystać w osadzeniach wstępnie przetrenowane wektory
    :return: obiekty DataFrame z dodatkowymi cechami
    """
    df["full_text"] = df["Title"].str.cat(df["body_text"], sep=" ", na_rep="")
    df = add_v1_features(df.copy())

    return df


def add_v1_features(df):
    """
    Dodanie pierwszych cech do wejściowego obiektu DataFrame.
    :param df: obiekt DataFrame z pytaniami
    :return: obiekt DataFrame z kolumnami zawierającymi dodane cechy
    """
    df["action_verb_full"] = (
        df["full_text"].str.contains("can", regex=False)
        | df["full_text"].str.contains("What", regex=False)
        | df["full_text"].str.contains("should", regex=False)
    )
    df["language_question"] = (
        df["full_text"].str.contains("punctuate", regex=False)
        | df["full_text"].str.contains("capitalize", regex=False)
        | df["full_text"].str.contains("abbreviate", regex=False)
    )
    df["question_mark_full"] = df["full_text"].str.contains("?", regex=False)
    df["text_len"] = df["full_text"].str.len()
    return df


def get_vectorized_inputs_and_label(df):
    """
    Złączenie obiektu DataFrame z wektorami tekstu.
    :param df: obiekt DataFrame z wyliczonymi cechami
    :return: połączony wektor zawierający cechy i tekst
    """
    vectorized_features = np.append(
        np.vstack(df["vectors"]),
        df[
            [
                "action_verb_full",
                "question_mark_full",
                "norm_text_len",
                "language_question",
            ]
        ],
        1,
    )
    label = df["Score"] > df["Score"].median()

    return vectorized_features, label


def get_feature_vector_and_label(df, feature_names):
    """
    Generowanie wektorów wejściowych i wyjściowych
    z wykorzystaniem wektora nazw cech.
    :param df: wejściowy obiekt DataFrame
    :param feature_names: nazwy kolumn z cechami (innych niż wektory)
    :return: tablice cech i etykiet
    """
    vec_features = vstack(df["vectors"])
    num_features = df[feature_names].astype(float)
    features = hstack([vec_features, num_features])
    labels = df["Score"] > df["Score"].median()
    return features, labels


def get_normalized_series(df, col):
    """
    Utworzenie znormalizowanej wersji kolumny.
    :param df: obiekt DataFrame
    :param col: nazwa kolumny
    :return: serie znormalizowane za pomocą Z-score
    """
    return (df[col] - df[col].mean()) / df[col].std()


def get_random_train_test_split(posts, test_size=0.3, random_state=40):
    """
    Podział danych treningowych i testowych zawartych w obiekcie DataFrame.
    Przyjęte założenie, że każde pytanie jest zapisane w osobnym wierszu obiektu.
    :param posts: wszystkie wpisy z etykietami
    :param test_size: proporcja zbioru testowego
    :param random_state: losowe ziarno
    """
    return train_test_split(
        posts, test_size=test_size, random_state=random_state
    )


def get_split_by_author(
    posts, author_id_column="OwnerUserId", test_size=0.3, random_state=40
):
    """
    Podział danych treningowych i testowych gwarantujący,
    że każdy autor zostanie umieszczony tylko w jednym podzbiorze.
    :param posts: wszystkie wpisy z etykietami
    :param author_id_column: nazwa kolumny zawierającej identyfikatory autorów
    :param test_size: proporcja zbioru testowego
    :param random_state: losowe ziarno
    """
    splitter = GroupShuffleSplit(
        n_splits=1, test_size=test_size, random_state=random_state
    )
    splits = splitter.split(posts, groups=posts[author_id_column])
    train_idx, test_idx = next(splits)
    return posts.iloc[train_idx, :], posts.iloc[test_idx, :]
