"""
Author: Alexandra DeLucia
"""
from sklearn.base import BaseEstimator
from sklearn.preprocessing import LabelEncoder
from sklearn.pipeline import Pipeline
from sklearn.model_selection import cross_validate, StratifiedShuffleSplit
from sklearn.metrics import precision_recall_fscore_support
import numpy as np
import pandas as pd
import os


class RandomGuessingClassifier(BaseEstimator):
    def __init__(self):
        super().__init__()
        self.class_names = None
        self.class_distribution = None

    def fit(self, X_train, y_train):
        """Find the distribution of targets. Assumes numerical labels"""
        self.class_names = np.unique(y_train)
        # Bins are [a, b) so we need an extra bin for the label counts to work
        bins = list(self.class_names) + [self.class_names.shape[0]]
        self.class_distribution, temp = np.histogram(y_train, bins=bins, density=True)

    def predict(self, X_test):
        rng = np.random.default_rng(42)
        predictions = [
            rng.choice(self.class_names, p=self.class_distribution) for i in range(len(X_test))
        ]
        return np.array(predictions)

    def predict_proba(self, X_test):
        probabilities = [
            self.class_distribution for i in range(len(X_test))
        ]
        return np.array(probabilities)

    @staticmethod
    def get_feature_importance(self):
        return [0]


if __name__ == "__main__":
    # Input data files are available in the read-only "../input/" directory
    # You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using "Save & Run All" 
    # You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session
    train_path, test_path = None, None
    for dirname, _, filenames in os.walk('/kaggle/input'):
        for filename in filenames:
            if "train" in filename:
                train_path = os.path.join(dirname, filename)
            elif "test" in filename:
                test_path = os.path.join(dirname, filename)
    
    # Read in train data
    train_df = pd.read_csv(train_path)

    # Transform the string labels into numerical
    label_encoder = LabelEncoder()
    train_df["target"] = label_encoder.fit_transform(train_df.label.values)

    # Fit our model
    clf = RandomGuessingClassifier()
    clf.fit(train_df.body, train_df.target)

    # Predict
    predictions = clf.predict(train_df.body)
    prec, rec, fscore, support = precision_recall_fscore_support(train_df.target, predictions, average="weighted")

    # Read in test data and predict
    test_df = pd.read_csv(test_path)
    test_df["predictions"] = clf.predict(test_df.body)
    test_df["prediction_labels"] = label_encoder.inverse_transform(test_df["predictions"])
    test_df[["id", "prediction_labels"]].to_csv("random_guessing.csv", index=False)
    
    print(f"Train scores:\nPrecision: {prec:.2}\tRecall: {rec:.2}\tFscore: {fscore:.2}")