{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0da926ac",
   "metadata": {},
   "source": [
    "# 04 \u2014 Vision Transformer (ViT-B/16) Transfer Learning\n",
    "\n",
    "Fine-tune a pretrained ViT-B/16 from `timm` for breed classification.\n",
    "\n",
    "### Two-Phase Training\n",
    "1. **Phase 1** (8 epochs): Freeze backbone, train custom head + warmup\n",
    "2. **Phase 2** (17 epochs): Unfreeze last 2 transformer blocks + norm, fine-tune\n",
    "\n",
    "### Requirements\n",
    "```bash\n",
    "pip install timm\n",
    "```\n",
    "\n",
    "### Expected Output\n",
    "- Best checkpoint \u2192 `ml/artifacts/checkpoints/vit_best.pth`\n",
    "- Training curves, confusion matrix, per-class F1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2454542f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys, os\n",
    "from pathlib import Path\n",
    "\n",
    "PROJECT_ROOT = Path(os.getcwd()).resolve()\n",
    "if 'notebooks' in str(PROJECT_ROOT):\n",
    "    PROJECT_ROOT = PROJECT_ROOT.parent.parent\n",
    "sys.path.insert(0, str(PROJECT_ROOT))\n",
    "print(f'Project root: {PROJECT_ROOT}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b800ca4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ml.src.utils.seed import set_seed\n",
    "from ml.src.utils.device import get_device\n",
    "from ml.src.utils.io import load_config\n",
    "from ml.src.utils.manifests import create_dataloaders\n",
    "from ml.src.data.transforms import get_train_transforms, get_eval_transforms\n",
    "from ml.src.models.vit import CattleViT\n",
    "from ml.src.training.trainer import Trainer\n",
    "\n",
    "set_seed(42)\n",
    "device = get_device()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "329b6a3c",
   "metadata": {},
   "source": [
    "## 1. Load Config & Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d50b49b",
   "metadata": {},
   "outputs": [],
   "source": [
    "config = load_config('vit', PROJECT_ROOT / 'ml' / 'configs')\n",
    "config['num_classes'] = 26\n",
    "\n",
    "img_size = config['image']['size']\n",
    "batch_size = config['training']['batch_size']\n",
    "arch = config['model']['architecture']\n",
    "\n",
    "print(f'Backbone: {arch[\"backbone\"]}')\n",
    "print(f'Batch size: {batch_size} (smaller due to ViT memory)')\n",
    "print(f'Warmup: {config[\"training\"].get(\"warmup_epochs\", 0)} epochs')\n",
    "print(f'Epochs: {config[\"training\"][\"num_epochs\"]}')\n",
    "print(f'LR: {config[\"training\"][\"learning_rate\"]}  Fine-tune LR: {config[\"training\"][\"fine_tune_lr\"]}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13bf7c17",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_transforms = get_train_transforms(img_size=img_size)\n",
    "eval_transforms = get_eval_transforms(img_size=img_size)\n",
    "\n",
    "dataloaders = create_dataloaders(\n",
    "    manifests_dir=PROJECT_ROOT / 'ml' / 'artifacts' / 'manifests',\n",
    "    data_root=PROJECT_ROOT / 'Cattle_Resized',\n",
    "    train_transform=train_transforms,\n",
    "    eval_transform=eval_transforms,\n",
    "    batch_size=batch_size,\n",
    "    num_workers=config['data'].get('num_workers', 4),\n",
    ")\n",
    "\n",
    "class_names = dataloaders['train'].dataset.classes\n",
    "print(f'Classes: {len(class_names)}')\n",
    "print(f'Train: {len(dataloaders[\"train\"].dataset)} | Val: {len(dataloaders[\"val\"].dataset)} | Test: {len(dataloaders[\"test\"].dataset)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40dcea94",
   "metadata": {},
   "source": [
    "## 2. Create Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f518445",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = CattleViT.from_config(config)\n",
    "\n",
    "trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "total = sum(p.numel() for p in model.parameters())\n",
    "print(f'Total parameters:     {total:,}')\n",
    "print(f'Trainable (Phase 1):  {trainable:,} ({100*trainable/total:.1f}%)')\n",
    "print(f'Frozen:               {total - trainable:,}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4e643d9",
   "metadata": {},
   "source": [
    "## 3. Two-Phase Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72d9d8d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    config=config,\n",
    "    dataloaders=dataloaders,\n",
    "    class_names=class_names,\n",
    "    device=device,\n",
    "    model_name='vit',\n",
    ")\n",
    "\n",
    "training_summary = trainer.train_with_phase_switch(\n",
    "    phase1_epochs=arch['unfreeze_after_epochs'],\n",
    "    phase2_lr=config['training']['fine_tune_lr'],\n",
    "    unfreeze_layers=arch['unfreeze_layers'],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7102ecaf",
   "metadata": {},
   "source": [
    "## 4. Evaluate on Test Set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "babe7c17",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_results = trainer.evaluate(split='test')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69443b69",
   "metadata": {},
   "source": [
    "## 5. Save Artifacts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36ffc567",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.save_artifacts()\n",
    "\n",
    "print('\\n=== ViT-B/16 Transfer Learning Summary ===')\n",
    "print(f'Best Val Loss:  {training_summary[\"best_val_loss\"]:.4f}')\n",
    "print(f'Best Val Acc:   {training_summary[\"best_val_accuracy\"]:.4f}')\n",
    "print(f'Test Accuracy:  {test_results[\"metrics\"][\"accuracy\"]:.4f}')\n",
    "print(f'Test Macro F1:  {test_results[\"metrics\"][\"macro_f1\"]:.4f}')\n",
    "print(f'Latency:        {test_results[\"latency\"][\"avg_ms\"]:.2f} ms')\n",
    "print(f'Model Size:     {test_results[\"model_size_mb\"]:.2f} MB')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82b3f7ed",
   "metadata": {},
   "source": [
    "## 6. Classification Report"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1cdab10b",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(test_results['metrics']['classification_report'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17e85d9e",
   "metadata": {},
   "source": [
    "---\n",
    "**\u2705 ViT Transfer Learning complete.** Proceed to `05_model_comparison.ipynb`."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Cattle Classifier (Python 3.12)",
   "language": "python",
   "name": "cattle-classifier"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}