# @title
# Install MatterGen into a repo-local Python 3.10 environment and anchor all later paths to this notebook's repo root.
from pathlib import Path
import os
import shutil
import subprocess
import sys

!pip install ase==3.25.0

def find_notebook_root():
    cwd = Path.cwd().resolve()
    search_roots = [cwd, *cwd.parents]
    explicit_roots = [
        Path('$HOME/../content'),
    ]
    for candidate in [*search_roots, *explicit_roots]:
        if (candidate / 'sota-crystals-comparison.ipynb').exists() and (candidate / 'README.md').exists():
            return candidate.resolve()
        nested = candidate / 'diffusion-model-tutorial'
        if (nested / 'sota-crystals-comparison.ipynb').exists() and (nested / 'README.md').exists():
            return nested.resolve()
    raise FileNotFoundError('Could not locate the diffusion-model-tutorial repo root from the current working directory.')


# Clone the diffusion-model-tutorial repository if it doesn't exist
repo_name = 'diffusion-model-tutorial'
repo_path_in_content = Path('/content') / repo_name
if not repo_path_in_content.exists():
    print(f'Cloning {repo_name} into /content/...')
    subprocess.run(['git', 'clone', 'https://github.com/Neutrino155/diffusion-model-tutorial.git', str(repo_path_in_content)], check=True)
    print(f'{repo_name} cloned successfully.')
else:
    print(f'Repository {repo_name} already exists in {repo_path_in_content}. Skipping clone.')


NOTEBOOK_ROOT = find_notebook_root()
os.chdir(NOTEBOOK_ROOT)
print('Notebook root:', NOTEBOOK_ROOT)

repo_dir = NOTEBOOK_ROOT / 'mattergen_repo'
venv_dir = repo_dir / '.venv'
MATTERGEN_REMOTE = 'https://github.com/microsoft/mattergen.git'
MATTERGEN_COMMIT = 'a245cf2b7538eea6d873e6430b0e30c56d26c60e'
cache_dir = Path('/tmp/uv-cache')
python_dir = Path('/tmp/uv-python')
cache_dir.mkdir(parents=True, exist_ok=True)
python_dir.mkdir(parents=True, exist_ok=True)

env = os.environ.copy()
env['UV_CACHE_DIR'] = str(cache_dir)
env['UV_PYTHON_INSTALL_DIR'] = str(python_dir)

uv_bin = shutil.which('uv')
if uv_bin is None:
    subprocess.run([sys.executable, '-m', 'pip', 'install', 'uv'], check=True)
    uv_bin = shutil.which('uv') or str(Path.home() / '.local/bin/uv')

if not repo_dir.exists():
    subprocess.run([
        'git',
        'clone',
        MATTERGEN_REMOTE,
        str(repo_dir),
    ], check=True)


def repo_has_commit(path: Path, commit: str) -> bool:
    return subprocess.run(
        ['git', '-C', str(path), 'cat-file', '-e', f'{commit}^{{commit}}'],
        check=False,
        capture_output=True,
    ).returncode == 0


if not repo_has_commit(repo_dir, MATTERGEN_COMMIT):
    subprocess.run([
        'git', '-C', str(repo_dir), 'fetch', '--depth', '1', 'origin', MATTERGEN_COMMIT
    ], check=True)
subprocess.run(['git', '-C', str(repo_dir), 'switch', '--detach', MATTERGEN_COMMIT], check=True)

if not venv_dir.exists():
    subprocess.run([
        uv_bin,
        'venv',
        str(venv_dir),
        '--python',
        '3.10',
    ], check=True, env=env)

print('Pinned MatterGen commit:', MATTERGEN_COMMIT)
print('Installing MatterGen into:', venv_dir.resolve())
subprocess.run([
    uv_bin,
    'pip',
    'install',
    '--python',
    str(venv_dir / 'bin' / 'python'),
    '-e',
    str(repo_dir),
], check=True, env=env)

subprocess.run([
    str(venv_dir / 'bin' / 'python'),
    '-c',
    'import mattergen; print(mattergen.__file__)',
], check=True, env=env)

# @title
import torch
import platform

print("Python:", platform.python_version())
print("PyTorch:", torch.__version__)
print("CUDA available:", torch.cuda.is_available())
if torch.cuda.is_available():
    print("GPU:", torch.cuda.get_device_name(0))

from pathlib import Path
import os
import shutil
import subprocess

from ase.io import read as ase_read

try:
    NOTEBOOK_ROOT
except NameError:
    NOTEBOOK_ROOT = Path.cwd().resolve()

MATTERGEN_FORCE_RERUN = False
mattergen_bin = (NOTEBOOK_ROOT / 'mattergen_repo' / '.venv' / 'bin' / 'mattergen-generate').resolve()
if not mattergen_bin.exists():
    raise FileNotFoundError('MatterGen CLI not found. Run the install cell first.')


def run_mattergen_generation(
    output_dir,
    *,
    pretrained_name,
    batch_size,
    num_batches,
    record_trajectories=False,
    properties=None,
    guidance=None,
    force=MATTERGEN_FORCE_RERUN,
):
    output_dir = Path(output_dir)
    if not output_dir.is_absolute():
        output_dir = NOTEBOOK_ROOT / output_dir
    output_dir = output_dir.resolve()
    extxyz_path = output_dir / 'generated_crystals.extxyz'
    expected_samples = batch_size * num_batches
    if extxyz_path.exists() and not force:
        try:
            existing_samples = len(ase_read(extxyz_path, index=':'))
        except Exception:
            existing_samples = None
        if existing_samples == expected_samples:
            print(f'Reusing existing MatterGen outputs in {output_dir} ({existing_samples} samples)')
            return output_dir
        print(f'Found {existing_samples} existing samples in {output_dir}; regenerating for {expected_samples} samples.')

    if output_dir.exists():
        shutil.rmtree(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    cmd = [
        str(mattergen_bin),
        str(output_dir),
        f'--pretrained-name={pretrained_name}',
        f'--batch_size={batch_size}',
        f'--num_batches={num_batches}',
        f'--record_trajectories={record_trajectories}',
    ]
    if properties is not None:
        cmd.append(f'--properties_to_condition_on={properties}')
    if guidance is not None:
        cmd.append(f'--diffusion_guidance_factor={guidance}')

    env = os.environ.copy()
    env['MPLCONFIGDIR'] = env.get('MPLCONFIGDIR', '/tmp/matplotlib-cache')
    Path(env['MPLCONFIGDIR']).mkdir(parents=True, exist_ok=True)

    print('Running:', ' '.join(cmd))
    subprocess.run(cmd, check=True, env=env)
    return output_dir


results_dir = run_mattergen_generation(
    'results_unconditional',
    pretrained_name='mattergen_base',
    batch_size=3,
    num_batches=2,
    record_trajectories=True,
)

mattergen_run_info = {
    'unconditional': {
        'display_name': 'Unconditional base',
        'pretrained_name': 'mattergen_base',
        'target_label': 'none',
        'condition_kind': 'unconditional',
        'output_dir': results_dir,
    }
}

print('\nDone. Files in output directory:')
for p in sorted(results_dir.iterdir()):
    print(' ', p.name)

from pathlib import Path

# These are real MatterGen conditional runs using the same fine-tuned band-gap checkpoint
# at two very different target values.
MATTERGEN_CONDITIONAL_SPECS = {
    'low_band_gap': {
        'display_name': 'Low band gap target',
        'pretrained_name': 'dft_band_gap',
        'properties': repr({'dft_band_gap': 0.2}),
        'guidance': 2.0,
        'batch_size': 2,
        'num_batches': 1,
        'target_label': 'dft_band_gap = 0.2 eV',
        'condition_kind': 'scalar target',
    },
    'high_band_gap': {
        'display_name': 'High band gap target',
        'pretrained_name': 'dft_band_gap',
        'properties': repr({'dft_band_gap': 4.0}),
        'guidance': 2.0,
        'batch_size': 2,
        'num_batches': 1,
        'target_label': 'dft_band_gap = 4.0 eV',
        'condition_kind': 'scalar target',
    },
}

for label, spec in MATTERGEN_CONDITIONAL_SPECS.items():
    out_dir = run_mattergen_generation(
        f'results_{label}',
        pretrained_name=spec['pretrained_name'],
        batch_size=spec['batch_size'],
        num_batches=spec['num_batches'],
        record_trajectories=True,
        properties=spec['properties'],
        guidance=spec['guidance'],
    )
    mattergen_run_info[label] = {
        **spec,
        'output_dir': out_dir,
    }

print('\nMatterGen run registry:')
for label, spec in mattergen_run_info.items():
    print(f"- {label}: {spec['display_name']} -> {spec['output_dir']}")

# Keep the notebook analysis stack aligned with the real MatterGen environment.
!python -m pip install --force-reinstall numpy==1.26.4 scipy==1.13.1 ase==3.25.0 matplotlib==3.8.4 pymatgen

from collections import Counter
from pathlib import Path
from html import escape
from statistics import mean

from ase.io import read
from IPython.display import HTML, display


def load_extxyz_frames(run_dir: Path):
    extxyz_path = run_dir / 'generated_crystals.extxyz'
    if not extxyz_path.exists():
        raise FileNotFoundError(f'Could not find MatterGen outputs in {run_dir}')
    return extxyz_path, read(extxyz_path, index=':')


def mattergen_row(label, atoms, local_index):
    volume = float(atoms.get_volume())
    mass = float(atoms.get_masses().sum())
    chemical_symbols = atoms.get_chemical_symbols()
    return {
        'label': label,
        'display_name': mattergen_run_info[label]['display_name'],
        'target_label': mattergen_run_info[label]['target_label'],
        'condition_kind': mattergen_run_info[label]['condition_kind'],
        'local_index': local_index,
        'formula': atoms.get_chemical_formula(),
        'elements': '-'.join(sorted(set(chemical_symbols))),
        'n_sites': len(atoms),
        'volume': volume,
        'density': mass / max(volume, 1e-12),
    }


def render_summary_table(rows):
    headers = ['run', 'target', 'samples', 'mean_n_sites', 'mean_volume', 'mean_density', 'top_formula']
    th_style = 'border:1px solid #ddd; padding:4px 8px; text-align:left; background:#eff6ff;'
    td_style = 'border:1px solid #ddd; padding:4px 8px; text-align:left;'
    body = []
    for row in rows:
        body.append(
            '<tr>'
            + ''.join(f'<td style="{td_style}">{escape(str(row[h]))}</td>' for h in headers)
            + '</tr>'
        )
    html = f'''
    <table style="border-collapse:collapse; min-width:760px; margin:0.75rem 0 1rem 0;">
      <thead><tr>{''.join(f'<th style="{th_style}">{escape(h)}</th>' for h in headers)}</tr></thead>
      <tbody>{''.join(body)}</tbody>
    </table>
    '''
    display(HTML(html))


mattergen_atoms = {}
mattergen_rows = []
summary_rows = []

for label, info in mattergen_run_info.items():
    extxyz_path, atoms_list = load_extxyz_frames(Path(info['output_dir']))
    mattergen_atoms[label] = atoms_list
    run_rows = [mattergen_row(label, atoms, i) for i, atoms in enumerate(atoms_list)]
    mattergen_rows.extend(run_rows)

    formula_counter = Counter(row['formula'] for row in run_rows)
    summary_rows.append({
        'run': info['display_name'],
        'target': info['target_label'],
        'samples': len(run_rows),
        'mean_n_sites': f"{mean(row['n_sites'] for row in run_rows):.2f}",
        'mean_volume': f"{mean(row['volume'] for row in run_rows):.2f}",
        'mean_density': f"{mean(row['density'] for row in run_rows):.3f}",
        'top_formula': formula_counter.most_common(1)[0][0],
    })
    print(f'Loaded {len(run_rows)} structures for {info["display_name"]} from {extxyz_path}')

print('\nMatterGen run summary:')
render_summary_table(summary_rows)
print('Reminder: the conditional target here is band gap, while density and volume are downstream structural diagnostics.')

import matplotlib.pyplot as plt
import numpy as np
from html import escape
from IPython.display import HTML, display
from ase.visualize.plot import plot_atoms

palette = {
    'unconditional': '#4c78a8',
    'low_band_gap': '#54a24b',
    'high_band_gap': '#e45756',
}
label_order = list(mattergen_run_info.keys())


def rows_for(label):
    return [row for row in mattergen_rows if row['label'] == label]


def render_rank_table(title, rows, headers):
    th_style = 'border:1px solid #ddd; padding:4px 8px; text-align:left; background:#eff6ff;'
    td_style = 'border:1px solid #ddd; padding:4px 8px; text-align:left;'

    def fmt(value):
        if isinstance(value, float):
            return f'{value:.3f}'
        return str(value)

    body = []
    for row in rows:
        body.append(
            '<tr>'
            + ''.join(f'<td style="{td_style}">{escape(fmt(row[h]))}</td>' for h in headers)
            + '</tr>'
        )
    html = f'''<div style="margin:0.5rem 0 0.25rem 0; font-weight:600;">{escape(title)}</div>
    <table style="border-collapse:collapse; min-width:760px; margin:0.35rem 0 1rem 0;">
      <thead><tr>{''.join(f'<th style="{th_style}">{escape(h)}</th>' for h in headers)}</tr></thead>
      <tbody>{''.join(body)}</tbody>
    </table>'''
    display(HTML(html))


fig, axes = plt.subplots(2, 3, figsize=(16, 10), facecolor='white')
rng = np.random.default_rng(0)
for x, label in enumerate(label_order, start=1):
    rows = rows_for(label)
    densities = [row['density'] for row in rows]
    volumes = [row['volume'] for row in rows]
    n_sites = [row['n_sites'] for row in rows]
    jitter = rng.normal(0.0, 0.04, size=len(rows)) if rows else np.array([])

    axes[0, 0].scatter(np.full(len(rows), x) + jitter, densities, s=60, alpha=0.85, color=palette.get(label, '#666666'), edgecolors='black', linewidths=0.3)
    axes[0, 1].scatter(np.full(len(rows), x) + jitter, volumes, s=60, alpha=0.85, color=palette.get(label, '#666666'), edgecolors='black', linewidths=0.3)
    axes[1, 2].bar(x, len({row['formula'] for row in rows}), color=palette.get(label, '#666666'), width=0.6)
    axes[1, 0].scatter(volumes, densities, s=90, alpha=0.85, label=mattergen_run_info[label]['display_name'], color=palette.get(label, '#666666'), edgecolors='black', linewidths=0.3)
    axes[0, 2].hist(densities, bins=4, alpha=0.55, color=palette.get(label, '#666666'), label=mattergen_run_info[label]['display_name'])
    axes[1, 1].hist(n_sites, bins=np.arange(min(n_sites) - 0.5, max(n_sites) + 1.5, 1.0) if n_sites else np.arange(0.5, 1.5), alpha=0.5, color=palette.get(label, '#666666'), label=mattergen_run_info[label]['display_name'])

for ax, ylabel, title in [
    (axes[0, 0], 'density (mass / ų proxy)', 'Density by MatterGen run'),
    (axes[0, 1], 'volume (ų)', 'Volume by MatterGen run'),
]:
    ax.set_xticks(range(1, len(label_order) + 1))
    ax.set_xticklabels([mattergen_run_info[label]['display_name'] for label in label_order], rotation=15, ha='right')
    ax.set_ylabel(ylabel)
    ax.set_title(title)

axes[0, 2].set_title('Density distributions across unconditional and conditioned runs')
axes[0, 2].set_xlabel('density')
axes[0, 2].set_ylabel('count')
axes[0, 2].legend(frameon=False)

axes[1, 0].set_title('Density vs volume across unconditional and conditioned runs')
axes[1, 0].set_xlabel('volume (ų)')
axes[1, 0].set_ylabel('density')
axes[1, 0].legend(frameon=False)

axes[1, 1].set_title('Atom-count distributions across MatterGen runs')
axes[1, 1].set_xlabel('n_sites')
axes[1, 1].set_ylabel('count')
axes[1, 1].legend(frameon=False)

axes[1, 2].set_xticks(range(1, len(label_order) + 1))
axes[1, 2].set_xticklabels([mattergen_run_info[label]['display_name'] for label in label_order], rotation=15, ha='right')
axes[1, 2].set_ylabel('unique formulas')
axes[1, 2].set_title('Formula diversity per MatterGen run')

plt.tight_layout()
plt.show()

print('Per-run averages:')
for label in label_order:
    rows = rows_for(label)
    print(
        f"- {mattergen_run_info[label]['display_name']}: "
        f"mean density={np.mean([row['density'] for row in rows]):.3f}, "
        f"mean volume={np.mean([row['volume'] for row in rows]):.2f}, "
        f"mean n_sites={np.mean([row['n_sites'] for row in rows]):.2f}"
    )


sorted_by_density = sorted(mattergen_rows, key=lambda row: row['density'])
render_rank_table(
    'Lowest-density MatterGen candidates across all runs',
    sorted_by_density[:4],
    ['display_name', 'target_label', 'formula', 'n_sites', 'volume', 'density'],
)
render_rank_table(
    'Highest-density MatterGen candidates across all runs',
    list(reversed(sorted_by_density[-4:])),
    ['display_name', 'target_label', 'formula', 'n_sites', 'volume', 'density'],
)


def show_mattergen_gallery(selected_rows, title, *, columns=3):
    selected_rows = list(selected_rows)
    if not selected_rows:
        print(f'No MatterGen structures available for {title}')
        return
    cols = min(columns, len(selected_rows))
    nrows = int(np.ceil(len(selected_rows) / cols))
    fig, axes = plt.subplots(nrows, cols, figsize=(4 * cols, 4 * nrows), squeeze=False, facecolor='white')
    for ax in axes.ravel():
        ax.axis('off')
    for ax, row in zip(axes.ravel(), selected_rows):
        atoms = mattergen_atoms[row['label']][row['local_index']]
        plot_atoms(atoms, ax=ax, rotation='20x,30y,0z', radii=0.35, show_unit_cell=2)
        ax.set_title(
            f"{row['display_name']}\n{row['formula']} | density={row['density']:.3f}\nvolume={row['volume']:.1f} ų",
            fontsize=8,
        )
        ax.set_axis_off()
    fig.suptitle(title, fontsize=13)
    plt.tight_layout()
    plt.show()


representative_rows = []
for label in label_order:
    representative_rows.extend(rows_for(label)[:2])
show_mattergen_gallery(
    representative_rows,
    'Representative MatterGen samples across the unconditional and band-gap-conditioned runs',
    columns=3,
)
show_mattergen_gallery(
    sorted_by_density[:3],
    'Lowest-density MatterGen candidates across all runs',
    columns=3,
)
show_mattergen_gallery(
    list(reversed(sorted_by_density[-3:])),
    'Highest-density MatterGen candidates across all runs',
    columns=3,
)

import matplotlib.pyplot as plt
from ase.io import read
from ase.visualize.plot import plot_atoms
from IPython.display import Image, display

preview_dir = (NOTEBOOK_ROOT / 'mattergen_figures').resolve()
preview_dir.mkdir(parents=True, exist_ok=True)

unconditional_extxyz = mattergen_run_info['unconditional']['output_dir'] / 'generated_crystals.extxyz'
if not unconditional_extxyz.exists():
    raise FileNotFoundError(f'Could not find {unconditional_extxyz}. Run the unconditional MatterGen cell first.')

unconditional_atoms = read(unconditional_extxyz, index=':')
selected_sample_index = 0
selected_atoms = unconditional_atoms[selected_sample_index]
preview_path = preview_dir / 'mattergen_unconditional_sample.png'

print(f'Showing unconditional sample {selected_sample_index} from {unconditional_extxyz}')
print('Formula:', selected_atoms.get_chemical_formula())
print('Number of atoms:', len(selected_atoms))
print(f'Volume: {selected_atoms.get_volume():.2f} ų')

fig, axes = plt.subplots(1, 2, figsize=(10, 4), squeeze=False, facecolor='white')
for ax, rotation, title in zip(
    axes[0],
    ['20x,30y,0z', '90x,0y,0z'],
    ['Perspective view', 'Side view'],
):
    plot_atoms(selected_atoms, ax, rotation=rotation, radii=0.35, show_unit_cell=2)
    ax.set_title(title, fontsize=10)
    ax.set_axis_off()
plt.tight_layout()
fig.savefig(preview_path, dpi=180, bbox_inches='tight', facecolor='white')
plt.close(fig)

display(Image(filename=str(preview_path)))

from pathlib import Path
import shutil
import zipfile

import matplotlib.pyplot as plt
import numpy as np
from ase.io import read, write
from ase.visualize.plot import plot_atoms
from IPython.display import Image, display

trajectory_dir = (NOTEBOOK_ROOT / 'unzipped_trajectories').resolve()
if trajectory_dir.exists():
    shutil.rmtree(trajectory_dir)
trajectory_dir.mkdir(parents=True, exist_ok=True)

zip_candidates = [
    mattergen_run_info['unconditional']['output_dir'] / 'generated_trajectories.zip',
    NOTEBOOK_ROOT / 'results_unconditional' / 'generated_trajectories.zip',
]
zip_file_path = next((p for p in zip_candidates if p.exists()), None)
if zip_file_path is None:
    raise FileNotFoundError(
        'Could not find a MatterGen trajectory archive. Run the unconditional generation cell with record_trajectories=True.'
    )

with zipfile.ZipFile(zip_file_path) as zf:
    zf.extractall(trajectory_dir)

trajectory_files = sorted(trajectory_dir.glob('gen_*.extxyz'))
if not trajectory_files:
    trajectory_files = sorted(trajectory_dir.glob('*.extxyz'))
if not trajectory_files:
    raise FileNotFoundError(f'Archive {zip_file_path} did not contain any .extxyz trajectory files.')

selected_trajectory = trajectory_files[0]
trajectory_frames = read(selected_trajectory, index=':')
if len(trajectory_frames) < 2:
    raise ValueError(f'Trajectory file {selected_trajectory.name} did not contain multiple frames.')

frame_ids = np.unique(np.linspace(0, len(trajectory_frames) - 1, num=min(5, len(trajectory_frames)), dtype=int))
preview_path = trajectory_dir / f'{selected_trajectory.stem}_preview.png'
gif_path = trajectory_dir / f'{selected_trajectory.stem}.gif'

fig, axes = plt.subplots(1, len(frame_ids), figsize=(4 * len(frame_ids), 4), squeeze=False, facecolor='white')
for ax, frame_idx in zip(axes[0], frame_ids):
    plot_atoms(trajectory_frames[frame_idx], ax, rotation='20x,30y,0z', radii=0.35, show_unit_cell=2)
    ax.set_title(f'frame {frame_idx + 1}/{len(trajectory_frames)}', fontsize=9)
    ax.set_axis_off()
plt.tight_layout()
fig.savefig(preview_path, dpi=180, bbox_inches='tight', facecolor='white')
plt.close(fig)

write(str(gif_path), trajectory_frames[::max(1, len(trajectory_frames) // 80)])

print(f'Extracted {len(trajectory_files)} trajectory files from {zip_file_path.name}')
print(f'Selected one trajectory: {selected_trajectory.name}')
print(f'Loaded {len(trajectory_frames)} frames from that single reverse-diffusion path')
print('First-frame formula:', trajectory_frames[0].get_chemical_formula())
print('Last-frame formula:', trajectory_frames[-1].get_chemical_formula())

display(Image(filename=str(preview_path)))
display(Image(filename=str(gif_path)))

import os

os.chdir(NOTEBOOK_ROOT)
print('Using notebook root:', os.getcwd())

import json
import os
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
from ase.visualize.plot import plot_atoms

print('Python:', sys.version)
print('Working dir:', os.getcwd())

cwd = Path.cwd().resolve()
candidate_roots = [
    cwd,
    cwd / 'diffusion-model-tutorial',
    Path('/home/uccabaa/OneDrive/Repositories/diffusion-model-tutorial'),
    Path('/mnt/c/Users/keith/OneDrive - University College London/Repositories/diffusion-model-tutorial'),
]
notebook_root = next((root for root in candidate_roots if (root / 'sota-crystals-comparison.ipynb').exists()), cwd)

CHEMELEON_DNG_REPO = (notebook_root / 'chemeleon_dng_repo').resolve()
CHEMELEON_DNG_REMOTE = 'https://github.com/hspark1212/chemeleon-dng.git'
CHEMELEON_DNG_COMMIT = '0d8da3a82a0c2211245a1b1394b599ca0545883c'
CHEMELEON_DNG_BACKUP = (notebook_root / 'chemeleon_dng_repo_incomplete_backup').resolve()
CHEMELEON_DNG_OUT = (notebook_root / 'results_chemeleon_dng').resolve()
CHEMELEON_DNG_OUT.mkdir(parents=True, exist_ok=True)
TMP_ROOT = Path(tempfile.gettempdir()).resolve()
CHEMELEON_DNG_VENV = TMP_ROOT / 'chemeleon-dng-venv'
CHEMELEON_DNG_PYTHON = CHEMELEON_DNG_VENV / 'bin' / 'python'
CHEMELEON_DNG_CLI = CHEMELEON_DNG_VENV / 'bin' / 'chemeleon-dng'
CHEMELEON_DNG_DEVICE = 'cuda'  # change to 'cuda' later if you want faster sampling on a GPU machine
CHEMELEON_DNG_DEMO_TIMESTEPS = 32  # shorter reverse process for notebook-speed demos
uv_cache_dir = TMP_ROOT / 'uv-cache-chemeleon-dng'
uv_cache_dir.mkdir(parents=True, exist_ok=True)
matplotlib_cache_dir = TMP_ROOT / 'matplotlib-cache-chemeleon-dng'
matplotlib_cache_dir.mkdir(parents=True, exist_ok=True)
base_env = os.environ.copy()
base_env['UV_CACHE_DIR'] = str(uv_cache_dir)
base_env['UV_LINK_MODE'] = 'copy'
base_env['PYTHONNOUSERSITE'] = '1'
base_env['MPLBACKEND'] = 'Agg'
base_env['MPLCONFIGDIR'] = str(matplotlib_cache_dir)


def repo_is_valid(path: Path) -> bool:
    return (path / 'pyproject.toml').exists() and (path / '.git').exists()

def repo_has_commit(path: Path, commit: str) -> bool:
    return subprocess.run(
        ['git', '-C', str(path), 'cat-file', '-e', f'{commit}^{{commit}}'],
        check=False,
        capture_output=True,
    ).returncode == 0


uv_bin = shutil.which('uv')
if uv_bin is None:
    subprocess.run([sys.executable, '-m', 'pip', 'install', 'uv'], check=True)
    uv_bin = shutil.which('uv')
if uv_bin is None:
    raise RuntimeError('Could not find uv after installation.')

if CHEMELEON_DNG_REPO.exists() and not repo_is_valid(CHEMELEON_DNG_REPO):
    if CHEMELEON_DNG_BACKUP.exists():
        shutil.rmtree(CHEMELEON_DNG_BACKUP)
    CHEMELEON_DNG_REPO.rename(CHEMELEON_DNG_BACKUP)
    print('Moved incomplete Chemeleon-DNG directory to:', CHEMELEON_DNG_BACKUP)
if not repo_is_valid(CHEMELEON_DNG_REPO):
    subprocess.run([
        'git',
        'clone',
        CHEMELEON_DNG_REMOTE,
        str(CHEMELEON_DNG_REPO),
    ], check=True)
else:
    print('Using existing Chemeleon-DNG clone:', CHEMELEON_DNG_REPO)

if not repo_has_commit(CHEMELEON_DNG_REPO, CHEMELEON_DNG_COMMIT):
    subprocess.run([
        'git', '-C', str(CHEMELEON_DNG_REPO), 'fetch', '--depth', '1', 'origin', CHEMELEON_DNG_COMMIT
    ], check=True)
subprocess.run(['git', '-C', str(CHEMELEON_DNG_REPO), 'switch', '--detach', CHEMELEON_DNG_COMMIT], check=True)


def import_check() -> bool:
    if not CHEMELEON_DNG_PYTHON.exists():
        return False
    result = subprocess.run(
        [
            str(CHEMELEON_DNG_PYTHON),
            '-c',
            'import chemeleon_dng, numpy, torch; print(chemeleon_dng.__file__); print(numpy.__version__); print(torch.__version__)',
        ],
        cwd=str(CHEMELEON_DNG_REPO),
        env=base_env,
        capture_output=True,
        text=True,
    )
    if result.returncode == 0:
        print(result.stdout.strip())
        return True
    if result.stdout:
        print(result.stdout)
    if result.stderr:
        print(result.stderr)
    return False


if not import_check():
    if CHEMELEON_DNG_VENV.exists():
        shutil.rmtree(CHEMELEON_DNG_VENV)
    subprocess.run([uv_bin, 'venv', str(CHEMELEON_DNG_VENV), '--python', '3.11'], check=True, env=base_env)
    subprocess.run([
        uv_bin,
        'pip',
        'install',
        '--python',
        str(CHEMELEON_DNG_PYTHON),
        '-e',
        str(CHEMELEON_DNG_REPO),
    ], check=True, env=base_env)
    if not import_check():
        raise RuntimeError('Chemeleon-DNG environment was created but the import check still failed.')

commit = subprocess.check_output(
    ['git', '-C', str(CHEMELEON_DNG_REPO), 'rev-parse', '--short', 'HEAD'],
    text=True,
).strip()
print('Chemeleon-DNG repo:', CHEMELEON_DNG_REPO)
print('Chemeleon-DNG commit:', commit)
print('Pinned Chemeleon-DNG commit:', CHEMELEON_DNG_COMMIT)
print('Chemeleon-DNG venv:', CHEMELEON_DNG_VENV)
print('Chemeleon-DNG Python:', CHEMELEON_DNG_PYTHON)
print('Chemeleon-DNG CLI:', CHEMELEON_DNG_CLI)
print('Default notebook device:', CHEMELEON_DNG_DEVICE)
print('Notebook demo timesteps:', CHEMELEON_DNG_DEMO_TIMESTEPS)


def run_chemeleon_dng_python(code: str, *, device: str = CHEMELEON_DNG_DEVICE):
    run_env = base_env.copy()
    run_env['MPLBACKEND'] = 'Agg'
    run_env['MPLCONFIGDIR'] = str(matplotlib_cache_dir)
    if device == 'cpu':
        run_env['CUDA_VISIBLE_DEVICES'] = ''
    subprocess.run(
        [str(CHEMELEON_DNG_PYTHON), '-c', code],
        cwd=str(CHEMELEON_DNG_REPO),
        env=run_env,
        check=True,
    )


def sample_chemeleon_dng(
    *,
    task: str,
    output_dir: Path,
    device: str = CHEMELEON_DNG_DEVICE,
    num_samples: int = 2,
    batch_size: int | None = None,
    formulas: list[str] | None = None,
    num_atom_distribution: str | list[int] | None = 'mp-20',
    reuse_existing: bool = True,
    demo_num_timesteps: int = CHEMELEON_DNG_DEMO_TIMESTEPS,
    model_path: str | None = None,
):
    output_dir = Path(output_dir).resolve()
    output_dir.mkdir(parents=True, exist_ok=True)
    expected_count = num_samples if task == 'dng' else num_samples * len(formulas or [])
    existing = sorted(output_dir.glob('sample_*.cif'))
    if reuse_existing and expected_count and len(existing) >= expected_count:
        print(f'Reusing {len(existing)} existing CIF files from {output_dir}')
        return existing[:expected_count]

    for old in output_dir.glob('sample_*.cif'):
        old.unlink()
    json_path = output_dir / 'generated_structures.json.gz'
    if json_path.exists():
        json_path.unlink()

    code_lines = [
        'from pathlib import Path',
        'from chemeleon_dng.diffusion.diffusion_module import DiffusionModule',
        'from chemeleon_dng.sample import DEFAULT_MODEL_PATH, get_checkpoint_path, sample_csp, sample_dng',
        f'task = {json.dumps(task)}',
        f'device = {json.dumps(device)}',
        f'output_dir = Path({json.dumps(str(output_dir))})',
        'output_dir.mkdir(parents=True, exist_ok=True)',
        f'demo_num_timesteps = {int(demo_num_timesteps)}',
        f'batch_size = {int(batch_size if batch_size is not None else num_samples)}',
        f'num_samples = {int(num_samples)}',
        f'model_path = {repr(model_path)}',
        'if model_path is None:',
        '    model_path = get_checkpoint_path(task, DEFAULT_MODEL_PATH)',
        'print(f"Using checkpoint path: {model_path}")',
        'dm = DiffusionModule.load_from_checkpoint(model_path, map_location=device, weights_only=False)',
        'print(f"Original timesteps: {dm.num_timesteps}")',
        'dm.num_timesteps = demo_num_timesteps',
        'print(f"Notebook demo timesteps: {dm.num_timesteps}")',
    ]
    if formulas is not None:
        code_lines.append(f'formulas = {json.dumps(list(formulas))}')
    if num_atom_distribution is not None:
        code_lines.append(f'num_atom_distribution = {json.dumps(num_atom_distribution)}')
    code_lines.extend([
        'if task == "dng":',
        '    sample_dng(dm=dm, num_atom_distribution=num_atom_distribution, num_samples=num_samples, batch_size=batch_size, output_path=output_dir)',
        'elif task == "csp":',
        '    sample_csp(dm=dm, formulas=formulas, num_samples=num_samples, batch_size=batch_size, output_path=output_dir)',
        'else:',
        '    raise ValueError(f"Unsupported task: {task}")',
        'print("Generated CIF files:", len(list(output_dir.glob("sample_*.cif"))))',
    ])
    run_chemeleon_dng_python('\n'.join(code_lines), device=device)
    cif_paths = sorted(output_dir.glob('sample_*.cif'))
    if not cif_paths:
        raise FileNotFoundError(f'Chemeleon-DNG sampling finished but wrote no CIF files in {output_dir}')
    return cif_paths


def show_atoms_gallery(atoms_list, title: str, save_path: Path, *, columns: int = 4, subtitles=None):
    atoms_list = list(atoms_list)
    subtitles = list(subtitles) if subtitles is not None else None
    if not atoms_list:
        print(f'No structures to display for {title}')
        return None
    cols = min(columns, len(atoms_list))
    rows = int(np.ceil(len(atoms_list) / cols))
    fig, axes = plt.subplots(rows, cols, figsize=(4 * cols, 4 * rows), squeeze=False, facecolor='white')
    for ax in axes.ravel():
        ax.axis('off')
    for idx, (ax, atoms) in enumerate(zip(axes.ravel(), atoms_list)):
        plot_atoms(atoms, ax=ax, radii=0.35, rotation=('10x,15y,0z'))
        ax.set_axis_off()
        subtitle = subtitles[idx] if subtitles and idx < len(subtitles) else f'{atoms.get_chemical_formula()} | {len(atoms)} atoms'
        ax.set_title(subtitle, fontsize=9)
    fig.suptitle(title, fontsize=13)
    plt.tight_layout()
    save_path = Path(save_path)
    save_path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(save_path, dpi=160, bbox_inches='tight')
    plt.show()
    return fig

from ase.io import read

# The first DNG call downloads the pretrained checkpoints into chemeleon_dng_repo/ckpts if needed.
# We keep this to a single batch so the quickstart stays notebook-friendly on CPU while still
# giving us enough structures for a small gallery.
dng_output_dir = CHEMELEON_DNG_OUT / 'dng_quickstart'
dng_cif_paths = sample_chemeleon_dng(
    task='dng',
    num_samples=8,
    batch_size=8,
    output_dir=dng_output_dir,
    device=CHEMELEON_DNG_DEVICE,
    num_atom_distribution='mp-20',
)
dng_atoms_list = [read(path) for path in dng_cif_paths]

print('Generated:', len(dng_atoms_list), 'structures')
for i, atoms in enumerate(dng_atoms_list):
    print(i, atoms.get_chemical_formula(), len(atoms), f'{atoms.get_volume():.2f}')

show_atoms_gallery(
    dng_atoms_list,
    'Chemeleon-DNG DNG quickstart',
    CHEMELEON_DNG_OUT / 'dng_quickstart_gallery.png',
)

import matplotlib.pyplot as plt
from html import escape
from IPython.display import HTML, display

# Add two explicit atom-count schedules so we can compare the default mp-20 prior
# with smaller-cell and larger-cell steering.
dng_small_output_dir = CHEMELEON_DNG_OUT / 'dng_small_cells'
dng_large_output_dir = CHEMELEON_DNG_OUT / 'dng_large_cells'

try:
    render_rank_table
except NameError:
    def render_rank_table(title, rows, headers):
        th_style = 'border:1px solid #ddd; padding:4px 8px; text-align:left; background:#eff6ff;'
        td_style = 'border:1px solid #ddd; padding:4px 8px; text-align:left;'

        def fmt(value):
            if isinstance(value, float):
                return f'{value:.3f}'
            return str(value)

        body = []
        for row in rows:
            body.append(
                '<tr>'
                + ''.join(f'<td style="{td_style}">{escape(fmt(row[h]))}</td>' for h in headers)
                + '</tr>'
            )
        html = f'''<div style="margin:0.5rem 0 0.25rem 0; font-weight:600;">{escape(title)}</div>
        <table style="border-collapse:collapse; min-width:760px; margin:0.35rem 0 1rem 0;">
          <thead><tr>{''.join(f'<th style="{th_style}">{escape(h)}</th>' for h in headers)}</tr></thead>
          <tbody>{''.join(body)}</tbody>
        </table>'''
        display(HTML(html))

dng_small_cif_paths = sample_chemeleon_dng(
    task='dng',
    num_samples=3,
    batch_size=3,
    output_dir=dng_small_output_dir,
    device=CHEMELEON_DNG_DEVICE,
    num_atom_distribution=[6, 6, 7],
)
dng_large_cif_paths = sample_chemeleon_dng(
    task='dng',
    num_samples=3,
    batch_size=3,
    output_dir=dng_large_output_dir,
    device=CHEMELEON_DNG_DEVICE,
    num_atom_distribution=[14, 14, 16],
)

dng_run_atoms = {
    'mp20_prior': dng_atoms_list,
    'small_cells': [read(path) for path in dng_small_cif_paths],
    'large_cells': [read(path) for path in dng_large_cif_paths],
}
dng_run_titles = {
    'mp20_prior': 'Default mp-20 prior',
    'small_cells': 'Small-cell schedule',
    'large_cells': 'Large-cell schedule',
}
dng_colors = {
    'mp20_prior': '#7aa2f7',
    'small_cells': '#9ece6a',
    'large_cells': '#f7768e',
}

dng_rows = []
for label, atoms_list in dng_run_atoms.items():
    for i, atoms in enumerate(atoms_list):
        volume = float(atoms.get_volume())
        dng_rows.append({
            'label': label,
            'display_name': dng_run_titles[label],
            'sample_id': i,
            'formula': atoms.get_chemical_formula(),
            'n_sites': len(atoms),
            'volume': volume,
            'density': float(atoms.get_masses().sum() / max(volume, 1e-12)),
        })

print(f'Loaded {len(dng_rows)} DNG structures across {len(dng_run_atoms)} steering settings')
for label, atoms_list in dng_run_atoms.items():
    rows = [row for row in dng_rows if row['label'] == label]
    print(
        f"- {dng_run_titles[label]}: mean n_sites={np.mean([row['n_sites'] for row in rows]):.2f}, "
        f"mean density={np.mean([row['density'] for row in rows]):.3f}, "
        f"unique formulas={len({row['formula'] for row in rows})}"
    )

fig, axes = plt.subplots(2, 3, figsize=(16, 10), facecolor='white')
for x, label in enumerate(dng_run_atoms, start=1):
    rows = [row for row in dng_rows if row['label'] == label]
    jitter = np.linspace(-0.08, 0.08, num=len(rows)) if len(rows) > 1 else np.array([0.0])
    axes[0, 0].scatter(np.full(len(rows), x) + jitter, [row['n_sites'] for row in rows], s=65, color=dng_colors[label], edgecolors='black', linewidths=0.3)
    axes[0, 1].scatter(np.full(len(rows), x) + jitter, [row['density'] for row in rows], s=65, color=dng_colors[label], edgecolors='black', linewidths=0.3)
    axes[1, 2].bar(x, len({row['formula'] for row in rows}), color=dng_colors[label], width=0.6)
    axes[1, 0].scatter([row['volume'] for row in rows], [row['density'] for row in rows], s=95, color=dng_colors[label], label=dng_run_titles[label], edgecolors='black', linewidths=0.3)
    axes[0, 2].hist([row['density'] for row in rows], bins=4, alpha=0.55, color=dng_colors[label], label=dng_run_titles[label])
    axes[1, 1].hist([row['n_sites'] for row in rows], bins=np.arange(5.5, 17.6, 1.0), alpha=0.5, color=dng_colors[label], label=dng_run_titles[label])

for ax, ylabel, title in [
    (axes[0, 0], 'n_sites', 'Atom-count control in DNG'),
    (axes[0, 1], 'density', 'Density by DNG steering setting'),
]:
    ax.set_xticks(range(1, len(dng_run_atoms) + 1))
    ax.set_xticklabels([dng_run_titles[label] for label in dng_run_atoms], rotation=15, ha='right')
    ax.set_ylabel(ylabel)
    ax.set_title(title)

axes[0, 2].set_title('Density distributions across DNG steering settings')
axes[0, 2].set_xlabel('density')
axes[0, 2].set_ylabel('count')
axes[0, 2].legend(frameon=False)

axes[1, 0].set_title('Density vs volume across DNG steering settings')
axes[1, 0].set_xlabel('volume')
axes[1, 0].set_ylabel('density')
axes[1, 0].legend(frameon=False)

axes[1, 1].set_title('Atom-count distributions across DNG steering settings')
axes[1, 1].set_xlabel('n_sites')
axes[1, 1].set_ylabel('count')
axes[1, 1].legend(frameon=False)

axes[1, 2].set_xticks(range(1, len(dng_run_atoms) + 1))
axes[1, 2].set_xticklabels([dng_run_titles[label] for label in dng_run_atoms], rotation=15, ha='right')
axes[1, 2].set_ylabel('unique formulas')
axes[1, 2].set_title('Formula diversity per DNG steering setting')

plt.tight_layout()
plt.show()

dng_sorted_by_density = sorted(dng_rows, key=lambda row: row['density'])
render_rank_table(
    'Lowest-density DNG candidates across all steering settings',
    dng_sorted_by_density[:3],
    ['display_name', 'formula', 'n_sites', 'volume', 'density'],
)
render_rank_table(
    'Highest-density DNG candidates across all steering settings',
    list(reversed(dng_sorted_by_density[-3:])),
    ['display_name', 'formula', 'n_sites', 'volume', 'density'],
)

show_atoms_gallery(
    dng_run_atoms['small_cells'][:3],
    'Chemeleon-DNG small-cell DNG steering',
    CHEMELEON_DNG_OUT / 'dng_small_cells_gallery.png',
    subtitles=[
        f"{atoms.get_chemical_formula()} | {len(atoms)} atoms | density={float(atoms.get_masses().sum() / max(atoms.get_volume(), 1e-12)):.3f}"
        for atoms in dng_run_atoms['small_cells'][:3]
    ],
    columns=3,
)
show_atoms_gallery(
    dng_run_atoms['large_cells'][:3],
    'Chemeleon-DNG large-cell DNG steering',
    CHEMELEON_DNG_OUT / 'dng_large_cells_gallery.png',
    subtitles=[
        f"{atoms.get_chemical_formula()} | {len(atoms)} atoms | density={float(atoms.get_masses().sum() / max(atoms.get_volume(), 1e-12)):.3f}"
        for atoms in dng_run_atoms['large_cells'][:3]
    ],
    columns=3,
)
show_atoms_gallery(
    [dng_run_atoms[row['label']][row['sample_id']] for row in dng_sorted_by_density[:3]],
    'Lowest-density DNG candidates across all steering settings',
    CHEMELEON_DNG_OUT / 'dng_low_density_gallery.png',
    subtitles=[
        f"{row['display_name']} | {row['formula']} | density={row['density']:.3f}"
        for row in dng_sorted_by_density[:3]
    ],
    columns=3,
)
show_atoms_gallery(
    [dng_run_atoms[row['label']][row['sample_id']] for row in list(reversed(dng_sorted_by_density[-3:]))],
    'Highest-density DNG candidates across all steering settings',
    CHEMELEON_DNG_OUT / 'dng_high_density_gallery.png',
    subtitles=[
        f"{row['display_name']} | {row['formula']} | density={row['density']:.3f}"
        for row in list(reversed(dng_sorted_by_density[-3:]))
    ],
    columns=3,
)

from ase.io import read

# CSP is the explicitly conditioned Chemeleon-DNG workflow.
# We keep multiple samples per formula, but fit them into one batch so the notebook stays lighter.
csp_targets = ['NaCl', 'LiMnO2']
CSP_SAMPLES_PER_FORMULA = 2
csp_samples = {}
for formula in csp_targets:
    out_dir = CHEMELEON_DNG_OUT / f'csp_{formula.lower()}'
    cif_paths = sample_chemeleon_dng(
        task='csp',
        formulas=[formula],
        num_samples=CSP_SAMPLES_PER_FORMULA,
        batch_size=2,
        output_dir=out_dir,
        device=CHEMELEON_DNG_DEVICE,
    )
    atoms_list = [read(path) for path in cif_paths]
    csp_samples[formula] = atoms_list
    print(f'{formula}: {len(atoms_list)} structures')
    show_atoms_gallery(
        atoms_list,
        f'Chemeleon-DNG CSP: {formula}',
        CHEMELEON_DNG_OUT / f'csp_{formula.lower()}_gallery.png',
        subtitles=[
            f"{atoms.get_chemical_formula()} | {len(atoms)} atoms | volume={atoms.get_volume():.1f} ų"
            for atoms in atoms_list
        ],
    )

import matplotlib.pyplot as plt
from ase.formula import Formula

# Compare the CSP outputs across target formulas and rank them with the same lightweight
# density screen we used for MatterGen.
csp_rows = []
for formula, atoms_list in csp_samples.items():
    target_reduced = Formula(formula).reduce()[0].format('hill')
    for i, atoms in enumerate(atoms_list):
        volume = float(atoms.get_volume())
        generated_reduced = Formula(atoms.get_chemical_formula()).reduce()[0].format('hill')
        csp_rows.append({
            'formula_target': formula,
            'sample_id': i,
            'generated_formula': atoms.get_chemical_formula(),
            'generated_reduced_formula': generated_reduced,
            'formula_match': generated_reduced == target_reduced,
            'n_sites': len(atoms),
            'volume': volume,
            'density': float(atoms.get_masses().sum() / max(volume, 1e-12)),
            'atoms': atoms,
        })

print(f'Loaded {len(csp_rows)} CSP structures across {len(csp_samples)} formulas')
for formula in csp_targets:
    subset = [row for row in csp_rows if row['formula_target'] == formula]
    print(
        f"- {formula}: mean density={np.mean([row['density'] for row in subset]):.3f}, "
        f"mean volume={np.mean([row['volume'] for row in subset]):.2f}, "
        f"formula-match fraction={np.mean([row['formula_match'] for row in subset]):.2f}"
    )

colors = {'NaCl': '#7aa2f7', 'LiMnO2': '#f7768e'}
fig, axes = plt.subplots(2, 3, figsize=(15, 9), facecolor='white')
for x, formula in enumerate(csp_targets, start=1):
    subset = [row for row in csp_rows if row['formula_target'] == formula]
    jitter = np.linspace(-0.06, 0.06, num=len(subset)) if len(subset) > 1 else np.array([0.0])
    axes[0, 0].scatter(np.full(len(subset), x) + jitter, [row['density'] for row in subset], s=65, color=colors[formula], edgecolors='black', linewidths=0.3)
    axes[0, 1].scatter(np.full(len(subset), x) + jitter, [row['volume'] for row in subset], s=65, color=colors[formula], edgecolors='black', linewidths=0.3)
    axes[1, 0].scatter([row['volume'] for row in subset], [row['density'] for row in subset], label=formula, s=90, color=colors[formula], edgecolors='black', linewidths=0.3)
    axes[1, 1].bar(x, np.mean([row['formula_match'] for row in subset]), color=colors[formula], width=0.6)
    axes[0, 2].hist([row['density'] for row in subset], bins=4, alpha=0.55, color=colors[formula], label=formula)
    axes[1, 2].scatter(np.full(len(subset), x) + jitter, [row['n_sites'] for row in subset], s=65, color=colors[formula], edgecolors='black', linewidths=0.3)

for ax, ylabel, title in [
    (axes[0, 0], 'density', 'Density by CSP target formula'),
    (axes[0, 1], 'volume', 'Volume by CSP target formula'),
    (axes[1, 2], 'n_sites', 'Atom count by CSP target formula'),
]:
    ax.set_xticks(range(1, len(csp_targets) + 1))
    ax.set_xticklabels(csp_targets)
    ax.set_ylabel(ylabel)
    ax.set_title(title)

axes[0, 2].set_title('Density distributions across CSP target formulas')
axes[0, 2].set_xlabel('density')
axes[0, 2].set_ylabel('count')
axes[0, 2].legend(frameon=False)

axes[1, 0].set_title('CSP density vs volume')
axes[1, 0].set_xlabel('volume')
axes[1, 0].set_ylabel('density')
axes[1, 0].legend(frameon=False)

axes[1, 1].set_xticks(range(1, len(csp_targets) + 1))
axes[1, 1].set_xticklabels(csp_targets)
axes[1, 1].set_ylim(0, 1.05)
axes[1, 1].set_ylabel('formula-match fraction')
axes[1, 1].set_title('Did the generated stoichiometry match the target?')

plt.tight_layout()
plt.show()

csp_sorted_by_density = sorted(csp_rows, key=lambda row: row['density'])
render_rank_table(
    'Lowest-density CSP candidates across both target formulas',
    csp_sorted_by_density[:4],
    ['formula_target', 'generated_formula', 'formula_match', 'n_sites', 'volume', 'density'],
)
render_rank_table(
    'Highest-density CSP candidates across both target formulas',
    list(reversed(csp_sorted_by_density[-4:])),
    ['formula_target', 'generated_formula', 'formula_match', 'n_sites', 'volume', 'density'],
)

perfect_matches = [row for row in csp_rows if row['formula_match']]
if perfect_matches:
    perfect_matches = sorted(perfect_matches, key=lambda row: row['density'], reverse=True)
    render_rank_table(
        'Densest formula-matching CSP candidates',
        perfect_matches[:4],
        ['formula_target', 'generated_formula', 'n_sites', 'volume', 'density'],
    )

show_atoms_gallery(
    [row['atoms'] for row in csp_sorted_by_density[:2]],
    'Lowest-density CSP candidates across both target formulas',
    CHEMELEON_DNG_OUT / 'csp_low_density_gallery.png',
    subtitles=[
        f"{row['formula_target']} target | {row['generated_formula']} | density={row['density']:.3f}"
        for row in csp_sorted_by_density[:2]
    ],
    columns=2,
)
show_atoms_gallery(
    [row['atoms'] for row in list(reversed(csp_sorted_by_density[-2:]))],
    'Highest-density CSP candidates across both target formulas',
    CHEMELEON_DNG_OUT / 'csp_high_density_gallery.png',
    subtitles=[
        f"{row['formula_target']} target | {row['generated_formula']} | density={row['density']:.3f}"
        for row in list(reversed(csp_sorted_by_density[-2:]))
    ],
    columns=2,
)