{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5a360ab5-0158-4584-b3ce-c55937ada352",
   "metadata": {},
   "outputs": [],
   "source": [
    "import xml.etree.ElementTree as ET\n",
    "import os\n",
    "import openai\n",
    "import json\n",
    "from tqdm import tqdm\n",
    "from pydantic import BaseModel\n",
    "from html.parser import HTMLParser"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0b39e2f3-57f7-47dd-be32-5b6f819945e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# target_locale='es'\n",
    "# target_locale='zh_CN'\n",
    "target_locales=['de','es','zh_CN']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0c7b35cb-9a96-4545-96c5-5aefc73459f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read my API key\n",
    "with open('/Users/pauly/.private/openai_api_key_itksnap','rt') as f:\n",
    "    os.environ['OPENAI_API_KEY'] = f.read().strip()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "983dfdc2-a379-48f2-8c3c-800df7c46ffe",
   "metadata": {},
   "outputs": [],
   "source": [
    "client = openai.OpenAI()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "13ff9a89-47df-4b6d-9538-dceb5a4ee589",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SingleTranslation(BaseModel):\n",
    "    original: str\n",
    "    translated: str\n",
    "\n",
    "class TranslationResponse(BaseModel):\n",
    "    translations: list[SingleTranslation]\n",
    "\n",
    "def translate(client, locale:str, input: list[str]):\n",
    "    response = client.responses.parse(\n",
    "      prompt={\n",
    "        \"id\": \"pmpt_686e17744b3c8190a634009a15a02eeb0258750c14426f72\",\n",
    "        \"version\": \"7\",\n",
    "        \"variables\": {\n",
    "          \"target_locale\": locale\n",
    "        }\n",
    "      },\n",
    "      input=json.dumps(input, ensure_ascii=False),\n",
    "      text_format = TranslationResponse\n",
    "    )\n",
    "    return response.output_parsed.translations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "28703cd7-3155-487f-aa08-204230dc7ee7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['Zoom Inspector', 'Control zoom factor in all slice views.']"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Extractor for HTML strings - only retains stuff that's not tags\n",
    "class HTMLDataExtractor(HTMLParser):\n",
    "    def __init__(self):\n",
    "        HTMLParser.__init__(self)\n",
    "        self.all_data = []\n",
    "        \n",
    "    def handle_starttag(self, tag, attrs):\n",
    "        pass\n",
    "\n",
    "    def handle_endtag(self, tag):\n",
    "        pass\n",
    "\n",
    "    def handle_data(self, data):\n",
    "        string = data.strip()\n",
    "        if len(string):\n",
    "            self.all_data.append(string)\n",
    "\n",
    "parser = HTMLDataExtractor()\n",
    "parser.feed('<html><head/><body><p><span style=\" font-weight:600;\">Zoom Inspector</span></p><p>Control zoom factor in all slice views. </p></body></html>')\n",
    "parser.all_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "96a1c114-1c14-4afe-b99e-c37f97c7e461",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[SingleTranslation(original='Zoom Inspector', translated='Zoom-Inspektor'),\n",
       " SingleTranslation(original='Control zoom factor in all slice views.', translated='Steuert den Zoomfaktor in allen Schnittansichten.')]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "translate(client, target_locales[0], parser.all_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a65b9e5c-77fd-4799-ac4b-46a4ed2bdcd2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Translator for HTML strings - only retains stuff that's not tags\n",
    "class HTMLDataTranslator(HTMLParser):\n",
    "    def __init__(self, token_list, token_index):\n",
    "        HTMLParser.__init__(self)\n",
    "        self.all_data = []\n",
    "        self.token_list = token_list\n",
    "        self.token_index = token_index\n",
    "        self.result = ''\n",
    "        self.success = True\n",
    "        \n",
    "    def handle_starttag(self, tag, attrs):\n",
    "        attr_str = ''.join(f' {k}=\"{v}\"' for k, v in attrs)\n",
    "        self.result += f\"<{tag}{attr_str}>\"\n",
    "\n",
    "    def handle_endtag(self, tag):\n",
    "        self.result += f\"</{tag}>\"\n",
    "\n",
    "    def handle_data(self, data):\n",
    "        string = data.strip()\n",
    "        if len(string):\n",
    "            if self.token_index >= len(self.token_list):\n",
    "                print(f'Replacement list size mismatch')\n",
    "                self.success = False\n",
    "                self.result += data\n",
    "            else:\n",
    "                t = self.token_list[self.token_index]\n",
    "                if t.original != string:\n",
    "                    print(f'Mismatch, original f{t.original} != f{string}')\n",
    "                    self.success = False\n",
    "                    self.result += data\n",
    "                else:\n",
    "                    self.result += data.replace(t.original, t.translated)\n",
    "                self.token_index += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "86e58a01-30c3-445e-970a-d83fd4bc3d42",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Translator:\n",
    "\n",
    "    def __init__(self, rootdir, locale):\n",
    "        self.locale = locale\n",
    "        self.tsfile = os.path.join(rootdir, f'GUI/Qt/Translations/itksnap_{locale}.ts')\n",
    "        self.tree = ET.parse(self.tsfile)\n",
    "        self.root = self.tree.getroot()\n",
    "\n",
    "        self.all_con = []\n",
    "        for context in self.root.findall('context'):\n",
    "            all_tran = []\n",
    "            for message in context.findall('message'):\n",
    "                source = message.find('source')\n",
    "                translation = message.find('translation')\n",
    "                if translation.attrib.get('type') != 'unfinished':\n",
    "                    continue\n",
    "                if 'qrc:/' in source.text:\n",
    "                    continue\n",
    "                htmlextr = HTMLDataExtractor()\n",
    "                htmlextr.feed(source.text)\n",
    "                all_tran.append({'message':message, 'text':source.text, 'strings': htmlextr.all_data})\n",
    "            if len(all_tran) > 0:\n",
    "                self.all_con.append({'context': context, 'tran': all_tran})\n",
    "\n",
    "    def translate(self):\n",
    "        for con in self.all_con:\n",
    "            all_strings = [x for xs in con['tran'] for x in xs['strings']]\n",
    "            if len(all_strings):\n",
    "                print(f\"Translating {con['context'].find('name').text} to {self.locale}\")\n",
    "                gpt = translate(client, self.locale, all_strings)\n",
    "                if gpt is None:\n",
    "                    continue\n",
    "\n",
    "                gpt_index = 0\n",
    "                print(f\"Updating {con['context'].find('name').text}\")\n",
    "                for tran in con['tran']:\n",
    "                    message = tran['message']\n",
    "                    source = message.find('source')\n",
    "                    translation = message.find('translation')\n",
    "                    #if translation.attrib.get('type') != 'unfinished':\n",
    "                    #    continue\n",
    "                    if 'qrc:/' in source.text:\n",
    "                        continue\n",
    "                    htmltran = HTMLDataTranslator(gpt, gpt_index)\n",
    "                    htmltran.feed(source.text)\n",
    "                    gpt_index = htmltran.token_index\n",
    "                    if htmltran.success:\n",
    "                        translation.text = htmltran.result\n",
    "                        translation.attrib.pop('type', None)\n",
    "\n",
    "    def save(self):\n",
    "        self.tree.write(self.tsfile, encoding='utf-8')\n",
    "        print(f'Wrote {self.tsfile}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d53df266-e57e-4504-9876-49ccdf59093b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Translating ViewPanel3D to de\n",
      "Updating ViewPanel3D\n",
      "Wrote /Users/pauly/tk/itksnap/itksnap/GUI/Qt/Translations/itksnap_de.ts\n",
      "Translating ViewPanel3D to es\n",
      "Updating ViewPanel3D\n",
      "Wrote /Users/pauly/tk/itksnap/itksnap/GUI/Qt/Translations/itksnap_es.ts\n",
      "Translating ViewPanel3D to zh_CN\n",
      "Updating ViewPanel3D\n",
      "Wrote /Users/pauly/tk/itksnap/itksnap/GUI/Qt/Translations/itksnap_zh_CN.ts\n"
     ]
    }
   ],
   "source": [
    "for locale in target_locales:\n",
    "    t = Translator('/Users/pauly/tk/itksnap/itksnap', locale)\n",
    "    t.translate()\n",
    "    t.save()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa0e5552-f18d-4928-ae7b-afcc983b3bab",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
