from sklearn.feature_extraction.text import CountVectorizer
from nltk.corpus import stopwords
import re
from math import log

main_dir = '/Users/jonathanestrella/Documents/MineriaDeDatos/data/'
file = main_dir+'archivo.txt'
d = main_dir+'archivo_prueba.txt'
C = ['H', 'M']

def ExtractVocabulary(D):
    stop_words = stopwords.words('spanish')
    file=D
    corpus = []
    with open(file, 'r', encoding='utf-8') as content_file:
        for line in content_file:
            line = line.rstrip().lower()
            words = re.findall('[a-záéíóúñ]+', line)
            text = ' '.join(word for word in words if word not in stop_words and len(word)>3)
            corpus.append(text)      
        vec = CountVectorizer()
        tdm = vec.fit_transform(corpus)
        vocabulary = vec.vocabulary_ 
        lista= list(vocabulary.keys())
    return(lista)

def CountDocsInClass(D, c):
    counter = 0
    with open(D, 'r', encoding='utf-8') as content_file:
        for line in content_file:
            if (line[0] == c):
                counter += 1
    return counter

def ConcaTextOfAllDocsInClass(D, c):
    text = ""
    with open(D, 'r', encoding = 'utf-8') as i_r:
        for line in i_r:
            if line[0] == c:
                text += line[1:]
        
    text = text.lower()
    text = text.replace('\n', ' ')
    return text

def CountTokensOfTerm(textc, t):
    val = textc.count(t)
    return val

def TrainMultinomialNB(C, D):
    V = ExtractVocabulary(D) 
    N = sum(1 for line in open (file, encoding='utf-8')) # Count docs
    
    prior = {'H': 0, 'M': 0}
    condprob = []
    
    for v in range(len(V)):
        condprob.append([])
    
    for c in C:
        Tct = []
        Nc = CountDocsInClass(D, c)
        prior[c] = Nc / N
        textc = ConcaTextOfAllDocsInClass(D, c)
        for i in range(len(V)):
            Tct.append(CountTokensOfTerm(textc, V[i]))
        
        for i in range(len(V)):    
            condprob[i].append((Tct[i] + 1) / (sum(Tct) + 1))
    
    return V, prior, condprob


def applyB(C,V,prior,contprob,d):
    W=[]
    score=[]
    
    with open(d,'r',encoding='utf-8') as content_file:
        text=content_file.read()
    text=text.lower()
    
    for p in V:
        cont=text.count(p)
        if cont>0:
            W.append(p)
    
    for c in range(len(C)):
        
        if prior[c]==0:
            score.append(0)
        else:
            score.append(log(prior[c]))
        
        for t in W:
            score[c]+=log(contprob[V.index(t)][c]) 
    
    # print(score)
    return(C[score.index(max(score))])


# PRUEBA
Voc, prior, condprob = TrainMultinomialNB(C, file)

prior_lista = [prior['H'], prior['M']]
result=applyB(C, Voc, prior_lista, condprob ,d)
print(result)


# codigo aún incompleto el pdf no se ve muy bien