{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "25cbbf6d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8d4ddcba8b47413ab08f7bbd1f096fa3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(Dropdown(description='Filter:', options=('Semiconductors', 'Insulators', 'Metals'), value='Semi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from IPython.display import clear_output, display\n",
    "import ipywidgets as widgets\n",
    "from datetime import datetime\n",
    "\n",
    "# Constants\n",
    "h = 6.626e-34\n",
    "m_e = 9.11e-31\n",
    "\n",
    "# Material categories\n",
    "semiconductors = {\n",
    "    \"Si\": {\"bandgap\": 1.12, \"E_v\": 0.0, \"E_c\": 1.12, \"m_n\": 0.26 * m_e, \"m_p\": 0.39 * m_e},\n",
    "    \"GaAs\": {\"bandgap\": 1.43, \"E_v\": 0.0, \"E_c\": 1.43, \"m_n\": 0.067 * m_e, \"m_p\": 0.48 * m_e},\n",
    "    \"Ge\": {\"bandgap\": 0.66, \"E_v\": 0.0, \"E_c\": 0.66, \"m_n\": 0.55 * m_e, \"m_p\": 0.37 * m_e},\n",
    "    \"InP\": {\"bandgap\": 1.35, \"E_v\": 0.0, \"E_c\": 1.35, \"m_n\": 0.08 * m_e, \"m_p\": 0.6 * m_e},\n",
    "    \"CdTe\": {\"bandgap\": 1.5, \"E_v\": 0.0, \"E_c\": 1.5, \"m_n\": 0.11 * m_e, \"m_p\": 0.35 * m_e},\n",
    "    \"ZnSe\": {\"bandgap\": 2.7, \"E_v\": 0.0, \"E_c\": 2.7, \"m_n\": 0.16 * m_e, \"m_p\": 0.6 * m_e},\n",
    "    \"GaN\": {\"bandgap\": 3.4, \"E_v\": 0.0, \"E_c\": 3.4, \"m_n\": 0.2 * m_e, \"m_p\": 0.6 * m_e},\n",
    "    \"AlN\": {\"bandgap\": 6.2, \"E_v\": 0.0, \"E_c\": 6.2, \"m_n\": 0.4 * m_e, \"m_p\": 0.8 * m_e},\n",
    "    \"InAs\": {\"bandgap\": 0.36, \"E_v\": 0.0, \"E_c\": 0.36, \"m_n\": 0.023 * m_e, \"m_p\": 0.41 * m_e},\n",
    "    \"InSb\": {\"bandgap\": 0.17, \"E_v\": 0.0, \"E_c\": 0.17, \"m_n\": 0.014 * m_e, \"m_p\": 0.45 * m_e},\n",
    "    \"ZnO\": {\"bandgap\": 3.37, \"E_v\": 0.0, \"E_c\": 3.37, \"m_n\": 0.24 * m_e, \"m_p\": 0.59 * m_e},\n",
    "    \"CdS\": {\"bandgap\": 2.42, \"E_v\": 0.0, \"E_c\": 2.42, \"m_n\": 0.21 * m_e, \"m_p\": 0.8 * m_e},\n",
    "}\n",
    "\n",
    "insulators = {\n",
    "    \"SiO2\": {\"bandgap\": 9.0, \"E_v\": 0.0, \"E_c\": 9.0, \"m_n\": 0.5 * m_e, \"m_p\": 0.5 * m_e},\n",
    "    \"Al2O3\": {\"bandgap\": 8.8, \"E_v\": 0.0, \"E_c\": 8.8, \"m_n\": 0.5 * m_e, \"m_p\": 0.5 * m_e},\n",
    "    \"HfO2\": {\"bandgap\": 5.6, \"E_v\": 0.0, \"E_c\": 5.6, \"m_n\": 0.5 * m_e, \"m_p\": 0.5 * m_e},\n",
    "    \n",
    "}\n",
    "\n",
    "metals = {\n",
    "    \"Metals\": {\"bandgap\": 0.0, \"E_v\": 0.0, \"E_c\": 0.0, \"m_n\": 1.0 * m_e, \"m_p\": 1.0 * m_e},\n",
    "}\n",
    "\n",
    "all_materials = {**semiconductors, **insulators, **metals}\n",
    "\n",
    "categories = {\n",
    "    \"Semiconductors\": semiconductors,\n",
    "    \"Insulators\": insulators,\n",
    "    \"Metals\": metals\n",
    "}\n",
    "\n",
    "# Generate DoS\n",
    "def D(E_v, E_c, m_n, m_p):\n",
    "    E_above = np.linspace(E_c, E_c + 1, 1000)\n",
    "    D_c = 8 * np.pi * m_n * np.sqrt(2 * m_n * (E_above - E_c)) / h**3\n",
    "    E_below = np.linspace(E_v - 1, E_v, 1000)\n",
    "    D_v = 8 * np.pi * m_p * np.sqrt(2 * m_p * (E_v - E_below)) / h**3\n",
    "    return D_c, D_v, E_above, E_below\n",
    "\n",
    "# Store all checkboxes\n",
    "checkboxes = {}\n",
    "for material in all_materials:\n",
    "    default_checked = material in list(semiconductors.keys())[:3]\n",
    "    checkboxes[material] = widgets.Checkbox(value=default_checked, description=material)\n",
    "\n",
    "# UI elements\n",
    "category_dropdown = widgets.Dropdown(\n",
    "    options=list(categories.keys()),\n",
    "    value=\"Semiconductors\",\n",
    "    description=\"Filter:\"\n",
    ")\n",
    "\n",
    "plot_button = widgets.Button(description=\"Plot Selected\")\n",
    "clear_button = widgets.Button(description=\"Clear Output\", button_style='danger')\n",
    "checkbox_container = widgets.VBox()\n",
    "\n",
    "def update_checkbox_display(category):\n",
    "    \"\"\"Display checkboxes from the selected category.\"\"\"\n",
    "    materials_in_cat = categories[category].keys()\n",
    "    checkbox_container.children = [checkboxes[mat] for mat in materials_in_cat]\n",
    "\n",
    "# Initial load\n",
    "update_checkbox_display(category_dropdown.value)\n",
    "\n",
    "# Plot handler\n",
    "def on_plot_click(b):\n",
    "    clear_output(wait=True)\n",
    "    display(ui)  # Redisplay UI after clearing\n",
    "    \n",
    "    fig, ax = plt.subplots()\n",
    "\n",
    "    for mat, cb in checkboxes.items():\n",
    "        if cb.value:\n",
    "            props = all_materials[mat]\n",
    "            D_c, D_v, E_above, E_below = D(props[\"E_v\"], props[\"E_c\"], props[\"m_n\"], props[\"m_p\"])\n",
    "            color = ax._get_lines.get_next_color()\n",
    "            ax.plot(D_c, E_above, label=mat, color=color)\n",
    "            ax.plot(D_v, E_below, color=color)\n",
    "            ax.axhline(props[\"E_c\"], color=color, linestyle='--', alpha=0.5)\n",
    "\n",
    "    yticks_pos = [0] + [all_materials[mat][\"E_c\"] for mat, cb in checkboxes.items() if cb.value]\n",
    "    yticks_lbl = [r'$E_v$'] + [f'$E_c$ ({mat})' for mat, cb in checkboxes.items() if cb.value]\n",
    "    ax.set_yticks(yticks_pos)\n",
    "    ax.set_yticklabels(yticks_lbl)\n",
    "    ax.set_xlabel(r'D(E) [cm$^{-3}$ eV$^{-1}$]')\n",
    "    ax.set_ylabel('E [eV]')\n",
    "    ax.axhline(0, color='black', linestyle='--', alpha=0.5, label=r'$E_v$')\n",
    "    # ax.legend()\n",
    "    ax.legend(borderpad=0.5, framealpha=1, edgecolor='#000', borderaxespad=0, fancybox=False, loc='upper right')\n",
    "    ax.grid(True)\n",
    "    plt.tight_layout()\n",
    "    plt.margins(x=0.0, y=0.0, tight=True)\n",
    "    \n",
    "    # Save last figure to be exported\n",
    "    global last_fig\n",
    "    last_fig = fig\n",
    "    \n",
    "    plt.show()\n",
    "\n",
    "# Export handlers\n",
    "def on_export(format):\n",
    "    if last_fig:\n",
    "        timestamp = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
    "        filename = f\"DoS_plot_{timestamp}.{format}\"\n",
    "        last_fig.savefig(filename, format=format)\n",
    "        print(f\"✅ Exported as: {filename}\")\n",
    "\n",
    "export_png_btn = widgets.Button(description=\"🖼️ Export PNG\")\n",
    "export_pdf_btn = widgets.Button(description=\"📄 Export PDF\")\n",
    "\n",
    "export_png_btn.on_click(lambda b: on_export(\"png\"))\n",
    "export_pdf_btn.on_click(lambda b: on_export(\"pdf\"))\n",
    "\n",
    "# Clear handler\n",
    "def on_clear_click(b):\n",
    "    clear_output(wait=True)\n",
    "    display(ui)\n",
    "\n",
    "# Update category\n",
    "category_dropdown.observe(lambda change: update_checkbox_display(change['new']), names='value')\n",
    "plot_button.on_click(on_plot_click)\n",
    "clear_button.on_click(on_clear_click)\n",
    "\n",
    "# Layout\n",
    "ui = widgets.VBox([\n",
    "    category_dropdown,\n",
    "    checkbox_container,\n",
    "    widgets.HBox([plot_button, clear_button, export_png_btn, export_pdf_btn])\n",
    "])\n",
    "\n",
    "last_fig = None  # Holds last plotted figure for export\n",
    "display(ui)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}