# [1] https://doi.org/10.1002/jcc.23481
#     Exploring transition state structures for intramolecular pathways
#     by the artificial force induced reaction method
#     Maeda, Morokuma et al, 2013
# [2] https://pubs.acs.org/doi/10.1021/ct200290m
#     Finding Reaction Pathways of Type A + B → X: Toward Systematic
#     Prediction of Reaction Mechanisms
#     Maeda, Morokuma, 2011

from dataclasses import dataclass
import itertools as it
import logging
from functools import reduce
import os
from pathlib import Path
from pprint import pformat
import shutil
import traceback
from typing import Callable, Dict, List, Tuple, Optional

import numpy as np
from numpy.typing import NDArray
from rmsd import kabsch_rmsd
from scipy.spatial.distance import pdist
from scipy.optimize import least_squares

from pysisyphus import logger as pysis_logger
from pysisyphus.calculators.AFIR import AFIR, CovRadiiSumZero
from pysisyphus.calculators import HardSphere
from pysisyphus.config import AFIR_RMSD_THRESH, OUT_DIR_DEFAULT
from pysisyphus.constants import BOHR2ANG
from pysisyphus.cos.NEB import NEB
from pysisyphus.drivers.opt import run_opt
from pysisyphus.elem_data import COVALENT_RADII as CR
from pysisyphus.Geometry import Geometry
from pysisyphus.helpers import pick_image_inds, check_for_end_sign
from pysisyphus.helpers_pure import to_sets
from pysisyphus.intcoords.helpers import get_bond_difference
from pysisyphus.intcoords.setup import get_pair_covalent_radii
from pysisyphus.intcoords.setup_fast import find_bonds
from pysisyphus.optimizers.FIRE import FIRE
from pysisyphus.xyzloader import make_xyz_str


logger = pysis_logger.getChild("afir")
logger.setLevel(logging.DEBUG)
file_handler = logging.FileHandler("afir.log", mode="w", delay=True)
logger.addHandler(file_handler)


AFIR_BOND_FACTOR = 1.2


@dataclass
class AFIRPath:
    atoms: tuple
    cart_coords: np.ndarray
    energies: np.ndarray
    forces: np.ndarray
    charge: int
    mult: int
    opt_is_converged: Optional[bool] = None
    gamma: Optional[float] = None
    path_indices: Optional[List[int]] = None

    def compatible(self, other):
        check = ("atoms", "charge", "mult")
        return all([getattr(self, name) == getattr(other, name) for name in check])

    def __add__(self, other):
        """This assumes that the first item in other.cart_coords is the same as the
        last item of self.cart_coords."""
        assert self.compatible(other)

        def conc(name):
            return np.concatenate(
                (getattr(self, name), getattr(other, name)[1:]), axis=0
            )

        cart_coords = conc("cart_coords")
        energies = conc("energies")
        forces = conc("forces")

        if self.path_indices is None:
            path_indices = [0] * len(self.cart_coords)
        else:
            path_indices = self.path_indices
        path_indices += [path_indices[-1] + 1] * (len(other.cart_coords) - 1)

        return AFIRPath(
            atoms=self.atoms,
            cart_coords=cart_coords,
            energies=energies,
            forces=forces,
            charge=self.charge,
            mult=self.mult,
            path_indices=path_indices,
        )

    def dump_trj(self, fn):
        geom = Geometry(self.atoms, self.cart_coords[0])
        xyzs = [geom.as_xyz(cart_coords=cc) for cc in self.cart_coords]
        with open(fn, "w") as handle:
            handle.write("\n".join(xyzs))


def analyze_afir_path(energies):
    energies = np.array(energies)
    max_ind = energies.argmax()

    local_minima = list()
    local_maxima = list()
    stationary_points = list()
    for i, en in enumerate(energies[1:-1], 1):
        prev_en = energies[i - 1]
        next_en = energies[i + 1]
        if is_minimum := prev_en > en < next_en:
            local_minima.append(i)
        if is_ts := prev_en < en > next_en:
            local_maxima.append(i)
        if is_minimum or is_ts:
            stationary_points.append(i)

    if max_ind in stationary_points:
        ts_ind = np.where(stationary_points == max_ind)[0][0]
        prev_min_ind = stationary_points[ts_ind - 1]
        next_min_ind = stationary_points[ts_ind + 1]
        sp_inds = [prev_min_ind, max_ind, next_min_ind]
    else:
        sp_inds = list()
    return sp_inds


##########################
#  Multi-component AFIR  #
#        Helper          #
##########################


def generate_random_union(geoms, offset=2.0, rng=None):
    """Unite fragments into one Geometry with random fragment orientations.

    Center, rotate and translate from origin acoording to approximate radius
    and an offset.
    Displace along +x, -x, +y, -y, +z, -z.

    Results for > 3 fragments don't look so pretty ;).
    """

    axis_inds = (0, 0, 1, 1, 2, 2)
    axis_inds_num = len(axis_inds)
    axis_translations = np.zeros(axis_inds_num)
    randomized = list()
    for i, geom in enumerate(geoms):
        i_mod = i % axis_inds_num
        geom = geom.copy()
        geom.center()
        geom.rotate(rng=rng)
        axis = axis_inds[i_mod]
        step_size = axis_translations[i_mod] + geom.approximate_radius() + offset
        axis_translations[i_mod] = step_size
        step = np.zeros(3)
        # Alternate between negative and positive direction along x/y/z
        step[axis] = (-1) ** i * step_size
        geom.coords3d += step[None, :]
        randomized.append(geom)
    union = reduce(lambda geom1, geom2: geom1 + geom2, randomized)
    return union


def generate_random_union_ref(geoms, rng=None, opt_kwargs=None):
    """Unite fragments into one Geometry with random fragment orientations."""

    geoms = [geom.copy() for geom in geoms]

    if rng is None:
        rng = np.random.default_rng()

    if opt_kwargs is None:
        opt_kwargs = {}

    # Random rotations
    for geom in geoms:
        geom.center()
        geom.rotate(rng=rng)

    # HardSphere optimization to fix overlapping fragments
    #
    # Set up fragment lists.
    fragments = list()
    for geom in geoms:
        geom_inds = np.arange(len(geom.atoms))
        try:
            geom_inds += fragments[-1][-1] + 1
        except IndexError:
            pass
        fragments.append(geom_inds.tolist())
    union = reduce(lambda geom1, geom2: geom1 + geom2, geoms)
    calc = HardSphere(union, frags=fragments, permutations=False, kappa=1.0)
    union.set_calculator(calc)
    _opt_kwargs = {
        "max_step": 0.2,
        "max_cycles": 500,
    }
    _opt_kwargs.update(opt_kwargs)
    opt = FIRE(union, **_opt_kwargs)
    opt.run()
    if not opt.is_converged:
        union = None
    else:
        # Remove HardSphere calculator
        union.clear()
        del union.calculator

    return union


def prepare_mc_afir(geoms, rng=None, **kwargs):
    """Wrapper for generate_random_union(_ref)."""
    union = generate_random_union(geoms, rng=rng, **kwargs)

    # Set up list of fragments
    i = 0
    fragments = list()
    for geom in geoms:
        atom_num = len(geom.atoms)
        fragments.append(np.arange(atom_num) + i)
        i += atom_num

    afir_kwargs = {
        "fragment_indices": fragments,
    }
    broken_bonds = []
    return union, afir_kwargs, broken_bonds


###########################
#  Single-component AFIR  #
#        Helper           #
###########################


def decrease_distance(coords3d, m, n, frac=0.8):
    c3d_m = coords3d[m]
    c3d_n = coords3d[n]
    dist_vec = c3d_n - c3d_m
    step = (1 - frac) / 2 * dist_vec
    c3d_new = coords3d.copy()
    c3d_new[m] += step
    c3d_new[n] -= step
    return c3d_new


def lstsqs_with_reference(coords3d, ref_coords3d, freeze_atoms=None):
    """Least-squares w.r.t. reference coordinates while keeping some
    atoms frozen."""

    if freeze_atoms is None:
        freeze_atoms = []
    else:
        freeze_atoms = list(freeze_atoms)
    ref_dists = pdist(ref_coords3d)

    mask = np.ones_like(coords3d[:, 0], dtype=bool)
    mask[freeze_atoms] = False
    # All atoms w/o the frozen atoms
    coords = coords3d[mask].flatten()
    x0 = coords

    coords_full = coords3d.copy()

    def fun(x):
        # Consider all distances, including distances to the fixed atoms 'm' and 'n'.
        coords_full[mask] = x.reshape(-1, 3)
        dists = pdist(coords_full)
        return dists - ref_dists

    res = least_squares(fun, x0)
    opt_coords = res.x
    coords_full[mask] = opt_coords.reshape(-1, 3)
    return res, coords_full


def weight_function(atoms, coords3d, i, j, p=6):
    cr_sum = sum([CR[atoms[k].lower()] for k in (i, j)])
    r_ij = np.linalg.norm(coords3d[i] - coords3d[j])
    omega = (cr_sum / r_ij) ** p
    return omega


def find_candidates(center, bond_sets):
    center_set = {
        center,
    }
    bonded_to_center = list()
    for bond in bond_sets:
        if center not in bond:
            continue
        bonded_to_center.append(*set(bond) - center_set)
    return bonded_to_center


def automatic_fragmentation(
    atoms, coords3d, frag1, frag2, cycles=2, p=6, bond_factor=1.25
):
    """Automatic fragmentation scheme as described in SC-AFIR paper [1]."""

    frag1 = set(frag1)
    frag2 = set(frag2)

    def w(m, n):
        """Shortcut for weight function"""
        return weight_function(atoms, coords3d, m, n, p=p)

    pairs = list(it.product(frag1, frag2))
    weights = [w(m, n) for m, n in pairs]
    max_weight = max(weights)

    bonds = find_bonds(atoms, coords3d, bond_factor=bond_factor)
    bond_sets = [set(bond) for bond in bonds.tolist()]

    def filter_candidates(candidates, partners, max_weight):
        to_keep = set()
        for candidate in candidates:
            for partner in partners:
                if candidate == partner:
                    break
                weight = w(candidate, partner)
                if weight > max_weight:
                    break
            else:
                to_keep.add(candidate)
        return to_keep

    def grow_fragment(frag1, frag2):
        # Find candidates that are bonded to atoms in frag1. Step 2 in [1].
        candidates = [find_candidates(m, bond_sets) for m in frag1]
        # Filter out candidates that are already contained in frag1
        candidates = [c for c in it.chain(*candidates) if c not in frag1]
        # Filter out candidates with weights that are too big. Step 3/4 in [1].
        candidates = set(filter_candidates(candidates, frag2, max_weight))
        return candidates

    for _ in range(cycles):
        f1_candidates = grow_fragment(frag1, frag2)
        f2_candidates = grow_fragment(frag2, frag1)

        # Step 5 in [1].
        f1_candidates = filter_candidates(f1_candidates, f2_candidates, max_weight)
        f2_candidates = filter_candidates(f2_candidates, f1_candidates, max_weight)

        # Step 6 in [1].
        frag1.update(f1_candidates)
        frag2.update(f2_candidates)
        assert frag1.isdisjoint(
            frag2
        ), "Overlapping fragments detected!"  # Sanity check
    return frag1, frag2


def prepare_sc_afir(geom, m, n, bond_factor=AFIR_BOND_FACTOR):
    """Create perturbed geometry, determine fragments and set AFIR calculator."""
    geom = geom.copy()
    atoms = geom.atoms
    org_coords3d = geom.coords3d.copy()

    def bond_sets(coords3d):
        bonds = find_bonds(atoms, coords3d, bond_factor=bond_factor)
        return set([frozenset(bond) for bond in bonds.tolist()])

    org_bond_sets = bond_sets(org_coords3d)

    # Move target atoms closer together along distance vector (decrease distance)
    decr_coords3d = decrease_distance(geom.coords3d, m, n)
    # Optimize remaining coordinates using least-squares, while keeping target
    # atom pair fixed.
    _, opt_coords3d = lstsqs_with_reference(decr_coords3d.copy(), geom.coords3d, (m, n))
    # Determine fragments, using the automated fragmentation
    frag1, frag2 = automatic_fragmentation(atoms, opt_coords3d, [m], [n])
    logger.debug(f"Fragments for target pair [{m}, {n}]: ({frag1}, {frag2})")
    fragment_indices = [list(frag) for frag in (frag1, frag2)]
    # Set lstsq-optimized coordinates and created wrapped calculator
    geom.coords3d = opt_coords3d

    opt_bond_sets = bond_sets(opt_coords3d)
    broken_bonds = org_bond_sets - opt_bond_sets

    afir_kwargs = {
        "fragment_indices": fragment_indices,
        # If 'complete_fragments' would be True, all remaining atom indices not
        # present in 'fragment_indices' would be assigned to a third fragment.
        "complete_fragments": False,
    }

    def set_atoms(inds, atom_type="X", mod_atoms=None):
        if mod_atoms is None:
            mod_atoms = list(atoms)
        for i in inds:
            mod_atoms[i] = atom_type
        return mod_atoms

    atoms_target = set_atoms((m, n))
    atoms_fragments = set_atoms(frag1)
    atoms_fragments = set_atoms(frag2, atom_type="Q", mod_atoms=atoms_fragments)

    atoms_coords3d = {
        "original": (atoms, org_coords3d),
        "original w/ target atoms": (atoms_target, org_coords3d),
        "decreased distance": (atoms, decr_coords3d),
        "decreased distance w/ target atoms": (atoms_target, decr_coords3d),
        "lstsq optimized": (atoms, opt_coords3d),
        "lstsq optimized w/ target atoms": (atoms_target, opt_coords3d),
        "lstsq optimized w/ fragments": (atoms_fragments, opt_coords3d),
    }
    trj = "\n".join(
        [
            make_xyz_str(atoms, BOHR2ANG * coords3d, comment=key)
            for key, (atoms, coords3d) in atoms_coords3d.items()
        ]
    )
    return geom, afir_kwargs, broken_bonds, trj


def determine_target_pairs(
    atoms: Tuple[str],
    coords3d: NDArray,
    min_: float = 1.25,
    max_: float = 5.0,
    active_atoms=None,
) -> List[Tuple[int]]:
    """Determine possible target m, n atom pairs for SC-AFIR calculations."""
    if active_atoms is None:
        active_atoms = range(len(atoms))
    active_atoms = set(active_atoms)

    pair_cov_radii = get_pair_covalent_radii(atoms)
    pair_dists = pdist(coords3d)
    quots = pair_dists / pair_cov_radii
    pair_inds = it.combinations(range(len(atoms)), 2)
    target_pairs = list()
    for pair_ind, quot in zip(pair_inds, quots):
        if (min_ <= quot <= max_) and (set(pair_ind) & active_atoms):
            target_pairs.append(pair_ind)
    return target_pairs


def determine_target_pairs_for_geom(geom: Geometry, **kwargs) -> List[Tuple[int]]:
    """Determine possible target m, n atom pairs for SC-AFIR calculations
    from geom."""
    target_pairs = determine_target_pairs(geom.atoms, geom.coords3d, **kwargs)
    return target_pairs


def coordinates_similar(
    test_coords3d: NDArray, ref_coords3d: List[NDArray], rmsd_thresh: float = 1e-2
) -> Tuple[bool, int]:
    # When the reference coordinates are an empty list.
    if len(ref_coords3d) == 0:
        return False, -1
    test_centered3d = test_coords3d - test_coords3d.mean(axis=0)[None, :]
    for i, rcoords3d in enumerate(ref_coords3d):
        ref_centered3d = rcoords3d - rcoords3d.mean(axis=0)[None, :]
        rmsd_ = kabsch_rmsd(test_centered3d, ref_centered3d)
        if rmsd_ <= rmsd_thresh:
            break
    else:
        return False, -1
    return True, i


def geom_similar(test_geom: Geometry, ref_geoms: List[Geometry], **kwargs) -> bool:
    return coordinates_similar(
        test_geom.coords3d, [geom.coords3d for geom in ref_geoms], **kwargs
    )


########################
#  Actual AFIR drivers #
########################


def opt_afir_path(geom, calc_getter, afir_kwargs, opt_kwargs=None, out_dir=None):
    """Minimize geometry with AFIR calculator."""
    if opt_kwargs is None:
        opt_kwargs = dict()
    if out_dir is None:
        out_dir = "."
    out_dir = Path(out_dir)

    actual_calc = calc_getter(out_dir=out_dir / OUT_DIR_DEFAULT)

    def afir_calc_getter():
        afir_calc = AFIR(actual_calc, out_dir=out_dir, **afir_kwargs)
        return afir_calc

    _opt_kwargs = {
        "dump": True,
        "out_dir": out_dir,
        "prefix": "afir",
        "max_cycles": 125,
        "overachieve_factor": 3,
        "hessian_update": "flowchart",
    }
    logger.debug(
        "\n".join(
            (
                "afir_kwargs:",
                "\t" + pformat(afir_kwargs),
                "opt_kwargs:",
                "\t" + pformat(_opt_kwargs),
            )
        )
    )
    _opt_kwargs.update(opt_kwargs)
    opt_result = run_opt(geom, afir_calc_getter, opt_key="rfo", opt_kwargs=_opt_kwargs)
    opt = opt_result.opt

    afir_path = AFIRPath(
        atoms=geom.atoms,
        cart_coords=np.array(opt.cart_coords),
        energies=np.array(opt.true_energies),
        forces=np.array(opt.true_forces),
        opt_is_converged=opt.is_converged,
        charge=actual_calc.charge,
        mult=actual_calc.mult,
        gamma=afir_kwargs["gamma"],
    )

    return afir_path


def run_afir_path(
    geom,
    calc_getter,
    out_dir,
    gamma_max,
    gamma_interval: Tuple[float, float],
    rng,
    ignore_bonds=None,
    bond_factor=AFIR_BOND_FACTOR,
    afir_kwargs=None,
    opt_kwargs=None,
):
    """Driver for AFIR minimizations with increasing gamma values."""
    if ignore_bonds is None:
        ignore_bonds = list()
    if afir_kwargs is None:
        afir_kwargs = dict()
    if opt_kwargs is None:
        opt_kwargs = {}

    if out_dir.exists():
        dir_contents = os.listdir(out_dir)
        for fn in dir_contents:
            fn = out_dir / fn
            try:
                os.remove(fn)
            except IsADirectoryError:
                shutil.rmtree(fn)
    else:
        os.mkdir(out_dir)

    # Decreasing the distance between two atoms in SC-AFIR may lead to broken
    # bonds for these two bonds, e.g., hydrogen atoms "that are left behind".
    # These bonds will be formed again when the AFIR function is minimized, but
    # we are not interested in these changes. So they can be ignored here.
    ignore_bonds = set([frozenset(bond) for bond in ignore_bonds])

    ref_calc = calc_getter()
    ref_energy = ref_calc.get_energy(geom.atoms, geom.cart_coords)["energy"]
    geom_backup = geom.copy()

    # By using (1.0, 1.0) as interval we can directly start at gamma_max, e.g.,
    # in SC-AFIR.
    gamma_low, gamma_high = gamma_interval
    assert gamma_high >= gamma_low
    gamma_spread = gamma_high - gamma_low
    gamma_0 = (gamma_low + (gamma_spread * rng.random(1)[0])) * gamma_max
    gamma_inrc = 0.1 * gamma_max

    afir_paths = list()
    best_afir_path = None
    lowest_barrier = None
    gamma = gamma_0

    # Minimize AFIR functions until gamma exceeds gamma_max.
    logger.info(f"New AFIR run with γ_max={gamma_max:.6f} au")
    while gamma <= gamma_max:
        gamma_ratio = gamma / gamma_max
        logger.info(f"AFIR run with γ={gamma:.6f} au, γ/γ_max={gamma_ratio: >6.2%}")
        _afir_kwargs = afir_kwargs.copy()
        _afir_kwargs["gamma"] = gamma
        try:
            afir_path = opt_afir_path(
                geom,
                calc_getter,
                afir_kwargs=_afir_kwargs,
                opt_kwargs=opt_kwargs,
                out_dir=out_dir,
            )
        # Can happen in SC-AFIR runs when fragments comprise only hydrogens.
        except CovRadiiSumZero:
            logger.warning("Sum of covalent radii is 0.0!")
            best_afir_path = None
            break
        except Exception:
            logger.error(f"Optimization crashed!\n{traceback.format_exc()}")
            best_afir_path = None
            break

        afir_paths.append(afir_path)
        true_energies = afir_path.energies

        # Check for changes in bond topology by comparing the result of the current
        # minimization to the initial geometry.
        formed, broken = get_bond_difference(geom_backup, geom, bond_factor=bond_factor)
        formed = to_sets(formed) - ignore_bonds
        broken = to_sets(broken) - ignore_bonds
        if formed or broken:
            max_energy = true_energies.max()
            barrier = max_energy - ref_energy

            # Always store the first path that leads to a change in bond topology.
            if lowest_barrier is None:
                best_afir_path = afir_path
                lowest_barrier = barrier
            # Break when a path with a lower or higher barrier is detected. This
            # should happen in the cycle after the first change in bond toplogy was
            # detected.
            if barrier < lowest_barrier:
                best_afir_path = afir_path
                lowest_barrier = barrier
                break
            elif barrier > lowest_barrier:
                break
        else:
            barrier = None

        # Update values for next cycle
        gamma += gamma_inrc

    # reduce(lambda ap1, ap2: ap1 + ap2, afir_paths).dump_trj("dumped_trj.xyz")

    # Construct TS guess from highest energy point along the AFIR path.
    if best_afir_path:
        guess_ind = best_afir_path.energies.argmax()
        guess_coords = best_afir_path.cart_coords[guess_ind]
        ts_guess = Geometry(geom.atoms, guess_coords)
        ts_guess.dump_xyz(out_dir / "ts_guess.xyz")
        afir_path_merged = reduce(lambda ap1, ap2: ap1 + ap2, afir_paths)
    else:
        ts_guess = None
        afir_path_merged = None
    return ts_guess, afir_path_merged


def relax_afir_path(atoms, cart_coords, calc_getter, images=15, out_dir=None):
    """Sample imagef from AFIR path and do COS relaxation."""
    image_inds = pick_image_inds(cart_coords, images=images)
    images = [Geometry(atoms, cart_coords[i]) for i in image_inds]

    # Relax last image
    opt_kwargs = {
        "dump": True,
        "prefix": "last",
        "out_dir": out_dir,
    }
    last_image = images[-1]
    last_image_backup = last_image.copy()
    run_opt(last_image, calc_getter, opt_key="rfo", opt_kwargs=opt_kwargs)
    _, broken = get_bond_difference(last_image, last_image_backup)
    if broken:
        return

    cos_kwargs = {}
    cos = NEB(images, **cos_kwargs)
    cos_opt_kwargs = {
        "align": True,
        "dump": True,
        "max_cycles": 30,
        "out_dir": out_dir,
    }
    run_opt(cos, calc_getter, opt_key="lbfgs", opt_kwargs=cos_opt_kwargs)


def run_mc_afir_paths(
    geoms: List,
    calc_getter: Callable,
    gamma_max: float,
    rng,
    N_max: int = 5,
    gamma_interval: Tuple[float, float] = (0.0, 1.0),
    afir_kwargs: Optional[Dict] = None,
    opt_kwargs: Optional[Dict] = None,
):
    if afir_kwargs is None:
        afir_kwargs = dict()

    N = 0
    N_0 = 0
    fmt = " >4d"
    while True:
        geom, _afir_kwargs, *_ = prepare_mc_afir(geoms, rng=rng)
        afir_kwargs = afir_kwargs.copy()
        afir_kwargs.update(_afir_kwargs)
        out_dir = Path(f"out_{N:03d}")

        ts_guess, afir_path = run_afir_path(
            geom,
            calc_getter,
            out_dir,
            gamma_max,
            gamma_interval,
            rng=rng,
            afir_kwargs=afir_kwargs,
            opt_kwargs=opt_kwargs,
        )
        ts_guess_is_new = yield N, ts_guess, afir_path
        if ts_guess_is_new:
            N_0 = N
        # Allow up to consecutive N_max failures
        if (N - N_0) > N_max:
            break
        logger.debug(f"{N_0=:{fmt}}, {N=:{fmt}}, {N-N_0=:{fmt}}, {N_max=:{fmt}}, ")
        N += 1


def run_sc_afir_paths(
    geom: Geometry,
    calc_getter: Callable,
    gamma_max: float,
    rng,
    N_max: int = 5,
    N_sample: int = 0,
    gamma_interval: Tuple[float, float] = (1.0, 1.0),  # Start directly with gamma_max
    afir_kwargs: Optional[Dict] = None,
    opt_kwargs: Optional[Dict] = None,
    target_pairs: Optional[List] = None,
):
    if target_pairs is None:
        target_pairs = determine_target_pairs_for_geom(geom)

    if afir_kwargs is None:
        afir_kwargs = dict()
    if opt_kwargs is None:
        opt_kwargs = dict()
    _opt_kwargs = opt_kwargs.copy()

    i = 0
    while len(target_pairs) > 0:
        m, n = target_pairs.pop(0)
        logger.info(f"Running SC-AFIR with target_pair ({m}, {n}).")
        # _afir_kwargs will contain the automatically determined fragments
        geom_mod, _afir_kwargs, broken_bonds, trj = prepare_sc_afir(geom, m, n)
        afir_kwargs = afir_kwargs.copy()
        afir_kwargs.update(_afir_kwargs)

        # _opt_kwargs.update({
        # "fragments": [[m, ], [n, ]],
        # "monitor_frag_dists": 5,
        # })

        out_dir = Path(f"out_{i:03d}")
        ts_guess, afir_path = run_afir_path(
            geom_mod,
            calc_getter,
            out_dir,
            gamma_max,
            gamma_interval,
            rng=rng,
            ignore_bonds=broken_bonds,
            afir_kwargs=afir_kwargs,
            opt_kwargs=_opt_kwargs,
        )
        ts_guess_is_new = yield i, ts_guess, afir_path
        i += 1
        # Here, TS optimization & IRC integration could take place, when the TS is new.


def run_afir_paths(
    afir_key,
    geoms,
    calc_getter,
    afir_kwargs=None,
    opt_kwargs=None,
    seed=None,
    N_sample=None,
    rmsd_thresh: float = AFIR_RMSD_THRESH,
    **kwargs,
):
    if seed is None:
        rng = np.random.default_rng()
        seed = rng.integers(1_000_000_000_000)
    logger.info(f"{seed=}")
    rng = np.random.default_rng(seed)

    if afir_key == "sc":
        assert (
            len(geoms) == 1
        ), f"Expected only 1 geometry for SC-AFIR, but got {len(geoms)}!."
        geoms = geoms[0]

    afir_funcs = {
        "mc": run_mc_afir_paths,
        "sc": run_sc_afir_paths,
    }
    afir_func = afir_funcs[afir_key]
    afir_coroutine = afir_func(
        geoms,
        calc_getter,
        rng=rng,
        afir_kwargs=afir_kwargs,
        opt_kwargs=opt_kwargs,
        **kwargs,
    )
    ts_guesses = list()
    afir_paths = list()
    stop_sign = "afir_stop"

    i, ts_guess, afir_path = next(afir_coroutine)
    while True:
        prefix = f"{i:03d}"
        logger.info(f"AFIR run {prefix}")
        # Check if we found enough TS guesses
        if N_sample and (len(ts_guesses) >= N_sample):
            break

        # Check similarity of TS guess to previously found TS guesses
        try:
            rmsds = [ts_guess.rmsd(guess) for guess in ts_guesses]
            min_rmsd = min(rmsds)
        except AttributeError:  # Raised when ts_guess is None
            min_rmsd = None
        except ValueError:
            min_rmsd = rmsd_thresh if (ts_guess is not None) else None

        if ts_guess_is_new := ((min_rmsd is not None) and (min_rmsd >= rmsd_thresh)):
            logger.debug(f"rmsds different enough! {min_rmsd=:.6f} au")
            ts_guesses.append(ts_guess)
            afir_paths.append(afir_path)
            afir_path.dump_trj(f"{prefix}_afir_path_trj.xyz")
            ts_guess.dump_xyz(f"{prefix}_ts_guess.xyz")
        elif min_rmsd:
            logger.debug(f"rmsds too similar! {min_rmsd=:.6f} au)")

        sign = check_for_end_sign(add_signs=(stop_sign,))
        if sign == stop_sign:
            break
        logger.info(f"{len(ts_guesses)=: >5d}")
        try:
            i, ts_guess, afir_path = afir_coroutine.send(ts_guess_is_new)
        except StopIteration:
            break

    return ts_guesses, afir_paths
