import vowpalwabbit
import pandas as pd
import re


x_train = pd.read_csv('train/in.tsv', header=None, sep='\t')
y_train = pd.read_csv('train/expected.tsv', header=None, sep='\t')

x_train = x_train.drop(1, axis=1)
x_train.columns = ['year', 'text']
y_train.columns = ['category']

data = pd.concat([x_train, y_train], axis=1)


model = vowpalwabbit.Workspace('--oaa 7 --ngram 3')

map_dict = {}

for i, x in enumerate(data['category'].unique()):
    map_dict[x] = i+1  
    
data['train_input'] = data.apply(lambda row: to_vw_format(row, map_dict), axis=1)


for example in data['train_input']:
    model.learn(example)



def to_vw_format(row, map_dict):
    text = row['text'].replace('\n', ' ').lower().strip()
    text = re.sub("[^a-zA-Z -']", '', text)
    year = row['year']
    try:
        category = map_dict[row['category']]
    except KeyError:
        category = ''

    vw_input = f"{category} | year:{year} text:{text}\n"

    return vw_input



### Read data          

data_dev = pd.read_csv('dev-0/in.tsv', header=None, sep='\t')
data_dev = data_dev.drop(1, axis=1)
data_dev.columns = ['year', 'text']
data_dev['train_input'] = data_dev.apply(lambda row: to_vw_format(row, map_dict), axis=1)


data_A = pd.read_csv('test-A/in.tsv', header=None, sep='\t')
data_A = data_A.drop(1, axis=1)
data_A.columns = ['year', 'text']
data_A['train_input'] = data_A.apply(lambda row: to_vw_format(row, map_dict), axis=1)

data_B = pd.read_csv('test-B/in.tsv', header=None, sep='\t')
data_B = data_B.drop(1, axis=1)
data_B.columns = ['year', 'text']
data_B['train_input'] = data_B.apply(lambda row: to_vw_format(row, map_dict), axis=1)



### Write predictions 

with open("dev-0/out.tsv", 'w', encoding='utf-8') as file:
    for test_example in data_dev['train_input']:
        prediction_dev = model.predict(test_example)
        text_prediction_dev = dict((value, key) for key, value in map_dict.items()).get(prediction_dev)
        file.write(str(text_prediction_dev) + '\n')


with open("test-A/out.tsv", 'w', encoding='utf-8') as file:
    for test_example in data_A['train_input']:
        prediction_A = model.predict(test_example)
        text_prediction_A = dict((value, key) for key, value in map_dict.items()).get(prediction_A)
        file.write(str(text_prediction_A) + '\n')


with open("test-B/out.tsv", 'w', encoding='utf-8') as file:
    for test_example in data_B['train_input']:
        prediction_B = model.predict(test_example)
        text_prediction_B = dict((value, key) for key, value in map_dict.items()).get(prediction_B)
        file.write(str(text_prediction_B) + '\n')