from collections import defaultdict
import operator

class DataMajorityBaseline(object):

    unknown_id = -1

    def __init__(self, lexicon):
        self.lexicon = lexicon
        self.counts = defaultdict(lambda: defaultdict(int))
    
    def fit(self, lemma_poses, frame_ids):
        for lemma_pos, frame_id in zip(lemma_poses, frame_ids):
            self.counts[lemma_pos][frame_id] += 1

    def predict(self, lemma_poses):
        """ Predicts the frame id"""
        predictions = []

        for lemma_pos in lemma_poses:
            prediction = self.predict_single(lemma_pos)
            predictions.append(prediction)

        return predictions

    def predict_single(self, lemma_pos):
        possible_frames = self.counts[lemma_pos]

        if possible_frames:
            # Sorted is used to always return the alphabetically smaller frame name on tie
            # instead of relying on internal implementation details of Python dicts
            prediction = max(sorted(possible_frames.items()), key=operator.itemgetter(1))[0]
            return prediction
        else:
            # Use a different unknown ID everytime
            self.unknown_id -= 1
            return self.unknown_id

class LexiconBaseline(object):

    unknown_id = -1

    def __init__(self, lexicon):
        self.lexicon = lexicon
        self.counts = defaultdict(int)

    def fit(self, lemma_poses, frame_ids):
        for frame_id in frame_ids:
            self.counts[frame_id] += 1                

    def predict(self, lemma_poses):
        """ Predicts the frame id"""
        predictions = []

        for lemma_pos in lemma_poses:
            available_frames = self.lexicon.get_available_frame_ids(lemma_pos)

            if len(available_frames):
                prediction = max(sorted(available_frames), key=lambda frame_id: self.counts[frame_id])
            else:
                prediction = self.unknown_id
                self.unknown_id -= 1

            predictions.append(prediction)

        return predictions            


class DataMajorityLexiconBaseline(object):

    def __init__(self, lexicon):
        self.dm_baseline = DataMajorityBaseline(lexicon)
        self.lexicon = lexicon

    def fit(self, lemma_poses, frame_names):
        self.dm_baseline.fit(lemma_poses, frame_names)

    def predict(self, lemma_poses):
        predictions = []

        for lemma_pos in lemma_poses:
            if self.lexicon.is_unambiguous(lemma_pos):
                lexicon_predictions = self.lexicon.get_available_frames(lemma_pos)
                assert len(lexicon_predictions) == 1
                prediction = self.lexicon.get_id(lexicon_predictions[0])
            else:
                prediction = self.dm_baseline.predict_single(lemma_pos)

            predictions.append(prediction)

        return predictions

def fit_baseline_from_labels(baseline, lexicon, lemma_pos_train, labels_train, lemma_pos_test, labels_test):
    lemma_poses_train = []
    frame_ids_train = []

    for lemma_pos, frame_name in zip(lemma_pos_train, labels_train):
        lemma_poses_train.append(lemma_pos)
        frame_ids_train.append(lexicon.get_id(frame_name))

    lemma_poses_test = []
    frame_ids_test = []        

    for lemma_pos, frame_name in zip(lemma_pos_test, labels_test):
        lemma_poses_test.append(lemma_pos)
        frame_ids_test.append(lexicon.get_id(frame_name))
    
    baseline.fit(lemma_poses_train, frame_ids_train)    

def fit_baseline_from_dataset(baseline, lexicon, dataset_train, dataset_test):
    fit_baseline_from_labels(baseline, lexicon, dataset_train.lemma_pos, dataset_train.labels, dataset_test.lemma_pos, dataset_test.labels)