from collections import defaultdict
import math
import pickle
import re

open_file=('test-A/out.tsv')

#---------------TRAIN START

#Prawdopodobienstwo wylosowania dokumentu
def calc_class_logprob(expected_path):
    paranormal_classcount=0
    skeptic_classcount=0
    with open(expected_path,encoding='utf-8') as f:
        for line in f:
            if '1' in line:
                paranormal_classcount += 1
            if '0' in line:
                skeptic_classcount += 1

    paranormal_prob = paranormal_classcount / (paranormal_classcount + skeptic_classcount)
    skeptic_prob = skeptic_classcount / (paranormal_classcount + skeptic_classcount)

    return math.log(paranormal_prob), math.log(skeptic_prob)


def calc_word_count(in_path, expected_path):
    word_counts = {'paranormal':defaultdict(int), 'skeptic': defaultdict(int)}
    with open(in_path,encoding='utf-8') as in_file, open(expected_path,encoding='utf-8')  as  expected_file:
        for line, exp in zip(in_file, expected_file):
            class_ = exp.rstrip('\n').replace(' ','')
            text, timestamp = line.rstrip('\n').split('\t')
            text = text.lower()
            text = re.sub(r'\\n+', " ", text)
            text = re.sub(r'http\S+', " ", text)
            text = re.sub(r'\/[a-z]\/', " ", text)
            text = re.sub(r'[^a-z]', " ", text)
            text = re.sub(r'\s{2,}', " ", text)
            text = re.sub(r'\W\w{1,3}\W|\A\w{1,3}\W', " ", text)
            text = re.sub(r'\W\w{1,3}\W|\A\w{1,3}\W', " ", text)
            text = re.sub(r'\W\w{1,3}\W|\A\w{1,3}\W', " ", text)
            text = re.sub(r'^\s', "", text)
            tokens = text.split(' ')
            for token in tokens:
                if class_ == '1':
                    word_counts['paranormal'][token] += 1
                elif class_ == '0':
                    word_counts['skeptic'][token] += 1
    return word_counts

def calc_word_logprobs(word_counts):
    total_skeptic = sum(word_counts['skeptic'].values()) + len(word_counts['skeptic'].keys())
    total_paranormal = sum(word_counts['paranormal'].values()) + len(word_counts['paranormal'].keys())

    word_logprobs= {'paranormal': {}, 'skeptic': {}}
    for class_ in word_counts.keys(): # sceptic paranormal
        for token, tokens in word_counts[class_].items():
            if class_ == 'skeptic':
                word_prob = (tokens+1)/total_skeptic
            else:
                word_prob = (tokens+1)/total_paranormal
            word_logprobs[class_][token] = math.log(word_prob)
    return word_logprobs

#--------------- TRAIN END

def main():
    paranomal_class_logprob, skeptic_class_logprob = calc_class_logprob("train/expected.tsv")
    word_counts=calc_word_count("train/in.tsv","train/expected.tsv")
    word_counts['paranormal'][''] = 0
    word_counts['skeptic'][''] = 0
    word_logprobs = calc_word_logprobs(word_counts)
    pickle.dump([paranomal_class_logprob, skeptic_class_logprob, word_logprobs], open('naive_base_model.pkl','wb'))

main()