{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Biochar Structure Visualization Demo\n",
    "\n",
    "This notebook demonstrates generating and visualizing biochar molecular structures.\n",
    "\n",
    "**Requirements (in addition to `biochar`):**\n",
    "```bash\n",
    "conda install -c conda-forge py3dmol jupyter matplotlib pandas\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "import py3Dmol\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom rdkit.Chem import Draw, MolToMolBlock\nfrom rdkit.Chem.Draw import rdMolDraw2D\nfrom IPython.display import display, SVG\n\nfrom biochar import generate_biochar, generate_surface, generate_biochar_series"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Helper: 3D viewer\n",
    "\n",
    "A small utility that builds a py3Dmol view from an RDKit mol + coords array.\n",
    "Using the coords array (rather than the mol's internal conformer) ensures the\n",
    "viewer always reflects the final post-processed positions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def coords_to_xyz(mol, coords):\n",
    "    \"\"\"Build an XYZ-format string from an RDKit mol and a coords array (Å).\"\"\"\n",
    "    lines = [str(mol.GetNumAtoms()), \"\"]\n",
    "    for atom, (x, y, z) in zip(mol.GetAtoms(), coords):\n",
    "        lines.append(f\"{atom.GetSymbol()}  {x:.4f}  {y:.4f}  {z:.4f}\")\n",
    "    return \"\\n\".join(lines)\n",
    "\n",
    "\n",
    "def view3d(mol, coords, width=700, height=450, style=\"stick\"):\n",
    "    \"\"\"Display an interactive 3D structure.\"\"\"\n",
    "    v = py3Dmol.view(width=width, height=height)\n",
    "    v.addModel(coords_to_xyz(mol, coords), \"xyz\")\n",
    "    if style == \"stick\":\n",
    "        v.setStyle({\"stick\": {\"radius\": 0.12},\n",
    "                    \"sphere\": {\"scale\": 0.25}})\n",
    "    elif style == \"sphere\":\n",
    "        v.setStyle({\"sphere\": {\"colorscheme\": \"Jmol\", \"scale\": 0.4}})\n",
    "    # Colour oxygen red so functional groups stand out\n",
    "    v.setStyle({\"elem\": \"O\"}, {\"sphere\": {\"color\": \"red\", \"scale\": 0.45},\n",
    "                                \"stick\": {\"color\": \"red\", \"radius\": 0.15}})\n",
    "    v.zoomTo()\n",
    "    return v"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1  Single molecule — 2D structure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mol, coords, gro, top, itp = generate_biochar(\n",
    "    target_num_carbons=50,\n",
    "    H_C_ratio=0.5,\n",
    "    O_C_ratio=0.1,\n",
    "    functional_groups={\"phenolic\": 3, \"ether\": 1},\n",
    "    output_directory=\"output/demo\",\n",
    "    basename=\"bc50\",\n",
    "    seed=42,\n",
    ")\n",
    "\n",
    "print(f\"{mol.GetNumAtoms()} atoms — {mol.GetNumHeavyAtoms()} heavy\")\n",
    "\n",
    "# 2D depiction (kekulized for clarity)\n",
    "drawer = rdMolDraw2D.MolDraw2DSVG(600, 380)\n",
    "drawer.drawOptions().addAtomIndices = False\n",
    "drawer.DrawMolecule(mol)\n",
    "drawer.FinishDrawing()\n",
    "display(SVG(drawer.GetDrawingText()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2  Single molecule — interactive 3D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "view3d(mol, coords).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3  Functional group comparison\n",
    "\n",
    "Side-by-side 2D depictions of the same carbon skeleton with different oxygen\n",
    "functional groups."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fg_configs = [\n",
    "    (\"Phenolic (–OH)\",  {\"phenolic\": 4}),\n",
    "    (\"Carboxyl (–COOH)\", {\"carboxyl\": 2}),\n",
    "    (\"Ether (–O–)\",     {\"ether\": 3}),\n",
    "    (\"Mixed\",           {\"phenolic\": 2, \"carboxyl\": 1, \"ether\": 1}),\n",
    "]\n",
    "\n",
    "mols, labels = [], []\n",
    "for label, fg in fg_configs:\n",
    "    m, *_ = generate_biochar(\n",
    "        target_num_carbons=38,\n",
    "        functional_groups=fg,\n",
    "        output_directory=\"output/demo/fg\",\n",
    "        basename=label.split()[0].lower(),\n",
    "        seed=7,\n",
    "    )\n",
    "    mols.append(m)\n",
    "    labels.append(label)\n",
    "\n",
    "img = Draw.MolsToGridImage(mols, molsPerRow=2, subImgSize=(500, 320),\n",
    "                           legends=labels)\n",
    "display(img)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4  Pentagon ring defects\n",
    "\n",
    "Topological disorder increases with `defect_fraction`; the structures become\n",
    "non-planar at higher fractions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "defect_mols, defect_coords, defect_labels = [], [], []\n",
    "for frac in [0.0, 0.10, 0.20]:\n",
    "    m, c, *_ = generate_biochar(\n",
    "        target_num_carbons=60,\n",
    "        defect_fraction=frac,\n",
    "        output_directory=\"output/demo/defects\",\n",
    "        basename=f\"defect_{int(frac*100):02d}pct\",\n",
    "        seed=99,\n",
    "    )\n",
    "    defect_mols.append(m)\n",
    "    defect_coords.append(c)\n",
    "    defect_labels.append(f\"defect_fraction={frac:.2f}\")\n",
    "\n",
    "display(Draw.MolsToGridImage(defect_mols, molsPerRow=3, subImgSize=(420, 280),\n",
    "                              legends=defect_labels))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 3D view of the 20% defect structure — notice the out-of-plane curvature\n",
    "view3d(defect_mols[2], defect_coords[2]).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5  Slit-pore surface\n",
    "\n",
    "Two parallel graphene-like sheets separated by a user-defined pore diameter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sheets, gro_surf, top_surf, itp_surf = generate_surface(\n",
    "    target_num_carbons=40,\n",
    "    functional_groups={\"phenolic\": 2, \"ether\": 1},\n",
    "    pore_diameter=10.0,\n",
    "    output_directory=\"output/demo/surface\",\n",
    "    basename=\"slit_10A\",\n",
    "    seed=42,\n",
    ")\n",
    "\n",
    "print(f\"Generated {len(sheets)} sheets\")\n",
    "for i, s in enumerate(sheets):\n",
    "    zmin, zmax = s.coords[:, 2].min(), s.coords[:, 2].max()\n",
    "    print(f\"  Sheet {i+1}: {s.mol.GetNumAtoms()} atoms, \"\n",
    "          f\"z ∈ [{zmin:.1f}, {zmax:.1f}] Å\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualise both sheets in one viewer, coloured differently\n",
    "SHEET_COLOURS = [\"#4C8EDA\", \"#E07B39\"]  # blue / orange\n",
    "\n",
    "v = py3Dmol.view(width=750, height=500)\n",
    "for i, sheet in enumerate(sheets):\n",
    "    v.addModel(coords_to_xyz(sheet.mol, sheet.coords), \"xyz\")\n",
    "    v.setStyle(\n",
    "        {\"model\": i},\n",
    "        {\"stick\": {\"color\": SHEET_COLOURS[i], \"radius\": 0.12},\n",
    "         \"sphere\": {\"color\": SHEET_COLOURS[i], \"scale\": 0.25}},\n",
    "    )\n",
    "    # Highlight oxygens in red regardless of sheet colour\n",
    "    v.setStyle(\n",
    "        {\"model\": i, \"elem\": \"O\"},\n",
    "        {\"sphere\": {\"color\": \"red\", \"scale\": 0.45},\n",
    "         \"stick\": {\"color\": \"red\", \"radius\": 0.15}},\n",
    "    )\n",
    "v.zoomTo()\n",
    "v.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6  Temperature series — van Krevelen diagram\n",
    "\n",
    "A van Krevelen plot (H/C vs O/C) is the standard way to characterise biochar\n",
    "across pyrolysis temperatures. Lower temperature → upper-right; higher\n",
    "temperature → lower-left."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "from biochar import BiocharGenerator, GeneratorConfig\n\ntemp_configs = [\n    {\"label\": \"300–400 °C\", \"target_num_carbons\": 70,  \"H_C_ratio\": 0.70, \"O_C_ratio\": 0.22, \"seed\": 1},\n    {\"label\": \"400–500 °C\", \"target_num_carbons\": 80,  \"H_C_ratio\": 0.60, \"O_C_ratio\": 0.16, \"seed\": 2},\n    {\"label\": \"500–600 °C\", \"target_num_carbons\": 90,  \"H_C_ratio\": 0.50, \"O_C_ratio\": 0.11, \"seed\": 3},\n    {\"label\": \"600–700 °C\", \"target_num_carbons\": 100, \"H_C_ratio\": 0.38, \"O_C_ratio\": 0.06, \"seed\": 4},\n    {\"label\": \"700–800 °C\", \"target_num_carbons\": 110, \"H_C_ratio\": 0.28, \"O_C_ratio\": 0.03, \"seed\": 5},\n]\n\nresults = []\nfor cfg in temp_configs:\n    gen = BiocharGenerator(GeneratorConfig(\n        target_num_carbons=cfg[\"target_num_carbons\"],\n        H_C_ratio=cfg[\"H_C_ratio\"],\n        O_C_ratio=cfg[\"O_C_ratio\"],\n        seed=cfg[\"seed\"],\n    ))\n    _, _, comp = gen.generate()\n    results.append({\n        \"label\": cfg[\"label\"],\n        \"H_C_target\": cfg[\"H_C_ratio\"],\n        \"O_C_target\": cfg[\"O_C_ratio\"],\n        \"H_C_actual\": comp.H_C_ratio,\n        \"O_C_actual\": comp.O_C_ratio,\n        \"n_C\": comp.num_carbons,\n    })"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(7, 5))\n",
    "\n",
    "cmap = plt.cm.plasma\n",
    "colours = [cmap(i / (len(results) - 1)) for i in range(len(results))]\n",
    "\n",
    "for r, c in zip(results, colours):\n",
    "    ax.scatter(r[\"O_C_actual\"], r[\"H_C_actual\"], s=120, color=c,\n",
    "               edgecolors=\"k\", linewidths=0.6, zorder=3)\n",
    "    ax.annotate(r[\"label\"], (r[\"O_C_actual\"], r[\"H_C_actual\"]),\n",
    "                textcoords=\"offset points\", xytext=(8, 4), fontsize=9)\n",
    "\n",
    "# Target values as ×\n",
    "for r, c in zip(results, colours):\n",
    "    ax.scatter(r[\"O_C_target\"], r[\"H_C_target\"], marker=\"x\", s=60,\n",
    "               color=c, linewidths=1.5, zorder=2)\n",
    "\n",
    "ax.set_xlabel(\"O/C atomic ratio\", fontsize=12)\n",
    "ax.set_ylabel(\"H/C atomic ratio\", fontsize=12)\n",
    "ax.set_title(\"Van Krevelen diagram — simulated biochar series\", fontsize=13)\n",
    "ax.legend(\n",
    "    handles=[\n",
    "        plt.scatter([], [], s=80, c=\"gray\", edgecolors=\"k\", label=\"Actual\"),\n",
    "        plt.scatter([], [], marker=\"x\", s=60, c=\"gray\", label=\"Target\"),\n",
    "    ],\n",
    "    fontsize=9,\n",
    ")\n",
    "ax.grid(True, alpha=0.3)\n",
    "ax.set_xlim(-0.02, 0.30)\n",
    "ax.set_ylim(0.0, 0.90)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7  Asymmetric slit-pore surface\n",
    "\n",
    "Different chemistry on each wall — e.g. one oxidised wall and one clean wall."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "asym_sheets, *_ = generate_surface(\n",
    "    pore_diameter=12.0,\n",
    "    sheet_overrides=[\n",
    "        {\"target_num_carbons\": 38, \"functional_groups\": {\"phenolic\": 4, \"carboxyl\": 1}, \"seed\": 10},\n",
    "        {\"target_num_carbons\": 38, \"functional_groups\": {\"ether\": 2}, \"seed\": 11},\n",
    "    ],\n",
    "    output_directory=\"output/demo/asymmetric\",\n",
    "    basename=\"asym_12A\",\n",
    ")\n",
    "\n",
    "# 2D depictions of each wall\n",
    "wall_mols   = [s.mol for s in asym_sheets]\n",
    "wall_labels = [\"Wall 1: phenolic + carboxyl\", \"Wall 2: ether\"]\n",
    "display(Draw.MolsToGridImage(wall_mols, molsPerRow=2, subImgSize=(500, 320),\n",
    "                              legends=wall_labels))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v2 = py3Dmol.view(width=750, height=500)\n",
    "for i, sheet in enumerate(asym_sheets):\n",
    "    v2.addModel(coords_to_xyz(sheet.mol, sheet.coords), \"xyz\")\n",
    "    v2.setStyle(\n",
    "        {\"model\": i},\n",
    "        {\"stick\": {\"color\": SHEET_COLOURS[i], \"radius\": 0.12},\n",
    "         \"sphere\": {\"color\": SHEET_COLOURS[i], \"scale\": 0.25}},\n",
    "    )\n",
    "    v2.setStyle(\n",
    "        {\"model\": i, \"elem\": \"O\"},\n",
    "        {\"sphere\": {\"color\": \"red\", \"scale\": 0.45},\n",
    "         \"stick\": {\"color\": \"red\", \"radius\": 0.15}},\n",
    "    )\n",
    "v2.zoomTo()\n",
    "v2.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}