{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e5815ee7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "41496f14",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import random\n",
    "from utils import *\n",
    "import pandas as pd\n",
    "from torch import nn\n",
    "import matplotlib.pyplot as plt\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import datasets, transforms"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4377730c",
   "metadata": {},
   "source": [
    "# Visualize Ordered Fashion MNIST dataset "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e177d26",
   "metadata": {},
   "source": [
    "### Labels\n",
    "\n",
    "| Label | Description  |\n",
    "|-------|-------------|\n",
    "| 0     | T-shirt/top |\n",
    "| 1     | Trouser     |\n",
    "| 2     | Pullover    |\n",
    "| 3     | Dress       |\n",
    "| 4     | Coat        |\n",
    "| 5     | Sandal      |\n",
    "| 6     | Shirt       |\n",
    "| 7     | Sneaker     |\n",
    "| 8     | Bag         |\n",
    "| 9     | Ankle boot  |\n",
    "\n",
    "\n",
    "### Natural language constraints\n",
    "\n",
    "If you wear a T-shirt/top, you will not be able to wear a Shirt, Dress, or another T-shirt/top.  \n",
    "If you wear Trousers, you will not be able to wear a Dress or another Trouser.  \n",
    "If you wear a Pullover, you will not be able to wear a Dress, T-shirt/top, Shirt, or another Pullover.  \n",
    "If you wear a Dress, you will not be able to wear a T-shirt/top, Shirt, Trouser, Pullover, or another Dress.  \n",
    "If you wear a Coat, you will not be able to wear a T-shirt/top, Shirt, Pullover, Dress, or another Coat.  \n",
    "If you wear Sandals, you will not be able to wear Sneakers, Trousers, Ankle boots, or another pair of Sandals.  \n",
    "If you wear a Shirt, you will not be able to wear a T-shirt/top, Dress, or another Shirt.  \n",
    "If you wear Sneakers, you will not be able to wear Sandals, Trousers, Ankle boots, or another pair of Sneakers.  \n",
    "If you wear a Bag, you will not be able to wear a T-shirt/top, Shirt, Dress, Pullover, Coat, or another Bag.  \n",
    "If you wear Ankle boots, you will not be able to wear Sandals, Trousers, Sneakers, or another pair of Ankle boots.  \n",
    "\n",
    "You must wear at least one of the following: T-shirt/top, Pullover, Shirt, or Dress.  \n",
    "You must wear at least one of the following: Trouser or Dress.  \n",
    "You must wear at least one of the following: Sandal, Sneaker, or Ankle boot.\n",
    "\n",
    "### LTLf constraints\n",
    "\n",
    "(G(t_shirt_top → ¬F shirt) ∧ G(t_shirt_top → ¬F dress) ∧ G(t_shirt_top → WXG ¬t_shirt_top))  \n",
    "∧ (G(Trouser → ¬F Dress) ∧ G(Trouser → WXG ¬Trouser))  \n",
    "∧ (G(Pullover → ¬F Dress) ∧ G(Pullover → ¬F T-shirt/top) ∧ G(Pullover → ¬F Shirt) ∧ G(Pullover → WXG ¬Pullover))  \n",
    "∧ (G(Dress → ¬F T-shirt/top) ∧ G(Dress → ¬F Shirt) ∧ G(Dress → ¬F Trouser) ∧ G(Dress → ¬F Pullover) ∧ G(Dress → WXG ¬Dress))  \n",
    "∧ (G(Coat → ¬F T-shirt/top) ∧ G(Coat → ¬F Shirt) ∧ G(Coat → ¬F Pullover) ∧ G(Coat → ¬F Dress) ∧ G(Coat → WXG ¬Coat))  \n",
    "∧ (G(Sandal → ¬F Sneaker) ∧ G(Sandal → ¬F Trouser) ∧ G(Sandal → ¬F Ankle boot) ∧ G(Sandal → WXG ¬Sandal))  \n",
    "∧ (G(Shirt → ¬F T-shirt/top) ∧ G(Shirt → ¬F Dress) ∧ G(Shirt → WXG ¬Shirt))  \n",
    "∧ (G(Sneaker → ¬F Sandal) ∧ G(Sneaker → ¬F Trouser) ∧ G(Sneaker → ¬F Ankle boot) ∧ G(Sneaker → WXG ¬Sneaker))  \n",
    "∧ (G(Bag → ¬F T-shirt/top) ∧ G(Bag → ¬F Shirt) ∧ G(Bag → ¬F Dress) ∧ G(Bag → ¬F Pullover) ∧ G(Bag → ¬F Coat) ∧ G(Bag → WXG ¬Bag))  \n",
    "∧ (G(Ankle boot → ¬F Sandal) ∧ G(Ankle boot → ¬F Trouser) ∧ G(Ankle boot → ¬F Sneaker) ∧ G(Ankle boot → WXG ¬Ankle boot))\n",
    "\n",
    "∧ (F(T-shirt/top ∨ Pullover ∨ Shirt ∨ Dress))  \n",
    "∧ (F(Trouser ∨ Dress))  \n",
    "∧ (F(Sandal ∨ Sneaker ∨ Ankle boot))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "d85fba8a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 25804 characters]",
      "text/plain": [
       "<Figure size 1200x400 with 5 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 22884 characters]",
      "text/plain": [
       "<Figure size 1200x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 20512 characters]",
      "text/plain": [
       "<Figure size 1200x400 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "clothing_items = [\"t_shirt_top\", \"trouser\", \"pullover\", \"dress\", \"coat\", \"sandal\", \"shirt\", \"sneaker\", \"bag\", \"ankle_boot\"]\n",
    "\n",
    "root = \"../data/ordered_fashion_mnist\"\n",
    "dataset = datasets.FashionMNIST(root=root, train=True, download=True)\n",
    "\n",
    "csv_path = \"../data/Ordered Fashion MNIST/train_ordered_fashion_mnist.csv\"\n",
    "df = pd.read_csv(csv_path)\n",
    "\n",
    "def display_images(indices, labels):\n",
    "    _, axes = plt.subplots(1, len(indices), figsize=(12, 4))\n",
    "    \n",
    "    if len(indices) == 1:\n",
    "        axes = [axes]\n",
    "\n",
    "    for ax, idx, label in zip(axes, indices, labels):\n",
    "        image = dataset.data[idx].numpy()\n",
    "        class_name = clothing_items[label]\n",
    "\n",
    "        ax.imshow(image, cmap=\"gray\")\n",
    "        ax.set_title(f\"Index: {idx}\\n{label} ({class_name})\")\n",
    "        ax.axis('off')\n",
    "    \n",
    "    plt.show()\n",
    "\n",
    "random_rows = random.sample(range(len(df)), 3)\n",
    "\n",
    "for index in random_rows:\n",
    "    row = df.iloc[index]\n",
    "    indices = eval(row['Indices'])\n",
    "    labels = eval(row['Labels'])\n",
    "    display_images(indices, labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69e37a03",
   "metadata": {},
   "source": [
    "# Automata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ea244d2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "state          distance  | t_shirt_top  trouser      pullover     dress        coat         sandal       shirt        sneaker      bag          ankle_boot   \n",
      "-------------------------------------------------------------------------------------------------------------------------------------------------------------\n",
      "1               2         | 8            5            4            6            2            3            8            3            2            3            \n",
      "2(Deadlock)     inf       | 2            2            2            2            2            2            2            2            2            2            \n",
      "3               1         | 2            2            2            7            2            2            2            2            2            2            \n",
      "4               2         | 2            6            2            2            11           2            2            2            10           2            \n",
      "5               2         | 9            2            6            2            2            14           9            14           2            14           \n",
      "6               1         | 2            2            2            2            13           7            2            7            12           7            \n",
      "7(Accepting)    0         | 2            2            2            2            17           2            2            2            16           2            \n",
      "8               2         | 2            9            4            2            11           2            2            2            10           2            \n",
      "9               1         | 2            2            6            2            13           15           2            15           12           15           \n",
      "10              2         | 2            12           2            2            2            2            2            2            2            2            \n",
      "11              2         | 2            13           2            2            2            2            2            2            10           2            \n",
      "12              1         | 2            2            2            2            2            16           2            16           2            16           \n",
      "13              1         | 2            2            2            2            2            17           2            17           12           17           \n",
      "14              1         | 15           2            7            2            2            2            15           2            2            2            \n",
      "15(Accepting)   0         | 2            2            7            2            17           2            2            2            16           2            \n",
      "16(Accepting)   0         | 2            2            2            2            2            2            2            2            2            2            \n",
      "17(Accepting)   0         | 2            2            2            2            2            2            2            2            16           2            \n"
     ]
    }
   ],
   "source": [
    "clothing_items = [\"t_shirt_top\", \"trouser\", \"pullover\", \"dress\", \"coat\", \"sandal\", \"shirt\", \"sneaker\", \"bag\", \"ankle_boot\"]\n",
    "\n",
    "automata = import_automata(\"../data/automata/ordered_fashion_mnist_automata.json\")\n",
    "print_automata(automata, clothing_items)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c95342a4",
   "metadata": {},
   "source": [
    "# Import Ordered Fashion MNIST dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "eff985c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset, train_custom_indices, train_custom_labels = import_ordered_fashion_mnist( \"../data/Ordered Fashion MNIST\", \"../data/Ordered Fashion MNIST/train_ordered_fashion_mnist.csv\", train=True)\n",
    "test_dataset, test_custom_indices, test_custom_labels = import_ordered_fashion_mnist( \"../data/Ordered Fashion MNIST\", \"../data/Ordered Fashion MNIST/test_ordered_fashion_mnist.csv\", train=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7cca95ec",
   "metadata": {},
   "source": [
    "# Experiments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4c02a741",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fashionConvNet(\n",
      "  (conv1): Sequential(\n",
      "    (0): Conv2d(1, 32, kernel_size=(2, 2), stride=(1, 1), padding=same)\n",
      "    (1): ReLU()\n",
      "    (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
      "  )\n",
      "  (conv2): Sequential(\n",
      "    (0): Conv2d(32, 64, kernel_size=(5, 5), stride=(1, 1))\n",
      "    (1): ReLU()\n",
      "    (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
      "  )\n",
      "  (flatten): Flatten(start_dim=1, end_dim=-1)\n",
      "  (fc1): Linear(in_features=1600, out_features=512, bias=True)\n",
      "  (drop1): Dropout2d(p=0.3, inplace=False)\n",
      "  (fc2): Linear(in_features=512, out_features=128, bias=True)\n",
      "  (drop2): Dropout2d(p=0.2, inplace=False)\n",
      "  (layerNorm): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
      "  (fc3): Linear(in_features=128, out_features=10, bias=True)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "convModel = fashionConvNet()\n",
    "print(convModel)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "180736e4",
   "metadata": {},
   "source": [
    "### run to train CNN from scratch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d330093",
   "metadata": {},
   "outputs": [],
   "source": [
    "#set_seed(0)\n",
    "\n",
    "#transform = transforms.Compose([\n",
    "#    transforms.ToTensor(),\n",
    "#    transforms.Normalize((0.5,), (0.5,))\n",
    "#])\n",
    "\n",
    "#train_data = datasets.FashionMNIST(\n",
    "#    root='../data/Ordered Fashion MNIST',\n",
    "#    train=True,\n",
    "#    download=True,\n",
    "#    transform=transform, \n",
    "#)\n",
    "\n",
    "#test_data = datasets.FashionMNIST(\n",
    "#    root='../data/Ordered Fashion MNIST',\n",
    "#    train=False,\n",
    "#    download=True,\n",
    "#    transform=transform\n",
    "#)\n",
    "\n",
    "#train_dLoader = DataLoader(\n",
    "#    dataset=train_data,\n",
    "#    batch_size=64,\n",
    "#    shuffle=True\n",
    "#)\n",
    "\n",
    "#test_dLoader = DataLoader(\n",
    "#    dataset=test_data,\n",
    "#    batch_size=64,\n",
    "#    shuffle=False\n",
    "#)\n",
    "\n",
    "#loss_fn = nn.CrossEntropyLoss()\n",
    "\n",
    "#optimizer = torch.optim.Adam(\n",
    "#    convModel.parameters(),\n",
    "#    lr=0.002\n",
    "#)\n",
    "\n",
    "#epochs = 7\n",
    "\n",
    "#for e in range(epochs):\n",
    "#    print(f\"\\n\\n Epoch {e+1} -------------------- \\n\")\n",
    "#    train_cnn_without_automata(train_dLoader, convModel, loss_fn, optimizer, device=\"cpu\")\n",
    "\n",
    "#torch.save(convModel.state_dict(), '../data/CNN/fashionConvNet_model_on_original_fmnist.pth')\n",
    "#print(\"Model saved!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36e888d1",
   "metadata": {},
   "source": [
    "### import CNN trained on Fashion MNIST"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "2bf28564",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "convModel.load_state_dict(torch.load('../data/CNN/fashionConvNet_model_on_original_fmnist.pth', map_location=torch.device('cpu')))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9fc3577",
   "metadata": {},
   "source": [
    "### CNN NOT guided by automata on Ordered Fashion MNIST"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "29e7aa1c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Accuracy:               91.03%\n",
      "Sequence Test Accuracy:      64.09%\n"
     ]
    }
   ],
   "source": [
    "#results may vary slightly if using a different OS and/or CPU model\n",
    "#to obtain the paper's results we used: Windows 11 and CPU Intel Core Ultra 7 165H\n",
    "\n",
    "accuracy_without_automata, sequence_accuracy_without_automata = test_cnn_without_automata(test_custom_indices, test_custom_labels, convModel, test_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f24ce0f",
   "metadata": {},
   "source": [
    "### CNN guided by automata on Ordered Fashion MNIST"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "862c85ce",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "k = 2:\taccuracy = 93.37%, sequence accuracy = 78.50%\n",
      "k = 3:\taccuracy = 93.99%, sequence accuracy = 79.01%\n",
      "k = 4:\taccuracy = 94.58%, sequence accuracy = 79.73%\n",
      "k = 5:\taccuracy = 94.79%, sequence accuracy = 79.94%\n",
      "k = 6:\taccuracy = 94.93%, sequence accuracy = 80.04%\n",
      "k = 7:\taccuracy = 95.17%, sequence accuracy = 80.25%\n",
      "k = 8:\taccuracy = 95.38%, sequence accuracy = 80.45%\n",
      "k = 9:\taccuracy = 95.54%, sequence accuracy = 80.66%\n",
      "k = 10:\taccuracy = 95.58%, sequence accuracy = 80.66%\n"
     ]
    }
   ],
   "source": [
    "#results may vary slightly if using a different OS and/or CPU model\n",
    "#to obtain the paper's results we used: Windows 11 and CPU Intel Core Ultra 7 165H\n",
    "\n",
    "accuracy_scores = []\n",
    "sequence_accuracy_scores = []\n",
    "\n",
    "for k in range(2, 11):\n",
    "    pred = test_cnn_with_automata(test_custom_indices, convModel, test_dataset, automata, k, device=\"cpu\")\n",
    "    accuracy_with_automata, sequence_accuracy_with_automata = ordered_fashion_mnist_accuracy(test_custom_labels, pred)\n",
    "    accuracy_scores.append(accuracy_with_automata)\n",
    "    sequence_accuracy_scores.append(sequence_accuracy_with_automata)\n",
    "    print(f\"k = {k}:\\taccuracy = {accuracy_with_automata:.2f}%, sequence accuracy = {sequence_accuracy_with_automata:.2f}%\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "4b4eef9e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 69056 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 71500 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "k_values = list(range(2, 11))\n",
    "line_width = 3\n",
    "font_properties = {'fontsize': 20, 'fontweight': 'bold'}\n",
    "\n",
    "# Line chart for Accuracy\n",
    "fig, ax = plt.subplots(figsize=(10, 6))\n",
    "ax.plot(k_values, accuracy_scores, marker='o', label='With TRIDENT', color='skyblue', linewidth=line_width, markersize=15)\n",
    "ax.axhline(y=accuracy_without_automata, color='salmon', linestyle='--', label='Without TRIDENT', linewidth=line_width)\n",
    "\n",
    "# Annotate each point\n",
    "for k, acc in zip(k_values, accuracy_scores):\n",
    "    ax.annotate(f'{acc:.2f}', xy=(k, acc), xytext=(0, 8), textcoords='offset points', ha='center', **font_properties)\n",
    "\n",
    "# Annotate horizontal line value\n",
    "center_k = (k_values[0] + k_values[-1]) / 2\n",
    "ax.text(center_k, accuracy_without_automata, f'{accuracy_without_automata:.2f}',\n",
    "        va='bottom', ha='center', fontsize=20, fontweight='bold')\n",
    "\n",
    "ax.set_xlabel('Number of beams', **font_properties)\n",
    "ax.set_ylabel('Accuracy', **font_properties)\n",
    "ax.set_title('Accuracy With vs Without TRIDENT', **font_properties)\n",
    "ax.set_ylim(90, 100)\n",
    "ax.grid(True)\n",
    "ax.legend(fontsize=20)\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "#fig.savefig(\"accuracy_comparison.pdf\")\n",
    "\n",
    "# Line chart for Sequence Accuracy\n",
    "fig, ax = plt.subplots(figsize=(10, 6))\n",
    "ax.plot(k_values, sequence_accuracy_scores, marker='o', label='With TRIDENT', color='mediumseagreen', linewidth=line_width, markersize=15)\n",
    "ax.axhline(y=sequence_accuracy_without_automata, color='tomato', linestyle='--', label='Without TRIDENT', linewidth=line_width)\n",
    "\n",
    "# Annotate each point\n",
    "for k, seq_acc in zip(k_values, sequence_accuracy_scores):\n",
    "    ax.annotate(f'{seq_acc:.2f}', xy=(k, seq_acc), xytext=(0, 8), textcoords='offset points', ha='center', **font_properties)\n",
    "\n",
    "# Annotate horizontal line value\n",
    "center_k = (k_values[0] + k_values[-1]) / 2\n",
    "ax.text(center_k, sequence_accuracy_without_automata, f'{sequence_accuracy_without_automata:.2f}',\n",
    "        va='bottom', ha='center', fontsize=20, fontweight='bold')\n",
    "\n",
    "ax.set_xlabel('Number of beams', **font_properties)\n",
    "ax.set_ylabel('Sequence Accuracy', **font_properties)\n",
    "ax.set_title('Sequence Accuracy With vs Without TRIDENT', **font_properties)\n",
    "ax.set_ylim(60, 100)\n",
    "ax.grid(True)\n",
    "ax.legend(fontsize=20)\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "#fig.savefig(\"sequence_accuracy_comparison.pdf\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
