{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6981eae7",
   "metadata": {},
   "source": [
    "# Programming Assistant"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f2e168c",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "27c69a83",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tkinter as tk\n",
    "from tkinter import scrolledtext, messagebox\n",
    "from PIL import ImageGrab\n",
    "import datetime\n",
    "import os\n",
    "import numpy as np\n",
    "import subprocess\n",
    "import easyocr\n",
    "import threading"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee7c16b1",
   "metadata": {},
   "source": [
    "## Configs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "d678ba76",
   "metadata": {},
   "outputs": [],
   "source": [
    "DOCUMENTS_PATH = os.path.expanduser(\"~/Documents/OCR_ScreenCaptures\")\n",
    "INPUT_TEXT_FILE = os.path.join(DOCUMENTS_PATH, \"screen_text.txt\")\n",
    "OUTPUT_CODE_FILE = os.path.join(DOCUMENTS_PATH, \"code_snippets.txt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f05ca831",
   "metadata": {},
   "source": [
    "## Screen Capture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "5e311658",
   "metadata": {},
   "outputs": [],
   "source": [
    "def textExtractionEasyOCR():\n",
    "    os.makedirs(DOCUMENTS_PATH, exist_ok=True)\n",
    "    \n",
    "    open(INPUT_TEXT_FILE, \"w\", encoding=\"utf-8\").close()\n",
    "    open(OUTPUT_CODE_FILE, \"w\", encoding=\"utf-8\").close()\n",
    "\n",
    "    timeStamp = datetime.datetime.now().strftime(\"%Y-%m-%d_%H-%M-%S\")\n",
    "    imagePath = os.path.join(DOCUMENTS_PATH, f\"screenshot_{timeStamp}.png\")\n",
    "    textPath = os.path.join(DOCUMENTS_PATH, \"screen_text.txt\")\n",
    "\n",
    "    screenshot = ImageGrab.grab()\n",
    "    screenshot.save(imagePath)\n",
    "\n",
    "    imageArray = np.array(screenshot)\n",
    "\n",
    "    reader = easyocr.Reader(['en'], gpu=False)\n",
    "    results = reader.readtext(imageArray)\n",
    "    results = sorted(results, key=lambda r: r[0][0][1])\n",
    "\n",
    "    lines = []\n",
    "    currentLineY = None\n",
    "    lineText = []\n",
    "\n",
    "    for bbox, text, conf in results:\n",
    "        y = bbox[0][1]\n",
    "        if currentLineY is None:\n",
    "            currentLineY = y\n",
    "\n",
    "        if abs(y - currentLineY) > 10:\n",
    "            lines.append(\" \".join(lineText))\n",
    "            lineText = [text]\n",
    "            currentLineY = y\n",
    "        else:\n",
    "            lineText.append(text)\n",
    "\n",
    "    if lineText:\n",
    "        lines.append(\" \".join(lineText))\n",
    "\n",
    "    extractedText = \"\\n\".join(lines)\n",
    "\n",
    "    with open(textPath, \"a\", encoding=\"utf-8\") as f:\n",
    "        f.write(f\"\\n--- Capture at {timeStamp} (EasyOCR) ---\\n\")\n",
    "        f.write(extractedText.strip() + \"\\n\")\n",
    "\n",
    "    return extractedText, imagePath, textPath"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9931d3e3",
   "metadata": {},
   "source": [
    "## Code Extraction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "aecb4aa1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def codeResemblance(text: str) -> bool:\n",
    "    codeIndicators = [\n",
    "        \"{\", \"}\", \";\", \"def \", \"class \", \"import \",\n",
    "        \"public \", \"private \", \"function \", \"=>\", \"(\", \")\", \"=\"\n",
    "    ]\n",
    "    return any(ind in text for ind in codeIndicators)\n",
    "\n",
    "\n",
    "def codeExtraction():\n",
    "    if not os.path.exists(INPUT_TEXT_FILE):\n",
    "        return \"\", \"[!] Input file not found.\"\n",
    "\n",
    "    with open(INPUT_TEXT_FILE, \"r\", encoding=\"utf-8\") as f:\n",
    "        text = f.read()\n",
    "\n",
    "    blocks = text.split(\"\\n\\n\")\n",
    "    codeBlocks = []\n",
    "\n",
    "    for block in blocks:\n",
    "        cleanedBlock = block.strip()\n",
    "        if not cleanedBlock:\n",
    "            continue\n",
    "        if codeResemblance(cleanedBlock):\n",
    "            codeBlocks.append(cleanedBlock)\n",
    "\n",
    "    if codeBlocks:\n",
    "        with open(OUTPUT_CODE_FILE, \"w\", encoding=\"utf-8\") as f:\n",
    "            for block in codeBlocks:\n",
    "                f.write(block + \"\\n\\n\")\n",
    "        return \"\\n\\n\".join(codeBlocks), f\"[+] Extracted {len(codeBlocks)} code block(s).\"\n",
    "    else:\n",
    "        return \"\", \"[!] No code detected in the input text.\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f731612c",
   "metadata": {},
   "source": [
    "# Model Selection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "df2f9fab",
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_OPTIONS = [\n",
    "    \"codellama:7b\",\n",
    "    \"codellama:13b\",\n",
    "    \"llama2:7b\",\n",
    "    \"llama2:13b\",\n",
    "    \"llama3:8b\"\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "740820c2",
   "metadata": {},
   "source": [
    "## Model Query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "1b83a17e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def queryCodellama(prompt: str, model: str = \"\"):\n",
    "    try:\n",
    "        result = subprocess.run(\n",
    "            [\"ollama\", \"run\", model],\n",
    "            input=prompt.encode(\"utf-8\"),\n",
    "            capture_output=True,\n",
    "            timeout=120\n",
    "        )\n",
    "        if result.returncode != 0:\n",
    "            return f\"[!] Error running model: {result.stderr.decode('utf-8')}\"\n",
    "        return result.stdout.decode(\"utf-8\")\n",
    "    except Exception as e:\n",
    "        return f\"[!] Exception running model: {str(e)}\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1570f65e",
   "metadata": {},
   "source": [
    "# Feedback and GUI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "adbafe33",
   "metadata": {},
   "outputs": [],
   "source": [
    "class OCRCodeAnalyzerApp:\n",
    "    def __init__(self, root):\n",
    "        self.root = root\n",
    "        root.title(\"OCR Code Analyzer\")\n",
    "\n",
    "        self.langLabel = tk.Label(root, text=\"Programming Language (optional):\")\n",
    "        self.langLabel.grid(row=0, column=0, sticky=\"w\", padx=5, pady=5)\n",
    "\n",
    "        self.langEntry = tk.Entry(root, width=30)\n",
    "        self.langEntry.grid(row=0, column=1, sticky=\"we\", padx=5, pady=5)\n",
    "\n",
    "        self.modelLabel = tk.Label(root, text=\"Select Model:\")\n",
    "        self.modelLabel.grid(row=0, column=2, sticky=\"w\", padx=5, pady=5)\n",
    "\n",
    "        self.selectedModel = tk.StringVar(root)\n",
    "        self.selectedModel.set(MODEL_OPTIONS[0])\n",
    "\n",
    "        self.modelDropdown = tk.OptionMenu(root, self.selectedModel, *MODEL_OPTIONS)\n",
    "        self.modelDropdown.grid(row=0, column=3, sticky=\"we\", padx=5, pady=5)\n",
    "\n",
    "        self.captureButton = tk.Button(root, text=\"Capture Screen and Analyze\", command=self.startAnalysis)\n",
    "        self.captureButton.grid(row=0, column=4, sticky=\"we\", padx=5, pady=5)\n",
    "\n",
    "        self.statusText = scrolledtext.ScrolledText(root, height=10, state='disabled')\n",
    "        self.statusText.grid(row=1, column=0, columnspan=5, sticky=\"nsew\", padx=5, pady=5)\n",
    "\n",
    "        self.feedbackLabel = tk.Label(root, text=\"AI Feedback:\")\n",
    "        self.feedbackLabel.grid(row=2, column=0, sticky=\"w\", padx=5, pady=(10,0))\n",
    "\n",
    "        self.feedbackText = scrolledtext.ScrolledText(root, height=20)\n",
    "        self.feedbackText.grid(row=3, column=0, columnspan=5, sticky=\"nsew\", padx=5, pady=5)\n",
    "\n",
    "        root.grid_rowconfigure(1, weight=1)\n",
    "        root.grid_rowconfigure(3, weight=3)\n",
    "        root.grid_columnconfigure(1, weight=1)\n",
    "        root.grid_columnconfigure(3, weight=1)\n",
    "\n",
    "    def log(self, message):\n",
    "        self.statusText.config(state='normal')\n",
    "        self.statusText.insert(tk.END, message + \"\\n\")\n",
    "        self.statusText.see(tk.END)\n",
    "        self.statusText.config(state='disabled')\n",
    "\n",
    "    def setFeedback(self, text):\n",
    "        self.feedbackText.delete(1.0, tk.END)\n",
    "        self.feedbackText.insert(tk.END, text)\n",
    "\n",
    "    def startAnalysis(self):\n",
    "        self.captureButton.config(state='disabled')\n",
    "        self.setFeedback(\"\")\n",
    "        self.statusText.config(state='normal')\n",
    "        self.statusText.delete(1.0, tk.END)\n",
    "        self.statusText.config(state='disabled')\n",
    "\n",
    "        thread = threading.Thread(target=self.runAnalysis)\n",
    "        thread.start()\n",
    "\n",
    "    def runAnalysis(self):\n",
    "        try:\n",
    "            lang = self.langEntry.get().strip()\n",
    "            model = self.selectedModel.get()\n",
    "            self.log(f\"Using model: {model}\")\n",
    "\n",
    "            self.log(\"Capturing screen and extracting text...\")\n",
    "            extractedText, imagePath, textPath = textExtractionEasyOCR()\n",
    "            self.log(f\"[+] Screenshot saved: {imagePath}\")\n",
    "            self.log(f\"[+] Text extracted and saved to: {textPath}\")\n",
    "\n",
    "            if not extractedText.strip():\n",
    "                self.log(\"[!] No text extracted from the screen.\")\n",
    "                self.captureButton.config(state='normal')\n",
    "                return\n",
    "\n",
    "            self.log(\"Extracting code blocks from text...\")\n",
    "            codeText, msg = codeExtraction()\n",
    "            self.log(msg)\n",
    "            if not codeText.strip():\n",
    "                self.log(\"[!] No code detected in the extracted text.\")\n",
    "                self.captureButton.config(state='normal')\n",
    "                return\n",
    "\n",
    "            lines = [line for line in codeText.splitlines() if line.strip()]\n",
    "            chunkSize = 50\n",
    "            chunks = [lines[i:i+chunkSize] for i in range(0, len(lines), chunkSize)]\n",
    "\n",
    "            feedbackList = []\n",
    "            totalChunks = len(chunks)\n",
    "\n",
    "            for i, chunkLines in enumerate(chunks):\n",
    "                chunkText = \"\\n\".join(chunkLines)\n",
    "\n",
    "                if lang:\n",
    "                    prompt = f\"\"\"\n",
    "The following text is extracted code from OCR. There may be some noise and incorrectly detected characters. Filter out any non-code noise. The programming language that you are analysing is {lang}.\n",
    "Before your feedback, state what the detected programming language is.\n",
    "Then, state if there are any improvements to be made or not. \n",
    "If there are, provide feedback on correctness and improvements. Do not feedback on syntax errors caused by potential noise. Make sure you state your reasoning.\n",
    "Do not generate any code at all, and do not provide any improved versions of already existing code.\n",
    "\n",
    "Code chunk:\n",
    "{chunkText}\n",
    "\"\"\"\n",
    "                else:\n",
    "                    prompt = f\"\"\"\n",
    "The following text is extracted code from OCR. There may be some noise and incorrectly detected characters. Filter out any non-code noise. The programming language is unknown; please detect it yourself.\n",
    "Before your feedback, state what the detected programming language is.\n",
    "Then, state if there are any improvements to be made or not. \n",
    "If there are, provide feedback on correctness and improvements. Do not feedback on syntax errors caused by potential noise. Make sure you state your reasoning.\n",
    "Do not generate any code at all, and do not provide any improved versions of already existing code.\n",
    "\n",
    "Code chunk:\n",
    "{chunkText}\n",
    "\"\"\"\n",
    "\n",
    "                self.log(f\"[+] Getting feedback for chunk {i+1}/{totalChunks}...\")\n",
    "                feedback = queryCodellama(prompt, model=model)\n",
    "                feedbackList.append(feedback)\n",
    "                self.log(f\"[+] Feedback for chunk {i+1}/{totalChunks} done.\")\n",
    "\n",
    "            finalFeedback = \"\\n\\n\".join(feedbackList)\n",
    "            self.setFeedback(finalFeedback)\n",
    "            self.log(\"Analysis complete.\")\n",
    "        except Exception as e:\n",
    "            self.log(f\"[!] Exception occurred: {str(e)}\")\n",
    "        finally:\n",
    "            self.captureButton.config(state='normal')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "09dd98ad",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using CPU. Note: This module is much faster with a GPU.\n",
      "C:\\Users\\mtii0\\AppData\\Local\\Packages\\PythonSoftwareFoundation.Python.3.13_qbz5n2kfra8p0\\LocalCache\\local-packages\\Python313\\site-packages\\torch\\utils\\data\\dataloader.py:666: UserWarning: 'pin_memory' argument is set as true but no accelerator is found, then device pinned memory won't be used.\n",
      "  warnings.warn(warn_msg)\n"
     ]
    }
   ],
   "source": [
    "def main():\n",
    "    root = tk.Tk()\n",
    "    app = OCRCodeAnalyzerApp(root)\n",
    "    root.geometry(\"900x700\")\n",
    "    root.mainloop()\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.13.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}