{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "\n",
    "import copy\n",
    "import nltk\n",
    "import pandas as pd\n",
    "import rapidfuzz\n",
    "import time\n",
    "\n",
    "from nltk.stem import WordNetLemmatizer\n",
    "from rapidfuzz.fuzz import  partial_ratio\n",
    "from rapidfuzz.utils import default_process\n",
    "\n",
    "\n",
    "wl = WordNetLemmatizer()\n",
    "\n",
    "glossary = pd.read_csv('mt-summit-corpora/glossary.tsv', sep='\\t', header=None, names=['source', 'result'])\n",
    "\n",
    "source_lemmatized = []\n",
    "for word in glossary['source']:\n",
    "    word = nltk.word_tokenize(word)\n",
    "    source_lemmatized.append(' '.join([wl.lemmatize(x) for x in word]))\n",
    "\n",
    "glossary['source_lem'] = source_lemmatized\n",
    "glossary = glossary[['source', 'source_lem', 'result']]\n",
    "glossary.set_index('source_lem')\n",
    "\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": true
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.194806501\n"
     ]
    }
   ],
   "source": [
    "# train_in_path = 'mt-summit-corpora/train/in.tsv'\n",
    "# train_expected_path = 'mt-summit-corpora/train/expected.tsv'\n",
    "\n",
    "train_in_path = 'mt-summit-corpora/dev-0/in.tsv'\n",
    "train_expected_path = 'mt-summit-corpora/dev-0/expected.tsv'\n",
    "\n",
    "\n",
    "file_pl = pd.read_csv(train_expected_path, sep='\\t', header=None, names=['text'])\n",
    "\n",
    "start_time = time.time_ns()\n",
    "file_lemmatized = []\n",
    "with open(train_in_path, 'r') as file:\n",
    "    for line in file:\n",
    "        if len(file_lemmatized) % 50000 == 0:\n",
    "            print(len(file_lemmatized), end='\\r')\n",
    "        line = nltk.word_tokenize(line)\n",
    "        file_lemmatized.append(' '.join([wl.lemmatize(x) for x in line]))\n",
    "\n",
    "stop = time.time_ns()\n",
    "timex = (stop - start_time) / 1000000000\n",
    "print(timex)\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "6.904260614\n"
     ]
    }
   ],
   "source": [
    "\n",
    "THRESHOLD = 88\n",
    "\n",
    "def is_injectable(sentence_pl, sequence):\n",
    "    sen = sentence_pl.split()\n",
    "    window_size = len(sequence.split())\n",
    "    maxx = 0\n",
    "    for i in range(len(sen) - window_size):\n",
    "        current = rapidfuzz.fuzz.ratio(' '.join(sen[i:i + window_size]), sequence)\n",
    "        if current > maxx:\n",
    "            maxx = current\n",
    "    if maxx >= THRESHOLD:\n",
    "        return True\n",
    "    else:\n",
    "        return False\n",
    "\n",
    "def get_injected(sentence, sequence, inject):\n",
    "    sen = sentence.split()\n",
    "    window_size = len(sequence.split())\n",
    "    maxx = 0\n",
    "    maxxi = 0\n",
    "    for i in range(len(sen) - window_size + 1):\n",
    "        current = rapidfuzz.fuzz.ratio(' '.join(sen[i:i + window_size]), sequence)\n",
    "        if current >= maxx:\n",
    "            maxx = current\n",
    "            maxxi = i\n",
    "    return ' '.join(sen[:maxxi + window_size]) + ' $' + inject + '$ ' + ' '.join(sen[maxxi + window_size:])\n",
    "\n",
    "glossary['source_lem'] = [' ' + str(default_process(x)) + ' ' for x in glossary['source_lem']]\n",
    "\n",
    "start_time = time.time_ns()\n",
    "en = []\n",
    "translation_line_counts = []\n",
    "for line, line_pl in zip(file_lemmatized, file_pl['text'].values.tolist()):\n",
    "    line = default_process(line)\n",
    "    line_pl = default_process(line_pl)\n",
    "    matchez = rapidfuzz.process.extract(query=line, choices=glossary['source_lem'], limit=5, score_cutoff=THRESHOLD, scorer=partial_ratio)\n",
    "    if len(matchez) > 0:\n",
    "        lines_added = 0\n",
    "        for match in matchez:\n",
    "            polish_translation = glossary.loc[lambda df: df['source_lem'] == match[0]]['result'].astype(str).values.flatten()[0]\n",
    "            if is_injectable(line_pl, polish_translation):\n",
    "                en.append(get_injected(line, match[0], polish_translation))\n",
    "                lines_added += 1\n",
    "        if lines_added == 0:\n",
    "            en.append(line)\n",
    "            lines_added = 1\n",
    "        translation_line_counts.append(lines_added)\n",
    "    else:\n",
    "        translation_line_counts.append(1)\n",
    "        en.append(line)\n",
    "\n",
    "\n",
    "stop = time.time_ns()\n",
    "timex = (stop - start_time) / 1000000000\n",
    "print(timex)\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "outputs": [],
   "source": [
    "\n",
    "translations = pd.read_csv(train_expected_path, sep='\\t', header=0, names=['text'])\n",
    "with open(train_expected_path + '.injected', 'w') as file_plx:\n",
    "    for line, translation_line_ct in zip(translations['text'].values.tolist(), translation_line_counts):\n",
    "        for i in range(translation_line_ct):\n",
    "            file_plx.write(line + '\\n')\n",
    "\n",
    "\n",
    "with open(train_in_path + '.injected', 'w') as file_en:\n",
    "    for e in en:\n",
    "        file_en.write(e + '\\n')"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}