import vowpalwabbit
import pandas as pd
import re


def to_vw_format(row, map_dict):
    text = row['text'].replace('\n', ' ').lower().strip()
    #text = re.sub("[^a-zA-Z0-9 -']", '', text)
    text = re.sub("[^a-zA-Z -']", '', text)
    text = re.sub(" +", ' ', 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


def predict_and_write(folder_name, model, map_dict):
    data = pd.read_csv(f'{folder_name}/in.tsv', header=None, sep='\t')
    data = data.drop(1, axis=1)
    data.columns = ['year', 'text']

    data['train_input'] = data.apply(lambda row: to_vw_format(row, map_dict), axis=1)

    with open(f"{folder_name}/out.tsv", 'w', encoding='utf-8') as file:
        for test_example in data['train_input']:
            prediction = model.predict(test_example)
            text_prediction = dict((value, key) for key, value in map_dict.items()).get(prediction)
            file.write(str(text_prediction) + '\n')


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

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)
map_dict = {}
for i, x in enumerate(data['category'].unique()):
    map_dict[x] = i+1   #0 nie może być
print(map_dict)
data['train_input'] = data.apply(lambda row: to_vw_format(row, map_dict), axis=1)
print(data.head(5))


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

predict_and_write('dev-0', model, map_dict)
predict_and_write('test-A', model, map_dict)
predict_and_write('test-B', model, map_dict)