{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Varying Composition Substitutions\n",
    "\n",
    "This guide demonstrates how to generate symmetry-inequivalent structures across multiple compositions in a single operation using `unique_structure_substitutions_by_composition`.\n",
    "\n",
    "## Overview\n",
    "\n",
    "The `unique_structure_substitutions_by_composition` function extends the basic substitution functionality by systematically exploring all possible compositions of your specified species.\n",
    "\n",
    "**Key differences from `unique_structure_substitutions`:**\n",
    "- Input: Provide a **list of species** instead of a fixed `site_distribution` dict\n",
    "- Output: Returns a **dictionary** mapping composition tuples to lists of structures\n",
    "- Explores **all possible compositions** (or a constrained range)\n",
    "\n",
    "This is particularly useful when you want to:\n",
    "- Survey structures across a composition range\n",
    "- Build phase diagrams\n",
    "- Screen materials with variable stoichiometry"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Basic Example: Binary Substitution\n",
    "\n",
    "Let's start with a simple example: a 4-site system where we substitute with two species (A and B)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Created structure with 4 sites\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "from pymatgen.core import Structure, Lattice\n",
    "from bsym.interface.pymatgen import unique_structure_substitutions_by_composition\n",
    "\n",
    "# Create a simple 2×2 square lattice\n",
    "coords = np.array([[0.0, 0.0, 0.0]])\n",
    "atom_list = ['Li']\n",
    "lattice = Lattice.from_parameters(a=1.0, b=1.0, c=1.0, alpha=90, beta=90, gamma=90)\n",
    "parent_structure = Structure(lattice, atom_list, coords) * [2, 2, 1]\n",
    "\n",
    "print(f\"Created structure with {len(parent_structure)} sites\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we'll generate all unique structures for substituting with species A and B:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated structures for 5 different compositions:\n",
      "  Composition (0, 4): 1 unique structure(s)\n",
      "  Composition (4, 0): 1 unique structure(s)\n",
      "  Composition (1, 3): 1 unique structure(s)\n",
      "  Composition (3, 1): 1 unique structure(s)\n",
      "  Composition (2, 2): 2 unique structure(s)\n"
     ]
    }
   ],
   "source": [
    "results = unique_structure_substitutions_by_composition(\n",
    "    parent_structure,\n",
    "    'Li',\n",
    "    ['A', 'B']  # List of species (order matters for composition tuples)\n",
    ")\n",
    "\n",
    "print(f\"Generated structures for {len(results)} different compositions:\")\n",
    "for composition, structures in results.items():\n",
    "    print(f\"  Composition {composition}: {len(structures)} unique structure(s)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Understanding the Output Format\n",
    "\n",
    "The function returns a dictionary where:\n",
    "- **Keys** are composition tuples: `(n_A, n_B)` representing the count of each species\n",
    "- **Values** are lists of Structure objects\n",
    "\n",
    "The order of species in the tuple matches the order in your species list.\n",
    "\n",
    "Let's examine a specific composition:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Composition (2A, 2B) has 2 unique structure(s)\n",
      "\n",
      "Structure 0:\n",
      "Full Formula (A2 B2)\n",
      "Reduced Formula: AB\n",
      "abc   :   2.000000   2.000000   1.000000\n",
      "angles:  90.000000  90.000000  90.000000\n",
      "pbc   :       True       True       True\n",
      "Sites (4)\n",
      "  #  SP      a    b    c\n",
      "---  ----  ---  ---  ---\n",
      "  0  A0+   0    0      0\n",
      "  1  A0+   0    0.5    0\n",
      "  2  B     0.5  0      0\n",
      "  3  B     0.5  0.5    0\n",
      "\n",
      "Degeneracy: 4\n"
     ]
    }
   ],
   "source": [
    "# Access structures with 2 A atoms and 2 B atoms\n",
    "composition_2_2 = results[(2, 2)]\n",
    "\n",
    "print(f\"Composition (2A, 2B) has {len(composition_2_2)} unique structure(s)\")\n",
    "print(f\"\\nStructure 0:\")\n",
    "print(composition_2_2[0])\n",
    "print(f\"\\nDegeneracy: {composition_2_2[0].number_of_equivalent_configurations}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example: Li-Na Binary System\n",
    "\n",
    "Let's look at a more realistic example with actual elements on a larger lattice:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Created structure with 9 sites\n"
     ]
    }
   ],
   "source": [
    "# Create a 3×3 lattice\n",
    "coords = np.array([[0.0, 0.0, 0.0]])\n",
    "atom_list = ['X']  # Placeholder to substitute\n",
    "lattice = Lattice.from_parameters(a=1.0, b=1.0, c=1.0, alpha=90, beta=90, gamma=90)\n",
    "parent_structure = Structure(lattice, atom_list, coords) * [3, 3, 1]\n",
    "\n",
    "print(f\"Created structure with {len(parent_structure)} sites\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of compositions explored: 10\n",
      "\n",
      "Composition summary:\n",
      "  Li0Na9: 1 unique, 1 total configurations\n",
      "  Li1Na8: 1 unique, 9 total configurations\n",
      "  Li2Na7: 2 unique, 36 total configurations\n",
      "  Li3Na6: 4 unique, 84 total configurations\n",
      "  Li4Na5: 5 unique, 126 total configurations\n",
      "  Li5Na4: 5 unique, 126 total configurations\n",
      "  Li6Na3: 4 unique, 84 total configurations\n",
      "  Li7Na2: 2 unique, 36 total configurations\n",
      "  Li8Na1: 1 unique, 9 total configurations\n",
      "  Li9Na0: 1 unique, 1 total configurations\n"
     ]
    }
   ],
   "source": [
    "# Generate Li-Na structures across all compositions\n",
    "li_na_results = unique_structure_substitutions_by_composition(\n",
    "    parent_structure,\n",
    "    'X',\n",
    "    ['Li', 'Na']\n",
    ")\n",
    "\n",
    "print(f\"Number of compositions explored: {len(li_na_results)}\")\n",
    "print(\"\\nComposition summary:\")\n",
    "for composition, structures in sorted(li_na_results.items()):\n",
    "    n_li, n_na = composition\n",
    "    total_configs = sum(s.number_of_equivalent_configurations for s in structures)\n",
    "    print(f\"  Li{n_li}Na{n_na}: {len(structures)} unique, {total_configs} total configurations\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Constraining Compositions with Bounds\n",
    "\n",
    "Often you don't want to explore all possible compositions. The `bounds` parameter lets you constrain the range:\n",
    "\n",
    "**Format:** `{'species_name': (min_count, max_count)}`\n",
    "\n",
    "For example, to explore only lithium-rich compositions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Li-rich compositions (6+ Li):\n",
      "  Li6Na3: 4 unique structures\n",
      "  Li7Na2: 2 unique structures\n",
      "  Li8Na1: 1 unique structures\n",
      "  Li9Na0: 1 unique structures\n"
     ]
    }
   ],
   "source": [
    "# Only compositions with at least 6 Li atoms\n",
    "li_rich_results = unique_structure_substitutions_by_composition(\n",
    "    parent_structure,\n",
    "    'X',\n",
    "    ['Li', 'Na'],\n",
    "    bounds={'Li': (6, 9)}  # 6-9 Li atoms\n",
    ")\n",
    "\n",
    "print(f\"Li-rich compositions (6+ Li):\")\n",
    "for composition in sorted(li_rich_results.keys()):\n",
    "    n_li, n_na = composition\n",
    "    print(f\"  Li{n_li}Na{n_na}: {len(li_rich_results[composition])} unique structures\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can constrain multiple species:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Constrained compositions:\n",
      "  Li5Na4: 5 unique structures\n",
      "  Li6Na3: 4 unique structures\n"
     ]
    }
   ],
   "source": [
    "# Compositions with 3-6 Li and 1-4 Na\n",
    "constrained_results = unique_structure_substitutions_by_composition(\n",
    "    parent_structure,\n",
    "    'X',\n",
    "    ['Li', 'Na'],\n",
    "    bounds={'Li': (3, 6), 'Na': (1, 4)}\n",
    ")\n",
    "\n",
    "print(f\"Constrained compositions:\")\n",
    "for composition in sorted(constrained_results.keys()):\n",
    "    n_li, n_na = composition\n",
    "    print(f\"  Li{n_li}Na{n_na}: {len(constrained_results[composition])} unique structures\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Three-Species Example\n",
    "\n",
    "The function works with any number of species. Here's an example with three:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated 12 compositions\n",
      "\n",
      "Example compositions:\n",
      "  Li0Na2K2: 2 unique structures\n",
      "  Li0Na3K1: 1 unique structures\n",
      "  Li0Na4K0: 1 unique structures\n",
      "  Li1Na1K2: 2 unique structures\n",
      "  Li1Na2K1: 2 unique structures\n"
     ]
    }
   ],
   "source": [
    "# Create a smaller structure for 3-species exploration\n",
    "small_structure = Structure(lattice, ['X'], [[0.0, 0.0, 0.0]]) * [2, 2, 1]\n",
    "\n",
    "# Li-Na-K ternary system\n",
    "ternary_results = unique_structure_substitutions_by_composition(\n",
    "    small_structure,\n",
    "    'X',\n",
    "    ['Li', 'Na', 'K'],\n",
    "    bounds={'Li': (0, 4), 'Na': (0, 4), 'K': (0, 2)}  # Limit K content\n",
    ")\n",
    "\n",
    "print(f\"Generated {len(ternary_results)} compositions\")\n",
    "print(\"\\nExample compositions:\")\n",
    "for i, (composition, structures) in enumerate(sorted(ternary_results.items())[:5]):\n",
    "    n_li, n_na, n_k = composition\n",
    "    print(f\"  Li{n_li}Na{n_na}K{n_k}: {len(structures)} unique structures\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using Progress Bars\n",
    "\n",
    "When running from a terminal or Python script, you can enable progress bars to monitor the computation:\n",
    "```python\n",
    "# Create a 5×5 lattice\n",
    "coords = np.array([[0.0, 0.0, 0.0]])\n",
    "atom_list = ['X']  # Placeholder to substitute\n",
    "lattice = Lattice.from_parameters(a=1.0, b=1.0, c=1.0, alpha=90, beta=90, gamma=90)\n",
    "parent_structure = Structure(lattice, atom_list, coords) * [5, 5, 1]\n",
    "\n",
    "# Enable progress bars\n",
    "results_with_progress = unique_structure_substitutions_by_composition(\n",
    "    parent_structure,\n",
    "    'X',\n",
    "    ['Li', 'Na'],\n",
    "    show_progress=True,\n",
    "    verbose=True\n",
    ")\n",
    "```\n",
    "\n",
    "This produces terminal output showing progress for each composition:\n",
    "```\n",
    "100%|████████████████████████| 1/1 [00:00<00:00, 1971.93 permutations/s, found=1]\n",
    "100%|█████████████████████| 25/25 [00:00<00:00, 50291.41 permutations/s, found=1]\n",
    "100%|██████████████████| 300/300 [00:00<00:00, 115196.48 permutations/s, found=5]\n",
    "100%|███████████████| 2300/2300 [00:00<00:00, 197986.64 permutations/s, found=19]\n",
    "100%|█████████████| 12650/12650 [00:00<00:00, 218314.01 permutations/s, found=88]\n",
    "100%|████████████| 53130/53130 [00:00<00:00, 240692.40 permutations/s, found=309]\n",
    "100%|██████████| 177100/177100 [00:00<00:00, 246178.45 permutations/s, found=975]\n",
    "...\n",
    "```\n",
    "\n",
    "Each progress bar corresponds to one composition being processed, showing the number of permutations evaluated and unique configurations found.\n",
    "\n",
    "**Note:** For Jupyter notebooks, use `show_progress='notebook'` to display interactive progress widgets instead of ASCII bars."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exporting Structures\n",
    "\n",
    "You might want to write structures to files, organised by composition:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Structures can be exported using structure.to(filename='...', fmt='cif')\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "# Example: save structures for each composition\n",
    "# (Commented out to avoid creating files in documentation)\n",
    "\n",
    "# output_dir = 'li_na_structures'\n",
    "# os.makedirs(output_dir, exist_ok=True)\n",
    "\n",
    "# for composition, structures in li_na_results.items():\n",
    "#     n_li, n_na = composition\n",
    "#     comp_dir = os.path.join(output_dir, f'Li{n_li}Na{n_na}')\n",
    "#     os.makedirs(comp_dir, exist_ok=True)\n",
    "#     \n",
    "#     for i, structure in enumerate(structures):\n",
    "#         filename = os.path.join(comp_dir, f'structure_{i}.cif')\n",
    "#         structure.to(filename=filename, fmt='cif')\n",
    "\n",
    "print(\"Structures can be exported using structure.to(filename='...', fmt='cif')\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Key Points\n",
    "\n",
    "- **`unique_structure_substitutions_by_composition`** explores multiple compositions in one call\n",
    "- Returns a **dictionary** with composition tuples as keys\n",
    "- Composition tuple order matches the species list order\n",
    "- Use **`bounds`** to constrain composition ranges\n",
    "- Each structure retains its `number_of_equivalent_configurations` attribute\n",
    "- Efficient for surveying composition space and building phase diagrams\n",
    "\n",
    "## When to Use This vs Fixed Composition\n",
    "\n",
    "**Use `unique_structure_substitutions_by_composition` when:**\n",
    "- You want to explore multiple compositions\n",
    "- Building a phase diagram or composition-property map\n",
    "- You don't know the optimal composition in advance\n",
    "\n",
    "**Use `unique_structure_substitutions` when:**\n",
    "- You know the exact composition you need\n",
    "- Performing sequential substitutions with intermediate analysis\n",
    "- You need to track `full_configuration_degeneracy` through multiple steps\n",
    "\n",
    "## Next Steps\n",
    "\n",
    "- For single-composition substitutions, see [Basic Substitutions](basic_substitutions.ipynb)\n",
    "- For sequential substitutions, see [Fixed Composition Substitutions](fixed_composition.ipynb)\n",
    "- To understand the composition enumeration algorithm, see [Composition Enumeration Theory](../theory/composition_enumeration.md)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}