import os
import pickle

import pandas as pd
import sklearn.svm as svm
from sklearn.calibration import CalibratedClassifierCV

from . import utility

LARGE_PUBLISHERS_LIMIT = 100  # minimum nb of records for a publisher to be modelled.
MODELS_REPO = os.path.join(utility.get_current_file_directory(), "repository")

class ModelRepository:
    """
    Container for the models available in the application.
    """

    def __init__(self):
        model_files = os.listdir(MODELS_REPO)
        model_files = (os.path.join(MODELS_REPO, model_file) for model_file in model_files)
        models = [LocationFilter.load(model_file) for model_file in model_files]
        self.models = {model.name: model for model in models}

    def __getitem__(self, item):
        return self.models[item]

    def list_models(self):
        return list(self.models.keys())


def form_feature_set(articles, locations, uri_index):
    """
    Constructs the dataframe of features for the articles provided.
    :param articles:
    :param locations:
    :param uri_index:
    :return:
    """
    pub_feature_set = []
    publishers = []
    for article in articles:
        for annotation in article["annotation_uri"]:
            if annotation in uri_index.keys():
                locs = uri_index[annotation]["locations"]
                for loc in locs:
                    publishers.append(utility.clean_publisher(article["publisher"]))
                    pub_feature_set.append(locations.loc[loc["id"], :])

    pub_feature_set = pd.DataFrame(pub_feature_set)
    pub_feature_set["publisher"] = publishers
    return pub_feature_set


def form_sample(dataset, locations, k=None):
    """
    generate a sample for training a model for a given publisher.
    We assume that all the records in dataset are labelled as positive.
    This function will generate a random sample of negative values of size k.
    k will default to the size of dataset if not provided.

    WARNING: the returned sample is not shuffled. All positive records will sit at the top of the frame.

    :param dataset: the frame of features for a publishers.
    :param locations: the frame of all locations features.
    :return: sample, a frame containing the records from dataset, plus new negative records sampled from location.
    """

    sample_size = k or dataset.shape[0]

    known_ids = set(dataset.index.values)
    dataset.drop("publisher", axis=1, inplace=True)
    dataset.reset_index(inplace=True, drop=True)

    available_location = locations.loc[~locations.index.isin(known_ids)]
    negative_sample = available_location.sample(sample_size).copy()
    negative_sample = negative_sample.reset_index(drop=True)
    negative_sample["Y"] = 0

    positive_sample = dataset.copy()
    positive_sample["Y"] = 1

    sample = pd.concat((positive_sample, negative_sample))
    return sample


class LocationFilter:
    """
    Instances of this class can predict the relevancy of a location for the articles generated by a given publisher.
    """
    def __init__(self, name, encoded_columns):
        self.name = name
        self.model = None
        self.encoded_columns = encoded_columns

    def fit(self, X, y):
        clf = svm.SVC()
        clf_sigmoid = CalibratedClassifierCV(clf, method='sigmoid')
        self.model = clf_sigmoid.fit(X, y)

    def process_location(self, location):
        """
        converts a location dictionary in a feature set consumable by the classifier.
        :param location:
        :return:
        """
        X = pd.DataFrame(columns=["lat", "lng", *self.encoded_columns], index=[0])
        X[:] = 0.0
        X["lat"] = location["lat"]
        X["lng"] = location["lng"]
        X[location["nuts_region"]] = 1
        return X

    def predict(self, X):

        return self.model.predict(X)

    def predict_proba(self, X):
        return self.model.predict_proba(X)[:, 1]

    def save(self, to_):
        path = os.path.join(to_, f"{self.name}_model.pickle")
        with open(path, 'wb') as fh:
            pickle.dump(self, fh)

    @classmethod
    def load(cls, from_):
        with open(from_, 'rb') as fh:
            obj = pickle.load(fh)
        return obj


def fit_all_filters():
    """
    Fits location filters for all the publishers and save the models to the filesystem.
    :return:
    """

    raw_articles, raw_annotations, raw_locations = utility.load_dataset()
    uri_index = utility.get_annotation_uri_index(raw_annotations, raw_locations)
    location_index = utility.get_location_index(raw_locations)
    article_to_loc_index = utility.get_article_to_loc_index(raw_articles, raw_locations, uri_index)

    locations = pd.DataFrame(raw_locations)[["id", "lat", "lng", "nuts_region"]]
    locations.set_index("id", inplace=True)
    locations, encoded_columns = utility.get_dummies(locations, "nuts_region", True)
    feature_set = form_feature_set(raw_articles, locations, uri_index)

    publishers = feature_set["publisher"].unique()
    publishers.sort()
    large_publishers = publishers[feature_set.groupby("publisher")["publisher"].count() > LARGE_PUBLISHERS_LIMIT]

    print(f"{len(large_publishers)} publishers with more than {LARGE_PUBLISHERS_LIMIT} records.")
    for publisher in large_publishers:
        print(f"fitting model for {publisher}")
        dataset = feature_set[feature_set["publisher"] == publisher]
        sample = form_sample(dataset, locations)
        location_filter = LocationFilter(publisher, encoded_columns)
        location_filter.fit(X=sample.drop("Y", axis=1), y=sample["Y"])
        location_filter.save(MODELS_REPO)


if __name__ == '__main__':
    fit_all_filters()
