{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Installation\n",
    "Available on PyPI as [`zanj`](https://pypi.org/project/zanj/)\n",
    "\n",
    "```\n",
    "pip install zanj\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "\n",
    "from muutils.json_serialize import (\n",
    "    SerializableDataclass,\n",
    "    serializable_dataclass,\n",
    "    serializable_field,\n",
    ")\n",
    "from zanj import ZANJ"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "# Usage\n",
    "\n",
    "## Saving a basic object\n",
    "\n",
    "Any `SerializableDataclass` of basic types can be saved as zanj:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "@serializable_dataclass\n",
    "class BasicZanj(SerializableDataclass):\n",
    "    a: str\n",
    "    q: int = 42\n",
    "    c: list[int] = serializable_field(default_factory=list)\n",
    "\n",
    "\n",
    "# initialize a zanj reader/writer\n",
    "zj = ZANJ()\n",
    "\n",
    "# create an instance\n",
    "instance: BasicZanj = BasicZanj(\"hello\", 42, [1, 2, 3])\n",
    "path: str = \"tests/junk_data/path_to_save_instance.zanj\"\n",
    "zj.save(instance, path)\n",
    "recovered: BasicZanj = zj.read(path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "type(recovered) = <class '__main__.BasicZanj'>\n",
      "os.path.getsize(path) = 3458\n"
     ]
    }
   ],
   "source": [
    "print(f\"{type(recovered) = }\")  # BasicZanj\n",
    "print(f\"{os.path.getsize(path) = }\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "ZANJ will intelligently handle nested serializable dataclasses, numpy arrays, pytorch tensors, and pandas dataframes: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "@serializable_dataclass\n",
    "class Complicated(SerializableDataclass):\n",
    "    name: str\n",
    "    arr1: np.ndarray\n",
    "    arr2: np.ndarray\n",
    "    iris_data: pd.DataFrame\n",
    "    brain_data: pd.DataFrame\n",
    "    container: list[BasicZanj]\n",
    "    torch_tensor: torch.Tensor"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For custom classes, you can specify a `serialization_fn` and `loading_fn` to handle the logic of converting to and from a json-serializable format:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "@serializable_dataclass\n",
    "class Complicated2(SerializableDataclass):\n",
    "    name: str\n",
    "    device: torch.device = serializable_field(\n",
    "        serialization_fn=lambda self: str(self.device),\n",
    "        loading_fn=lambda data: torch.device(data[\"device\"]),\n",
    "    )"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that `loading_fn` takes the dictionary of the whole class -- this is in case you've stored data in multiple fields of the dict which are needed to reconstruct the object."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Saving Models\n",
    "\n",
    "First, define a configuration class for your model. This class will hold the parameters for your model and any associated objects (like losses and optimizers). The configuration class should be a subclass of `SerializableDataclass` and use the `serializable_field` function to define fields that need special serialization.\n",
    "\n",
    "Here's an example that defines a configuration for a simple neural network:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from zanj.torchutil import ConfiguredModel, set_config_class\n",
    "\n",
    "\n",
    "@serializable_dataclass\n",
    "class MyNNConfig(SerializableDataclass):\n",
    "    input_dim: int\n",
    "    hidden_dim: int\n",
    "    output_dim: int\n",
    "\n",
    "    # store the activation function by name, reconstruct it by looking it up in torch.nn\n",
    "    act_fn: torch.nn.Module = serializable_field(\n",
    "        serialization_fn=lambda x: x.__name__,\n",
    "        loading_fn=lambda x: getattr(torch.nn, x[\"act_fn\"]),\n",
    "    )\n",
    "\n",
    "    # same for the loss function\n",
    "    loss_kwargs: dict = serializable_field(default_factory=dict)\n",
    "    loss_factory: torch.nn.modules.loss._Loss = serializable_field(\n",
    "        default_factory=lambda: torch.nn.CrossEntropyLoss,\n",
    "        serialization_fn=lambda x: x.__name__,\n",
    "        loading_fn=lambda x: getattr(torch.nn, x[\"loss_factory\"]),\n",
    "    )\n",
    "    loss = property(lambda self: self.loss_factory(**self.loss_kwargs))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, define your model class. It should be a subclass of `ConfiguredModel`, and use the `set_config_class` decorator to associate it with your configuration class. The `__init__` method should take a single argument, which is an instance of your configuration class. You must also call the superclass `__init__` method with the configuration instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "@set_config_class(MyNNConfig)\n",
    "class MyNN(ConfiguredModel[MyNNConfig]):\n",
    "    def __init__(self, config: MyNNConfig):\n",
    "        # call the superclass init!\n",
    "        # this will store the model in the zanj_model_config field\n",
    "        super().__init__(config)\n",
    "\n",
    "        # whatever you want here\n",
    "        self.net = torch.nn.Sequential(\n",
    "            torch.nn.Linear(config.input_dim, config.hidden_dim),\n",
    "            config.act_fn(),\n",
    "            torch.nn.Linear(config.hidden_dim, config.output_dim),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.net(x)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can now create instances of your model, save them to disk, and load them back into memory:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/miv/projects/tools/ZANJ/.venv/lib/python3.13/site-packages/muutils/json_serialize/serializable_dataclass.py:857: UserWarning: Type mismatch in fields of MyNNConfig:\n",
      "act_fn:\texpected cls_type_hints[k] = <class 'torch.nn.modules.module.Module'>, but got value getattr(output, k) = <class 'torch.nn.modules.activation.ReLU'>, type(getattr(output, k)) = <class 'type'>\n",
      "loss_factory:\texpected cls_type_hints[k] = <class 'torch.nn.modules.loss._Loss'>, but got value getattr(output, k) = <class 'torch.nn.modules.loss.CrossEntropyLoss'>, type(getattr(output, k)) = <class 'type'>\n",
      "  on_typecheck_mismatch.process(\n"
     ]
    }
   ],
   "source": [
    "config = MyNNConfig(\n",
    "    input_dim=10,\n",
    "    hidden_dim=20,\n",
    "    output_dim=2,\n",
    "    act_fn=torch.nn.ReLU,\n",
    "    loss_kwargs=dict(reduction=\"mean\"),\n",
    ")\n",
    "\n",
    "# create your model from the config, and save\n",
    "model = MyNN(config)\n",
    "fname = \"tests/junk_data/path_to_save_model.zanj\"\n",
    "ZANJ().save(model, fname)\n",
    "# load by calling the class method `read()`\n",
    "loaded_model = MyNN.read(fname)\n",
    "# zanj will actually infer the type of the object in the file\n",
    "# -- and will warn you if you don't have the correct package installed\n",
    "loaded_another_way = ZANJ().read(fname)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "type(loaded_model) = <class '__main__.MyNN'>\n",
      "x.shape = torch.Size([10])\n",
      "out_1 = tensor([ 0.0295, -0.4743], grad_fn=<ViewBackward0>), out_2 = tensor([ 0.0295, -0.4743], grad_fn=<ViewBackward0>), out_3 = tensor([ 0.0295, -0.4743], grad_fn=<ViewBackward0>)\n"
     ]
    }
   ],
   "source": [
    "print(f\"{type(loaded_model) = }\")\n",
    "x = torch.randn(config.input_dim)\n",
    "print(f\"{x.shape = }\")\n",
    "out_1 = model(x)\n",
    "out_2 = loaded_model(x)\n",
    "out_3 = loaded_another_way(x)\n",
    "\n",
    "print(f\"{out_1 = }, {out_2 = }, {out_3 = }\")\n",
    "assert torch.allclose(out_1, out_2)\n",
    "assert torch.allclose(out_1, out_3)"
   ]
  }
 ],
 "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"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
