import sys
import os
import numpy as np
import numpy.random as npr
import pandas as pd
import random
from nltk.stem import WordNetLemmatizer
import nltk
nltk.download('wordnet')

# Module file with functions that you fill in so that they can be
# called by the notebook.  This file should be in the same
# directory/folder as the notebook when you test with the notebook.

# You can modify anything here, add "helper" functions, more imports,
# and so on, as long as the notebook runs and produces a meaningful
# result (though not necessarily a good classifier).  Document your
# code with reasonable comments.

# Function for Part 1
def preprocess(inputfile):    
    col = ['Sentence Number','Word','POS','Named Entity']
    text = []
    lemma = WordNetLemmatizer().lemmatize
    
    for line in inputfile.readlines()[1:]:
        line = line.strip('\n')
        token = line.split('\t')
        
        #We remove the first column because it'll be indexed in the DataFrame
        token = token [1:]
        token[1] = token[1].lower()
        
        #We need to specify which POS tag we want to lemmatize (Adjective, Noun, Adverb and Verb)
        if token[1].isalpha():
            if token[2].startswith('J'):
                token[1] = lemma(token[1],pos='a')
            elif token[2].startswith('N'):
                token[1] = lemma(token[1],pos='n')
            elif token[2].startswith('R'):
                token[1] = lemma(token[1],pos='r')
            elif token[2].startswith('V'):
                token[1] = lemma(token[1],pos='v')
            text.append(token)
        elif token[2] == '.':
            text.append(token)
        elif token[1] == '!':
            text.append(token)
        elif token[1] == '?':
            text.append(token)
    
    processed_text = pd.DataFrame(text, columns = col)
    
    return processed_text

# Code for part 2
class Instance:
    def __init__(self, neclass, features):
        self.neclass = neclass
        self.features = features

    def __str__(self):
        return "Class: {} Features: {}".format(self.neclass, self.features)

    def __repr__(self):
        return str(self)

def create_instances(data):
    instances = []
    sentence_num = list(data['Sentence Number'])
    word = list(data['Word'])
    pos = list(data['POS'])
    named_entity = list(data['Named Entity'])
    sentences = list(zip(sentence_num, word, named_entity))
    start_end = ['<s>','</s>']
    
    for current, line in enumerate(sentences):
        if line[2].startswith('B'):
            neclass = line[2][2:]
            features = []
            
            #Previous words
            for i in range(5, 0, -1):
                #If part of the same sentence
                if sentences[current - i][0] == line[0]:
                    features.append(sentences[current-i][1])
                #Specify new sentence
                else:
                    features.append(start_end[0])
            
            #End of the neclass
            count = 1
            while sentences[current+count][2].startswith('I'):
                count = count + 1
            
            #Following words
            min = 1 + count
            max = 6 + count
            if (current + max) < len(sentences):
                for j in range(min, max):
                    if sentences[current+j][0] == line[0]:
                        features.append(sentences[current+j][1])
                    else:
                        features.append(start_end[1])
            
            instances.append(Instance(neclass, features))
            
    return instances

# Code for part 3
def create_table(instances):
    
    df_col = [instance.neclass for instance in instances]
    df_matrix=[]
 
    for instance in instances:
        df_dict = {}

        for word in instance.features:
            if word not in df_dict:
                df_dict[word] = 1
            else:
                df_dict[word] += 1

        df_matrix.append(df_dict)
    
    df = pd.DataFrame(df_matrix)
    df.insert(0, 'neClass', df_col)
    #We fill with 0 to avoid N/A
    df = df.fillna(0)
    
    return df

def ttsplit(bigdf):
    
    dfnames = bigdf.groupby("neClass", group_keys=False)
    df_train = bigdf.sample(frac=0.8).reset_index()
    df_test = bigdf.drop(df_train.index).reset_index()
        
    return df_train.drop('neClass', axis=1).to_numpy(), df_train['neClass'], df_test.drop('neClass', axis=1).to_numpy(), df_test['neClass']

# Code for part 5
def confusion_matrix(truth, predictions):
    from sklearn.metrics import confusion_matrix
    
    neClass = ['art', 'eve', 'geo', 'gpe', 'nat', 'org', 'per', 'tim']
    cm = confusion_matrix(truth, predictions, labels=neClass)
    df = pd.DataFrame(cm, index=neClass, columns=neClass)
    
    return df

# Code for bonus part B
def bonusb(filename):
    pass
