<file_sep>class KshellDataStructureError(Exception):
    """
    Raise when KSHELL data file has unexpected structure / syntax.
    """<file_sep># kshell-utilities
Handy utilities for processing nuclear shell model data from `KSHELL`. Se the [KSHELL repository](https://github.com/GaffaSnobb/kshell) for installation and usage instructions for `KSHELL`.

## Installation
Install from the PyPi repository
``` bash
pip install kshell-utilities
```
Or, for the very latest version, clone this repository to your downloads directory. cd to the root directory of this repository and run
```
pip install .
```

## Usage
Please see the *How to use the output from KSHELL* section of the Wiki in the `KSHELL` repository for an introduction on how to use `kshell-utilities`: https://github.com/GaffaSnobb/kshell/wiki/How-to-use-the-output-from-KSHELL

## Credits
KSHELL is created by <NAME> https://arxiv.org/abs/1310.5431. Code in this repository is built upon tools created by <NAME>: https://github.com/jorgenem/kshell_public.
<file_sep>from __future__ import annotations
import time, sys, ast, warnings
from fractions import Fraction
from typing import TextIO
import numpy as np
from .kshell_exceptions import KshellDataStructureError
from .data_structures import (
    Interaction, Partition, OrbitalParameters, Configuration
)
from .parameters import (
    spectroscopic_conversion, shell_model_order, flags
)
from .partition_tools import (
    _calculate_configuration_parity, _sanity_checks, configuration_energy
)

def load_interaction(
    filename_interaction: str,
    interaction: Interaction,
):
    interaction.name = filename_interaction
    with open(filename_interaction, "r") as infile:
        """
        Extract information from the interaction file about the orbitals
        in the model space.
        """
        for line in infile:
            """
            Example
            -------
            ! GXPF1A  pf-shell 
            ! <NAME>, <NAME>, <NAME>, and <NAME>, 
            !   Eur. Phys. J. A 25, Suppl. 1, 499 (2005). 
            !
            ! default input parameters 
            !namelist eff_charge = 1.5, 0.5
            !namelist orbs_ratio = 2, 3, 4, 6, 7, 8
            !
            ! model space
            4   4    20  20
            ...
            """
            if line[0] != "!":
                tmp = line.split()
                interaction.model_space_proton.n_orbitals = int(tmp[0])
                interaction.model_space_neutron.n_orbitals = int(tmp[1])
                interaction.model_space.n_orbitals = (
                    interaction.model_space_proton.n_orbitals + interaction.model_space_neutron.n_orbitals
                )
                interaction.n_core_protons = int(tmp[2])
                interaction.n_core_neutrons = int(tmp[3])
                break

        for line in infile:
            """
            Example
            -------
            4   4    20  20
            1       0   3   7  -1    !  1 = p 0f_7/2
            2       1   1   3  -1    !  2 = p 1p_3/2
            3       0   3   5  -1    !  3 = p 0f_5/2
            4       1   1   1  -1    !  4 = p 1p_1/2
            5       0   3   7   1    !  5 = n 0f_7/2
            6       1   1   3   1    !  6 = n 1p_3/2
            7       0   3   5   1    !  7 = n 0f_5/2
            8       1   1   1   1    !  8 = n 1p_1/2
            ! interaction
            ...
            """
            if line[0] == "!": break
            idx, n, l, j, tz = [int(i) for i in line.split("!")[0].split()]
            idx -= 1
            nucleon = "p" if tz == -1 else "n"
            name = f"{n}{spectroscopic_conversion[l]}{j}"
            tmp_orbital = OrbitalParameters(
                idx = idx,
                n = n,
                l = l,
                j = j,
                tz = tz,
                nucleon = nucleon,
                name = f"{nucleon}{name}",
                parity = (-1)**l,
                order = shell_model_order[name],
                ho_quanta = 2*n + l
            )
            interaction.model_space.orbitals.append(tmp_orbital)
            interaction.model_space.major_shell_names.add(shell_model_order[name].major_shell_name)

            if tz == -1:
                interaction.model_space_proton.orbitals.append(tmp_orbital)
                interaction.model_space_proton.major_shell_names.add(shell_model_order[name].major_shell_name)
            elif tz == +1:
                interaction.model_space_neutron.orbitals.append(tmp_orbital)
                interaction.model_space_neutron.major_shell_names.add(shell_model_order[name].major_shell_name)
            else:
                msg = f"Valid values for tz are -1 and +1, got {tz=}"
                raise ValueError(msg)
            
        for line in infile:
            """
            Example
            -------
            ! GCLSTsdpfsdgix5pn.int
            !   p-n formalism
                24      0
            ...
            """
            if line[0] != "!":
                tmp = line.split()
                if int(tmp[1]) != 0: raise NotImplementedError
                interaction.n_spe = int(tmp[0])
                break

        for line in infile:
            """
            Example
            -------
            1   1     -8.62400
            2   2     -5.67930
            3   3     -1.38290
            4   4     -4.13700
            5   5     -8.62400
            6   6     -5.67930
            7   7     -1.38290
            8   8     -4.13700
            518   1  42  -0.30000
            ...
            """
            tmp = line.split()
            if len(tmp) != 3: break
            interaction.spe.append(float(tmp[2]))

        try:
            interaction.n_tbme = int(tmp[0])
            interaction.fmd_mass = int(tmp[2])
            interaction.fmd_power = float(tmp[3])
        except IndexError:
            """
            I dont really know what this is yet.
            """
            msg = "Interactions with no mass dependence have not yet been implemented."
            raise NotImplementedError(msg)

        for line in infile:
            """
            NOTE: This way of structuring the TBMEs is taken from
            espe.py.
            """
            i0, i1, i2, i3, j, tbme = line.split()
            i0 = int(i0) - 1
            i1 = int(i1) - 1
            i2 = int(i2) - 1
            i3 = int(i3) - 1
            j = int(j)
            tbme = float(tbme)

            if (i0, i1, i2, i3, j) in interaction.tbme:
                """
                I dont yet understand why I should check the TBME value
                when I already know that the indices are the same. This
                is how it is done in espe.py (why check > 1.e-3 and not
                < 1.e-3?):
                if (i,j,k,l,J) in vtb:
                    if abs( v - vtb[i,j,k,l,J] )>1.e-3:
                            print( 'WARNING duplicate TBME', i+1,j+1,k+1,l+1,J,v,vtb[(i,j,k,l,J)] )

                """
                warnings.warn(f"Duplicate TBME! {i0 + 1}, {i1 + 1}, {i2 + 1}, {i3 + 1}, {j}, {tbme}, {interaction.tbme[(i0, i1, i2, i3, j)]}")

            interaction.tbme[(i0, i1, i2, i3, j)] = tbme
            s01 = (-1)**((interaction.model_space.orbitals[i0].j + interaction.model_space.orbitals[i1].j)/2 - j + 1)
            s23 = (-1)**((interaction.model_space.orbitals[i2].j + interaction.model_space.orbitals[i3].j)/2 - j + 1)
            
            if i0 != i1:
                interaction.tbme[(i1, i0, i2, i3, j)] = tbme*s01
            if i2 != i3:
                interaction.tbme[(i0, i1, i3, i2, j)] = tbme*s23
            if (i0 != i1) and (i2 != i3):
                interaction.tbme[(i1, i0, i3, i2, j)] = tbme*s01*s23
            
            if (i0, i1) != (i2, i3):
                interaction.tbme[(i2, i3, i0, i1, j)] = tbme
                if i0 != i1:
                    interaction.tbme[(i2, i3, i1, i0, j)] = tbme*s01
                if i2 != i3:
                    interaction.tbme[(i3, i2, i0, i1, j)] = tbme*s23
                if (i0 != i1) and (i2!=i3):
                    interaction.tbme[(i3, i2, i1, i0, j)] = tbme*s01*s23

    assert len(interaction.spe) == interaction.n_spe
    # assert len(interaction.tbme) == interaction.n_tbme

    interaction.vm = np.zeros((interaction.model_space.n_orbitals, interaction.model_space.n_orbitals), dtype=float)
    for i0 in range(interaction.model_space.n_orbitals):
        """
        Non-diagonal. TODO: Make a better description when I figure out
        what vm is.
        """
        for i1 in range(interaction.model_space.n_orbitals):
            j_min = abs(interaction.model_space.orbitals[i0].j - interaction.model_space.orbitals[i1].j)//2
            j_max =    (interaction.model_space.orbitals[i0].j + interaction.model_space.orbitals[i1].j)//2
            
            skip = 2 if (i0 == i1) else 1
            
            v: float = 0.0
            d: int = 0
            for j in range(j_min, j_max + 1, skip):
                """
                Using j_max + 1, not j_max + skip, because when i0 == i1
                both nucleons are in the same orbital and they cannot
                both have the same angular momentum z component.

                skip = 2 when i0 == i1 because only even numbered j are
                allowed when identical particles are in the same
                orbital.
                """
                try:
                    tbme = interaction.tbme[(i0, i1, i0, i1, j)]
                except KeyError:
                    """
                    I am unsure if this should be allowed at all and
                    should rather raise an exception.
                    """
                    warnings.warn(f"TBME entry not found! ({i0 + 1}, {i1 + 1}, {i0 + 1}, {i1 + 1}, {j})")
                    tbme = 0.0
                
                degeneracy = 2*j + 1
                v += degeneracy*tbme
                d += degeneracy
            
            interaction.vm[i0, i1] = v/d

    interaction.model_space.n_major_shells = len(interaction.model_space.major_shell_names)
    interaction.model_space_proton.n_major_shells = len(interaction.model_space_proton.major_shell_names)
    interaction.model_space_neutron.n_major_shells = len(interaction.model_space_neutron.major_shell_names)

    if not all(orb.idx == i for i, orb in enumerate(interaction.model_space.orbitals)):
        """
        Make sure that the list indices are the same as the orbit
        indices.
        """
        msg = (
            "The orbitals in the model space are not indexed correctly!"
        )
        raise KshellDataStructureError(msg)

def load_partition(
    filename_partition: str,
    interaction: Interaction,
    partition_proton: Partition,
    partition_neutron: Partition,
    partition_combined: Partition,
) -> str:
    header: str = ""
    with open(filename_partition, "r") as infile:
        # truncation_info: str = infile.readline()    # Eg. hw trucnation,  min hw = 0 ,   max hw = 1
        # hw_min, hw_max = [int(i.split("=")[1].strip()) for i in truncation_info.split(",")[1:]] # NOTE: No idea what happens if no hw trunc is specified.
        for line in infile:
            """
            Extract the information from the header before partitions
            are specified. Example:

            # hw trucnation,  min hw = 0 ,   max hw = 1
            # partition file of gs8.snt  Z=20  N=31  parity=-1
            20 31 -1
            # num. of  proton partition, neutron partition
            86 4
            # proton partition
            ...
            """
            if "#" not in line:
                tmp = [int(i) for i in line.split()]

                try:
                    """
                    For example:
                    20 31 -1
                    """
                    n_valence_protons, n_valence_neutrons, parity_partition = tmp

                    partition_proton.parity = parity_partition
                    partition_neutron.parity = parity_partition
                    partition_combined.parity = parity_partition
                    interaction.model_space.n_valence_nucleons = n_valence_protons + n_valence_neutrons
                    interaction.model_space_proton.n_valence_nucleons = n_valence_protons
                    interaction.model_space_neutron.n_valence_nucleons = n_valence_neutrons
                except ValueError:
                    """
                    For example:
                    86 4
                    """
                    n_proton_configurations, n_neutron_configurations = tmp
                    infile.readline()   # Skip header.
                    break

            header += line

        ho_quanta_min: int = +1000  # The actual number of harmonic oscillator quanta will never be smaller or larger than these values.
        ho_quanta_max: int = -1000
        for line in infile:
            """
            Extract proton configurations.
            """
            if "# neutron partition" in line: break

            configuration = [int(i) for i in line.split()[1:]]

            parity_tmp = _calculate_configuration_parity(
                configuration = configuration,
                model_space = interaction.model_space_proton.orbitals
            )
            if   parity_tmp == -1: partition_proton.n_existing_negative_configurations += 1
            elif parity_tmp == +1: partition_proton.n_existing_positive_configurations += 1

            assert len(interaction.model_space_proton.orbitals) == len(configuration)

            ho_quanta_tmp = sum([   # The number of harmonic oscillator quanta for each configuration.
                n*orb.ho_quanta for n, orb in zip(configuration, interaction.model_space_proton.orbitals)
            ])
            ho_quanta_min = min(ho_quanta_min, ho_quanta_tmp)
            ho_quanta_max = max(ho_quanta_max, ho_quanta_tmp)
            
            partition_proton.configurations.append(
                Configuration(
                    configuration = configuration,
                    parity = parity_tmp,
                    ho_quanta = ho_quanta_tmp,
                    energy = None,
                    is_original = True,
                )
            )
        partition_proton.ho_quanta_min_this_parity = ho_quanta_min
        partition_proton.ho_quanta_max_this_parity = ho_quanta_max
        ho_quanta_min: int = +1000  # Reset for neutrons.
        ho_quanta_max: int = -1000
        
        for line in infile:
            """
            Extract neutron configurations.
            """
            if "# partition of proton and neutron" in line: break

            configuration = [int(i) for i in line.split()[1:]]

            parity_tmp = _calculate_configuration_parity(
                configuration = configuration,
                model_space = interaction.model_space_neutron.orbitals
            )
            if   parity_tmp == -1: partition_neutron.n_existing_negative_configurations += 1
            elif parity_tmp == +1: partition_neutron.n_existing_positive_configurations += 1

            assert len(interaction.model_space_neutron.orbitals) == len(configuration)

            ho_quanta_tmp = sum([   # The number of harmonic oscillator quanta for each configuration.
                n*orb.ho_quanta for n, orb in zip(configuration, interaction.model_space_neutron.orbitals)
            ])
            ho_quanta_min = min(ho_quanta_min, ho_quanta_tmp)
            ho_quanta_max = max(ho_quanta_max, ho_quanta_tmp)
            
            partition_neutron.configurations.append(
                Configuration(
                    configuration = configuration,
                    parity = parity_tmp,
                    ho_quanta = ho_quanta_tmp,
                    energy = None,
                    is_original = True,
                )
            )
        partition_neutron.ho_quanta_min_this_parity = ho_quanta_min
        partition_neutron.ho_quanta_max_this_parity = ho_quanta_max
        ho_quanta_min: int = +1000  # Reset for combined.
        ho_quanta_max: int = -1000
        n_combined_configurations = int(infile.readline())

        for line in infile:
            """
            Extract the combined pn configurations.
            """
            proton_idx, neutron_idx = line.split()
            proton_idx = int(proton_idx) - 1
            neutron_idx = int(neutron_idx) - 1
            parity_tmp = partition_proton.configurations[proton_idx].parity*partition_neutron.configurations[neutron_idx].parity
            assert parity_partition == parity_tmp

            if   parity_tmp == -1: partition_combined.n_existing_negative_configurations += 1
            elif parity_tmp == +1: partition_combined.n_existing_positive_configurations += 1

            ho_quanta_tmp = (
                partition_proton.configurations[proton_idx].ho_quanta + 
                partition_neutron.configurations[neutron_idx].ho_quanta
            )
            ho_quanta_min = min(ho_quanta_min, ho_quanta_tmp)
            ho_quanta_max = max(ho_quanta_max, ho_quanta_tmp)

            energy = configuration_energy(
                interaction = interaction,
                proton_configuration = partition_proton.configurations[proton_idx],
                neutron_configuration = partition_neutron.configurations[neutron_idx],
            )

            partition_combined.configurations.append(
                Configuration(
                    configuration = [proton_idx, neutron_idx],
                    parity = parity_partition,
                    ho_quanta = ho_quanta_tmp,
                    energy = energy,
                    is_original = True,
                )
            )
        partition_combined.ho_quanta_min_this_parity = ho_quanta_min
        partition_combined.ho_quanta_max_this_parity = ho_quanta_max

        partition_combined.ho_quanta_min = min(ho_quanta_min, partition_combined.ho_quanta_min_opposite_parity)
        partition_combined.ho_quanta_max = max(ho_quanta_max, partition_combined.ho_quanta_max_opposite_parity)

    energies = [configuration.energy for configuration in partition_combined.configurations]
    partition_combined.min_configuration_energy = min(energies)
    partition_combined.max_configuration_energy = max(energies)
    partition_combined.max_configuration_energy_original = partition_combined.max_configuration_energy

    _sanity_checks(
        partition_proton = partition_proton,
        partition_neutron = partition_neutron,
        partition_combined = partition_combined,
        interaction = interaction,
    )
    assert len(partition_proton.configurations) == n_proton_configurations
    assert len(partition_neutron.configurations) == n_neutron_configurations
    assert len(partition_combined.configurations) == n_combined_configurations

    assert (
        partition_proton.n_existing_negative_configurations + partition_proton.n_existing_positive_configurations + 
        partition_proton.n_new_negative_configurations + partition_proton.n_new_positive_configurations
    ) == n_proton_configurations
    assert (
        partition_neutron.n_existing_negative_configurations + partition_neutron.n_existing_positive_configurations + 
        partition_neutron.n_new_negative_configurations + partition_neutron.n_new_positive_configurations
    ) == n_neutron_configurations
    assert (
        partition_combined.n_existing_negative_configurations + partition_combined.n_existing_positive_configurations + 
        partition_combined.n_new_negative_configurations + partition_combined.n_new_positive_configurations
    ) == n_combined_configurations

    return header

def _parity_string_to_integer(parity: str):
    if parity == "+":
        res = 1
    elif parity == "-":
        res = -1
    else:
        msg = f"Invalid parity read from file. Got: '{parity}'."
        raise KshellDataStructureError(msg)

    return res

def _load_energy_levels(infile: TextIO) -> tuple[list, int]:
    """
    Load excitation energy, spin and parity into a list of structure:
    levels = [[energy, spin, parity], ...].
    
    Parameters
    ----------
    infile : TextIO
        The KSHELL summary file at the starting position of the level
        data.

    Returns
    -------
    levels : list
        List of level data.
        
    negative_spin_counts : int
        The number of negative spin levels encountered.

    Example
    -------
    Energy levels

    N    J prty N_Jp    T     E(MeV)  Ex(MeV)  log-file

    1   5/2 +     1   3/2    -16.565    0.000  log_O19_sdpf-mu_m1p.txt 
    2   3/2 +     1   3/2    -15.977    0.588  log_O19_sdpf-mu_m1p.txt 
    3   1/2 +     1   3/2    -15.192    1.374  log_O19_sdpf-mu_m1p.txt 
    4   9/2 +     1   3/2    -13.650    2.915  log_O19_sdpf-mu_m1p.txt 
    5   7/2 +     1   3/2    -13.267    3.298  log_O19_sdpf-mu_m1p.txt 
    6   5/2 +     2   3/2    -13.074    3.491  log_O19_sdpf-mu_m1p.txt
    """
    levels = []
    negative_spin_counts = 0
    for _ in range(3): infile.readline()
    for line in infile:
        try:
            tmp = line.split()
            
            if tmp[1] == "-1":
                """
                -1 spin states in the KSHELL data file indicates
                bad states which should not be included.
                """
                negative_spin_counts += 1  # Debug.
                continue
            
            parity = 1 if tmp[2] == "+" else -1
            energy = float(tmp[5])
            spin = 2*float(Fraction(tmp[1]))
            idx = int(tmp[3])
            levels.append([energy, spin, parity, idx])
        except IndexError:
            """
            End of energies.
            """
            break

    return levels, negative_spin_counts

def _load_transition_probabilities_old(infile: TextIO) -> tuple[list, int]:
    """
    For summary files with old syntax (pre 2021-11-24).

    Parameters
    ----------
    infile : TextIO
        The KSHELL summary file at the starting position of either of
        the transition probability sections.

    Returns
    -------
    transitions : list
        List of transition data.
        
    negative_spin_counts : int
        The number of negative spin levels encountered.
    """
    negative_spin_counts = 0
    transitions = []
    for _ in range(2): infile.readline()
    for line in infile:
        try:
            """
            Example of possible lines in file:
            J_i    Ex_i     J_f    Ex_f   dE        B(M1)->         B(M1)<- 
            2+(11) 18.393 2+(10) 17.791 0.602 0.1(  0.0) 0.1( 0.0)
            3/2+( 1) 0.072 5/2+( 1) 0.000 0.071 0.127( 0.07) 0.084( 0.05)
            2+(10) 17.791 2+( 1) 5.172 12.619 0.006( 0.00) 0.006( 0.00)
            3+( 8) 19.503 2+(11) 18.393 1.111 0.000( 0.00) 0.000( 0.00)
            1+( 7) 19.408 2+( 9) 16.111 3.297 0.005( 0.00) 0.003( 0.00)
            5.0+(60) 32.170  4.0+(100) 31.734  0.436    0.198( 0.11)    0.242( 0.14)
            4.0-( 3)  3.191  3.0+(10)  3.137  0.054      0.0(  0.0)      0.0(  0.0)
            0.0+(46)', '47.248', '1.0+(97)', '45.384', '1.864', '23.973(13.39)', '7.991(', '4.46)
            """
            tmp = line.split()
            len_tmp = len(tmp)
            case_ = None # Used for identifying which if-else case reads wrong.
            
            # Location of initial parity is common for all cases.
            parity_idx = tmp[0].index("(") - 1 # Find index of initial parity.
            parity_initial = 1 if tmp[0][parity_idx] == "+" else -1
            parity_initial_symbol = tmp[0][parity_idx]
            
            # Location of initial spin is common for all cases.
            spin_initial = float(Fraction(tmp[0][:parity_idx]))
            
            if (tmp[1][-1] != ")") and (tmp[3][-1] != ")") and (len_tmp == 9):
                """
                Example:
                J_i    Ex_i     J_f    Ex_f   dE        B(M1)->         B(M1)<- 
                2+(11)   18.393  2+(10)    17.791  0.602    0.1(    0.0)    0.1(    0.0)
                5.0+(60) 32.170  4.0+(100) 31.734  0.436    0.198( 0.11)    0.242( 0.14)
                """
                case_ = 0
                E_gamma = float(tmp[4])
                Ex_initial = float(tmp[1])
                reduced_transition_prob_decay = float(tmp[5][:-1])
                reduced_transition_prob_excite = float(tmp[7][:-1])
                parity_final_symbol = tmp[2].split("(")[0][-1]
                spin_final = float(Fraction(tmp[2].split(parity_final_symbol)[0]))
                Ex_final = float(tmp[3])
                idx_initial = int(tmp[0].split("(")[1].split(")")[0])
                idx_final = int(tmp[2].split("(")[1].split(")")[0])

            elif (tmp[1][-1] != ")") and (tmp[3][-1] == ")") and (len_tmp == 10):
                """
                Example:
                J_i    Ex_i     J_f    Ex_f   dE        B(M1)->         B(M1)<- 
                2+(10) 17.791 2+( 1) 5.172 12.619 0.006( 0.00) 0.006( 0.00)
                """
                case_ = 1
                E_gamma = float(tmp[5])
                Ex_initial = float(tmp[1])
                reduced_transition_prob_decay = float(tmp[6][:-1])
                reduced_transition_prob_excite = float(tmp[8][:-1])
                parity_final_symbol = tmp[2].split("(")[0][-1]
                # spin_final = float(Fraction(tmp[2][:-2]))
                spin_final = float(Fraction(tmp[2].split(parity_final_symbol)[0]))
                Ex_final = float(tmp[4])
                idx_initial = int(tmp[0].split("(")[1].split(")")[0])
                idx_final = int(tmp[3][0])
            
            elif (tmp[1][-1] == ")") and (tmp[4][-1] != ")") and (len_tmp == 10):
                """
                Example:
                J_i    Ex_i     J_f    Ex_f   dE        B(M1)->         B(M1)<- 
                3+( 8)   19.503 2+(11)    18.393 1.111 0.000( 0.00) 0.000( 0.00)
                1.0+( 1) 5.357  0.0+(103) 0.000  5.357 0.002( 0.00) 0.007( 0.00)
                4.0-( 3)  3.191  3.0+(10)  3.137  0.054      0.0(  0.0)      0.0(  0.0)
                """
                case_ = 2
                E_gamma = float(tmp[5])
                Ex_initial = float(tmp[2])
                reduced_transition_prob_decay = float(tmp[6][:-1])
                reduced_transition_prob_excite = float(tmp[8][:-1])
                parity_final_symbol = tmp[3].split("(")[0][-1]
                spin_final = float(Fraction(tmp[3].split(parity_final_symbol)[0]))
                Ex_final = float(tmp[4])
                idx_initial = int(tmp[1][0])
                idx_final = int(tmp[3].split("(")[1].split(")")[0])

            elif (tmp[1][-1] == ")") and (tmp[4][-1] == ")") and (len_tmp == 11):
                """
                Example:
                J_i    Ex_i     J_f    Ex_f   dE        B(M1)->         B(M1)<- 
                1+( 7) 19.408 2+( 9) 16.111 3.297 0.005( 0.00) 0.003( 0.00)
                """
                case_ = 3
                E_gamma = float(tmp[6])
                Ex_initial = float(tmp[2])
                reduced_transition_prob_decay = float(tmp[7][:-1])
                reduced_transition_prob_excite = float(tmp[9][:-1])
                parity_final_symbol = tmp[3].split("(")[0][-1]
                # spin_final = float(Fraction(tmp[3][:-2]))
                spin_final = float(Fraction(tmp[3].split(parity_final_symbol)[0]))
                Ex_final = float(tmp[5])
                idx_initial = int(tmp[1][0])
                idx_final = int(tmp[4][0])

            elif (tmp[5][-1] == ")") and (tmp[2][-1] == ")") and (len_tmp == 8):
                """
                Example:
                J_i    Ex_i     J_f    Ex_f   dE        B(M1)->         B(M1)<- 
                0.0+(46) 47.248  1.0+(97) 45.384  1.864   23.973(13.39)    7.991( 4.46)
                """
                case_ = 4
                E_gamma = float(tmp[4])
                Ex_initial = float(tmp[1])
                reduced_transition_prob_decay = float(tmp[5].split("(")[0])
                reduced_transition_prob_excite = float(tmp[6][:-1])
                parity_final_symbol = tmp[2].split("(")[0][-1]
                spin_final = float(Fraction(tmp[2].split(parity_final_symbol)[0]))
                Ex_final = float(tmp[3])
                idx_initial = int(tmp[0].split("(")[1].split(")")[0])
                idx_final = int(tmp[2].split("(")[1].split(")")[0])

            else:
                msg = "ERROR: Structure not accounted for!"
                msg += f"\n{line=}"
                raise KshellDataStructureError(msg)

            if parity_final_symbol == "+":
                parity_final = 1
            elif parity_final_symbol == "-":
                parity_final = -1
            else:
                msg = f"Could not properly read the final parity! {case_=}"
                raise KshellDataStructureError(msg)

            if (spin_final == -1) or (spin_initial == -1):
                """
                -1 spin states in the KSHELL data file indicates
                bad states which should not be included.
                """
                negative_spin_counts += 1  # Debug.
                continue

            # reduced_transition_prob_decay_list.append([
            #     2*spin_initial, parity_initial, Ex_initial, 2*spin_final,
            #     parity_final, Ex_final, E_gamma, reduced_transition_prob_decay,
            #     reduced_transition_prob_excite
            # ])
            transitions.append([
                2*spin_initial, parity_initial, idx_initial, Ex_initial,
                2*spin_final, parity_final, idx_final, Ex_final, E_gamma,
                reduced_transition_prob_decay, reduced_transition_prob_excite
            ])

        except ValueError as err:
            """
            One of the float conversions failed indicating that
            the structure of the line is not accounted for.
            """
            msg = "\n" + err.__str__() + f"\n{case_=}" + f"\n{line=}"
            raise KshellDataStructureError(msg)

        except IndexError:
            """
            End of probabilities.
            """
            break

    return transitions, negative_spin_counts
    
def _load_transition_probabilities(infile: TextIO) -> tuple[list, int]:
    """
    For summary files with new syntax (post 2021-11-24).

    Parameters
    ----------
    infile : TextIO
        The KSHELL summary file at the starting position of either of
        the transition probability sections.

    Returns
    -------
    transitions : list
        List of transition data.

    negative_spin_counts : int
        The number of negative spin levels encountered.

    Example
    -------
    B(E2)  ( > -0.0 W.u.)  mass = 50    1 W.u. = 10.9 e^2 fm^4
    e^2 fm^4 (W.u.)
    J_i  pi_i idx_i Ex_i    J_f  pi_f idx_f Ex_f      dE         B(E2)->         B(E2)->[wu]     B(E2)<-         B(E2)<-[wu]
    5    +    1     0.036   6    +    1     0.000     0.036     70.43477980      6.43689168     59.59865983      5.44660066
    4    +    1     0.074   6    +    1     0.000     0.074     47.20641983      4.31409897     32.68136758      2.98668391
    """
    negative_spin_counts = 0
    transitions = []
    for _ in range(2): infile.readline()
    for line in infile:
        line_split = line.split()
        if not line_split: break
        
        spin_initial = float(Fraction(line_split[0]))
        parity_initial = _parity_string_to_integer(line_split[1])
        idx_initial = int(line_split[2])
        Ex_initial = float(line_split[3])

        spin_final = float(Fraction(line_split[4]))
        parity_final = _parity_string_to_integer(line_split[5])
        idx_final = int(line_split[2])
        Ex_final = float(line_split[7])

        E_gamma = float(line_split[8])
        reduced_transition_prob_decay = float(line_split[9])
        reduced_transition_prob_excite = float(line_split[11])

        if (spin_final < 0) or (spin_initial < 0):
            """
            -1 spin states in the KSHELL data file indicates
            bad states which should not be included.
            """
            negative_spin_counts += 1  # Debug.
            continue

        # reduced_transition_prob_decay_list.append([
        #     2*spin_initial, parity_initial, Ex_initial, 2*spin_final,
        #     parity_final, Ex_final, E_gamma, reduced_transition_prob_decay,
        #     reduced_transition_prob_excite
        # ])
        transitions.append([
            2*spin_initial, parity_initial, idx_initial, Ex_initial,
            2*spin_final, parity_final, idx_final, Ex_final, E_gamma,
            reduced_transition_prob_decay, reduced_transition_prob_excite
        ])

    return transitions, negative_spin_counts

def _generic_loader(arg_list: list) -> tuple[list, int]:
    """
    Constructed for parallel loading, but can be used in serial as well.
    """
    fname, condition, loader, thread_idx = arg_list
    
    if flags["parallel"]:
        print(f"Thread {thread_idx} loading {condition} values...")
    else:
        print(f"Loading {condition} values...")
        
    load_time = time.perf_counter()
    
    with open(fname, "r") as infile:
        for line in infile:
            if condition in line:
                ans = loader(infile)
                break
        else:
            ans = [], 0
    
    load_time = time.perf_counter() - load_time
    
    if not ans[0]:
        print(f"No {condition} transitions found in {fname}")
    else:
        print(f"Thread {thread_idx} finished loading {condition} values in {load_time:.2f} s")

    return ans

def _load_transition_probabilities_jem(infile: TextIO) -> tuple[list, int]:
    """
    JEM has modified the summary files from KSHELL with a slightly
    different syntax. This function reads that syntax. Note also that
    these summary files have 2*J, not J.

    Parameters
    ----------
    infile : TextIO
        The KSHELL summary file at the starting position of either of
        the transition probability sections.

    Returns
    -------
    transitions : list
        List of transition data.

    negative_spin_counts : int
        The number of negative spin levels encountered.

    Example
    -------
    B(M1)  larger than 1e-08 mu_N^2
    2Ji        Ei      2Jf        Ef       Ex            B(M1)->         B(M1)<- 
    2 - (   1) -35.935   0 - (   8) -35.583   0.352      0.00428800      0.01286400
    0 - (   8) -35.583   2 - (   2) -35.350   0.233      0.45171030      0.15057010
    0 - (   8) -35.583   2 - (   3) -34.736   0.847      0.04406500      0.01468830
    """
    negative_spin_counts = 0
    transitions = []

    infile.readline()   # Skip header line.
    for line in infile:
        line_split = line.split()
        if not line_split: break

        spin_initial = int(line_split[0])
        parity_initial = _parity_string_to_integer(line_split[1])
        idx_initial = int(line_split[3].strip(")"))
        Ex_initial = float(line_split[4])
        spin_final = int(line_split[5])
        parity_final = _parity_string_to_integer(line_split[6])
        idx_final = int(line_split[8].strip(")"))
        Ex_final = float(line_split[9])
        E_gamma = float(line_split[10])
        reduced_transition_prob_decay = float(line_split[11])
        reduced_transition_prob_excite = float(line_split[12])

        if (spin_final < 0) or (spin_initial < 0):
            """
            -1 spin states in the KSHELL data file indicates
            bad states which should not be included.
            """
            negative_spin_counts += 1  # Debug.
            continue

        transitions.append([
            spin_initial, parity_initial, idx_initial, Ex_initial,
            spin_final, parity_final, idx_final, Ex_final, E_gamma,
            reduced_transition_prob_decay, reduced_transition_prob_excite
        ])
    
    return transitions, negative_spin_counts
<file_sep>from __future__ import annotations
import os, sys, multiprocessing, hashlib, ast, time, re
from fractions import Fraction
from typing import Union, Callable, Tuple, Iterable
from itertools import chain
import numpy as np
import numba
import matplotlib.pyplot as plt
import seaborn as sns
from .collect_logs import collect_logs
from .kshell_exceptions import KshellDataStructureError
from .parameters import atomic_numbers, flags
from .general_utilities import (
    level_plot, level_density, gamma_strength_function_average, porter_thomas,
    isotope
)
from .loaders import (
    _generic_loader, _load_energy_levels, _load_transition_probabilities,
    _load_transition_probabilities_old, _load_transition_probabilities_jem
)

def _generate_unique_identifier(path: str) -> str:
    """
    Generate a unique identifier based on the shell script and the
    save_input file from KSHELL.

    Parameters
    ----------
    path : str
        The path to a summary file or a directory with a summary file.
    """
    shell_file_content = ""
    save_input_content = ""
    msg = "Not able to generate unique identifier!"
    if os.path.isfile(path):
        """
        If a file is specified, extract the directory from the path.
        """
        directory = path.rsplit("/", 1)[0]
        if directory == path:
            """
            Example: path is 'summary.txt'
            """
            directory = "."
    
    elif os.path.isdir(path):
        directory = path

    for elem in os.listdir(directory):
        """
        Loop over all elements in the directory and find the shell
        script and save_input file.
        """
        try:
            if elem.endswith(".sh"):
                with open(f"{directory}/{elem}", "r") as infile:
                    shell_file_content += infile.read()
                    
            elif "save_input_ui.txt" in elem:
                with open(f"{directory}/{elem}", "r") as infile:
                    save_input_content += infile.read()
        except UnicodeDecodeError:
            msg = f"Skipping {elem} for tmp file unique identifier due to UnicodeDecodeError."
            msg += " Are you sure this file is supposed to be in this directory?"
            print(msg)
            continue

    if (shell_file_content == "") and (save_input_content == ""):
        print(msg)

    return hashlib.sha1((shell_file_content + save_input_content).encode()).hexdigest()

class ReadKshellOutput:
    """
    Read `KSHELL` data files and store the values as instance
    attributes.

    Attributes
    ----------
    levels : np.ndarray
        Array containing energy, spin, and parity for each excited
        state. [[E, 2*spin, parity, idx], ...]. idx counts how many
        times a state of that given spin and parity has occurred. The
        first 0+ state will have an idx of 1, the second 0+ will have an
        idx of 2, etc.

    transitions_BE1 : np.ndarray
        Transition data for BE1 transitions. Structure:
        NEW:
        [2*spin_initial, parity_initial, idx_initial, Ex_initial,
        2*spin_final, parity_final, idx_final, Ex_final, E_gamma,
        B(.., i->f), B(.., f<-i)]
        OLD NEW:
        [2*spin_initial, parity_initial, Ex_initial, 2*spin_final,
        parity_final, Ex_final, E_gamma, B(.., i->f), B(.., f<-i)]
        OLD:
        Mx8 array containing [2*spin_final, parity_initial, Ex_final,
        2*spin_initial, parity_initial, Ex_initial, E_gamma, B(.., i->f)].

    transitions_BM1 : np.ndarray
        Transition data for BM1 transitions. Same structure as BE1.

    transitions_BE2 : np.ndarray
        Transition data for BE2 transitions. Same structure as BE1.
    """
    def __init__(self, path: str, load_and_save_to_file: bool, old_or_new: str):
        """
        Parameters
        ----------
        path : string
            Path of `KSHELL` output file directory, or path to a
            specific `KSHELL` data file.

        load_and_save_to_file : bool
            Toggle saving data as `.npy` files on / off. If `overwrite`,
            saved `.npy` files are overwritten.

        old_or_new : str
            Choose between old and new summary file syntax. All summary
            files generated pre 2021-11-24 use old style.
            New:
            J_i  pi_i idx_i Ex_i    J_f  pi_f idx_f Ex_f      dE         B(E2)->         B(E2)->[wu]     B(E2)<-         B(E2)<-[wu]
            5    +    1     0.036   6    +    1     0.000     0.036     70.43477980      6.43689168     59.59865983      5.44660066
            Old:
            J_i    Ex_i     J_f    Ex_f   dE        B(M1)->         B(M1)<- 
            2+(11) 18.393 2+(10) 17.791 0.602 0.1(  0.0) 0.1( 0.0)
        """

        self.path = path.rstrip("/")    # Just in case, prob. not necessary.
        self.load_and_save_to_file = load_and_save_to_file
        self.old_or_new = old_or_new
        # Some attributes might not be altered, depending on the input file.
        # self.fname_ptn = None     # I'll fix .pnt reading when the functionality is actually needed.
        # self.fname_summary = None
        # self.proton_partition = None
        # self.neutron_partition = None
        # self.truncation = None
        self.nucleus = None
        self.interaction = None
        self.levels = None
        self.transitions_BM1 = None
        self.transitions_BE2 = None
        self.transitions_BE1 = None
        self.npy_path = "tmp"   # Directory for storing .npy files.
        # Debug.
        self.negative_spin_counts = np.array([0, 0, 0, 0])  # The number of skipped -1 spin states for [levels, BM1, BE2, BE1].

        if isinstance(self.load_and_save_to_file, str) and (self.load_and_save_to_file != "overwrite"):
            msg = "Allowed values for 'load_and_save_to_file' are: 'True', 'False', 'overwrite'."
            msg += f" Got '{self.load_and_save_to_file}'."
            raise ValueError(msg)

        if os.path.isdir(self.path):
            """
            If input 'path' is a directory. Look for summaries.
            """
            summaries = [i for i in os.listdir(self.path) if "summary" in i]

            if not summaries:
                """
                No summaries found, call collect_logs.
                """
                self.fname_summary = collect_logs(
                    path = self.path,
                    old_or_new = "both"
                )
                self.path_summary = f"{self.path}/{self.fname_summary}"

            elif len(summaries) > 1:
                msg = f"Several summaries found in {self.path}. Please specify path"
                msg += " directly to the summary file you want to load."
                raise RuntimeError(msg)

            else:
                self.fname_summary = summaries[0]   # Just filename.
                self.path_summary = f"{self.path}/{summaries[0]}"    # Complete path (maybe relative).

        elif os.path.isfile(self.path):
            """
            'path' is a single file, not a directory.
            """
            self.fname_summary = path.split("/")[-1]       # Just filename.
            self.path_summary = self.path    # Complete path (maybe relative).

        else:
            msg = f"{self.path} is not a file or a directory!"
            tmp_directory = self.path.rsplit('/', 1)[0]
            if os.path.isdir(tmp_directory):
                summaries = [i for i in os.listdir(tmp_directory) if "summary" in i]
                if summaries:
                    msg += f" {len(summaries)} summary files were found in {tmp_directory}."
                    msg += f" {summaries}"
                
                logs = [i for i in os.listdir(tmp_directory) if i.startswith("log_")]
                if logs and (not summaries):
                    msg += f" Logs found in '{tmp_directory}'. Set path to '{tmp_directory}'"
                    msg += " to collect logs and load the summary file."

            raise RuntimeError(msg)

        self.base_fname = self.fname_summary.split(".")[0] # Base filename for .npy tmp files.
        self.unique_id = _generate_unique_identifier(self.path) # Unique identifier for .npy files.
        self._extract_info_from_summary_fname()
        self._read_summary()
        self.mixing_pairs_BM1_BE2 = self._mixing_pairs(B_left="M1", B_right="E2")
        
        self.ground_state_energy = self.levels[0, 0]
        
        try:
            self.A = int("".join(filter(str.isdigit, self.nucleus)))
            self.Z, self.N = isotope(
                name = "".join(filter(str.isalpha, self.nucleus)).lower(),
                A = self.A
            )
        except ValueError:
            """
            Prob. because the summary filename does not contain the name
            of the isotope.
            """
            self.A = None
            self.Z = None
            self.N = None
        
        self.check_data()

    def _extract_info_from_ptn_fname(self):
        """
        Extract nucleus and model space name.
        """
        fname_split = self.fname_ptn.split("/")[-1]
        fname_split = fname_split.split("_")
        self.nucleus = fname_split[0]
        self.interaction = fname_split[1]

    def _read_ptn(self):
        """
        Read `KSHELL` partition file (.ptn) and extract proton
        partition, neutron partition, and particle-hole truncation data.
        Save as instance attributes.
        """

        line_number = 0
        line_number_inner = 0
        self.truncation = []

        with open(self.fname_ptn, "r") as infile:
            for line in infile:
                line_number += 1
                
                if line.startswith("# proton partition"):
                    for line_inner in infile:
                        """
                        Read until next '#'.
                        """
                        line_number_inner += 1
                        if line_inner.startswith("#"):
                            line = line_inner
                            break
                    
                    self.proton_partition = np.loadtxt(
                        fname = self.fname_ptn,
                        skiprows = line_number,
                        max_rows = line_number_inner
                    )
                    line_number += line_number_inner
                    line_number_inner = 0
                
                if line.startswith("# neutron partition"):
                    for line_inner in infile:
                        """
                        Read until next '#'.
                        """
                        line_number_inner += 1
                        if line_inner.startswith("#"):
                            line = line_inner
                            break
                    
                    self.neutron_partition = np.loadtxt(
                        fname = self.fname_ptn,
                        skiprows = line_number,
                        max_rows = line_number_inner
                    )
                    line_number += line_number_inner
                    line_number_inner = 0

                if line.startswith("# particle-hole truncation"):
                    for line_inner in infile:
                        """
                        Loop over all particle-hole truncation lines.
                        """
                        line_number += 1
                        line_inner_split = line_inner.split()

                        if (len(line_inner_split) < 2):
                            """
                            Condition will probably not get fulfilled.
                            Safety precaution due to indexing in this
                            loop.
                            """
                            break

                        if (line_inner_split[1]).startswith("["):
                            """
                            '[' indicates that 'line_inner' is still
                            containing truncation information.
                            """
                            for colon_index, elem in enumerate(line_inner_split):
                                """
                                Find the index of the colon ':' to
                                decide the orbit numbers and occupation
                                numbers.
                                """
                                if (elem == ":"): break

                            occupation = [int(occ) for occ in line_inner_split[colon_index + 1:]]   # [min, max].
                            orbit_numbers = "".join(line_inner_split[1:colon_index])
                            orbit_numbers = orbit_numbers.replace("[", "")
                            orbit_numbers = orbit_numbers.replace("]", "")
                            orbit_numbers = orbit_numbers.replace(" ", "")  # This can prob. be removed because of the earlier split.
                            orbit_numbers = orbit_numbers.split(",")
                            orbit_numbers = [int(orbit) for orbit in orbit_numbers]
                            
                            for orbit in orbit_numbers:
                                self.truncation.append((orbit, occupation))
                        
                        else:
                            """
                            Line does not contain '[' and thus does not
                            contain truncation information.
                            """
                            break

    def _extract_info_from_summary_fname(self):
        """
        Extract nucleus and model space name.
        """
        fname_split = self.fname_summary.split("/")[-1]  # Remove path.
        fname_split = fname_split.split(".")[0] # Remove .txt.
        fname_split = fname_split.split("_")
        self.nucleus = fname_split[1]
        self.interaction = fname_split[2]

    def _read_summary(self):
        """
        Read energy level data, transition probabilities and transition
        strengths from `KSHELL` output files.

        Raises
        ------
        KshellDataStructureError
            If the `KSHELL` file has unexpected structure / syntax.
        """
        # npy_path = "tmp"
        # base_fname = self.path.split("/")[-1][:-4]
        # unique_id = _generate_unique_identifier(self.path)

        if self.load_and_save_to_file:
            try:
                os.mkdir(self.npy_path)
            except FileExistsError:
                pass

            with open(f"{self.npy_path}/README.txt", "w") as outfile:
                msg = "This directory contains binary numpy data of KSHELL summary data."
                msg += " The purpose is to speed up subsequent runs which use the same summary data."
                msg += " It is safe to delete this entire directory if you have the original summary text file, "
                msg += "though at the cost of having to read the summary text file over again which may take some time."
                msg += " The ksutil.loadtxt parameter load_and_save_to_file = 'overwrite' will force a re-write of the binary numpy data."
                outfile.write(msg)
        
        levels_fname = f"{self.npy_path}/{self.base_fname}_levels_{self.unique_id}.npy"
        transitions_BM1_fname = f"{self.npy_path}/{self.base_fname}_transitions_BM1_{self.unique_id}.npy"
        transitions_BE2_fname = f"{self.npy_path}/{self.base_fname}_transitions_BE2_{self.unique_id}.npy"
        transitions_BE1_fname = f"{self.npy_path}/{self.base_fname}_transitions_BE1_{self.unique_id}.npy"
        debug_fname = f"{self.npy_path}/{self.base_fname}_debug_{self.unique_id}.npy"

        fnames = [
            levels_fname, transitions_BE2_fname, transitions_BM1_fname,
            transitions_BE1_fname, debug_fname
        ]

        if self.load_and_save_to_file != "overwrite":
            """
            Do not load files if overwrite parameter has been passed.
            """
            if all([os.path.isfile(fname) for fname in fnames]) and self.load_and_save_to_file:
                """
                If all files exist, load them. If any of the files do
                not exist, all will be generated.
                """
                self.levels = np.load(file=levels_fname, allow_pickle=True)
                self.transitions_BM1 = np.load(file=transitions_BM1_fname, allow_pickle=True)
                self.transitions_BE2 = np.load(file=transitions_BE2_fname, allow_pickle=True)
                self.transitions_BE1 = np.load(file=transitions_BE1_fname, allow_pickle=True)
                self.debug = np.load(file=debug_fname, allow_pickle=True)
                msg = "Summary data loaded from .npy!"
                msg += " Use loadtxt parameter load_and_save_to_file = 'overwrite'"
                msg += " to re-read data from the summary file."
                print(msg)
                return

        parallel_args = [
            [self.path_summary, "Energy", "replace_this_entry_with_loader", 0],
            [self.path_summary, "B(E1)", "replace_this_entry_with_loader", 1],
            [self.path_summary, "B(M1)", "replace_this_entry_with_loader", 2],
            [self.path_summary, "B(E2)", "replace_this_entry_with_loader", 3],
        ]

        if self.old_or_new == "new":
            parallel_args[0][2] = _load_energy_levels
            parallel_args[1][2] = _load_transition_probabilities
            parallel_args[2][2] = _load_transition_probabilities
            parallel_args[3][2] = _load_transition_probabilities

        elif self.old_or_new == "old":
            parallel_args[0][2] = _load_energy_levels
            parallel_args[1][2] = _load_transition_probabilities_old
            parallel_args[2][2] = _load_transition_probabilities_old
            parallel_args[3][2] = _load_transition_probabilities_old

        elif self.old_or_new == "jem":
            parallel_args[0][2] = _load_energy_levels
            parallel_args[1][2] = _load_transition_probabilities_jem
            parallel_args[2][2] = _load_transition_probabilities_jem
            parallel_args[3][2] = _load_transition_probabilities_jem

        if flags["parallel"]:
            with multiprocessing.Pool() as pool:
                pool_res = pool.map(_generic_loader, parallel_args)
                self.levels, self.negative_spin_counts[0] = pool_res[0]
                self.transitions_BE1, self.negative_spin_counts[1] = pool_res[1]
                self.transitions_BM1, self.negative_spin_counts[2] = pool_res[2]
                self.transitions_BE2, self.negative_spin_counts[3] = pool_res[3]
        else:
            self.levels, self.negative_spin_counts[0] = _generic_loader(parallel_args[0])
            self.transitions_BE1, self.negative_spin_counts[1] = _generic_loader(parallel_args[1])
            self.transitions_BM1, self.negative_spin_counts[2] = _generic_loader(parallel_args[2])
            self.transitions_BE2, self.negative_spin_counts[3] = _generic_loader(parallel_args[3])

        self.levels = np.array(self.levels)
        self.transitions_BE1 = np.array(self.transitions_BE1)
        self.transitions_BM1 = np.array(self.transitions_BM1)
        self.transitions_BE2 = np.array(self.transitions_BE2)
        self.debug = "DEBUG\n"
        self.debug += f"skipped -1 states in levels: {self.negative_spin_counts[0]}\n"
        self.debug += f"skipped -1 states in BE1: {self.negative_spin_counts[1]}\n"
        self.debug += f"skipped -1 states in BM1: {self.negative_spin_counts[2]}\n"
        self.debug += f"skipped -1 states in BE2: {self.negative_spin_counts[3]}\n"
        self.debug = np.array(self.debug)

        if self.old_or_new == "jem":
            """
            'jem style' summary syntax lists all initial and final
            excitation energies in transitions as absolute values.
            Subtract the ground state energy to get the relative
            energies to match the newer KSHELL summary file syntax.
            """
            msg = "The issue of E_final > E_initial must be figured out before"
            msg += " JEM style syntax can be used!"
            raise NotImplementedError(msg)
            E_gs = abs(self.levels[0, 0])   # Can prob. just use ... -= E_gs
            try:
                self.transitions_BM1[:, 3] = E_gs - np.abs(self.transitions_BM1[:, 3])
                self.transitions_BM1[:, 7] = E_gs - np.abs(self.transitions_BM1[:, 7])
            except IndexError:
                """
                No BM1 transitions.
                """
                pass
            try:
                self.transitions_BE1[:, 3] = E_gs - np.abs(self.transitions_BE1[:, 3])
                self.transitions_BE1[:, 7] = E_gs - np.abs(self.transitions_BE1[:, 7])
            except IndexError:
                """
                No BE1 transitions.
                """
                pass
            try:
                self.transitions_BE2[:, 3] = E_gs - np.abs(self.transitions_BE2[:, 3])
                self.transitions_BE2[:, 7] = E_gs - np.abs(self.transitions_BE2[:, 7])
            except IndexError:
                """
                No BE2 transitions.
                """
                pass

            self.levels[:, 1] /= 2  # JEM style syntax has 2*J already. Without this correction it would be 4*J.

        if self.load_and_save_to_file:
            np.save(file=levels_fname, arr=self.levels, allow_pickle=True)
            np.save(file=transitions_BM1_fname, arr=self.transitions_BM1, allow_pickle=True)
            np.save(file=transitions_BE2_fname, arr=self.transitions_BE2, allow_pickle=True)
            np.save(file=transitions_BE1_fname, arr=self.transitions_BE1, allow_pickle=True)
            np.save(file=debug_fname, arr=self.debug, allow_pickle=True)

    def level_plot(self,
        include_n_levels: int = 1000,
        filter_spins: Union[None, list] = None,
        filter_parity: Union[None, str] = None,
        color: Union[None, str] = "black"
        ):
        """
        Wrapper method to include level plot as an attribute to this
        class. Generate a level plot for a single isotope. Angular
        momentum on the x axis, energy on the y axis.

        Parameters
        ----------
        include_n_levels : int
            The maximum amount of states to plot for each spin. Default
            set to a large number to indicate ≈ no limit.

        filter_spins : Union[None, list]
            Which spins to include in the plot. If `None`, all spins are
            plotted. Defaults to `None`

        color : Union[None, str]
            Set the color of the level lines.
        """
        level_plot(
            levels = self.levels,
            include_n_levels = include_n_levels,
            filter_spins = filter_spins,
            filter_parity = filter_parity,
            color = color,
        )

    def level_density_plot(self,
            bin_width: Union[int, float] = 0.2,
            include_n_levels: Union[None, int] = None,
            filter_spins: Union[None, int, list] = None,
            filter_parity: Union[None, str, int] = None,
            return_counts: bool = False,
            E_min: Union[float, int] = 0,
            E_max: Union[float, int] = np.inf,
            plot: bool = True,
            save_plot: bool = False
        ):
        """
        Wrapper method to include level density plotting as
        an attribute to this class. Generate the level density with the
        input bin size.

        Parameters
        ----------
        See level_density in general_utilities.py for parameter
        information.
        """
        bins, density = level_density(
            levels = self.levels,
            bin_width = bin_width,
            include_n_levels = include_n_levels,
            filter_spins = filter_spins,
            filter_parity = filter_parity,
            return_counts = return_counts,
            E_min = E_min,
            E_max = E_max,
            plot = plot,
            save_plot = save_plot
        )

        return bins, density

    def nld(self,
        bin_width: Union[int, float] = 0.2,
        include_n_levels: Union[None, int] = None,
        filter_spins: Union[None, int, list] = None,
        filter_parity: Union[None, str, int] = None,
        E_min: Union[float, int] = 0,
        E_max: Union[float, int] = np.inf,
        return_counts: bool = False,
        plot: bool = True,
        save_plot: bool = False
        ):
        """
        Wrapper method to level_density_plot.
        """
        return self.level_density_plot(
            bin_width = bin_width,
            include_n_levels = include_n_levels,
            filter_spins = filter_spins,
            filter_parity = filter_parity,
            E_min = E_min,
            E_max = E_max,
            return_counts = return_counts,
            plot = plot,
            save_plot = save_plot
        )

    def gamma_strength_function_average_plot(self,
        bin_width: Union[float, int] = 0.2,
        Ex_min: Union[float, int] = 5,
        Ex_max: Union[float, int] = 50,
        multipole_type: str = "M1",
        prefactor_E1: Union[None, float] = None,
        prefactor_M1: Union[None, float] = None,
        prefactor_E2: Union[None, float] = None,
        initial_or_final: str = "initial",
        partial_or_total: str = "partial",
        include_only_nonzero_in_average: bool = True,
        include_n_levels: Union[None, int] = None,
        filter_spins: Union[None, list] = None,
        filter_parities: str = "both",
        return_n_transitions: bool = False,
        plot: bool = True,
        save_plot: bool = False
        ):
        """
        Wrapper method to include gamma ray strength function
        calculations as an attribute to this class. Includes saving
        of GSF data to .npy files.

        Parameters
        ----------
        See gamma_strength_function_average in general_utilities.py
        for parameter descriptions.
        """
        transitions_dict = {
            "M1": self.transitions_BM1,
            "E2": self.transitions_BE2,
            "E1": self.transitions_BE1
        }
        is_loaded = False
        gsf_unique_string = f"{bin_width}{Ex_min}{Ex_max}{multipole_type}"
        gsf_unique_string += f"{prefactor_E1}{prefactor_M1}{prefactor_E2}"
        gsf_unique_string += f"{initial_or_final}{partial_or_total}{include_only_nonzero_in_average}"
        gsf_unique_string += f"{include_n_levels}{filter_spins}{filter_parities}"
        gsf_unique_id = hashlib.sha1((gsf_unique_string).encode()).hexdigest()
        gsf_fname = f"{self.npy_path}/{self.base_fname}_gsf_{gsf_unique_id}_{self.unique_id}.npy"
        bins_fname = f"{self.npy_path}/{self.base_fname}_gsfbins_{gsf_unique_id}_{self.unique_id}.npy"
        n_transitions_fname = f"{self.npy_path}/{self.base_fname}_gsfntransitions_{gsf_unique_id}_{self.unique_id}.npy"
        
        fnames = [gsf_fname, bins_fname]
        if return_n_transitions:
            fnames.append(n_transitions_fname)
        
        if all([os.path.isfile(fname) for fname in fnames]) and self.load_and_save_to_file and (self.load_and_save_to_file != "overwrite"):
            """
            If all these conditions are met, all arrays will be loaded
            from file. If any of these conditions are NOT met, all
            arrays will be re-calculated.
            """
            gsf = np.load(file=gsf_fname, allow_pickle=True)
            bins = np.load(file=bins_fname, allow_pickle=True)
            if return_n_transitions:
                n_transitions = np.load(file=n_transitions_fname, allow_pickle=True)
            
            msg = f"{self.nucleus} {multipole_type} GSF data loaded from .npy!"
            print(msg)
            is_loaded = True

        else:
            tmp = gamma_strength_function_average(
                levels = self.levels,
                transitions = transitions_dict[multipole_type],
                bin_width = bin_width,
                Ex_min = Ex_min,
                Ex_max = Ex_max,
                multipole_type = multipole_type,
                prefactor_E1 = prefactor_E1,
                prefactor_M1 = prefactor_M1,
                prefactor_E2 = prefactor_E2,
                initial_or_final = initial_or_final,
                partial_or_total = partial_or_total,
                include_only_nonzero_in_average = include_only_nonzero_in_average,
                include_n_levels = include_n_levels,
                filter_spins = filter_spins,
                filter_parities = filter_parities,
                return_n_transitions = return_n_transitions,
                # plot = plot,
                # save_plot = save_plot
            )
            if return_n_transitions:
                bins, gsf, n_transitions = tmp
            else:
                bins, gsf = tmp

        if self.load_and_save_to_file and not is_loaded:
            np.save(file=gsf_fname, arr=gsf, allow_pickle=True)
            np.save(file=bins_fname, arr=bins, allow_pickle=True)
            
            if return_n_transitions:
                np.save(file=n_transitions_fname, arr=n_transitions, allow_pickle=True)

        if plot:
            unit_exponent = 2*int(multipole_type[-1]) + 1
            fig, ax = plt.subplots()
            ax.plot(bins, gsf, label=multipole_type.upper(), color="black")
            ax.legend()
            ax.grid()
            ax.set_xlabel(r"E$_{\gamma}$ [MeV]")
            ax.set_ylabel(f"$\gamma$SF [MeV$^-$$^{unit_exponent}$]")
            if save_plot:
                fname = f"gsf_{multipole_type}.png"
                print(f"GSF saved as '{fname}'")
                fig.savefig(fname=fname, dpi=300)
            plt.show()

        if return_n_transitions:
            return bins, gsf, n_transitions
        else:
            return bins, gsf

    def gsf(self,
        bin_width: Union[float, int] = 0.2,
        Ex_min: Union[float, int] = 5,
        Ex_max: Union[float, int] = 50,
        multipole_type: str = "M1",
        prefactor_E1: Union[None, float] = None,
        prefactor_M1: Union[None, float] = None,
        prefactor_E2: Union[None, float] = None,
        initial_or_final: str = "initial",
        partial_or_total: str = "partial",
        include_only_nonzero_in_average: bool = True,
        include_n_levels: Union[None, int] = None,
        filter_spins: Union[None, list] = None,
        filter_parities: str = "both",
        return_n_transitions: bool = False,
        plot: bool = True,
        save_plot: bool = False
        ):
        """
        Alias for gamma_strength_function_average_plot. See that
        docstring for details.
        """
        return self.gamma_strength_function_average_plot(
            bin_width = bin_width,
            Ex_min = Ex_min,
            Ex_max = Ex_max,
            multipole_type = multipole_type,
            prefactor_E1 = prefactor_E1,
            prefactor_M1 = prefactor_M1,
            prefactor_E2 = prefactor_E2,
            initial_or_final = initial_or_final,
            partial_or_total = partial_or_total,
            include_only_nonzero_in_average = include_only_nonzero_in_average,
            include_n_levels = include_n_levels,
            filter_spins = filter_spins,
            filter_parities = filter_parities,
            return_n_transitions = return_n_transitions,
            plot = plot,
            save_plot = save_plot
        )

    def porter_thomas(self, multipole_type: str, **kwargs):
        """
        Wrapper for general_utilities.porter_thomas. See that docstring
        for details.

        Parameters
        ----------
        multipole_type : str
            Choose the multipolarity of the transitions. 'E1', 'M1',
            'E2'.
        """
        transitions_dict = {
            "E1": self.transitions_BE1,
            "M1": self.transitions_BM1,
            "E2": self.transitions_BE2,
        }
        
        return porter_thomas(transitions_dict[multipole_type], **kwargs)

    def porter_thomas_Ei_plot(self,
        Ei_range_min: float = 5,
        Ei_range_max: float = 9,
        Ei_values: Union[list, None] = None,
        Ei_bin_width: float = 0.2,
        BXL_bin_width: float = 0.1,
        multipole_type: Union[str, list] = "M1",
        set_title: bool = True
        ):
        """
        Porter-Thomas analysis of the reduced transition probabilities
        for different initial excitation energies. Produces a figure
        very similar to fig. 3.3 in JEM PhD thesis:
        http://urn.nb.no/URN:NBN:no-79895.

        Parameters
        ----------
        Ei_range_min : float
            Minimum value of the initial energy range. Three equally
            spaced intervals will be chosen from this range. Be sure to
            choose this value to be above the discrete region. MeV.
        
        Ei_range_max : float
            Maximum value of the initial energy range. Three equally
            spaced intervals will be chosen from this range. The neutron
            separation energy is a good choice. MeV.

        Ei_values : Union[list, None]
            List of initial energies to be used. If None, the
            initial energies will be chosen from the Ei_range_min
            and Ei_range_max. Values in a bin around Ei_values of size
            Ei_bin_width will be used. Max 3 values allowed. MeV.

        Ei_bin_width : float
            Width of the initial energy bins. MeV.

        BXL_bin_width : float
            Width of the BXL bins when the BXL/mean(BXL) values are
            counted. Unitless.

        multipole_type : str
            Choose the multipolarity of the transitions. 'E1', 'M1',
            'E2'.

        set_title: bool
            Toggle figure title on / off. Defaults to True (on).
        """

        if Ei_values is None:
            """
            Defaults to a range defined by Ei_range_min and Ei_range_max.
            """
            Ei_values = np.linspace(Ei_range_min, Ei_range_max, 3)
        
        if len(Ei_values) > 3:
            raise ValueError("Ei_values must be a list of length <= 3.")

        if isinstance(multipole_type, str):
            multipole_type = [multipole_type]
        
        elif isinstance(multipole_type, list):
            pass
        
        else:
            msg = f"multipole_type must be str or list. Got {type(multipole_type)}."
            raise TypeError(msg)

        colors = ["blue", "royalblue", "lightsteelblue"]
        Ei_range = np.linspace(Ei_range_min, Ei_range_max, 4)
        
        if len(multipole_type) == 1:
            fig, axd = plt.subplot_mosaic(
                [['upper'], ['middle'], ['lower']],
                gridspec_kw = dict(height_ratios=[1, 1, 0.7]),
                figsize = (6.4, 8),
                constrained_layout = True,
                sharex = True
            )

            for Ei, color in zip(Ei_values, colors):
                """
                Calculate in a bin size of 'Ei_bin_width' around given Ei
                values.
                """
                bins, counts, chi2 = self.porter_thomas(
                    multipole_type = multipole_type[0],
                    Ei = Ei,
                    BXL_bin_width = BXL_bin_width,
                    Ei_bin_width = Ei_bin_width,
                    return_chi2 = True
                )
                idx = np.argmin(np.abs(bins - 10))  # Slice the arrays at approx 10.
                bins = bins[:idx]
                counts = counts[:idx]
                chi2 = chi2[:idx]
                axd["upper"].step(
                    bins,
                    counts,
                    label = r"$E_i = $" + f"{Ei:.2f}" + r" $\pm$ " + f"{Ei_bin_width/2:.2f} MeV",
                    color = color
                )
        
            axd["upper"].plot(
                bins,
                chi2,
                color = "tab:green",
                label = r"$\chi_{\nu = 1}^2$"
            )
            axd["upper"].legend(loc="upper right")
            axd["upper"].set_ylabel(r"Normalised counts")

            for i, color in enumerate(colors):
                """
                Calculate in the specified range of Ei values.
                """
                bins, counts, chi2 = self.porter_thomas(
                    multipole_type = multipole_type[0],
                    Ei = [Ei_range[i], Ei_range[i+1]],
                    BXL_bin_width = BXL_bin_width,
                    return_chi2 = True
                )
                
                idx = np.argmin(np.abs(bins - 10))
                bins = bins[:idx]
                counts = counts[:idx]
                chi2 = chi2[:idx]
                
                axd["middle"].step(
                    bins,
                    counts,
                    color = color,
                    label = r"$E_i = $" + f"[{Ei_range[i]:.2f}, {Ei_range[i+1]:.2f}] MeV"
                )
                axd["lower"].step(
                    bins,
                    counts/chi2,
                    color = color,
                    label = r"($E_i = $" + f"[{Ei_range[i]:.2f}, {Ei_range[i+1]:.2f}] MeV)" + r"$/\chi_{\nu = 1}^2$",
                )

            axd["middle"].plot(bins, chi2, color="tab:green", label=r"$\chi_{\nu = 1}^2$")
            axd["middle"].legend(loc="upper right")
            axd["middle"].set_ylabel(r"Normalised counts")

            axd["lower"].hlines(y=1, xmin=bins[0], xmax=bins[-1], linestyle="--", color="black")
            axd["lower"].set_xlabel(r"$B(M1)/\langle B(M1) \rangle$")
            axd["lower"].legend(loc="upper left")
            axd["lower"].set_ylabel(r"Relative error")
            axd["lower"].set_xlabel(
                r"$B(" + f"{multipole_type[0]}" + r")/\langle B(" + f"{multipole_type[0]}" + r") \rangle$"
            )
            if set_title:
                axd["upper"].set_title(
                    f"{self.nucleus_latex}, {self.interaction}, " + r"$" + f"{multipole_type[0]}" + r"$"
                )
            fig.savefig(fname=f"{self.nucleus}_porter_thomas_Ei_{multipole_type[0]}.png", dpi=300)
        
        elif len(multipole_type) == 2:
            fig, axd = plt.subplot_mosaic([
                    ['upper left', 'upper right'],
                    ['middle left', 'middle right'],
                    ['lower left', 'lower right']
                ],
                gridspec_kw = dict(height_ratios=[1, 1, 0.7]),
                figsize = (10, 8),
                constrained_layout = True,
                sharex = True
            )

            for multipole_type_, loc in zip(multipole_type, ["left", "right"]):
                for Ei, color in zip(Ei_values, colors):
                    """
                    Calculate in a bin size of 'Ei_bin_width' around given Ei
                    values.
                    """
                    bins, counts, chi2 = self.porter_thomas(
                        multipole_type = multipole_type_,
                        Ei = Ei,
                        BXL_bin_width = BXL_bin_width,
                        Ei_bin_width = Ei_bin_width,
                        return_chi2 = True
                    )
                    idx = np.argmin(np.abs(bins - 10))  # Slice the arrays at approx 10.
                    bins = bins[:idx]
                    counts = counts[:idx]
                    chi2 = chi2[:idx]
                    axd["upper " + loc].step(
                        bins,
                        counts,
                        label = r"$E_i = $" + f"{Ei:.2f}" + r" $\pm$ " + f"{Ei_bin_width/2:.2f} MeV",
                        color = color
                    )
            
                axd["upper " + loc].plot(
                    bins,
                    chi2,
                    color = "tab:green",
                    label = r"$\chi_{\nu = 1}^2$"
                )

                for i, color in enumerate(colors):
                    """
                    Calculate in the specified range of Ei values.
                    """
                    bins, counts, chi2 = self.porter_thomas(
                        multipole_type = multipole_type_,
                        Ei = [Ei_range[i], Ei_range[i+1]],
                        BXL_bin_width = BXL_bin_width,
                        return_chi2 = True
                    )
                    
                    idx = np.argmin(np.abs(bins - 10))
                    bins = bins[:idx]
                    counts = counts[:idx]
                    chi2 = chi2[:idx]
                    
                    axd["middle " + loc].step(
                        bins,
                        counts,
                        color = color,
                        label = r"$E_i = $" + f"[{Ei_range[i]:.2f}, {Ei_range[i+1]:.2f}] MeV"
                    )
                    axd["lower " + loc].step(
                        bins,
                        counts/chi2,
                        color = color,
                        label = r"($E_i = $" + f"[{Ei_range[i]:.2f}, {Ei_range[i+1]:.2f}] MeV)" + r"$/\chi_{\nu = 1}^2$",
                    )

                axd["middle " + loc].plot(bins, chi2, color="tab:green", label=r"$\chi_{\nu = 1}^2$")

                axd["lower " + loc].hlines(y=1, xmin=bins[0], xmax=bins[-1], linestyle="--", color="black")
                # axd["lower " + loc].set_xlabel(r"$B(M1)/\langle B(M1) \rangle$")
                axd["lower " + loc].set_xlabel(
                    r"$B(" + f"{multipole_type_}" + r")/\langle B(" + f"{multipole_type_}" + r") \rangle$"
                )
                if set_title:
                    axd["upper " + loc].set_title(
                        f"{self.nucleus_latex}, {self.interaction}, " + r"$" + f"{multipole_type_}" + r"$"
                    )
            
            axd["lower left"].legend(loc="upper left")
            axd["middle left"].legend(loc="upper right")
            axd["middle left"].set_ylabel(r"Normalised counts")
            axd["lower left"].set_ylabel(r"Relative error")
            axd["upper left"].legend(loc="upper right")
            axd["upper left"].set_ylabel(r"Normalised counts")
            
            axd["upper right"].set_yticklabels([])
            axd["upper right"].set_ylim(axd["upper left"].get_ylim())

            axd["middle right"].set_yticklabels([])
            axd["middle right"].set_ylim(axd["middle left"].get_ylim())

            axd["lower right"].set_yticklabels([])
            axd["lower right"].set_ylim(axd["lower left"].get_ylim())
            fig.savefig(fname=f"{self.nucleus}_porter_thomas_Ei_{multipole_type[0]}_{multipole_type[1]}.png", dpi=300)

        else:
            msg = "Only 1 or 2 multipole types may be given at the same time!"
            msg += f" Got {len(multipole_type)}."
            raise ValueError(msg)
        
        plt.show()

    def _porter_thomas_j_plot_calculator(self,
        Ex_min: float,
        Ex_max: float,
        j_lists: Union[list, None],
        BXL_bin_width: float,
        multipole_type: str,
        ):
        """
        Really just a wrapper to self.porter_thomas with j checks.

        Parameters
        ----------
        Ex_min : float
            Minimum value of the initial energy. MeV.
        
        Ex_max : float
            Maximum value of the initial energy. MeV.

        j_lists : Union[list, None]
            Either a list of j values to compare, a list of lists of j
            values to compare, or None where all j values available
            will be used.

        BXL_bin_width : float
            Width of the BXL bins when the BXL/mean(BXL) values are
            counted. Unitless.

        multipole_type : str
            Choose the multipolarity of the transitions. 'E1', 'M1',
            'E2'.
        """        
        if isinstance(j_lists, list):
            if not j_lists:
                msg = "Please provide a list of j values or a list of lists of j values."
                raise ValueError(msg)

        else:
            msg = f"j_lists must be a list. Got {type(j_lists)}."
            raise TypeError(msg)
        
        if all(isinstance(j, list) for j in j_lists):
            """
            All entries in j_lists are lists.
            """
            pass
        
        elif any(isinstance(j, list) for j in j_lists):
            """
            Only some of the entries are lists. The case where all
            entries are lists will be captured by the previous check.
            """
            msg = "j_lists cant contain a mix of lists and numbers!"
            raise TypeError(msg)

        else:
            """
            None of the entries are lists. Combine all numbers as a
            single list inside j_lists.
            """
            if all(isinstance(j, (int, float)) for j in j_lists):
                j_lists = [j_lists]
            else:
                msg = "All entries in j_lists must either all be lists or all be numbers!"
                raise TypeError(msg)

        if (j_lists_len := len(j_lists)) > 3:
            msg = f"j_lists cannot contain more than 3 ranges of j values. Got {j_lists_len}."
            raise ValueError(msg)

        if Ex_min > Ex_max:
            msg = "Ex_min cannot be larger than Ex_max!"
            raise ValueError(msg)

        if (Ex_min < 0) or (Ex_max < 0):
            msg = "Ex_min and Ex_max cannot be negative!"
            raise ValueError(msg)
        
        binss = []
        countss = []
        chi2s = []
        for j_list in j_lists:
            """
            Calculate for the j values in j_list (note: not in j_lists).
            """
            bins, counts, chi2 = self.porter_thomas(
                multipole_type = multipole_type,
                j_list = j_list,
                Ei = [Ex_min, Ex_max],
                BXL_bin_width = BXL_bin_width,
                return_chi2 = True
            )
            idx = np.argmin(np.abs(bins - 10))  # Slice the arrays at approx 10.
            # bins = bins[:idx]
            # counts = counts[:idx]
            # chi2 = chi2[:idx]
            binss.append(bins[:idx])
            countss.append(counts[:idx])
            chi2s.append(chi2[:idx])

        return binss, countss, chi2s

    def porter_thomas_j_plot(self,
        Ex_min: float = 5,
        Ex_max: float = 9,
        j_lists: Union[list, None] = None,
        BXL_bin_width: float = 0.1,
        multipole_type: Union[str, list] = "M1",
        include_relative_difference: bool = True,
        set_title: bool = True
        ):
        """
        Porter-Thomas analysis of the reduced transition probabilities
        for different angular momenta.

        Parameter
        ---------
        multipole_type : Union[str, list]
            Choose the multipolarity of the transitions. 'E1', 'M1',
            'E2'. Accepts a list of max 2 multipolarities.

        set_title : bool
            Toggle plot title on / off.

        See the docstring of _porter_thomas_j_plot_calculator for the
        rest of the descriptions.
        """
        if j_lists is None:
            j_list_default = True
        else:
            j_list_default = False

        transitions_dict = {
            "E1": self.transitions_BE1,
            "M1": self.transitions_BM1,
            "E2": self.transitions_BE2,
        }

        colors = ["blue", "royalblue", "lightsteelblue"]
        if isinstance(multipole_type, str):
            multipole_type = [multipole_type]
        
        elif isinstance(multipole_type, list):
            pass
        
        else:
            msg = f"multipole_type must be str or list. Got {type(multipole_type)}."
            raise TypeError(msg)

        if len(multipole_type) == 1:
            if j_list_default:
                """
                Default j_lists values.
                """
                j_lists = []
                for elem in np.unique(transitions_dict[multipole_type[0]][:, 0]):
                    j_lists.append([elem/2])

                j_lists = j_lists[:3]   # _porter_thomas_j_plot_calculator supports max. 3 lists of j values.

            if include_relative_difference:
                fig, axd = plt.subplot_mosaic(
                    [['upper'], ['lower']],
                    gridspec_kw = dict(height_ratios=[1, 0.5]),
                    figsize = (6.4, 8),
                    constrained_layout = True,
                    sharex = True
                )
            else:
                fig, axd = plt.subplot_mosaic(
                    [['upper']],
                    gridspec_kw = dict(height_ratios=[1]),
                    # figsize = (6.4, 8),
                    constrained_layout = True,
                    sharex = True
                )

            binss, countss, chi2s = self._porter_thomas_j_plot_calculator(
                Ex_min = Ex_min,
                Ex_max = Ex_max,
                j_lists = j_lists,
                BXL_bin_width = BXL_bin_width,
                multipole_type = multipole_type[0],
            )

            for bins, counts, chi2, j_list, color in zip(binss, countss, chi2s, j_lists, colors):
                axd["upper"].step(
                    bins,
                    counts,
                    label = r"$j_i = $" + f"{j_list}",
                    color = color
                )
                if include_relative_difference:
                    axd["lower"].step(
                        bins,
                        counts/chi2,
                        color = color,
                        label = r"($j_i = $" + f"{j_list})" + r"$/\chi_{\nu = 1}^2$",
                    )
        
            axd["upper"].plot(
                bins,
                chi2,
                color = "tab:green",
                label = r"$\chi_{\nu = 1}^2$"
            )
            axd["upper"].legend(loc="upper right")
            axd["upper"].set_ylabel(r"Normalised counts")
            if set_title:
                axd["upper"].set_title(
                    f"{self.nucleus_latex}, {self.interaction}, " + r"$" + f"{multipole_type[0]}" + r"$"
                )

            if include_relative_difference:
                axd["lower"].hlines(y=1, xmin=bins[0], xmax=bins[-1], linestyle="--", color="black")
                axd["lower"].legend(loc="upper left")
                axd["lower"].set_ylabel(r"Relative error")
                axd["lower"].set_xlabel(
                    r"$B(" + f"{multipole_type[0]}" + r")/\langle B(" + f"{multipole_type[0]}" + r") \rangle$"
                )
            else:
                axd["upper"].set_xlabel(
                    r"$B(" + f"{multipole_type[0]}" + r")/\langle B(" + f"{multipole_type[0]}" + r") \rangle$"
                )

            fig.savefig(fname=f"{self.nucleus}_porter_thomas_j_{multipole_type[0]}.png", dpi=300)

        elif len(multipole_type) == 2:
            if j_list_default:
                """
                Default j_lists values.
                """
                j_lists = []
                for elem in np.unique(transitions_dict[multipole_type[0]][:, 0]):
                    j_lists.append([elem/2])

                j_lists = j_lists[:3]   # _porter_thomas_j_plot_calculator supports max. 3 lists of j values.

            fig, axd = plt.subplot_mosaic(
                [['upper left', 'upper right'], ['lower left', 'lower right']],
                gridspec_kw = dict(height_ratios=[1, 0.5]),
                figsize = (10, 8),
                constrained_layout = True,
                sharex = True
            )
            binss, countss, chi2s = self._porter_thomas_j_plot_calculator(
                Ex_min = Ex_min,
                Ex_max = Ex_max,
                j_lists = j_lists,
                BXL_bin_width = BXL_bin_width,
                multipole_type = multipole_type[0],
            )

            for bins, counts, chi2, j_list, color in zip(binss, countss, chi2s, j_lists, colors):
                axd["upper left"].step(
                    bins,
                    counts,
                    label = r"$j_i = $" + f"{j_list}",
                    color = color
                )
                axd["lower left"].step(
                    bins,
                    counts/chi2,
                    color = color,
                    label = r"($j_i = $" + f"{j_list})" + r"$/\chi_{\nu = 1}^2$",
                )
        
            axd["upper left"].plot(
                bins,
                chi2,
                color = "tab:green",
                label = r"$\chi_{\nu = 1}^2$"
            )
            axd["upper left"].legend(loc="upper right")
            axd["upper left"].set_ylabel(r"Normalised counts")

            axd["lower left"].hlines(y=1, xmin=bins[0], xmax=bins[-1], linestyle="--", color="black")
            axd["lower left"].set_xlabel(
                r"$B(" + f"{multipole_type[0]}" + r")/\langle B(" + f"{multipole_type[0]}" + r") \rangle$"
            )
            axd["lower left"].legend(loc="upper left")
            axd["lower left"].set_ylabel(r"Relative error")
            if set_title:
                axd["upper left"].set_title(
                    f"{self.nucleus_latex}, {self.interaction}, " + r"$" + f"{multipole_type[0]}" + r"$"
                )

            if j_list_default:
                """
                Default j_lists values.
                """
                j_lists = []
                for elem in np.unique(transitions_dict[multipole_type[1]][:, 0]):
                    j_lists.append([elem/2])

                j_lists = j_lists[:3]   # _porter_thomas_j_plot_calculator supports max. 3 lists of j values.

            binss, countss, chi2s = self._porter_thomas_j_plot_calculator(
                Ex_min = Ex_min,
                Ex_max = Ex_max,
                j_lists = j_lists,
                BXL_bin_width = BXL_bin_width,
                multipole_type = multipole_type[1],
            )

            for bins, counts, chi2, j_list, color in zip(binss, countss, chi2s, j_lists, colors):
                axd["upper right"].step(
                    bins,
                    counts,
                    label = r"$j_i = $" + f"{j_list}",
                    color = color
                )
                axd["lower right"].step(
                    bins,
                    counts/chi2,
                    color = color,
                    label = r"($j_i = $" + f"{j_list})" + r"$/\chi_{\nu = 1}^2$",
                )
        
            axd["upper right"].plot(
                bins,
                chi2,
                color = "tab:green",
                label = r"$\chi_{\nu = 1}^2$"
            )
            # axd["upper right"].legend(loc="upper right")
            axd["upper right"].set_yticklabels([])
            axd["upper right"].set_ylim(axd["upper left"].get_ylim())

            axd["lower right"].hlines(y=1, xmin=bins[0], xmax=bins[-1], linestyle="--", color="black")
            axd["lower right"].set_xlabel(
                r"$B(" + f"{multipole_type[1]}" + r")/\langle B(" + f"{multipole_type[1]}" + r") \rangle$"
            )
            # axd["lower right"].legend(loc="upper left")
            axd["lower right"].set_yticklabels([])
            axd["lower right"].set_ylim(axd["lower left"].get_ylim())
            if set_title:
                axd["upper right"].set_title(
                    f"{self.nucleus_latex}, {self.interaction}, " + r"$" + f"{multipole_type[1]}" + r"$"
                )
            fig.savefig(fname=f"{self.nucleus}_porter_thomas_j_{multipole_type[0]}_{multipole_type[1]}.png", dpi=300)
        else:
            msg = "Only 1 or 2 multipole types may be given at the same time!"
            msg += f" Got {len(multipole_type)}."
            raise ValueError(msg)

        plt.show()

    def angular_momentum_distribution_plot(self,
        bin_width: float = 0.2,
        E_min: float = 5,
        E_max: float = 10,
        j_list: Union[None, int, float, Iterable] = None,
        filter_parity: Union[None, int, str] = None,
        plot: bool = True,
        # single_spin_plot: Union[None, int, float, Iterable] = None,
        save_plot: bool = True,
        set_title: bool = True
    ):
        """
        Plot the angular momentum distribution of the levels.

        Parameters
        ----------
        bin_width : float
            Width of the energy bins. MeV.

        E_min : float
            Minimum value of the energy range.

        E_max : float
            Maximum value of the energy range.

        j_list : Union[None, int, float, Iterable]
            Filter the levels by their angular momentum. If None,
            all levels are plotted.

        filter_parity : Union[None, int, str]
            Filter the levels by their parity. If None, all levels
            are plotted.

        plot : bool
            If True, the plot will be shown.
        
        """
        if not isinstance(filter_parity, (type(None), int, str)):
            msg = f"'filter_parity' must be of type: None, int, str. Got {type(j_list)}."
            raise TypeError(msg)

        if isinstance(filter_parity, str):
            valid_filter_parity = ["+", "-"]
            if filter_parity not in valid_filter_parity:
                msg = f"Valid parity filters are: {valid_filter_parity}."
                raise ValueError(msg)
            
            filter_parity = 1 if (filter_parity == "+") else -1

        if j_list is None:
            """
            If no angular momentum filter, then include all angular
            momenta in the data set.
            """
            angular_momenta = np.unique(self.levels[:, 1])/2
        else:
            if isinstance(j_list, (float, int)):
                angular_momenta = [j_list]
            
            elif isinstance(j_list, Iterable) and not isinstance(j_list, str):
                angular_momenta = j_list
            
            else:
                msg = f"'j_list' must be of type: None, Iterable, int, float. Got {type(j_list)}."
                raise TypeError(msg)
        
        bins = np.arange(E_min, E_max, bin_width)
        n_bins = len(bins)
        n_angular_momenta = len(angular_momenta)
        densities = np.zeros((n_bins, n_angular_momenta))

        for i in range(n_angular_momenta):
            """
            Calculate the nuclear level density for each angular
            momentum.

            NOTE: Each column is the density for a unique angular
            momentum. Consequently, rows are the energy axis.
            """
            bins_tmp, densities_tmp, counts_tmp = level_density(
                levels = self.levels,
                bin_width = bin_width,
                filter_spins = angular_momenta[i],
                filter_parity = filter_parity,
                E_min = E_min,
                E_max = E_max,
                return_counts = True,
                plot = False,
                save_plot = False
            )
            assert len(bins_tmp) <= n_bins, "There are more NLD bins than bins in the expected range!"
            for b1, b2 in zip(bins_tmp, bins):
                """
                Check that the returned bins match the expected bin
                range. Note that bins_tmp might be shorter than
                bins, but never longer.
                """
                assert b1 == b2, "NLD bins do not match the expected bins!"

            """
            Check that the total number of levels returned by the NLD
            function is actually the correct number of levels as counted
            with the following stright-forward counter:
            """
            E_tmp = self.levels[self.levels[:, 1] == 2*angular_momenta[i]]  # Extract levels of correct angular momentum.
            if filter_parity is not None:
                E_tmp = E_tmp[E_tmp[:, 2] == filter_parity]
            mask_lower = (E_tmp[:, 0] - self.ground_state_energy) >= E_min  # Keep only levels larger or equal to E_min.
            mask_upper = (E_tmp[:, 0] - self.ground_state_energy) < E_max   # Keep only levels lower than E_max.
            mask_total = np.logical_and(mask_lower, mask_upper)             # Combine the two masks with and.
            n_levels_in_range = sum(mask_total) # The number of Trues is the number of levels within the energy range.
            n_counts_tmp = sum(counts_tmp)
            msg = f"The number of levels from the NLD in the range [{E_min, E_max})"
            msg += " does not match the expected value!"
            msg += f" {n_levels_in_range = }, {n_counts_tmp = }"
            success = n_levels_in_range == n_counts_tmp
            assert success, msg

            """
            If E_max exceeds the max energy in the data set,
            level_density will adjust E_max to the max energy value. In
            this function, the bin range should be the same for all
            angular momenta, meaning that we might have to add some
            extra zeros to the end of the densities array.
            """
            n_missing_values = n_bins - len(bins_tmp)
            densities[:, i] = np.concatenate((densities_tmp, np.zeros(n_missing_values)))

            if flags["debug"]:
                print("-----------------------------------")
                print("angular_momentum_distribution debug")
                if all(densities_tmp == 0):
                    print("No levels in this energy range!")
                print(f"{angular_momenta[i] = }")
                print(f"{filter_parity = }")
                print(f"{bin_width = }")
                print(f"{E_min = }")
                print(f"{E_max = }")
                print(f"densities extended with {n_missing_values} zeros")
                print("-----------------------------------")

        if filter_parity is None:
            exponent = r"$^{\pm}$"  # For exponent in yticklabels.
            parity_str = "+-"       # For filename in savefig.
        elif filter_parity == 1:
            exponent = r"$^{+}$"
            parity_str = "+"
        elif filter_parity == -1:
            exponent = r"$^{-}$"
            parity_str = "-"

        if plot:
            xticklabels = []
            for i in bins:
                """
                Loop for generating x tick labels.
                
                Loop over the bins of the first angular momentum entry.
                Note that the bins are equal for all the different
                angular momenta (all columns in 'bins' are equal). Round
                all energies to 1 decimal when that decimal is != 0. If
                that decimal is zero, round to 0 decimals.
                """
                if (tmp := int(i)) == i:
                    """
                    Eg.
                    >>> 1 == 1.0
                    True
                    
                    (round to 0 decimals)
                    """
                    xticklabels.append(tmp)
                else:
                    """
                    Round to 1 decimal.
                    """
                    xticklabels.append(round(i, 1))
            
            fig, ax = plt.subplots(figsize=(7, 6.4))
            sns.heatmap(
                data = densities.T[-1::-1],
                linewidth = 0.5,
                annot = True,
                cmap = 'gray',
                xticklabels = xticklabels,
                fmt = ".0f",
                ax = ax
            )

            ax.set_yticklabels(np.flip([f"{Fraction(i)}" + exponent for i in angular_momenta]), rotation=0)
            ax.set_xlabel(r"$E$ [MeV]")
            ax.set_ylabel(r"$j$ [$\hbar$]")
            if set_title:
                ax.set_title(f"{self.nucleus_latex}, {self.interaction}")
            cbar = ax.collections[0].colorbar
            cbar.ax.set_ylabel(r"NLD [MeV$^{-1}$]", rotation=90)

            if save_plot:
                fig.savefig(f"{self.nucleus}_j{parity_str}_distribution_heatmap.png", dpi=300)
        
        if plot:
            plt.show()

        return bins, densities

    def B_distribution(self,
        partial_or_total: str,
        multipole_type: str = "M1",
        filter_spins: Union[None, list] = None,
        filter_parity: Union[None, int] = None,
        filter_indices: Union[None, int, list] = None,
        plot: bool = True,
    ) -> np.ndarray:
        """
        Plot a histogram of the distribution of B values.

        Parameters
        ----------
        partial_or_total : str
            If total, then all partial B values will be summed per
            level. If partial, then the distribution of all partial B
            values will be generated.

        multipole_type : str
            Choose the multipolarity of the transitions. 'E1', 'M1',
            'E2'.
        
        filter_spins : Union[None, list]
            Filter the levels by their angular momentum. If None,
            all levels are included.

        filter_parity : Union[None, int]
            Filter the levels by their parity. If None, both parities
            are included.

        plot : bool
            If True, the plot will be shown.

        Returns
        -------
        total_B : np.ndarray
            The sum over every partial B value for each level.
        """
        total_time = time.perf_counter()

        is_loaded = False
        B_unique_string = f"{multipole_type}{filter_spins}{filter_parity}{filter_indices}{partial_or_total}"
        B_unique_id = hashlib.sha1((B_unique_string).encode()).hexdigest()
        B_fname = f"{self.npy_path}/{self.base_fname}_Bdist_{B_unique_id}_{self.unique_id}.npy"

        if os.path.isfile(B_fname) and self.load_and_save_to_file:
            total_B = np.load(file=B_fname, allow_pickle=True)
            msg = f"{self.nucleus} {multipole_type} B distribution data loaded from .npy!"
            print(msg)
            is_loaded = True

        else:
            transitions_dict = {
                "E1": self.transitions_BE1,
                "M1": self.transitions_BM1,
                "E2": self.transitions_BE2,
            }
            transitions = transitions_dict[multipole_type]

            if filter_spins is None:
                initial_j = np.unique(transitions[:, 0])
            else:
                initial_j = [2*j for j in filter_spins]
            
            if filter_parity is None:
                initial_pi = [-1, 1]
            else:
                initial_pi = [filter_parity]

            if filter_indices is None:
                initial_indices = np.unique(transitions[:, 2]).astype(int)
            elif isinstance(filter_indices, list):
                initial_indices = [int(i) for i in filter_indices]
            elif isinstance(filter_indices, int):
                initial_indices = [filter_indices]
            
            total_B = []    # The sum of every partial B value for each level.
            idxi_masks = []
            pii_masks = []
            ji_masks = []

            mask_time = time.perf_counter()
            for idxi in initial_indices:
                idxi_masks.append(transitions[:, 2] == idxi)

            for pii in initial_pi:
                pii_masks.append(transitions[:, 1] == pii)

            for ji in initial_j:
                ji_masks.append(transitions[:, 0] == ji)
            mask_time = time.perf_counter() - mask_time

            for pii in pii_masks:
                for idxi in idxi_masks:
                    for ji in ji_masks:
                        mask = np.logical_and(ji, np.logical_and(pii, idxi))
                        # total_B.append(np.sum(transitions[mask][:, 9]))   # 9 is B decay
                        total_B.append(transitions[mask][:, 9])   # 9 is B decay

            if partial_or_total == "total":
                """
                Sum partial B values to get total B values.
                """
                for i in range(len(total_B)):
                    total_B[i] = sum(total_B[i])
            
            elif partial_or_total == "partial":
                """
                Keep a 1D list of partial B values.
                """
                total_B = list(chain.from_iterable(total_B))

        total_B = np.asarray(total_B)
        
        if self.load_and_save_to_file and not is_loaded:
            np.save(file=B_fname, arr=total_B, allow_pickle=True)
        
        total_time = time.perf_counter() - total_time

        if flags["debug"]:
            if not is_loaded: print(f"B_distribution {mask_time = :.4f}")
            print(f"B_distribution {total_time = :.4f}")

        if plot:
            plt.hist(total_B, bins=100, color="black")
            plt.xlabel(r"$B(" + f"{multipole_type}" + r")$")
            plt.show()

        return total_B

    def _brink_axel_j_calculator(self,
        bin_width: Union[float, int],
        Ex_min: Union[float, int],
        Ex_max: Union[float, int],
        multipole_type: str,
        j_list: list,
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
        """
        Calculate the GSF as a function of Eg for all Ei and ji, as well
        as the GSF individually for all ji in j_list.
        """
        bins_all_j, gsf_all_j = self.gsf(
            bin_width = bin_width,
            Ex_min = Ex_min,
            Ex_max = Ex_max,
            multipole_type = multipole_type,
            return_n_transitions = False,
            plot = False,
        )
        n_j = len(j_list)
        n_gsf_all_j = len(gsf_all_j)
        gsf = np.zeros((n_gsf_all_j, n_j))
        bins = np.zeros((n_gsf_all_j, n_j))   # Eg rows, j cols.

        for i in range(n_j):
            """
            Calculate the GSF individually for each ji.
            """
            bins_one_j, gsf_one_j = self.gsf(
                bin_width = bin_width,
                Ex_min = Ex_min,
                Ex_max = Ex_max,
                multipole_type = multipole_type,
                filter_spins = [j_list[i]],
                return_n_transitions = False,
                plot = False,
            )
            gsf[:, i] = gsf_one_j
            bins[:, i] = bins_one_j

        return bins, gsf, bins_all_j, gsf_all_j

    def brink_axel_j(self,
        bin_width: Union[float, int] = 0.2,
        Ex_min: Union[float, int] = 5,
        Ex_max: Union[float, int] = 50,
        multipole_type: Union[list, str] = "M1",
        j_list: Union[list, None] = None,
        set_title: bool = True
    ):
        """
        Plot a comparison of the GSF averaged over all ji and GSFs for
        each individual ji.
        """
        if j_list is None:
            j_list_default = True
        else:
            j_list_default = False
            
        transitions_dict = {
            "E1": self.transitions_BE1,
            "M1": self.transitions_BM1,
            "E2": self.transitions_BE2,
        }

        if isinstance(multipole_type, str):
            multipole_type = [multipole_type]
        
        n_multipole_types = len(multipole_type)

        fig, ax = plt.subplots(
            nrows = n_multipole_types,
            ncols = 1,
            figsize = (6.4, 4.8*n_multipole_types)
        )
        if not isinstance(ax, np.ndarray):
            ax = [ax]

        upper_ylim = -np.inf
        lower_ylim = np.inf
        for i in range(n_multipole_types):
            if j_list_default:
                """
                Choose all available angular momenta as default.
                """
                j_list = np.unique(transitions_dict[multipole_type[i]][:, 0])/2
                j_list.sort()   # Just in case.

            bins, gsf, bins_all_j, gsf_all_j = self._brink_axel_j_calculator(
                    bin_width = bin_width,
                    Ex_min = Ex_min,
                    Ex_max = Ex_max,
                    multipole_type = multipole_type[i],
                    j_list = j_list,
                )

            ax[i].plot(bins_all_j, gsf_all_j, color="black", label=r"All $j_i$")
            ax[i].plot(bins, gsf, color="black", alpha=0.2)
            ax[i].plot([0], [0], color="black", alpha=0.2, label=r"Single $j_i$")  # Dummy for legend.
            ax[i].set_yscale('log')
            ax[i].set_ylabel(r"GSF [MeV$^{-3}$]")
            ax[i].legend()
            if set_title:
                ax[i].set_title(
                    f"{self.nucleus_latex}, {self.interaction}, " + r"$" + f"{multipole_type[i]}" + r"$"
                )
            else:
                ax[i].set_title(
                    r"$" + f"{multipole_type[i]}" + r"$"
                )

            lower_ylim = min(ax[i].get_ylim()[0], lower_ylim)
            upper_ylim = max(ax[i].get_ylim()[1], upper_ylim)

        for i in range(n_multipole_types):
            """
            Set both lower ylim to the lowest of the two and both upper
            ylims to the largest of the two.
            """
            ax[i].set_ylim((lower_ylim, upper_ylim))
        
        ax[-1].set_xlabel(r"E$_{\gamma}$ [MeV]")
        fig.savefig(
            fname = f"{self.nucleus}_brink-axel_ji_{'_'.join(multipole_type)}.png",
            dpi = 300
        )
        plt.show()

    def primary_matrix(self,
        bin_width: Union[float, int] = 0.2,
        Ex_min: Union[float, int] = 0,
        Ex_max: Union[float, int] = 50,
        multipole_type: str = "M1",
        plot: bool = True,
    ):
        """
        Create a Ex Eg primary matrix like the ones which are calculated
        with the Oslo method.

        Parameters
        ----------
        bin_width : float, optional
            Width of the bins in MeV, by default 0.2
        
        Ex_min : float, optional
            Minimum excitation energy of initial level in a transition,
            by default 0 MeV.

        Ex_max : float, optional
            Maximum excitation energy of initial level in a transition,
            by default 50 MeV.
        
        multipole_type : str
            Multipolarity of the transitions to be considered, by
            default "M1".

        plot : bool, optional
            Whether to plot the matrix, by default True.

        Returns
        -------
        Ex_range : np.ndarray
            The range of the initial level energies.

        Eg_range : np.ndarray
            The range of the gamma energies.

        B_matrix : np.ndarray
            The primary matrix.
        """
        transitions_dict = {
            "E1": self.transitions_BE1,
            "M1": self.transitions_BM1,
            "E2": self.transitions_BE2,
        }

        n_bins = int(np.ceil((Ex_max - Ex_min)/bin_width))
        B_matrix = np.zeros((n_bins, n_bins), dtype=float)    # rows: Ei, cols: Eg
        Ex_range = np.linspace(Ex_min, Ex_max, n_bins)
        Eg_range = np.linspace(Ex_min, Ex_max, n_bins)
        
        transitions = transitions_dict[multipole_type]
        Ex_masks = []
        Eg_masks = []
        for i in range(n_bins-1):
            """
            Generate the masks once. It is not necessary to generate the
            Eg masks over and over again while iterating over the Ex
            masks.
            """
            Ex_masks.append(np.logical_and(
                transitions[:, 3] >= Ex_range[i],
                transitions[:, 3] < Ex_range[i+1],
            ))
        for i in range(n_bins-1):
            """
            Using separate loops allows using different bin widths for
            Ex and Eg, however, that has not yet been implemented.
            """
            Eg_masks.append(np.logical_and(
                transitions[:, 8] >= Eg_range[i],
                transitions[:, 8] < Eg_range[i+1],
            ))
        for i in range(n_bins-1):
            """
            Loop over each Ex Eg mask pair and create a combined mask.
            Take the mean of the B values for each combined mask.
            """
            for j in range(i+1):
                mask = np.logical_and(
                    Ex_masks[i],
                    Eg_masks[j],
                )
                B_slice = transitions[mask][:, 9]
                if B_slice.size == 0: continue  # Skip if there are no B values in the mask.
                
                B_matrix[i, j] = np.mean(B_slice)

        if plot:
            fig, ax = plt.subplots()
            im = ax.pcolormesh(Ex_range, Eg_range, B_matrix, cmap="jet", norm="log")
            ax.set_title(r"$\langle B($" + f"{multipole_type}" + r"$) \rangle$")
            ax.set_xlabel(r"$E_{\gamma}$")
            ax.set_ylabel(r"$E_{x}$")
            fig.colorbar(im)
            plt.show()

        return Ex_range, Eg_range, B_matrix

    def _mixing_pairs(self,
        B_left: str = "M1",
        B_right: str = "E2",
    ):
        """
        NOTE: Currently hard-coded for BM1_BE2!
        """
        if B_left != "M1":#not in ["M1", "E2"]:
            # msg = f"'B_left' must be 'M1' or 'E2', got: {B_left}"
            # raise ValueError(msg)
            raise ValueError
        
        if B_right != "E2":# not in ["M1", "E2"]:
            # msg = f"'B_right' must be 'M1' or 'E2', got: {B_right}"
            # raise ValueError(msg)
            raise ValueError
        
        transitions_dict = {
            "M1": self.transitions_BM1,
            "E2": self.transitions_BE2,
            "E1": self.transitions_BE1
        }
        transitions_left = transitions_dict[B_left]
        transitions_right = transitions_dict[B_right]
        
        mixing_pairs_fname = f"{self.npy_path}/{self.base_fname}_mixing_pairs_B{B_left}_B{B_right}_{self.unique_id}.npy"

        if os.path.isfile(mixing_pairs_fname) and self.load_and_save_to_file and (self.load_and_save_to_file != "overwrite"):
            print(f"Mixing pairs loaded from .npy!")
            mixing_pairs = np.load(file=mixing_pairs_fname, allow_pickle=True)
            return mixing_pairs        
        
        @numba.njit
        def calculate_mixing_pairs(
            possible_j: list[int],
            possible_indices: list[int],
            transitions_BM1: np.ndarray,
            transitions_BE2: np.ndarray,
        ):
            mixing_pairs = []

            for ji in possible_j:
                ji_slice_BM1 = transitions_BM1[transitions_BM1[:, 0] == ji]
                ji_slice_BE2 = transitions_BE2[transitions_BE2[:, 0] == ji]

                if (not ji_slice_BM1.size) or (not ji_slice_BE2.size): return

                for jf in possible_j:
                    # if not j_allowed(ji=ji, jf=jf): continue
                    jf_slice_BM1 = ji_slice_BM1[ji_slice_BM1[:, 4] == jf]
                    jf_slice_BE2 = ji_slice_BE2[ji_slice_BE2[:, 4] == jf]

                    if (not jf_slice_BM1.size) or (not jf_slice_BE2.size): continue

                    for pi in [-1, 1]:
                        """
                        [2*spin_initial, parity_initial, idx_initial, Ex_initial, 2*spin_final,
                        parity_final, idx_final, Ex_final, E_gamma, B(.., i->f), B(.., f<-i)]
                        """
                        pi_i_slice_BM1 = jf_slice_BM1[jf_slice_BM1[:, 1] == pi]
                        pi_i_slice_BE2 = jf_slice_BE2[jf_slice_BE2[:, 1] == pi]

                        if (not pi_i_slice_BM1.size) or (not pi_i_slice_BE2.size): continue
                        
                        pi_f_slice_BM1 = pi_i_slice_BM1[pi_i_slice_BM1[:, 5] == pi]
                        pi_f_slice_BE2 = pi_i_slice_BE2[pi_i_slice_BE2[:, 5] == pi]

                        if (not pi_f_slice_BM1.size) or (not pi_f_slice_BE2.size): continue

                        for idx_i in possible_indices:
                            idx_i_slice_BM1 = pi_f_slice_BM1[pi_f_slice_BM1[:, 2] == idx_i]
                            idx_i_slice_BE2 = pi_f_slice_BE2[pi_f_slice_BE2[:, 2] == idx_i]

                            if (not idx_i_slice_BM1.size) or (not idx_i_slice_BE2.size): continue

                            for idx_f in possible_indices:
                                idx_f_slice_BM1 = idx_i_slice_BM1[idx_i_slice_BM1[:, 2] == idx_f]
                                idx_f_slice_BE2 = idx_i_slice_BE2[idx_i_slice_BE2[:, 2] == idx_f]

                                if (not idx_f_slice_BM1.size) or (not idx_f_slice_BE2.size): continue

                                for transition_BM1 in idx_f_slice_BM1:
                                    for transition_BE2 in idx_f_slice_BE2:
                                        if (transition_BM1[3] == transition_BE2[3]) and (transition_BM1[7] == transition_BE2[7]):
                                            mixing_pairs.append([transition_BM1, transition_BE2])
                
            return mixing_pairs

        mixing_pairs_time = time.perf_counter()
        possible_j = np.unique(transitions_left[:, 0])
        possible_indices = np.unique(transitions_left[:, 2])
        assert np.all(possible_indices == np.unique(transitions_left[:, 6]))    # Check that indices of the initial and final levels are the same range.
        
        mixing_pairs = calculate_mixing_pairs(
            transitions_BM1 = transitions_left,
            transitions_BE2 = transitions_right,
            possible_j = possible_j,
            possible_indices = possible_indices,
        )
        mixing_pairs.sort(key=lambda tup: tup[0][8])    # Sort wrt. gamma energy.
        mixing_pairs = np.array(mixing_pairs)
        np.save(file=mixing_pairs_fname, arr=mixing_pairs, allow_pickle=True)
        mixing_pairs_time = time.perf_counter() - mixing_pairs_time
        print(f"{mixing_pairs_time = :.3f} s")
        return mixing_pairs

    def mixing_ratio(self,
        bin_width: float = 0.2,
        plot: bool = True,
        save_plot: bool = False,
    ):
        """
        Calculate the ratio of T(E2)/(T(E2) + T(M1)), aka. how large the
        E2 contribution is. Currently hard-coded for this specific
        ratio.

        Parameters
        ----------
        bin_width : float
            The ratios are sorted by gamma energy and averaged over
            the ratio values in a gamma energy bin of bin_with.
        """
        E_min = self.mixing_pairs_BM1_BE2[0, 0, 8]  # BM1 and BE2 has to be at the exact same gamma energies so it doesnt matter which one we take E_min and E_max from.
        E_max = self.mixing_pairs_BM1_BE2[-1, 0, 8]

        bins = np.arange(E_min, E_max + bin_width, bin_width)
        n_bins = len(bins)
        ratios = np.zeros(n_bins - 1)

        for i in range(n_bins - 1):
            mask_1 = self.mixing_pairs_BM1_BE2[:, 0, 8] >= bins[i]
            mask_2 = self.mixing_pairs_BM1_BE2[:, 0, 8] < bins[i + 1]
            mask_3 = np.logical_and(mask_1, mask_2)

            # M1_mean = self.mixing_pairs_BM1_BE2[:, 0, 9][mask_3].mean()
            # E2_mean = self.mixing_pairs_BM1_BE2[:, 1, 9][mask_3].mean()

            M1_mean = (1.76e13*(self.mixing_pairs_BM1_BE2[:, 0, 8][mask_3]**3)*self.mixing_pairs_BM1_BE2[:, 0, 9][mask_3]).mean()
            E2_mean = (1.22e09*(self.mixing_pairs_BM1_BE2[:, 1, 8][mask_3]**5)*self.mixing_pairs_BM1_BE2[:, 1, 9][mask_3]).mean()

            ratios[i] = E2_mean/(E2_mean + M1_mean)
        
        if plot:
            fig, ax = plt.subplots()
            ax.step(bins[:-1], ratios, color="black")
            # ax.legend()
            ax.grid()
            ax.set_ylabel(r"TE2/(TE2 + TM1)")
            ax.set_xlabel(r"$E_{\gamma}$ [MeV]")
            if save_plot:
                fname = f"mixing_ratio_E2_M1.png"
                print(f"Mixing ratio plot saved as '{fname}'")
                fig.savefig(fname=fname, dpi=300)
            plt.show()

        return bins[:-1], ratios

    @property
    def help(self):
        """
        Generate a list of instance attributes without magic and private
        methods.

        Returns
        -------
        help_list : list
            A list of non-magic instance attributes.
        """
        help_list = []
        for elem in dir(self):
            if not elem.startswith("_"):   # Omit magic and private methods.
                help_list.append(elem)
        
        return help_list

    @property
    def parameters(self) -> dict:
        """
        Get the KSHELL parameters from the shell file.

        Returns
        -------
        : dict
            A dictionary of KSHELL parameters.
        """
        path = self.path
        if os.path.isfile(path):
            path = path.rsplit("/", 1)[0]
        return get_parameters(path)

    @property
    def nucleus_latex(self):
        m = re.search(r"\d+$", self.nucleus)
        A = m.group()
        X = self.nucleus[:m.span()[0]]
        return r"$^{" + f"{A}" + r"}$" + f"{X}"

    def check_data(self):
        """
        Check that the 'levels' data is sorted in order of increasing
        energy.
        """
        msg = "'levels' should be sorted in order of increasing energy,"
        msg += " but it isnt! Check the data!"
        energies = self.levels[:, 0]
        success = np.all(energies[:-1] <= energies[1:])
        assert success, msg

        """
        Check that the ground state energy is correct.
        """
        msg = "'ground_state_energy' is not the lowest energy in 'levels'"
        msg += " as it should be! Check the data!"
        success = self.ground_state_energy == np.min(self.levels[:, 0])
        assert success, msg

        """
        Insert more checks under...
        """

def loadtxt(
    path: str,
    load_and_save_to_file: Union[bool, str] = True,
    old_or_new = "new"
    ) -> ReadKshellOutput:
    """
    Wrapper for using ReadKshellOutput class as a function.

    Parameters
    ----------
    path : str
        Path to summary file or path to directory with summary / log
        files.

    load_and_save_to_file : Union[bool, str]
        Toggle saving data as `.npy` files on / off. If 'overwrite',
        saved `.npy` files are overwritten.

    old_or_new : str
        Choose between old and new summary file syntax. All summary
        files generated pre 2021-11-24 use old style.
        New:
        J_i  pi_i idx_i Ex_i    J_f  pi_f idx_f Ex_f      dE         B(E2)->         B(E2)->[wu]     B(E2)<-         B(E2)<-[wu]
        5    +    1     0.036   6    +    1     0.000     0.036     70.43477980      6.43689168     59.59865983      5.44660066
        Old:
        J_i    Ex_i     J_f    Ex_f   dE        B(M1)->         B(M1)<- 
        2+(11) 18.393 2+(10) 17.791 0.602 0.1(  0.0) 0.1( 0.0)

    Returns
    -------
    res : ReadKshellOutput
        Instances with data from `KSHELL` data file as attributes.
    """
    loadtxt_time = time.perf_counter()  # Debug.

    if old_or_new not in (old_or_new_allowed := ["old", "new", "jem"]):
        msg = f"'old_or_new' argument must be in {old_or_new_allowed}!"
        msg += f" Got '{old_or_new}'."
        raise ValueError(msg)

    res = ReadKshellOutput(path, load_and_save_to_file, old_or_new)

    loadtxt_time = time.perf_counter() - loadtxt_time
    if flags["debug"]:
        print(f"{loadtxt_time = } s")

    return res

def _get_timing_data(path: str):
    """
    Get timing data from KSHELL log files.

    Parameters
    ----------
    path : str
        Path to log file.

    Examples
    --------
    Last 10 lines of log_Ar30_usda_m0p.txt:
    ```
          total      20.899         2    10.44928   1.0000
    pre-process       0.029         1     0.02866   0.0014
        operate       3.202      1007     0.00318   0.1532
     re-orthog.      11.354       707     0.01606   0.5433
  thick-restart       0.214        12     0.01781   0.0102
   diag tri-mat       3.880       707     0.00549   0.1857
           misc       2.220                         0.1062

           tmp        0.002       101     0.00002   0.0001
    ```
    """
    if "log" not in path:
        msg = f"Unknown log file name! Got '{path}'"
        raise KshellDataStructureError(msg)

    if not os.path.isfile(path):
        raise FileNotFoundError(path)

    res = os.popen(f'tail -n 20 {path}').read()    # Get the final 10 lines.
    res = res.split("\n")
    total = None
    
    if "_tr_" not in path:
        """
        KSHELL log.
        """
        for elem in res:
            tmp = elem.split()
            try:
                if tmp[0] == "total":
                    total = float(tmp[1])
                    break
            except IndexError:
                continue
        
    elif "_tr_" in path:
        """
        Transit log.
        """
        for elem in res:
            tmp = elem.split()
            try:
                if tmp[0] == "total":
                    total = float(tmp[3])
                    break
            except IndexError:
                continue

    if total is None:
        msg = f"Not able to extract timing data from '{path}'!"
        raise KshellDataStructureError(msg)
    
    return total

def _get_memory_usage(path: str) -> Union[float, None]:
    """
    Get memory usage from KSHELL log files.

    Parameters
    ----------
    path : str
        Path to a single log file.

    Returns
    -------
    total : float, None
        Memory usage in GB or None if memory usage could not be read.
    """
    total = None
    
    if "tr" not in path:
        """
        KSHELL log.
        """
        with open(path, "r") as infile:
            for line in infile:
                if line.startswith("Total Memory for Lanczos vectors:"):
                    try:
                        total = float(line.split()[-2])
                    except ValueError:
                        msg = f"Error reading memory usage from '{path}'."
                        msg += f" Got '{line.split()[-2]}'."
                        raise KshellDataStructureError(msg)
                    break
        
    elif "tr" in path:
        """
        Transit log. NOTE: Not yet implemented.
        """
        return 0

    if total is None:
        msg = f"Not able to extract memory data from '{path.split('/')[-1]}'!"
        raise KshellDataStructureError(msg)
    
    return total

def _sortkey(filename):
    """
    Key for sorting filenames based on angular momentum and parity.
    Example filename: 'log_Sc44_GCLSTsdpfsdgix5pn_j0n.txt'
    (angular momentum  = 0). 
    """
    tmp = filename.split("_")[-1]
    tmp = tmp.split(".")[0]
    # parity = tmp[-1]
    spin = int(tmp[1:-1])
    # return f"{spin:03d}{parity}"    # Examples: 000p, 000n, 016p, 016n
    return spin

def _get_data_general(
    path: str,
    func: Callable,
    plot: bool
    ):
    """
    General input handling for timing data and memory data.

    Parameters
    ----------
    path : str
        Path to a single log file or path to a directory of log files.

    func : Callable
        _get_timing_data or _get_memory_usage.
    """
    total_negative = []
    total_positive = []
    filenames_negative = []
    filenames_positive = []
    if os.path.isfile(path):
        return func(path)
    
    elif os.path.isdir(path):
        for elem in os.listdir(path):
            """
            Select only log files in path.
            """
            tmp = elem.split("_")
            try:
                if ((tmp[0] == "log") or (tmp[1] == "log")) and elem.endswith(".txt"):
                    tmp = tmp[-1].split(".")
                    parity = tmp[0][-1]
                    if parity == "n":
                        filenames_negative.append(elem)
                    elif parity == "p":
                        filenames_positive.append(elem)
            except IndexError:
                continue
        
        filenames_negative.sort(key=_sortkey)
        filenames_positive.sort(key=_sortkey)

        for elem in filenames_negative:
            total_negative.append(func(f"{path}/{elem}"))
        for elem in filenames_positive:
            total_positive.append(func(f"{path}/{elem}"))
        
        if plot:
            xticks_negative = ["sum"] + [str(Fraction(_sortkey(i)/2)) for i in filenames_negative]
            xticks_positive = ["sum"] + [str(Fraction(_sortkey(i)/2)) for i in filenames_positive]
            sum_total_negative = sum(total_negative)
            sum_total_positive = sum(total_positive)
            
            fig0, ax0 = plt.subplots(ncols=1, nrows=2)
            fig1, ax1 = plt.subplots(ncols=1, nrows=2)

            bars = ax0[0].bar(
                xticks_negative,
                [sum_total_negative/60/60] + [i/60/60 for i in total_negative],
                color = "black",
            )
            ax0[0].set_title("negative")
            for rect in bars:
                height = rect.get_height()
                ax0[0].text(
                    x = rect.get_x() + rect.get_width() / 2.0,
                    y = height,
                    s = f'{height:.3f}',
                    ha = 'center',
                    va = 'bottom'
                )
            
            bars = ax1[0].bar(
                xticks_negative,
                [sum_total_negative/sum_total_negative] + [i/sum_total_negative for i in total_negative],
                color = "black",
            )
            ax1[0].set_title("negative")
            for rect in bars:
                height = rect.get_height()
                ax1[0].text(
                    x = rect.get_x() + rect.get_width() / 2.0,
                    y = height,
                    s = f'{height:.3f}',
                    ha = 'center',
                    va = 'bottom'
                )
            
            bars = ax0[1].bar(
                xticks_positive,
                [sum_total_positive/60/60] + [i/60/60 for i in total_positive],
                color = "black",
            )
            ax0[1].set_title("positive")
            for rect in bars:
                height = rect.get_height()
                ax0[1].text(
                    x = rect.get_x() + rect.get_width() / 2.0,
                    y = height,
                    s = f'{height:.3f}',
                    ha = 'center',
                    va = 'bottom'
                )

            bars = ax1[1].bar(
                xticks_positive,
                [sum_total_positive/sum_total_positive] + [i/sum_total_positive for i in total_positive],
                color = "black",
            )
            ax1[1].set_title("positive")
            for rect in bars:
                height = rect.get_height()
                ax1[1].text(
                    x = rect.get_x() + rect.get_width() / 2.0,
                    y = height,
                    s = f'{height:.3f}',
                    ha = 'center',
                    va = 'bottom'
                )

            fig0.text(x=0.02, y=0.5, s="Time [h]", rotation="vertical")
            fig0.text(x=0.5, y=0.02, s="Angular momentum")
            fig1.text(x=0.02, y=0.5, s="Norm. time", rotation="vertical")
            fig1.text(x=0.5, y=0.02, s="Angular momentum")
            plt.show()

        return sum(total_negative) + sum(total_positive)

    else:
        msg = f"'{path}' is neither a file nor a directory!"
        raise FileNotFoundError(msg)

def get_timing_data(path: str, plot: bool = False) -> float:
    """
    Wrapper for _get_timing_data. Input a single log filename and get
    the timing data. Input a path to a directory several log files and
    get the summed timing data. In units of seconds.

    Parameters
    ----------
    path : str
        Path to a single log file or path to a directory of log files.

    Returns
    -------
    : float
        The summed times for all input log files.
    """
    return _get_data_general(path, _get_timing_data, plot)

def get_memory_usage(path: str) -> float:
    """
    Wrapper for _get_memory_usage. Input a single log filename and get
    the memory data. Input a path to a directory several log files and
    get the summed memory data. In units of GB.

    Parameters
    ----------
    path : str
        Path to a single log file or path to a directory of log files.

    Returns
    -------
    : float
        The summed memory usage for all input log files.
    """
    return _get_data_general(path, _get_memory_usage, False)

def get_parameters(path: str, verbose: bool = True) -> dict:
    """
    Extract the parameters which are fed to KSHELL throught the shell
    script.

    Parameters
    ----------
    path : str
        Path to a KSHELL work directory.

    Returns
    -------
    res : dict
        A dictionary where the keys are the parameter names and the
        values are the corresponding values.
    """
    res = {}
    shell_filename = None
    if os.path.isdir(path):
        for elem in os.listdir(path):
            if elem.endswith(".sh"):
                shell_filename = f"{path}/{elem}"
                break
    else:
        print("Directly specifying path to .sh file not yet implemented!")

    if shell_filename is None:
        if verbose:
            msg = f"No .sh file found in path '{path}'!"
            print(msg)

        return res
    
    with open(shell_filename, "r") as infile:
        for line in infile:
            if line.startswith(r"&input"):
                break
        
        for line in infile:
            if line.startswith(r"&end"):
                """
                End of parameters.
                """
                break
            
            tmp = line.split("=")
            key = tmp[0].strip()
            value = tmp[1].strip()

            try:
                value = ast.literal_eval(value)
            except ValueError:
                """
                Cant convert strings. Keep them as strings.
                """
                pass
            except SyntaxError:
                """
                Cant convert Fortran booleans (.true., .false.). Keep
                them as strings.
                """
                pass
            
            res[key] = value

    return res

def inspect_log(
    path: str
):
    """
    Load log file(s) for inspection.

    Parameters
    ----------
    path : str
        Path to directory of log files or directly to a log file.
    """
    allowed_multipolarities_list = []
    if os.path.isdir(path):
        """
        List all log files in 'path' and let the user decide which one
        to inspect.
        """
        content = [i for i in os.listdir(path) if "log_" in i]
        content.sort()
        content = np.array(content)
        
        if not content.size:
            msg = f"No log files found in '{path}'!"
            raise FileNotFoundError(msg)

        for i, elem in enumerate(content):
            """
            Make log file names more readable and calculate allowed
            multipolarities (E1, M1, E2) based on angular momentum and
            parity (if transition log file).
            """
            if "_tr_" in elem:
                """
                Transition log file.
                """
                elem = elem.rstrip(".txt")
                elem = elem.split("_")[-2:]
                elem[0] = elem[0].lstrip("j")
                elem[1] = elem[1].lstrip("j")
                
                delta_parity = elem[0][-1] != elem[1][-1]   # Change in parity.

                j_i = int(elem[0][:-1])/2   # Initial total angular momentum.
                j_f = int(elem[1][:-1])/2   # Final.

                if delta_parity:
                    j_max = 1   # Because KSHELL doesnt support M2.
                else:
                    j_max = 2   # Because KSHELL supports max. E2.
                
                allowed_j = np.arange(max(np.abs(j_i - j_f), 1), min(j_i + j_f, j_max) + 1, 1, dtype=int)   # Possible j couplings.

                tmp_multipolarity = []
                if delta_parity:
                    """
                    If there is a change in parity: Even numbered
                    magnetic and odd numbered electric.
                    """
                    for j in allowed_j:
                        if j%2 == 0:
                            tmp_multipolarity.append(f"M{j}")
                        else:
                            tmp_multipolarity.append(f"E{j}")

                else:
                    """
                    If there is not a change in parity: Even numbered
                    electric and odd numbered magnetic.
                    """
                    for j in allowed_j:
                        if j%2 == 0:
                            tmp_multipolarity.append(f"E{j}")
                        else:
                            tmp_multipolarity.append(f"M{j}")

                allowed_multipolarities_list.append(tmp_multipolarity)
                
                elem = " <-> ".join(elem)
            else:
                elem = elem.rstrip(".txt")
                elem = elem.split("_")[-1]
                elem = elem.lstrip("j")
                allowed_multipolarities_list.append([])

            print(f"{i}: {elem}")

        allowed_multipolarities_list = np.array(allowed_multipolarities_list, dtype=object) # Numpy arrays support indexing with lists of indices.

        while True:
            """
            Prompt user for a single index or indices separated by
            comma. The latter is converted to a list and used as indices
            to the correct path and accompanying allowed multipolarity.
            """
            choice = input("Choose a log file for inspection: ")
            try:
                choice = ast.literal_eval(choice)
            except (SyntaxError, ValueError):
                print("Input must be an index or comma separated indices!")
                continue
            try:
                if isinstance(choice, int):
                    path_log = [content[choice]]
                    allowed_multipolarities = [allowed_multipolarities_list[choice]]
                elif isinstance(choice, Iterable):
                    choice = list(choice)
                    path_log = content[choice]
                    allowed_multipolarities = allowed_multipolarities_list[choice]
                else:
                    continue
                
                break

            except IndexError:
                print(f"Input indices cannot be larger than {len(content) - 1}.")
                continue

        print(f"{path_log} chosen")
        path_log = [f"{path}/{log}" for log in path_log]    # Prepend entire path (relative or absolute depending on user input).

    elif os.path.isfile(path):
        """
        If path to a single log file is given.
        """
        path_log = [path]
        elem = path_log[0].split("/")[-1]
        elem = elem.rstrip(".txt")
        elem = elem.split("_")[-2:]
        elem[0] = elem[0].lstrip("j")
        elem[1] = elem[1].lstrip("j")
        
        delta_parity = elem[0][-1] != elem[1][-1]

        j_i = int(elem[0][:-1])/2   # Initial total angular momentum.
        j_f = int(elem[1][:-1])/2   # Final.

        if delta_parity:
            j_max = 1   # Because KSHELL doesnt support M2.
        else:
            j_max = 2   # Because KSHELL supports max. E2.
        
        allowed_j = np.arange(max(np.abs(j_i - j_f), 1), min(j_i + j_f, j_max) + 1, 1, dtype=int)   # Possible j couplings.

        tmp_multipolarity = []
        if delta_parity:
            """
            If there is a change in parity: Even numbered
            magnetic and odd numbered electric.
            """
            for j in allowed_j:
                if j%2 == 0:
                    tmp_multipolarity.append(f"M{j}")
                else:
                    tmp_multipolarity.append(f"E{j}")

        else:
            """
            If there is not a change in parity: Even numbered
            electric and odd numbered magnetic.
            """
            for j in allowed_j:
                if j%2 == 0:
                    tmp_multipolarity.append(f"E{j}")
                else:
                    tmp_multipolarity.append(f"M{j}")

        allowed_multipolarities = [tmp_multipolarity]
    
    else:
        msg = "Input 'path' must be a file or directory!"
        raise FileNotFoundError(msg)

    for i, log in enumerate(path_log):
        if "_tr_" not in log:
            msg = "Inspection implemented only for transition log files (for now)."
            msg += f"\n{log = }"
            raise NotImplementedError(msg)
        
        for multipolarity in allowed_multipolarities[i]:
            j_f = []
            idx_f = []
            E_f = []
            j_i = []
            idx_i = []
            E_i = []
            E_x = []
            mred = []
            B_decay = []
            B_excite = []
            mom = []

            with open(log, "r") as infile:
                for line in infile:
                    """
                    Find where 'multipolarity' starts in the log file.
                    """
                    if f"{multipolarity} transition" not in line:
                        continue
                    else:
                        break
                
                infile.readline()   # Skip a line.

                for line in infile:
                    try:
                        tmp_j_f, tmp_idx_f, tmp_E_f, tmp_j_i, tmp_idx_i, tmp_E_i, \
                        tmp_E_x, tmp_mred, tmp_B_decay, tmp_B_excite, tmp_mom = \
                            [ast.literal_eval(i) for i in line.split()]

                        j_f.append(tmp_j_f)
                        idx_f.append(tmp_idx_f)
                        E_f.append(tmp_E_f)
                        j_i.append(tmp_j_i)
                        idx_i.append(tmp_idx_i)
                        E_i.append(tmp_E_i)
                        E_x.append(tmp_E_x)
                        mred.append(tmp_mred)
                        B_decay.append(tmp_B_decay)
                        B_excite.append(tmp_B_excite)
                        mom.append(tmp_mom)
                    except ValueError:
                        break
            
            j_f = np.array(j_f)
            idx_f = np.array(idx_f)
            E_f = np.array(E_f)
            j_i = np.array(j_i)
            idx_i = np.array(idx_i)
            E_i = np.array(E_i)
            E_x = np.array(E_x)
            mred = np.array(mred)
            B_decay = np.array(B_decay)
            B_excite = np.array(B_excite)
            mom = np.array(mom)

            print(f"{multipolarity} inspection summary of {log.split('/')[-1]}")
            print(f"{np.mean(B_decay) = }")
            print(f"{np.min(B_decay) = }")
            print(f"{np.max(B_decay) = }")
            print(f"{sum(B_decay) = }")
            print(f"{sum(B_decay == 0) = }")
            print(f"{sum(B_decay != 0) = }")
            print(f"{len(B_decay) = }")
            print()
<file_sep>from __future__ import annotations
"""
usage: count_dim.py foo.snt bar.ptn
count M-scheme dimension
Author: <NAME>
Modified by: https://github.com/GaffaSnobb
"""
from functools import lru_cache
import sys, math, time, multiprocessing, os
from typing import TextIO, Tuple

def _readline_sk(
    fp: TextIO,
    ) -> list:
    """
    Skips lines starting with '#' or '!'. Splits and returns all other
    lines.

    Parameters
    ----------
    fp:
        TextIOWrapper of .snt or .ptn file.

    Returns
    -------
    arr:
        The current line in 'fp' split by whitespace.
    """
    arr = fp.readline().split()
    while arr[0][0] in ['!','#']:
        arr = fp.readline().split()
    return arr
    
def read_snt(
    model_space_filename: str
    ) -> Tuple[list, list, list, list, list, list]:
    """
    Parameters
    ----------
    model_space_filename : str
        The filename of the .snt model space file.

    Returns
    -------
    orbits_proton_neutron : list
        The number of proton and neutron valence orbitals.
        TODO: Double check.

    core_protons_neutrons : list
        The number of protons and neutrons in the core.

    norb : list
        List of n orbital numbers.
    
    lorb : list
        List of l orbital numbers.

    jorb : list
        List of j orbital numbers.

    itorb : list
        ?
    """
    fp = open(model_space_filename,'r')
    arr = _readline_sk(fp)
    orbits_proton_neutron = [ int(arr[0]), int(arr[1]) ]
    core_protons_neutrons = [ int(arr[2]), int(arr[3]) ]

    norb, lorb, jorb, itorb = [], [], [], []
    for _ in range(sum(orbits_proton_neutron)):
        arr = _readline_sk(fp)
        norb.append( int(arr[1]))
        lorb.append( int(arr[2]))
        jorb.append( int(arr[3]))
        itorb.append(int(arr[4]))

    fp.close()
    return orbits_proton_neutron, core_protons_neutrons, norb, lorb, jorb, itorb

def read_ptn(
    partition_filename: str
    ) -> Tuple[tuple, int, list, list, list]:
    """
    Read partition file (.ptn) and extract information.

    Parameters
    ----------
    partition_filename : str
        The filename of the .ptn partition file.
    
    Returns
    -------
    valence_protons_neutrons : tuple
        A tuple containing the number of valence protons and neutrons.
        Example: (#p, #n).

    parity : int
        The parity of the configuration. There are separate files for
        the configurations of +1 and -1 parity.

    proton_partition : list
        The proton occupation for the different orbitals.

    neutron_partition : list
        The neutron occupation for the different orbitals.

    total_partition : list
        The total proton and neutron configuration (?).
    """
    fp = open(partition_filename, 'r')
    arr = _readline_sk(fp)
    valence_protons_neutrons, parity = (int(arr[0]), int(arr[1])),  int(arr[2])

    arr = _readline_sk(fp)
    n_idp, n_idn = int(arr[0]), int(arr[1])

    proton_partition = []
    for _ in range(n_idp):
        arr = _readline_sk(fp)
        proton_partition.append(tuple([int(x) for x in arr[1:]]))

    neutron_partition = []
    for _ in range(n_idn):
        arr = _readline_sk(fp)
        neutron_partition.append(tuple([int(x) for x in arr[1:]]))


    arr = _readline_sk(fp)
    n_idpn = int(arr[0])

    total_partition = []
    for _ in range(n_idpn):
        arr = _readline_sk(fp)
        total_partition.append((int(arr[0])-1, int(arr[1])-1))

    fp.close()
    
    return valence_protons_neutrons, parity, proton_partition, neutron_partition, total_partition

def _mp_add(mp1,mp2):
    for (m, p), v in mp2.items():
        mp1[m,p] = mp1.get((m,p), 0) + v

def _mp_product(mp1,mp2):
    mp = {}
    for (m1, p1), v1 in mp1.items():
        for (m2, p2), v2 in mp2.items():
            mp[m1+m2, p1*p2] = mp.get((m1+m2, p1*p2), 0) + v1*v2
    return mp

def _mps_product(mps):
    while len(mps) != 1:
        mp1 = mps.pop(0)
        mp2 = mps.pop(0)
        mps.append( _mp_product(mp1, mp2) )
    return mps[0]

def _set_dim_singlej( jorb ):
    # set dimension for each single-j orbit as a function of (j, occupation, jz)
    dim_jnm = {}
    for j in set( jorb ):
        dim_jnm[j] = [ {} for x in range(j+2) ]
        for mbit in range(2**(j+1)):
            n, m = 0, 0
            for i in range(j+1):
                n += (mbit&2**i)>>i
                m += ((mbit&2**i)>>i)*(2*i-j)
            dim_jnm[j][n][m] = dim_jnm[j][n].get(m,0) + 1
    return dim_jnm

def _parallel(args):
    dim_mp_parallel = {}
    i, dim_idp_mp, dim_idn_mp = args
    # print(i) # Process number for debug.
    _mp_add( dim_mp_parallel, _mp_product(dim_idp_mp, dim_idn_mp) )
    return dim_mp_parallel

# @lru_cache(maxsize=None, typed=False)
def count_dim(
    model_space_filename: str,
    parity: None | int,
    proton_partition: None | list[list[int]],
    neutron_partition: None | list[list[int]],
    total_partition: None | list[list[int]],
    partition_filename: str | None = None,
    print_dimensions: bool = True,
    debug: bool = False,
    ):
    """ 
    Product dimension calculation is parallelized. Some timing data for
    a .ptn file of approximately 1 million lines:
    
    TIMING (parallel):
    -------
    ```
    where                         abs. time  rel. time
    timing_read_ptn               0.7270s    0.0326
    timing_read_snt               0.0005s    0.0000
    timing_set_dim_singlej        0.0016s    0.0001
    timing_proton_partition_loop  0.0004s    0.0000
    timing_neutron_partition_loop 0.0040s    0.0002
    timing_product_dimension      21.5805s    0.9671
    timing_data_gather            0.0002s    0.0000
    timing_total                  22.3142s    1.0000
    ```
    TIMING (serial):
    -------
    ```
    where                         abs. time  rel. time
    timing_read_ptn               0.7295s    0.0108
    timing_read_snt               0.0004s    0.0000
    timing_set_dim_singlej        0.0016s    0.0000
    timing_proton_partition_loop  0.0004s    0.0000
    timing_neutron_partition_loop 0.0035s    0.0001
    timing_product_dimension      67.0947s    0.9892
    timing_data_gather            0.0003s    0.0000
    timing_total                  67.8304s    1.0000
    ```
    Parameters
    ----------
    model_space_filename : str
        The filename of the .snt file which contains the model space
        information. Example 'usda.snt'.

    partition_filename : str
        The filename of the .ptn file which contains the proton and
        neutron configuration, with truncation information if applied.

    print_dimensions : bool
        For toggling print on / off.

    debug : bool
        For toggling debug print on / off.
    """
    if partition_filename is None and (proton_partition is None or neutron_partition is None or total_partition is None or parity is None):
        msg = (
            "If 'partition_filename' is not supplied then 'proton_partition',"
            " 'neutron_partition', 'total_partition', and 'parity' must be defined."
        )
        raise ValueError(msg)

    timing_total = time.time()
    timing_read_snt = time.time()
    orbits_proton_neutron, core_protons_neutrons, norb, lorb, jorb, itorb = \
        read_snt(model_space_filename)
    timing_read_snt = time.time() - timing_read_snt
    
    timing_read_ptn = time.time()
    if partition_filename is not None:
        valence_protons_neutrons, parity, proton_partition, neutron_partition, total_partition = \
            read_ptn(partition_filename)

    timing_read_ptn = time.time() - timing_read_ptn

    timing_set_dim_singlej = time.time()
    dim_jnm = _set_dim_singlej( jorb )
    timing_set_dim_singlej = time.time() - timing_set_dim_singlej

    timing_proton_partition_loop = time.time()
    # dimension for proton
    dim_idp_mp = []
    for ptn in proton_partition:
        """
        Loop over all proton partitions.

        Ex: proton_partition=[(2, 6, 2), (3, 5, 2), (3, 6, 1), (4, 4, 2), (4, 5, 1), (4, 6, 0)]
        """
        mps = []
        for i,n in enumerate(ptn):
            p = (-1)**(lorb[i]*n) 
            j = jorb[i]
            mps.append( dict( ( (m, p), d ) for m,d in dim_jnm[j][n].items() ) )
        dim_idp_mp.append(_mps_product(mps))

    timing_proton_partition_loop = time.time() - timing_proton_partition_loop

    timing_neutron_partition_loop = time.time()
    # dimension for neutron
    dim_idn_mp = []
    for ptn in neutron_partition:
        mps = []
        for i,n in enumerate(ptn):
            p = (-1)**( lorb[ orbits_proton_neutron[0]+i ] * n )
            j = jorb[ orbits_proton_neutron[0]+i ]
            mps.append( dict( ( (m, p), d ) for m,d in dim_jnm[j][n].items() ) )
        dim_idn_mp.append( _mps_product(mps) )

    timing_neutron_partition_loop = time.time() - timing_neutron_partition_loop

    """
    Product dimensions of protons and neutrons. Parallelization gives
    some penalty for small partition files, but a large speed-up for
    huge partition files.
    """
    timing_product_dimension = time.time()
    dim_mp = {}
    with multiprocessing.Pool() as pool:
        list_of_dicts = pool.map(
            _parallel,
            # [(dim_idp_mp[idp], dim_idn_mp[idn]) for idp, idn in total_partition]
            [(i, dim_idp_mp[idx[0]], dim_idn_mp[idx[1]]) for i, idx in enumerate(total_partition)]
        )
    for dict_ in list_of_dicts:
        for key in dict_:
            try:
                dim_mp[key] += dict_[key]
            except KeyError:
                dim_mp[key] = dict_[key]
    
    # dim_mp = {}
    # for idp, idn in total_partition:
    #     _mp_add( dim_mp, _mp_product(dim_idp_mp[idp], dim_idn_mp[idn]) )
    
    timing_product_dimension = time.time() - timing_product_dimension
    
    timing_data_gather = time.time()
    M, mdim, jdim, mpow, jpow = [], [], [], [], []  # Might be unnecessary to make all of these lists.

    for m in range( max([x[0] for x in dim_mp]), -1, -2 ):
        M.append(m)
        mdim_tmp = dim_mp[m, parity]
        jdim_tmp = dim_mp[m, parity] - dim_mp.get((m+2, parity), 0)
        mdim.append(mdim_tmp)
        jdim.append(jdim_tmp)
        mpow.append(int( math.log10(mdim_tmp) ) if mdim_tmp != 0 else 0)
        jpow.append(int( math.log10(jdim_tmp) ) if jdim_tmp != 0 else 0)

    if print_dimensions:
        print("      2*M        M-scheme dim.          J-scheme dim.")
        for i in range(len(M)):
            msg = f"dim. {M[i]:5d}{mdim[i]:21d}{jdim[i]:21d}"
            msg += f"   {mdim[i]/(10**mpow[i]):4.2f}x10^{mpow[i]:2d}"
            msg += f"  {jdim[i]/(10**jpow[i]):4.2f}x10^{jpow[i]:2d}"
            print(msg)
    timing_data_gather = time.time() - timing_data_gather
    timing_total = time.time() - timing_total

    if debug:
        print("TIMING:")
        print("-------")
        print("where                         abs. time  rel. time")
        print(f"timing_read_ptn               {timing_read_ptn:.4f}s    {timing_read_ptn/timing_total:.4f}")
        print(f"timing_read_snt               {timing_read_snt:.4f}s    {timing_read_snt/timing_total:.4f}")
        print(f"timing_set_dim_singlej        {timing_set_dim_singlej:.4f}s    {timing_set_dim_singlej/timing_total:.4f}")
        print(f"timing_proton_partition_loop  {timing_proton_partition_loop:.4f}s    {timing_proton_partition_loop/timing_total:.4f}")
        print(f"timing_neutron_partition_loop {timing_neutron_partition_loop:.4f}s    {timing_neutron_partition_loop/timing_total:.4f}")
        print(f"timing_product_dimension      {timing_product_dimension:.4f}s    {timing_product_dimension/timing_total:.4f}")
        print(f"timing_data_gather            {timing_data_gather:.4f}s    {timing_data_gather/timing_total:.4f}")
        print(f"timing_total                  {timing_total:.4f}s    {timing_total/timing_total:.4f}")

    return M, mdim, jdim

def input_choice(content, content_type):
    if (n := len(content)) > 1:
        for i in range(n):
            print(f"{content[i]} ({i}), ", end="")

        while True:
            choice = input(": ")
            try:
                choice = int(choice)
                filename = content[choice]
                break
            except (ValueError, IndexError):
                continue

    elif n == 1:
        filename = content[0]

    else:
        print(f"No {content_type} file in this directory. Exiting...")
        sys.exit()

    return filename

def handle_input():
    try:
        model_space_filename, partition_filename = sys.argv[1:3]
    except ValueError:
        """
        Ask for input if none is given in the command line, or choose
        the only combination of .ptn and .snt.
        """
        dir_content = os.listdir()
        snt_content = [elem for elem in dir_content if elem.endswith(".snt")]
        ptn_content = [elem for elem in dir_content if elem.endswith(".ptn")]

        model_space_filename = input_choice(snt_content, ".snt")
        partition_filename = input_choice(ptn_content, ".ptn")

    count_dim(model_space_filename, partition_filename)

if __name__ == "__main__":
    handle_input()
<file_sep>"""
Author: <NAME> (?)
Modified by: <NAME>
"""
import sys, os, warnings
from typing import List, Tuple
from fractions import Fraction
from math import pi

#weisskopf_threshold = 1.0 # threshold to show in W.u.
weisskopf_threshold = -0.001

n_jnp = {}

E_gs = 0.0

def parity_integer_to_string(i: int) -> str:
    """
    Convert 1 to '+' and -1 to '-'.

    Parameters
    ----------
    i : int
        Parity in integer representation.

    Returns
    -------
    : str
        The string representation of the parity.
    """
    if i == 1: return '+'
    else: return '-'

def weisskopf_unit(multipole_type: str, mass: int) -> Tuple[float, str]:
    """
    Generate the Weisskopf unit_weisskopf for input multipolarity and mass.
    Ref. Bohr and Mottelson, Vol.1, p. 389.

    Parameters
    ----------
    multipole_type : str
        The electromagnetic character and angular momentum of the gamma
        radiation. Examples: 'E1', 'M1', 'E2'.

    mass : int
        Mass of the nucleus.

    Returns
    -------
    B_weisskopf : float
        Reduced transition probability in the Weisskopf estimate.

    unit_weisskopf : str
        The accompanying unit.
    """
    l = int(multipole_type[1:])
    if multipole_type[0].upper() == "E":
        B_weisskopf = 1.2**(2*l)/(4*pi)*(3/(l + 3))**2*mass**(2*l/3)  
        unit_weisskopf = f"e^2 fm^{str(2*l)}"
    
    elif multipole_type[0].upper() == "M":
        B_weisskopf = 10/pi*1.2**(2*l - 2)*(3/(l + 3))**2*mass**((2*l - 2)/3) 
        if l == 1: 
            unit_weisskopf = 'mu_N^2  '
        else:
            unit_weisskopf = f"mu_N^2 fm^{str(2*l - 2)}"
    else:
        msg = f"Got invalid multipole type: '{multipole_type}'."
        raise ValueError(msg)
    
    return B_weisskopf, unit_weisskopf
    
def read_energy_logfile(filename: str, E_data: dict):
    """
    Extract the energy, spin, and parity for each eigenstate and arrange
    the data in a dictionary, E_data, where the keys are the energies
    and the values are tuples of
    (log filename, spin, parity, eigenstate number, tt).

    The transition logs will not be read by this function as they do not
    contain the energy level information. Only the KSHELL logs will be
    read.

    Parameters
    ----------
    filename : str
        The log filename.
    """
    with open(filename, "r") as infile:
        while True:
            line = infile.readline()
            if not line: break
            if len(line) >= 11 and line[:11] != "H converged": continue
            if len(line) >= 14 and line[:14] != "H bn converged": continue
            while True:
                line = infile.readline()
                if not line: break
                if line[6:10] == '<H>:':
                    """
                    Example:
                    -------------------------------------------------
                    1  <H>:  -391.71288  <JJ>:    -0.00000  J:  0/2  prty -1     <-- this line will be read
                        <Hcm>:     0.00022  <TT>:     6.00000  T:  4/2              <-- this line will be read
                    <p Nj>  5.944  3.678  1.489  3.267  0.123  0.456  0.043        <-- this line will be skipped
                    <n Nj>  5.993  3.902  1.994  5.355  0.546  0.896  0.314        <-- this line will be skipped
                    hw:  1:1.000                                                   <-- this line will be skipped
                    -------------------------------------------------
                    NOTE: All this substring indexing seems to be easily
                    replaceable with a simple split!
                    """
                    n_eig = int(line[:5])       # Eigenvalue number. 1, 2, 3, ...
                    energy = float(line[11:22])
                    spin = int(line[45:48])     # 2*spin actually.
                    parity = int(line[57:59])
                    parity = parity_integer_to_string(parity)
                    while energy in E_data: energy += 0.000001  # NOTE: To separate energies close together? Else keys may be identical!
                    while True:
                        line = infile.readline()
                        if line[42:45] != ' T:': continue
                        tt = int(line[45:48])
                        E_data[ energy ] = (filename, spin, parity, n_eig, tt)
                        break

def spin_to_string(spin: int) -> str:
    """
    Divide spin by 2 and represent integer results as integers, and
    fraction results as fractions.

    NOTE: Tempting to just do str(Fraction(spin/2)) on the negatives too
    and just let them be negative fractions.

    Parameters
    ----------
    spin : int
        Two times the actual spin (for easier integer representation).

    Returns
    -------
    res : str
        The actual spin (input divided by 2) represented correctly as
        an integer or as a fraction.
    """
    if spin < 0:
        """
        For the invalid -1 cases. Spin cannot actually be < 0.
        """
        res = str(spin)
    else:
        res = str(Fraction(spin/2))

    return res

def read_transit_logfile_old(
    filename: str,
    multipole_type: str
    ):
    """
    Extract transit information from transit logfile. Old syntax style,
    pre 2021-11-24.

    Parameters
    ----------
    filename : str
        Filename of the log file.

    multipole_type : str
        The electromagnetic character and angular momentum of the gamma
        radiation. Examples: 'E1', 'M1', 'E2'.

    Raises
    ------
    Exception:
        If mass != mass_save. Unsure why mass_save is here at all.

    ValueError:
        If the conversion of log file values to int or float cannot be
        done. This indicates wrong log syntax.
    """
    out_e = {}
    mass_save = 0           # NOTE: Unclear what this is used for.
    with open(filename, "r") as infile:
        for line in infile:
            """
            Fetch mass number, wave function filenames and parities.
            Loop breaks when the line with parity information is found.
            """
            line_split = line.split()
            if len(line_split) == 0: continue   # Skip empty lines.
            if 'mass=' in line:
                """
                Fetch mass number. Example:
                N. of valence protons and neutrons =  15 19   mass= 50   n,z-core     8    8
                """
                n = line.index('mass=')
                mass = int(line[n + 5:n + 8])
                if not mass_save:
                    mass_save = mass
                
                if mass_save != mass: 
                    msg = f"ERROR  mass: {mass=}, {mass_save=}"
                    raise RuntimeError(msg)
                
                B_weisskopf, unit_weisskopf = weisskopf_unit(multipole_type, mass)
                continue

            if line_split[0] == 'fn_load_wave_l': 
                filename_wavefunction_left = line_split[2]
                continue
            
            if line_split[0] == 'fn_load_wave_r': 
                filename_wavefunction_right = line_split[2]
                continue
            
            if f"{multipole_type} transition" in line:
                """
                Example:
                 E2 transition  e^2 fm^4  eff_charge=  1.3500  0.3500 parity -1 -1
                   2Jf   idx  Ef        2Ji   idx  Ei          Ex        Mred.           B(EM )->        B(EM)<-         Mom.
                   4     1   -387.729   4     1   -387.729     0.000    -20.43244863     83.49699137     83.49699137    -15.48643011
                ...
                """
                parity_final = parity_integer_to_string(int(line_split[-2]))
                parity_initial = parity_integer_to_string(int(line_split[-1]))
                
                if filename_wavefunction_left == filename_wavefunction_right:
                    is_diag = True
                else:
                    is_diag = False
                break

            continue    # Included for readability.

        infile.readline()    # Skip table header.
        for line in infile:
            """
            Extract transition data from log_*_tr_*.txt. Example:

            NEW (higher decimal precision and only whitespace between values):
            E2 transition  e^2 fm^4  eff_charge=  1.3500  0.3500 parity -1 -1
            2Jf   idx  Ef        2Ji   idx  Ei          Ex        Mred.           B(EM )->        B(EM)<-         Mom.
            4     1   -387.729   4     1   -387.729     0.000    -20.43244863     83.49699137     83.49699137    -15.48643011
            4     1   -387.729   6     2   -387.461     0.267     10.55639593     22.28749899     15.91964213      0.00000000
            4     1   -387.729   8     3   -387.196     0.532     14.53838975     42.27295529     23.48497516      0.00000000
            4     1   -387.729   6     4   -387.080     0.649    -17.34631937     60.17895916     42.98497083      0.00000000
            4     1   -387.729   8     5   -386.686     1.042    -11.24379628     25.28459094     14.04699497      0.00000000
            ...

            OLD:
            ...
            4(   2) -200.706 6(   4) -198.842   1.864    0.0147    0.0000    0.0000    0.0000
            4(   2) -200.706 6(  10) -197.559   3.147   -0.0289    0.0002    0.0001    0.0000
            ...
            """
            line_split = line.split()
            if not line_split: break    # End file read when blank lines are encountered.
            if line.startswith("pn="):
                """
                jonkd: I had to add this because 'pn' showed up in the
                middle of the 'log_Ni56_gxpf1a_tr_m0p_m0p.txt' file
                after trying (unsuccessfully) to run on Fram. Example:

                ...
                4(   2) -200.706 6(   4) -198.842   1.864    0.0147    0.0000    0.0000    0.0000
                pn= 1   # of mbits=            286
                4(   2) -200.706 6(  10) -197.559   3.147   -0.0289    0.0002    0.0001    0.0000
                ...
                """
                continue
            
            try:
                spin_final = int(line[:2])
                idx_1 = int(line[3:7])
                spin_initial = int(line[17:19])
                idx_2 = int(line[20:24])
                dE = float(line[34:42]) # Gamma energy.
                E_final = float(line[8:17]) - E_gs
                E_initial = float(line[25:34]) - E_gs
                B_decay = float(line[52:62])
                B_excite = float(line[62:72])
            except ValueError as err:
                msg = f"\n{err.__str__()}"
                msg += "\nThis might be due to wrong log file syntax."
                msg += " Try using old_or_new='new' or 'both' as argument"
                msg += " to collect_logs."
                msg += f"\n{filename = }"
                msg += f"\n{line = }"
                raise ValueError(msg)

            B_weisskopf_decay  = B_decay/B_weisskopf
            B_weisskopf_excite = B_excite/B_weisskopf
            if (spin_final == spin_initial) and (idx_1 == idx_2): continue
            if is_diag and (dE < 0.0): continue
            if (B_weisskopf_decay < weisskopf_threshold): continue
            if (B_weisskopf_excite < weisskopf_threshold): continue
            if abs(E_final) < 1e-3: E_final = 0.
            if abs(E_initial) < 1e-3: E_initial = 0.
            
            idx_1 = n_jnp[ (spin_final, parity_final, idx_1) ]
            idx_2 = n_jnp[ (spin_initial, parity_initial, idx_2) ]
                
            if dE > 0:
                out = f"{spin_to_string(spin_initial):4s} "
                out += f"{parity_initial:1s} "
                out += f"{idx_2:4d} "
                out += f"{E_initial:9.3f}   "
                out += f"{spin_to_string(spin_final):4s} "
                out += f"{parity_final:1s} "
                out += f"{idx_1:4d} "
                out += f"{E_final:9.3f} "
                out += f"{dE:9.3f} "
                out += f"{B_excite:15.8f} "
                out += f"{B_weisskopf_excite:15.8f} "
                out += f"{B_decay:15.8f} "
                out += f"{B_weisskopf_decay:15.8f}\n"
                key = E_initial + E_final * 1e-5 + spin_initial *1e-10 + idx_2*1e-11 + spin_final*1e-13 + idx_1*1e-14
            else:
                """
                NOTE: What is this option used for? In what case is the
                excitation energy negative?
                """
                out = f"{spin_to_string(spin_final):4s} "
                out += f"{parity_final:1s} "
                out += f"{idx_1:4d} "
                out += f"{E_final:9.3f}   "
                out += f"{spin_to_string(spin_initial):4s} "
                out += f"{parity_initial:1s} "
                out += f"{idx_2:4d} "
                out += f"{E_initial:9.3f} "
                out += f"{-dE:9.3f} "
                out += f"{B_decay:15.8f} "
                out += f"{B_weisskopf_decay:15.8f} "
                out += f"{B_excite:15.8f} "
                out += f"{B_weisskopf_excite:15.8f}\n"
                key = E_final + E_initial * 1e-5 + spin_final *1.e-10 + idx_1*1.e-11 + spin_initial*1.e-12 + idx_2*1.e-14
            out_e[key] = out

    return unit_weisskopf, out_e, mass_save   

def read_transit_logfile(filename: str, multipole_type: str):
    """
    Extract transit information from transit logfile.

    Parameters
    ----------
    filename : str
        Filename of the log file.

    multipole_type : str
        The electromagnetic character and angular momentum of the gamma
        radiation. Examples: 'E1', 'M1', 'E2'.

    Raises
    ------
    Exception:
        If mass != mass_save. Unsure why mass_save is here at all.

    ValueError:
        If the conversion of log file values to int or float cannot be
        done. This indicates wrong log syntax.
    """
    out_e = {}
    mass_save = 0           # NOTE: Unclear what this is used for.
    unit_weisskopf = None   # Stays None if the logfile is invalid.
    with open(filename, "r") as infile:
        for line in infile:
            """
            Fetch mass number, wave function filenames and parities.
            Loop breaks when the line with parity information is found.
            """
            line_split = line.split()
            if len(line_split) == 0: continue   # Skip empty lines.
            if 'mass=' in line:
                """
                Fetch mass number. Example:
                N. of valence protons and neutrons =  15 19   mass= 50   n,z-core     8    8
                """
                n = line.index('mass=')
                mass = int(line[n + 5:n + 8])
                if not mass_save:
                    mass_save = mass
                
                if mass_save != mass: 
                    msg = f"ERROR  mass: {mass=}, {mass_save=}"
                    raise RuntimeError(msg)
                
                B_weisskopf, unit_weisskopf = weisskopf_unit(multipole_type, mass)
                continue

            if line_split[0] == 'fn_load_wave_l': 
                filename_wavefunction_left = line_split[2]
                continue
            
            if line_split[0] == 'fn_load_wave_r': 
                filename_wavefunction_right = line_split[2]
                continue
            
            if f"{multipole_type} transition" in line:
                """
                Example:
                 E2 transition  e^2 fm^4  eff_charge=  1.3500  0.3500 parity -1 -1
                   2Jf   idx  Ef        2Ji   idx  Ei          Ex        Mred.           B(EM )->        B(EM)<-         Mom.
                   4     1   -387.729   4     1   -387.729     0.000    -20.43244863     83.49699137     83.49699137    -15.48643011
                ...
                """
                parity_final = parity_integer_to_string(int(line_split[-2]))
                parity_initial = parity_integer_to_string(int(line_split[-1]))
                
                if filename_wavefunction_left == filename_wavefunction_right:
                    is_diag = True
                else:
                    is_diag = False
                break

            continue    # Included for readability.

        infile.readline()    # Skip table header.
        for line in infile:
            """
            Extract transition data from log_*_tr_*.txt. Example:

            NEW (higher decimal precision and only whitespace between values):
            E2 transition  e^2 fm^4  eff_charge=  1.3500  0.3500 parity -1 -1
            2Jf   idx  Ef        2Ji   idx  Ei          Ex        Mred.           B(EM )->        B(EM)<-         Mom.
            4     1   -387.729   4     1   -387.729     0.000    -20.43244863     83.49699137     83.49699137    -15.48643011
            4     1   -387.729   6     2   -387.461     0.267     10.55639593     22.28749899     15.91964213      0.00000000
            4     1   -387.729   8     3   -387.196     0.532     14.53838975     42.27295529     23.48497516      0.00000000
            4     1   -387.729   6     4   -387.080     0.649    -17.34631937     60.17895916     42.98497083      0.00000000
            4     1   -387.729   8     5   -386.686     1.042    -11.24379628     25.28459094     14.04699497      0.00000000
            ...

            OLD:
            ...
            4(   2) -200.706 6(   4) -198.842   1.864    0.0147    0.0000    0.0000    0.0000
            4(   2) -200.706 6(  10) -197.559   3.147   -0.0289    0.0002    0.0001    0.0000
            ...
            """
            line_split = line.split()
            if not line_split: break    # End file read when blank lines are encountered.
            if line.startswith("pn="):
                """
                jonkd: I had to add this because 'pn' showed up in the
                middle of the 'log_Ni56_gxpf1a_tr_m0p_m0p.txt' file
                after trying (unsuccessfully) to run on Fram. Example:

                ...
                4(   2) -200.706 6(   4) -198.842   1.864    0.0147    0.0000    0.0000    0.0000
                pn= 1   # of mbits=            286
                4(   2) -200.706 6(  10) -197.559   3.147   -0.0289    0.0002    0.0001    0.0000
                ...
                """
                continue
            try:
                spin_final   = int(line_split[0])
                idx_1        = int(line_split[1])
                E_final      = float(line_split[2]) - E_gs
                spin_initial = int(line_split[3])
                idx_2        = int(line_split[4])
                E_initial    = float(line_split[5]) - E_gs
                dE           = float(line_split[6])
                Mred         = float(line_split[7])
                B_decay      = float(line_split[8])
                B_excite     = float(line_split[9])
                Mom          = float(line_split[10])
            except ValueError as err:
                msg = f"\n{err.__str__()}"
                msg += "\nThis might be due to wrong log file syntax."
                msg += " Try using old_or_new='old' or 'both' as argument to"
                msg += " collect_logs."
                msg += f"\n{filename = }"
                msg += f"\n{line = }"
                raise ValueError(msg)

            B_weisskopf_decay  = B_decay/B_weisskopf
            B_weisskopf_excite = B_excite/B_weisskopf

            if (spin_final == spin_initial) and (idx_1 == idx_2): continue
            if is_diag and (dE < 0.0): continue
            if (B_weisskopf_decay < weisskopf_threshold): continue
            if (B_weisskopf_excite < weisskopf_threshold): continue
            if abs(E_final) < 1e-3: E_final = 0.
            if abs(E_initial) < 1e-3: E_initial = 0.
            
            idx_1 = n_jnp[ (spin_final, parity_final, idx_1) ]
            idx_2 = n_jnp[ (spin_initial, parity_initial, idx_2) ]
                
            if dE > 0.0:
                out = f"{spin_to_string(spin_initial):4s} "
                out += f"{parity_initial:1s} "
                out += f"{idx_2:4d} "
                out += f"{E_initial:9.3f}   "
                out += f"{spin_to_string(spin_final):4s} "
                out += f"{parity_final:1s} "
                out += f"{idx_1:4d} "
                out += f"{E_final:9.3f} "
                out += f"{dE:9.3f} "
                out += f"{B_excite:15.8f} "
                out += f"{B_weisskopf_excite:15.8f} "
                out += f"{B_decay:15.8f} "
                out += f"{B_weisskopf_decay:15.8f}\n"
                key = E_initial + E_final * 1e-5 + spin_initial *1e-10 + idx_2*1e-11 + spin_final*1e-13 + idx_1*1e-14
            else:
                """
                NOTE: What is this option used for? In what case is the
                excitation energy negative?
                """
                out = f"{spin_to_string(spin_final):4s} "
                out += f"{parity_final:1s} "
                out += f"{idx_1:4d} "
                out += f"{E_final:9.3f}   "
                out += f"{spin_to_string(spin_initial):4s} "
                out += f"{parity_initial:1s} "
                out += f"{idx_2:4d} "
                out += f"{E_initial:9.3f} "
                out += f"{-dE:9.3f} "
                out += f"{B_decay:15.8f} "
                out += f"{B_weisskopf_decay:15.8f} "
                out += f"{B_excite:15.8f} "
                out += f"{B_weisskopf_excite:15.8f}\n"
                key = E_final + E_initial * 1e-5 + spin_final *1.e-10 + idx_1*1.e-11 + spin_initial*1.e-12 + idx_2*1.e-14
            out_e[key] = out

    return unit_weisskopf, out_e, mass_save

def check_multipolarities(path: str="."):
    """
    Check whether E1, M1, and E2 values are present in transit logfiles.
    Currently only checks that the correct headers are there, not the
    actual values.

    Parameters
    ----------
    path : str
        Path to directory with transit logfiles.
    """
    print("logfile                                 E1        M1        E2        Is it correct?")
    print("------------------------------------------------------------------------------------")
    # print("log_V50_sdpf-mu_tr_j12n_j16n.txt     True   False  True")
    for elem in os.listdir(path):
        E1 = False
        should_it_have_E1 = False
        M1 = False
        should_it_have_M1 = False
        E2 = False
        should_it_have_E2 = False

        if elem.startswith("log_") and elem.endswith(".txt") and ("_tr_" in elem):
            tmp = elem[:-4]
            tmp = tmp.split("_")
            orbital_1_parity = tmp[-2][-1]
            orbital_2_parity = tmp[-1][-1]
            orbital_1_spin = int(tmp[-2][1:-1])/2
            orbital_2_spin = int(tmp[-1][1:-1])/2
            # spin_diff = abs(orbital_1_spin - orbital_2_spin)
            spin_range = list(range(
                int(abs(orbital_1_spin - orbital_2_spin)),
                int(orbital_1_spin + orbital_2_spin + 1),
                1
            ))

            if orbital_1_parity == orbital_2_parity:
                if (1 in spin_range) and (2 not in spin_range):
                    should_it_have_E1 = False
                    should_it_have_M1 = True
                    should_it_have_E2 = False
                elif (1 in spin_range) and (2 in spin_range):
                    should_it_have_E1 = False
                    should_it_have_M1 = True
                    should_it_have_E2 = True
                elif (1 not in spin_range) and (2 in spin_range):
                    should_it_have_E1 = False
                    should_it_have_M1 = False
                    should_it_have_E2 = True
                else:
                    msg = "No conditions met! (parity change)"
                    raise RuntimeError(msg)

            elif orbital_1_parity != orbital_2_parity:
                if (1 in spin_range) and (2 not in spin_range):
                    should_it_have_E1 = True
                    should_it_have_M1 = False
                    should_it_have_E2 = False
                elif (1 in spin_range) and (2 in spin_range):
                    should_it_have_E1 = True
                    should_it_have_M1 = False
                    should_it_have_E2 = False
                elif (1 not in spin_range) and (2 in spin_range):
                    """
                    This condition should never be met since KSHELL
                    will never calculate these transitions.
                    """
                    should_it_have_E1 = False
                    should_it_have_M1 = False
                    should_it_have_E2 = False
                    msg = "This condition should not ever be met!"
                    raise RuntimeError(msg)
                else:
                    msg = "No conditions met! (no parity change)"
                    raise RuntimeError(msg)
            else:
                msg = "No conditions met! (could not decide parity change)"
                raise RuntimeError(msg)

            with open(f"{path}/{elem}", "r") as infile:
                for line in infile:
                    if "E1 transition" in line:
                        E1 = True
                    elif "M1 transition" in line:
                        M1 = True
                    elif "E2 transition" in line:
                        E2 = True
                
                if (E1 == should_it_have_E1) and (M1 == should_it_have_M1) and (E2 == should_it_have_E2):
                    correct = True
                else:
                    correct = False
                print(f"{elem:40s}{str(E1):10s}{str(M1):10s}{str(E2):10s}{correct}")

def collect_logs(path: str=".", old_or_new: str="new"):
    """
    Collect energy and transition data from all log files in 'path'.

    Parameters
    ----------
    path : str
        Path to directory with log files, or parth to single energy log
        file, if you wanna do that for some reason.

    old_or_new : str
        Choose between old or new log file syntax. Pre or post
        2021-11-24. Summary file syntax is of new type, regardless.
        The option "both" is for the rare case of log files of mixed
        syntax in the same directory.

    Raises
    ------
    RuntimeError:
        If no energy log files are found in 'path'.

    ValueError:
        If 'old_or_new' is of invalid value.
    """
    allowed_old_or_new = ["new", "old", "both"]
    old_or_new = old_or_new.lower()
    if old_or_new not in allowed_old_or_new:
        msg = f"old_or_new must be in {allowed_old_or_new}. Got {old_or_new}."
        raise ValueError(msg)

    energy_log_files = []
    transit_log_files = []
    isotopes = []
    for elem in os.listdir(path):
        # if elem.startswith("log_") and elem.endswith(".txt"):
        if ("log_" in elem) and elem.endswith(".txt"):
            tmp = elem.split("_")
            isotope_index = tmp.index("log") + 1    # Isotope name is always after 'log'.
            # if (tmp := elem.split("_")[1].lower()) not in isotopes:
            if tmp[isotope_index] not in isotopes:
                isotopes.append(tmp[isotope_index])
            if not "_tr_" in elem:
                energy_log_files.append(f"{path}/{elem}")
            elif "_tr_" in elem:
                transit_log_files.append(f"{path}/{elem}")
    
    if len(isotopes) > 1:
        print(f"Log files for different isotopes have been found in {path}")
        msg = f"Found: {isotopes}. Your choice: "
        while True:
            choice = input(msg)
            if choice in isotopes:
                break
        
        # Remove all log files not of type 'choice'.
        energy_log_files = [i for i in energy_log_files if choice in i.split("/")[-1].lower()]
        transit_log_files = [i for i in transit_log_files if choice in i.split("/")[-1].lower()]

    if len(energy_log_files) == 0:
        msg = f"No energy log files in path '{path}'."
        raise FileNotFoundError(msg)

    if len(transit_log_files) == 0:
        msg = f"No transit log files in path '{path}', only energy log files."
        warnings.warn(msg, RuntimeWarning)

    E_data = {} # E_data[energy] = (log filename, spin, parity, eigenstate number, tt).
    spin_parity_occurrences = {}    # Count the occurrences of each (spin, parity) pair.
    multipole_types = ["E1", "M1", "E2"]
    
    print("Loading energy log files...")
    for i, log_file in enumerate(energy_log_files):
        progress_message = f"Reading file {i + 1:3d} of {len(energy_log_files)},"
        progress_message += f" '{log_file.split('/')[-1]}'"
        print(progress_message)
        read_energy_logfile(log_file, E_data)

    energies = E_data.keys()
    if len(energies) == 0:
        msg = "No energy data has been read from energy logs!"
        raise RuntimeError(msg)

    energies = sorted(energies)
    for energy in energies:
        """
        What does this loop actually do...?
        """
        filename, spin, parity, n_eig, tt = E_data[energy]
        spin_parity = (spin, parity)
        try:
            """
            Count the number of each (spin, parity) occurrence.
            """
            spin_parity_occurrences[spin_parity] += 1
        except KeyError:
            """
            Create initial value if key has not yet occurred.
            """
            spin_parity_occurrences[spin_parity] = 1
        n_jnp[ (spin, parity, n_eig) ] = spin_parity_occurrences[spin_parity]
        E_data[energy] = filename, spin, parity, spin_parity_occurrences[spin_parity], tt
    
    global E_gs
    E_gs = energies[0]

    counter = 0
    filename_without_path = energy_log_files[0].split("/")[-1]
    isotope = filename_without_path.split("_")[1]
    model_space = filename_without_path.split("_")[2]
    while True:
        """
        Create unique summary filename.
        """
        summary_filename = f"summary_{isotope}_{model_space}_{counter:03d}"
        if os.path.isfile(f"{path}/{summary_filename}.txt"):
            counter += 1
        else:
            summary_filename += ".txt"
            break
    
    with open(f"{path}/{summary_filename}", "w") as outfile:
        outfile.write("\n Energy levels\n")
        outfile.write('\n    N   J     prty N_Jp T        E(MeV)    Ex(MeV)  log-file\n\n')
        for i, energy in enumerate(energies):
            filename, spin, parity, n_eig, tt = E_data[energy]
            out = f"{i + 1:5d}   "
            out += f"{spin_to_string(spin):5s} "
            out += f"{parity:1s} "
            out += f"{n_eig:5d}   "
            out += f"{spin_to_string(tt):3s} "
            out += f"{energy:10.3f} "
            out += f"{energy - E_gs:10.3f}   "
            out += f"{filename}\n"
            outfile.write(out)
        outfile.write("\n")

        if len(transit_log_files) > 0:
            print("\nLoading transit log files...")
            for multipole_type in multipole_types:
                output_e = {}
                
                for i, filename in enumerate(transit_log_files):
                    progress_message = f"Reading file {i + 1:3d} of {len(transit_log_files)},"
                    progress_message += f" multipole: {multipole_type},"
                    progress_message += f" '{filename.split('/')[-1]}'"
                    print(progress_message)
                    if old_or_new == "new":
                        unit_weisskopf, out_e, mass = read_transit_logfile(filename, multipole_type)
                    elif old_or_new == "old":
                        unit_weisskopf, out_e, mass = read_transit_logfile_old(filename, multipole_type)
                    elif old_or_new == "both":
                        try:
                            unit_weisskopf, out_e, mass = read_transit_logfile_old(filename, multipole_type)
                        except ValueError:
                            """
                            If a ValueError is raised inside
                            read_transit_logfile_old, it is due to the
                            log file being new syntax styled.
                            """
                            unit_weisskopf, out_e, mass = read_transit_logfile(filename, multipole_type)
                    
                    if unit_weisskopf is None:
                        """
                        Incomplete and invalid logfile.
                        """
                        print(f"Incomplete logfile '{filename}'. Skipping...")
                        continue
                    
                    output_e.update(out_e)
                
                B_weisskopf, unit_weisskopf = weisskopf_unit(multipole_type, mass)
                outfile.write(f"B({multipole_type})  ( > {weisskopf_threshold:.1f} W.u.)  mass = {mass}    1 W.u. = {B_weisskopf:.1f} {unit_weisskopf}")
                outfile.write(f"\n{unit_weisskopf} (W.u.)")
                outfile.write(f"\nJ_i  pi_i idx_i Ex_i    J_f  pi_f idx_f Ex_f      dE         B({multipole_type})->         B({multipole_type})->[wu]     B({multipole_type})<-         B({multipole_type})<-[wu]\n")

                for _, out in sorted(output_e.items()):
                    outfile.write(out)
                outfile.write("\n\n")

    return summary_filename<file_sep>from __future__ import annotations
from .data_structures import OrbitalOrder

GS_FREE_PROTON = 5.585
GS_FREE_NEUTRON = -3.826
flags = {"debug": False, "parallel": True}

def debug_mode(switch):
    if isinstance(switch, bool):
        flags["debug"] = switch
    else:
        print(f"Invalid debug switch '{switch}'")

def latex_plot():
    import matplotlib.pyplot as plt
    plt.rcParams.update({
        "text.usetex": True,
        "font.family": "serif",
        "font.serif": ["roman"],
        "legend.fontsize": 14,
        "xtick.labelsize": 15,
        "ytick.labelsize": 15,
        "axes.labelsize": 14,
        "axes.titlesize": 15,
        "figure.titlesize": 15
    })

spectroscopic_conversion: dict[int, str] = {
    0: "s", 1: "p", 2: "d", 3: "f", 4: "g", 5: "h"
}

atomic_numbers = {
    "oxygen": 8, "fluorine": 9, "neon": 10, "sodium": 11, "magnesium": 12,
    "aluminium": 13, "silicon": 14, "phosphorus": 15, "sulfur": 16,
    "chlorine": 17, "argon": 18
}

atomic_numbers_reversed = {
    8: 'oxygen', 9: 'fluorine', 10: 'neon', 11: 'sodium', 12: 'magnesium',
    13: 'aluminium', 14: 'silicon', 15: 'phosphorus', 16: 'sulfur',
    17: 'chlorine', 18: 'argon'
}

elements = {
    'h': 1, 'he': 2, 'li': 3, 'be': 4, 'b': 5, 'c': 6, 'n': 7, 'o': 8, 'f': 9,
    'ne': 10, 'na': 11, 'mg': 12, 'al': 13, 'si': 14, 'p': 15, 's': 16,
    'cl': 17, 'ar': 18, 'k': 19, 'ca': 20, 'sc': 21, 'ti': 22, 'v': 23,
    'cr': 24, 'mn': 25, 'fe': 26, 'co': 27, 'ni': 28, 'cu': 29, 'zn': 30,
    'ga': 31, 'ge': 32, 'as': 33, 'se': 34, 'br': 35, 'kr': 36, 'rb': 37,
    'sr': 38, 'y': 39, 'zr': 40, 'nb': 41, 'mo': 42, 'tc': 43, 'ru': 44,
    'rh': 45, 'pd': 46, 'ag': 47, 'cd': 48, 'in': 49, 'sn': 50, 'sb': 51,
    'te': 52, 'i': 53, 'xe': 54, 'cs': 55, 'ba': 56, 'la': 57, 'ce': 58,
    'pr': 59, 'nd': 60, 'pm': 61, 'sm': 62, 'eu': 63, 'gd': 64, 'tb': 65,
    'dy': 66, 'ho': 67, 'er': 68, 'tm': 69, 'yb': 70, 'lu': 71, 'hf': 72,
    'ta': 73, 'w': 74, 're': 75, 'os': 76, 'ir': 77, 'pt': 78, 'au': 79,
    'hg': 80, 'tl': 81, 'pb': 82, 'bi': 83, 'po': 84, 'at': 85, 'rn': 86,
    'fr': 87, 'ra': 88, 'ac': 89, 'th': 90, 'pa': 91, 'u': 92, 'np': 93,
    'pu': 94, 'am': 95, 'cm': 96, 'bk': 97, 'cf': 98, 'es': 99, 'fm': 100,
    'md': 101, 'no': 102, 'lr': 103, 'rf': 104, 'db': 105, 'sg': 106,
    'bh': 107, 'hs': 108, 'mt': 109, 'ds': 110, 'rg': 111, 'cn': 112,
    'nh': 113, 'fl': 114, 'mc': 115, 'lv': 116, 'ts': 117, 'og': 118
}

recommended_quenching_factors = {
    "GCLSTsdpfsdgix5pn.snt": f"0.75*GS_FREE = {round(0.75*GS_FREE_PROTON, 3), round(0.75*GS_FREE_NEUTRON, 3)}",
    "gs8.snt": f"0.75*GS_FREE = {round(0.75*GS_FREE_PROTON, 3), round(0.75*GS_FREE_NEUTRON, 3)}",
    "jun45.snt": f"0.7*GS_FREE = {round(0.7*GS_FREE_PROTON, 3), round(0.7*GS_FREE_NEUTRON, 3)}",
    "gxpf1a.snt": f"0.9*GS_FREE = {round(0.9*GS_FREE_PROTON, 3), round(0.9*GS_FREE_NEUTRON, 3)}",
    "gxpf1.snt": f"0.9*GS_FREE = {round(0.9*GS_FREE_PROTON, 3), round(0.9*GS_FREE_NEUTRON, 3)}",
    "sdpf-mu.snt": f"0.9*GS_FREE = {round(0.9*GS_FREE_PROTON, 3), round(0.9*GS_FREE_NEUTRON, 3)}",
    "sn100pn.snt": f"0.7*GS_FREE = {round(0.7*GS_FREE_PROTON, 3), round(0.7*GS_FREE_NEUTRON, 3)}"
}

# shell_model_order: dict[str, int] = {   # Standard shell order for spherical nuclei.
#     "0s1": 0,
#     "0p3": 1,
#     "0p1": 2,
#     "0d5": 3,
#     "1s1": 4,
#     "0d3": 5,
#     "0f7": 6,
#     "1p3": 7,
#     "0f5": 8,
#     "1p1": 9,
#     "0g9": 10,
#     "1d5": 11,
#     "0g7": 12,
#     "1d3": 13,
#     "2s1": 14,
# }

shell_model_order: dict[str, OrbitalOrder] = {   # Standard shell order for spherical nuclei.
    "0s1": OrbitalOrder(idx=0,  major_shell_idx=0, major_shell_name="s"),
    "0p3": OrbitalOrder(idx=1,  major_shell_idx=1, major_shell_name="p"),
    "0p1": OrbitalOrder(idx=2,  major_shell_idx=1, major_shell_name="p"),
    "0d5": OrbitalOrder(idx=3,  major_shell_idx=2, major_shell_name="sd"),
    "1s1": OrbitalOrder(idx=4,  major_shell_idx=2, major_shell_name="sd"),
    "0d3": OrbitalOrder(idx=5,  major_shell_idx=2, major_shell_name="sd"),
    "0f7": OrbitalOrder(idx=6,  major_shell_idx=3, major_shell_name="pf"),
    "1p3": OrbitalOrder(idx=7,  major_shell_idx=3, major_shell_name="pf"),
    "0f5": OrbitalOrder(idx=8,  major_shell_idx=3, major_shell_name="pf"),
    "1p1": OrbitalOrder(idx=9,  major_shell_idx=3, major_shell_name="pf"),
    "0g9": OrbitalOrder(idx=10, major_shell_idx=4, major_shell_name="sdg"),
    "1d5": OrbitalOrder(idx=11, major_shell_idx=4, major_shell_name="sdg"),
    "0g7": OrbitalOrder(idx=12, major_shell_idx=4, major_shell_name="sdg"),
    "1d3": OrbitalOrder(idx=13, major_shell_idx=4, major_shell_name="sdg"),
    "2s1": OrbitalOrder(idx=14, major_shell_idx=4, major_shell_name="sdg"),
}
major_shell_order: dict[str, int] = {
    "s"  : 0,
    "p"  : 1,
    "sd" : 2,
    "pf" : 3,
    "sdg": 4,
}
<file_sep>import kshell_utilities.kshell_utilities
from kshell_utilities.general_utilities import create_spin_parity_list

def test_spin_parity_list():
    """
    Check that all spin-parity pairs appear in spin_parity_list.

    TODO: Loop over BE2 values too?
    """
    res = kshell_utilities.loadtxt(
        path = "summary_O19_sdpf-mu.txt",
        load_and_save_to_file = False,
        old_or_new = "old"
    )
    n_transitions = len(res.transitions_BM1[:, 0])
    spins = res.levels[:, 1]
    parities = res.levels[:, 2]
    spin_parity_list = create_spin_parity_list(spins, parities)
    index_list = []
    
    for transition_idx in range(n_transitions):
        """
        list.index raises ValueError if x is not in list.
        """
        spin_initial = int(res.transitions_BM1[transition_idx, 0])
        parity_initial = int(res.transitions_BM1[transition_idx, 1])
        index = spin_parity_list.index([spin_initial, parity_initial])

        if index not in index_list:
            index_list.append(index)

    index_list.sort()
    for i in range(len(index_list)-1):
        msg = "Warning! Not all spin parity pairs have been used!"
        msg += f" {index_list[i+1]} != {index_list[i] + 1}"
        assert index_list[i+1] == index_list[i] + 1, msg
    
if __name__ == "__main__":
    test_spin_parity_list()<file_sep>from __future__ import annotations
from dataclasses import dataclass, field
import numpy as np

@dataclass
class Skips:
    n_proton_skips: int = 0
    n_neutron_skips: int = 0
    n_parity_skips: int = 0
    n_ho_skips: int = 0
    n_monopole_skips: int = 0

@dataclass
class OrbitalOrder:
    """
    For storing information about the general classification of the
    shell model orbitals.
    """
    idx: int    # NOTE: Indices are in the "standard shell model order", not the order of the interaction file.
    major_shell_idx: int
    major_shell_name: str

@dataclass
class OrbitalParameters:
    """
    For storing parameters of the model space orbitals.
    """
    idx: int        # Index of the orbital from the interaction file.
    n: int          # The "principal quantum number".
    l: int          # Orbital angular momentum.
    j: int          # Total angular momentum.
    tz: int         # Isospin.
    nucleon: str    # 'p' or 'n'.
    name: str       # Ex: 'p 0d5/2'.
    parity: int
    order: OrbitalOrder
    ho_quanta: int  # Harmonic oscillator quanta (2*n + l) of the orbital (little bit unsure if this is a good name). Used to limit the possible combinations of pn configuration combinations.

    def __str__(self):
        return self.name

@dataclass
class ModelSpace:
    orbitals: list[OrbitalParameters] = field(default_factory=list)
    major_shell_names: set[str] = field(default_factory=set)
    n_major_shells: int = 0
    n_orbitals: int = 0
    n_valence_nucleons: int = 0

@dataclass
class Interaction:
    model_space: ModelSpace = field(default_factory=ModelSpace)
    model_space_proton: ModelSpace = field(default_factory=ModelSpace)
    model_space_neutron: ModelSpace = field(default_factory=ModelSpace)
    name: str = ""
    n_core_protons: int = 0
    n_core_neutrons: int = 0
    n_spe: int = 0
    spe: list[float] = field(default_factory=list)  # Single-particle energies.
    n_tbme: int = 0
    # tbme: list[list[int | float]] = field(default_factory=list) # Two-body matrix elements.
    tbme: dict[tuple[int, int, int, int, int], float] = field(default_factory=dict) # Two-body matrix elements.
    fmd_mass: int = 0   # Dont know what this is yet.
    fmd_power: float = 0    # The exponent of fmd_mass.
    vm: np.ndarray = field(default_factory=lambda: np.zeros(shape=0, dtype=float))  # Dont know what this is yet.

@dataclass
class Configuration:
    """
    Terminology:
        - "Occupation" refers to the number of particles occupying one
        orbital. Ex: 1 is an occupation.
        
        - "Configuration" refers to a set of orbitals with a given
        occupation. Ex: [0, 1, 1, 0, 0, 0] is a configuration.

        - "Partition" refers to a set of configurations with a given
        number of particles. Ex: [[0, 1, 1, 0, 0, 0], [0, 1, 0, 1, 0, 0]]
        is a partition.

    Parameters
    ----------    
    parity : int
        The parity of the configuration.

    configuration : list[int]
        The configuration as a list of occupations. Ex: [0, 1, 1, 0, 0, 0].

    ho_quanta : int
        The sum of the harmonic oscillator quanta of all the particles
        in the configuration. Used to limit the possible combinations of
        pn configuration combinations when using hw truncation.

    energy : float | None
        The energy of the combined pn configuration. For just proton
        configurations and just neutron configurations the energy has
        not yet been properly defined and is therefore set to None.

    is_original : bool
        If True, the configuration existed in the original partition
        file. If the configuration is new, then False. NOTE: Snce the
        total configurations are re-calculated, they are set to
        is_original = True if both proton and neutron configuration
        which it is built of of is is_original = True. There might
        however be a problem with this, namely that if the user changes
        the max H.O. quanta, then previously disallowed combinations of
        the original proton and neutron configurations might now be
        allowed, and these new total configurations will subsequently be
        is_original = True even though they are not.
    """
    parity: int
    configuration: list[int]
    ho_quanta: int
    energy: float | None
    is_original: bool

    def __lt__(self, other):
        if isinstance(other, Configuration):
            return self.configuration < other.configuration
        return NotImplemented

    def __le__(self, other):
        if isinstance(other, Configuration):
            return self.configuration <= other.configuration
        return NotImplemented

    def __gt__(self, other):
        if isinstance(other, Configuration):
            return self.configuration > other.configuration
        return NotImplemented

    def __ge__(self, other):
        if isinstance(other, Configuration):
            return self.configuration >= other.configuration
        return NotImplemented

    def __eq__(self, other):
        if isinstance(other, Configuration):
            return self.configuration == other.configuration
        return NotImplemented

    def __ne__(self, other):
        if isinstance(other, Configuration):
            return self.configuration != other.configuration
        return NotImplemented

@dataclass
class Partition:
    """
    Parameters
    ----------
    parity : int
        The parity of the partition file.

    Properties
    ----------
    n_configurations : int
        The number of configurations.
    """
    configurations: list[Configuration] = field(default_factory=list)
    parity: int = 0
    n_existing_positive_configurations: int = 0
    n_existing_negative_configurations: int = 0
    n_new_positive_configurations: int = 0
    n_new_negative_configurations: int = 0
    ho_quanta_min_opposite_parity: int = +1000
    ho_quanta_max_opposite_parity: int = -1000
    ho_quanta_min_this_parity: int = +1000
    ho_quanta_max_this_parity: int = -1000
    ho_quanta_min: int = +1000
    ho_quanta_max: int = -1000
    min_configuration_energy: float = 0.0   # Is this also the ground state energy?
    max_configuration_energy: float = 0.0
    max_configuration_energy_original: float = 0.0

    @property
    def n_configurations(self) -> int:
        expected = (
            self.n_existing_negative_configurations + self.n_existing_positive_configurations +
            self.n_new_negative_configurations + self.n_new_positive_configurations
        )
        calculated = len(self.configurations)
        assert expected == calculated
        return calculated
    
    @property
    def n_new_configurations(self) -> int:
        return self.n_new_negative_configurations + self.n_new_positive_configurations
    
    @property
    def n_existing_configurations(self) -> int:
        return self.n_existing_negative_configurations + self.n_existing_positive_configurations
    
    def clear(self):
        self.configurations.clear()
        self.n_existing_positive_configurations = 0
        self.n_existing_negative_configurations = 0
        self.n_new_positive_configurations = 0
        self.n_new_negative_configurations = 0
        self.ho_quanta_min_this_parity = +1000
        self.ho_quanta_max_this_parity = -1000
<file_sep>import os
from kshell_utilities.data_structures import OrbitalParameters
from kshell_utilities.partition_editor import partition_editor
from kshell_utilities.other_tools import HidePrint

def test_partition_editor():
    """
    Add one proton and one neutron configuration to Ni67_gs8_n.ptn with
    a 1hw truncation and check that the new partition file is
    constructed correctly.
    """
    def input_wrapper_test(arg: str):

        if (arg == "Add new proton configuration? (y/n): "):
            return "y"
        
        if (arg == "Add another proton configuration? (y/n): "):
            return "n"
        
        if (arg == "p 0d5/2 (remaining: 20): "):
            return "0"
        
        if (arg == "p 0d3/2 (remaining: 20): "):
            return "0"
        
        if (arg == "p 1s1/2 (remaining: 20): "):
            return "0"
        
        if (arg == "p 0f7/2 (remaining: 20): "):
            return "0"
        
        if (arg == "p 0f5/2 (remaining: 20): "):
            return "6"
        
        if (arg == "p 1p3/2 (remaining: 14): "):
            return "2"
        
        if (arg == "p 1p1/2 (remaining: 12): "):
            return "2"
        
        if (arg == "p 0g9/2 (remaining: 10): "):
            return "2"
        
        if (arg == "p 0g7/2 (remaining: 8): "):
            return "2"
        
        if (arg == "p 1d5/2 (remaining: 6): "):
            return "2"
        
        if (arg == "p 1d3/2 (remaining: 4): "):
            return "2"
        
        if (arg == "p 2s1/2 (remaining: 2): "):
            return "2"
        
        if (arg == "Add new neutron configuration? (y/n): "):
            return "y"
        
        if (arg == "n 0d5/2 (remaining: 31): "):
            return "6"
        
        if (arg == "n 0d3/2 (remaining: 25): "):
            return "4"
        
        if (arg == "n 1s1/2 (remaining: 21): "):
            return "2"
        
        if (arg == "n 0f7/2 (remaining: 19): "):
            return "8"
        
        if (arg == "n 0f5/2 (remaining: 11): "):
            return "6"
        
        if (arg == "n 1p3/2 (remaining: 5): "):
            return "2"
        
        if (arg == "n 1p1/2 (remaining: 3): "):
            return "1"
        
        if (arg == "n 0g9/2 (remaining: 2): "):
            return "2"
        
        if (arg == "n 0g7/2 (remaining: 0): "):
            return "0"
        
        if (arg == "n 1d5/2 (remaining: 0): "):
            return "0"
        
        if (arg == "n 1d3/2 (remaining: 0): "):
            return "0"
        
        if (arg == "n 2s1/2 (remaining: 0): "):
            return "0"
        
        if (arg == "Add another neutron configuration? (y/n): "):
            return "n"

        msg = f"'{arg}' is not accounted for in the testing procedure!"
        raise RuntimeError(msg)

    filename_partition_edited = "tmp_partition_editor_output_can_be_deleted_anytime.ptn"
    filename_partition_original = "Ni67_gs8_n_test.ptn"
    
    partition_editor(
        filename_interaction = "gs8_test.snt",
        filename_partition = filename_partition_original,
        filename_partition_edited = filename_partition_edited,
        input_wrapper = input_wrapper_test,
        is_interactive = False,
    )

    with open(filename_partition_edited, "r") as infile_edited, open(filename_partition_original, "r") as infile_original:
        for i, (line_edited, line_original) in enumerate(zip(infile_edited, infile_original)):
            """
            Loop over the proton configurations, neutron configurations,
            and metadata and check that they are as expected.
            """
            if line_edited != line_original:
                if line_edited == " 87 5\n":
                    """
                    The edited file will have one extra proton and
                    neutron configuration.
                    """
                    continue

                if line_edited == "    87     0  0  0  0  6  2  2  2  2  2  2  2\n":
                    """
                    The edited file will have this as an extra proton
                    configuration. That extra line has its own check in
                    this block.
                    """
                    assert infile_edited.readline() == "# neutron partition\n"
                    continue

                if line_edited == "     5     6  4  2  8  6  2  1  2  0  0  0  0\n":
                    """
                    The edited file will have this as an extra neutron
                    configuration. That extra line has its own check in
                    this block.
                    """
                    assert infile_edited.readline() == "# partition of proton and neutron\n"
                    continue

                if line_edited == "435\n":
                    """
                    The edited file will have more proton-neutron
                    configurations than the original because of the
                    extra proton and neutron configurations.
                    """
                    break

                msg = (
                    f"Error on line {i+1}. {line_edited = }, {line_original = }"
                )
                assert False, msg

        msg = "Incorrect proton-neutron indices: "
        for proton_index in range(1, 87+1):
            """
            Check that the proton-neutron indices are calculated
            correctly.
            """
            for neutron_index in range(1, 5+1):
                calculated = infile_edited.readline().split()
                expected = [str(proton_index), str(neutron_index)]
                assert calculated == expected, f"{msg}{calculated = }, {expected = }"

        msg = (
            "There are still untested lines in the edited file!"
            " It should be exhausted now!"
        )
        assert not infile_edited.readline(), msg

    os.remove(filename_partition_edited)    # Will be removed only if there is no AssertionError.

if __name__ == "__main__":
    with HidePrint():
        test_partition_editor()<file_sep>from __future__ import annotations
import sys, time, warnings
from typing import Union, Tuple, Optional
from fractions import Fraction
import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import chi2
from .parameters import flags, elements
# warnings.filterwarnings("error")    # To catch warnings with try except.
# from scipy.optimize import curve_fit

def isotope(name: str, A: int):
    protons = elements[name]
    neutrons = A - protons
    return protons, neutrons

def create_spin_parity_list(
    spins: np.ndarray,
    parities: np.ndarray
    ) -> list:
    """
    Pair up input spins and parities in a list of lists.

    Parameters
    ----------
    spins : np.ndarray
        Array of spins for each energy level.

    parities : np.ndarray
        Array of corresponding parities for each energy level.

    Returns
    -------
    spins_parities : list
        A nested list of spins and parities [[spin, parity], ...] sorted
        with respect to the spin. N is the number of unique spins in
        'spins'.

    Examples
    --------
    Example list:
    ``` python
    [[1, +1], [3, +1], [5, +1], [7, +1], [9, +1], [11, +1], [13, +1]]
    ```
    """
    spin_parity_list = []
    for i in range(len(spins)):
        if (tmp := [int(spins[i]), int(parities[i])]) in spin_parity_list:
            continue
        spin_parity_list.append(tmp)

    return spin_parity_list

def div0(numerator, denominator):
    """
    Suppress ZeroDivisionError, set x/0 to 0, and set inf, -inf and nan
    to 0. Author <NAME>.

    Examples
    --------
    >>> div0([1, 1, 1], [1, 2, 0])
    array([1. , 0.5, 0. ])
    """
    with np.errstate(divide='ignore', invalid='ignore'):
        res = np.true_divide(numerator, denominator)
        res[~np.isfinite(res)] = 0    # -inf inf NaN
    return res

def gamma_strength_function_average(
    levels: np.ndarray,
    transitions: np.ndarray,
    bin_width: Union[float, int],
    Ex_min: Union[float, int],
    Ex_max: Union[float, int],
    multipole_type: str,
    prefactor_E1: Union[None, float] = None,
    prefactor_M1: Union[None, float] = None,
    prefactor_E2: Union[None, float] = None,
    initial_or_final: str = "initial",
    partial_or_total: str = "partial",
    include_only_nonzero_in_average: bool = True,
    include_n_levels: Union[None, int] = None,
    filter_spins: Union[None, list] = None,
    filter_parities: str = "both",
    return_n_transitions: bool = False,
    # plot: bool = False,
    # save_plot: bool = False
    ) -> Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]]:
    """
    Calculate the gamma strength function averaged over total angular
    momenta, parities, and initial excitation energies.
    
    Author: <NAME>.
    Modified by: GaffaSnobb.
    
    TODO: Figure out the pre-factors.
    TODO: Use numpy.logical_or to filter levels and transitions to avoid
    using many if statements in the loops (less readable, though!).
    TODO: Make res.transitions_BXL.ji, res.transitions_BXL.pii, etc.
    class attributes (properties).

    Parameters
    ----------
    levels : np.ndarray
        Array containing energy, spin, and parity for each excited
        state. [[E, 2*spin, parity, idx], ...]. idx counts how many
        times a state of that given spin and parity has occurred. The
        first 0+ state will have an idx of 1, the second 0+ will have an
        idx of 2, etc.

    transitions : np.ndarray
        Array containing transition data for the specified
        multipolarity.
        OLD:
        Mx8 array containing [2*spin_final, parity_initial, Ex_final,
        2*spin_initial, parity_initial, Ex_initial, E_gamma, B(.., i->f)]

        OLD NEW:
        [2*spin_initial, parity_initial, Ex_initial, 2*spin_final,
        parity_final, Ex_final, E_gamma, B(.., i->f), B(.., f<-i)]

        NEW:
        [2*spin_initial, parity_initial, idx_initial, Ex_initial,
        2*spin_final, parity_final, idx_final, Ex_final, E_gamma,
        B(.., i->f), B(.., f<-i)]

    bin_width : Union[float, int]
        The width of the energy bins. A bin width of 0.2 contains 20
        states of uniform spacing of 0.01.

    Ex_min : Union[float, int]
        Lower limit for initial level excitation energy, usually in MeV.

    Ex_max : Union[float, int]
        Upper limit for initial level excitation energy, usually in MeV.

    multipole_type : str
        Choose whether to calculate for 'E1', 'M1' or 'E2'. NOTE:
        Currently only M1 and E1 is implemented.

    prefactor_E1 : Union[None, float]
        E1 pre-factor from the definition of the GSF. Defaults to a
        standard value if None.

    prefactor_M1 : Union[None, float]
        M1 pre-factor from the definition of the GSF. Defaults to a
        standard value if None.

    prefactor_E2 : Union[None, float]
        E2 pre-factor from the definition of the GSF. Defaults to a
        standard value if None.

    initial_or_final : str
        Choose whether to use the energy of the initial or final state
        for the transition calculations. NOTE: This may be removed in
        a future release since the correct alternative is to use the
        initial energy.

    partial_or_total : str
        Choose whether to use the partial level density
        rho(E_i, J_i, pi_i) or the total level density rho(E_i) for
        calculating the gamma strength function. Note that the partial
        level density, the default value, is probably the correct
        alternative. Using the total level density will introduce an
        arbitrary scaling depending on how many (J, pi) combinations
        were included in the calculations.

        This argument is included for easy comparison between the two
        densities. See the appendix of PhysRevC.98.064321 for details.

    include_only_nonzero_in_average : bool
        If True (default) only non-zero values are included in the final
        averaging of the gamma strength function. The correct
        alternative is to use only the non-zero values, so setting this
        parameter to False should be done with care.

    include_n_levels : Union[None, int]
        The number of states per spin to include. Example:
        include_n_levels = 100 will include only the 100 lowest laying
        states for each spin.

    filter_spins : Union[None, list]
        Which spins to include in the GSF. If None, all spins are
        included. TODO: Make int valid input too.

    filter_parities : str
        Which parities to include in the GSF. 'both', '+', '-' are
        allowed.

    return_n_transitions : bool
        Count the number of transitions, as a function of gamma energy,
        involved in the GSF calculation and return this number as a
        third return value. For calculating Porter-Thomas fluctuations
        in the GSF by

            r(E_gamma) = sqrt(2/n(E_gamma))

        where n is the number of transitions for each gamma energy, used
        to calculate the GSF. The value n is called n_transitions_array
        in the code. See for example DOI: 10.1103/PhysRevC.98.054303 for
        details.

    plot : bool
        Toogle plotting on / off.

    save_plot : bool
        Toogle saving of plot (as .png with dpi=300) on / off.

    Variables
    ---------
    Ex : np.ndarray 
        The excitation energy of all levels.

    Ex_initial : np.ndarray
        The excitation energy of the initial state of a transition.

    spins : np.ndarray
        The spins of all levels.

    parities : np.ndarray
        The parities of all levels.

    Returns
    -------
    bins : np.ndarray
        The bins corresponding to gSF_ExJpiavg (x values for plot).
        
    gSF_ExJpiavg : np.ndarray
        The gamma strength function.
    """
    skip_counter = {    # Debug.
        "Transit: Energy range": 0,
        "Transit: Number of levels": 0,
        "Transit: Parity": 0,
        "Level density: Energy range": 0,
        "Level density: Number of levels": 0,
        "Level density: Parity": 0
    }
    total_gsf_time = time.perf_counter()

    allowed_filter_parities = ["+", "-", "both"]
    if filter_parities not in allowed_filter_parities:
        msg = f"filter_parities must be {allowed_filter_parities}"
        raise TypeError(msg)
    if filter_parities == "both":
        filter_parities = [-1, +1]
    elif filter_parities == "-":
        filter_parities = [-1]
    elif filter_parities == "+":
        filter_parities = [+1]

    if include_n_levels is None:
        include_n_levels = np.inf   # Include all states.

    if (Ex_min < 0) or (Ex_max < 0):
        msg = "Ex_min and Ex_max cannot be negative!"
        raise ValueError(msg)

    if Ex_max < Ex_min:
        msg = "Ex_max cannot be smaller than Ex_min!"
        raise ValueError(msg)

    prefactors: dict[str, float] = {   # Factor for converting from B(XL) to GSF
        "M1": 11.5473e-9, # [1/(mu_N**2*MeV**2)].
        # "E1": 1.047e-6,
        "E1": 3.4888977e-7,
        "E2": 0.80632e-12,   # PhysRevC.90.064321
    }
    if prefactor_E1 is not None:
        """
        Override the E1 prefactor.
        """
        prefactors["E1"] = prefactor_E1
    
    if prefactor_M1 is not None:
        prefactors["M1"] = prefactor_M1
    
    if prefactor_E2 is not None:
        prefactors["E2"] = prefactor_E2
    
    prefactor = prefactors[multipole_type]

    # Extract data to a more readable form:
    n_transitions = len(transitions[:, 0])
    n_levels = len(levels[:, 0])
    E_ground_state = levels[0, 0] # Read out the absolute ground state energy so we can get relative energies later.

    try:
        Ex, spins, parities, level_counter = np.copy(levels[:, 0]), levels[:, 1], levels[:, 2], levels[:, 3]
    
    except IndexError as err:
        msg = f"{err.__str__()}\n"
        msg += "Error probably due to old tmp files. Use loadtxt parameter"
        msg += " load_and_save_to_file = 'overwrite' (once) to re-read data from the"
        msg += " summary file and generate new tmp files."
        raise Exception(msg) from err
    
    if initial_or_final == "initial":
        Ex_initial_or_final = np.copy(transitions[:, 3])   # To avoid altering the raw data.
        spin_initial_or_final_idx = 0
        parity_initial_or_final_idx = 1
    
    elif initial_or_final == "final":
        Ex_initial_or_final = np.copy(transitions[:, 7])   # To avoid altering the raw data.
        spin_initial_or_final_idx = 4
        parity_initial_or_final_idx = 5
        msg = "Using final states for the energy limits is not correct"
        msg += " and should only be used for comparison with the correct"
        msg += " option which is using initial states for the energy limits."
        warnings.warn(msg, RuntimeWarning)
    
    else:
        msg = "'initial_or_final' must be either 'initial' or 'final'."
        msg += f" Got {initial_or_final}"
        raise ValueError(msg)

    if abs(Ex_initial_or_final[0]) > 10:
        """
        Adjust energies relative to the ground state energy if they have
        not been adjusted already. The ground state energy is usually
        minus a few tens of MeV and above, so checking absolute value
        above 10 MeV is probably safe. Cant check for equality to zero
        since the initial state will never be zero.
        NOTE: Just check if the value is negative instead?
        """
        Ex_initial_or_final -= E_ground_state

    if Ex[0] != 0:
        """
        Adjust energies relative to the ground state energy if they have
        not been adjusted already.
        """
        Ex -= E_ground_state

    if (Ex_actual_max := np.max(Ex)) < Ex_max:
        msg = "Requested max excitation energy is greater than the largest"
        msg += " excitation energy in the data file."
        msg += f" Changing Ex_max from {Ex_max} to {Ex_actual_max}."
        Ex_max = Ex_actual_max
        print(msg)

    """
    Find index of first and last bin (lower bin edge) where we put
    counts. It's important to not include the other Ex bins in the
    averaging later, because they contain zeros which will pull the
    average down.
    
    Bin alternatives:
    bin_array = np.linspace(0, bin_width*n_bins, n_bins + 1) # Array of lower bin edge energy values
    bin_array_middle = (bin_array[0: -1] + bin_array[1:])/2 # Array of middle bin values
    """
    Ex_min_idx = int(Ex_min/bin_width)
    Ex_max_idx = int(Ex_max/bin_width)
    n_bins = int(np.ceil(Ex_max/bin_width)) # Make sure the number of bins cover the whole Ex region.
    # Ex_max = bin_width*n_bins # Adjust Ex_max to match the round-off in the bin width. NOTE: Unsure if this is needed.

    """
    B_pixel_sum[Ex_final_idx, E_gamma_idx, spin_parity_idx] contains the
    summed reduced transition probabilities for all transitions
    contained within the Ex_final_idx bin, E_gamma_idx bin, and
    spin_parity_idx bin. B_pixel_counts counts the number of transitions
    within the same bins.
    """
    spin_parity_list = create_spin_parity_list(spins, parities) # To create a unique index for every [spin, parity] pair.
    n_unique_spin_parity_pairs = len(spin_parity_list)
    B_pixel_sum = np.zeros((n_bins, n_bins, n_unique_spin_parity_pairs))     # Summed B(..) values for each pixel.
    B_pixel_count = np.zeros((n_bins, n_bins, n_unique_spin_parity_pairs))   # The number of transitions.
    rho_ExJpi = np.zeros((n_bins, n_unique_spin_parity_pairs))  # (Ex, Jpi) matrix to store level density
    gSF = np.zeros((n_bins, n_bins, n_unique_spin_parity_pairs))    
    n_transitions_array = np.zeros(n_bins, dtype=int)  # Count the number of transitions per gamma energy bin.
    transit_gsf_time = time.perf_counter()
    
    for transition_idx in range(n_transitions):
        """
        Iterate over all transitions in the transitions matrix and add
        up all reduced transition probabilities and the number of
        transitions in the correct bins.
        """
        if (Ex_initial_or_final[transition_idx] < Ex_min) or (Ex_initial_or_final[transition_idx] >= Ex_max):
            """
            Check if transition is within min max limits, skip if not.
            """
            skip_counter["Transit: Energy range"] += 1   # Debug.
            continue

        idx_initial = transitions[transition_idx, 2]
        idx_final = transitions[transition_idx, 6]

        if (idx_initial > include_n_levels) or (idx_final > include_n_levels):
            """
            Include only 'include_n_levels' number of levels. Defaults
            to np.inf (include all).
            """
            skip_counter["Transit: Number of levels"] += 1   # Debug.
            continue

        spin_initial = transitions[transition_idx, 0]/2
        spin_final = transitions[transition_idx, 4]/2

        if filter_spins is not None:
            # if (spin_initial not in filter_spins) or (spin_final not in filter_spins):
            if spin_initial not in filter_spins:
                """
                Skip transitions to or from levels of total angular momentum
                not in the filter list.
                """
                try:
                    skip_counter[f"Transit: j init: {spin_initial}"] += 1
                except KeyError:
                    skip_counter[f"Transit: j init: {spin_initial}"] = 1
                continue

        parity_initial = transitions[transition_idx, 1]
        parity_final = transitions[transition_idx, 5]

        if (parity_initial not in filter_parities) or (parity_final not in filter_parities):
            """
            Skip initial or final parities which are not in the filter
            list. NOTE: Might be wrong to filter on the final parity.
            """
            skip_counter["Transit: Parity"] += 1
            continue

        # Get bin index for E_gamma and Ex. Indices are defined with respect to the lower bin edge.
        E_gamma_idx = int(transitions[transition_idx, 8]/bin_width)
        Ex_initial_or_final_idx = int(Ex_initial_or_final[transition_idx]/bin_width)
        n_transitions_array[E_gamma_idx] += 1    # Count the number of transitions involved in this GSF (Porter-Thomas fluctuations).
        """
        transitions : np.ndarray
            OLD:
            Mx8 array containing [2*spin_final, parity_initial, Ex_final,
            2*spin_initial, parity_initial, Ex_initial, E_gamma, B(.., i->f)]
            OLD NEW:
            [2*spin_initial, parity_initial, Ex_initial, 2*spin_final,
            parity_final, Ex_final, E_gamma, B(.., i->f), B(.., f<-i)]
            NEW:
            [2*spin_initial, parity_initial, idx_initial, Ex_initial,
            2*spin_final, parity_final, idx_final, Ex_final, E_gamma,
            B(.., i->f), B(.., f<-i)]
        """
        spin_initial_or_final = int(transitions[transition_idx, spin_initial_or_final_idx])  # Superfluous int casts?
        parity_initial_or_final = int(transitions[transition_idx, parity_initial_or_final_idx])
        spin_parity_idx = spin_parity_list.index([spin_initial_or_final, parity_initial_or_final])

        try:
            """
            Add B(..) value and increment transition count,
            respectively. NOTE: Hope to remove this try-except by
            implementing suitable input checks to this function. Note to
            the note: Will prob. not be removed to keep the ability to
            compare initial and final.
            """
            B_pixel_sum[Ex_initial_or_final_idx, E_gamma_idx, spin_parity_idx] += \
                transitions[transition_idx, 9]
            B_pixel_count[Ex_initial_or_final_idx, E_gamma_idx, spin_parity_idx] += 1
        except IndexError as err:
            """
            NOTE: This error usually occurs because Ex_max is set to
            limit Ex_final instead of Ex_initial. If so, E_gamma might
            be larger than Ex_max and thus be placed in a B_pixel
            outside of the allocated scope. This error has a larger
            probability of occuring if Ex_max is set to a low value
            because then the probability of 
                
                E_gamma = Ex_initial - Ex_final

            is larger.
            """
            msg = f"{err.__str__()}\n"
            msg += f"{Ex_initial_or_final_idx=}, {E_gamma_idx=}, {spin_parity_idx=}, {transition_idx=}\n"
            msg += f"{B_pixel_sum.shape=}\n"
            msg += f"{transitions.shape=}\n"
            msg += f"{Ex_max=}\n"
            msg += f"2*spin_final: {transitions[transition_idx, 4]}\n"
            msg += f"parity_initial: {transitions[transition_idx, 1]}\n"
            msg += f"Ex_final: {transitions[transition_idx, 7]}\n"
            msg += f"2*spin_initial: {transitions[transition_idx, 0]}\n"
            msg += f"parity_initial: {transitions[transition_idx, 1]}\n"
            msg += f"Ex_initial: {transitions[transition_idx, 3]}\n"
            msg += f"E_gamma: {transitions[transition_idx, 8]}\n"
            msg += f"B(.., i->f): {transitions[transition_idx, 9]}\n"
            msg += f"B(.., f<-i): {transitions[transition_idx, 10]}\n"
            raise Exception(msg) from err

    transit_gsf_time = time.perf_counter() - transit_gsf_time
    level_density_gsf_time = time.perf_counter()
    
    for levels_idx in range(n_levels):
        """
        Calculate the level density for each (Ex, spin_parity) pixel.
        """
        if Ex[levels_idx] >= Ex_max:
            """
            Skip if level is outside range. Only upper limit since
            decays to levels below the lower limit are allowed.
            """
            skip_counter["Level density: Energy range"] += 1
            continue
        
        if level_counter[levels_idx] > include_n_levels:
            """
            Include only 'include_n_levels' number of levels. Defaults
            to np.inf (include all).
            """
            skip_counter["Level density: Number of levels"] += 1
            continue

        if filter_spins is not None:
            if (spin_tmp := levels[levels_idx, 1]/2) not in filter_spins:
                """
                Skip levels of total angular momentum not in the filter
                list.
                """
                try:
                    skip_counter[f"Level density: j: {spin_tmp}"] += 1
                except KeyError:
                    skip_counter[f"Level density: j: {spin_tmp}"] = 1
                continue

        Ex_idx = int(Ex[levels_idx]/bin_width)

        spin_parity_idx = \
            spin_parity_list.index([spins[levels_idx], parities[levels_idx]])
        
        rho_ExJpi[Ex_idx, spin_parity_idx] += 1

    level_density_gsf_time = time.perf_counter() - level_density_gsf_time

    if partial_or_total == "total":
        """
        Use the total level density, rho(E_i), instead of the partial
        level density, rho(E_i, J_i, pi_i). Sum over all (J_i, pi_i)
        pairs and then copy these summed values to all columns in
        rho_ExJpi.
        """
        tmp_sum = rho_ExJpi.sum(axis=1)

        for i in range(rho_ExJpi.shape[1]):
            """
            All columns in rho_ExJpi will be identical. This is for
            compatibility with the following for loop.
            """
            rho_ExJpi[:, i] = tmp_sum

        msg = "Using the total level density is not correct and"
        msg += " should only be used when comparing with the correct"
        msg += " alternative which is using the partial level density."
        warnings.warn(msg, RuntimeWarning)

    rho_ExJpi /= bin_width # Normalize to bin width, to get density in MeV^-1.
    gsf_time = time.perf_counter()
    for spin_parity_idx in range(n_unique_spin_parity_pairs):
        """
        Calculate gamma strength functions for each [Ex, E_gamma,
        spin_parity] individually using the partial level density for
        each [Ex, spin_parity].
        """
        for Ex_idx in range(n_bins):
            gSF[Ex_idx, :, spin_parity_idx] = \
                prefactor*rho_ExJpi[Ex_idx, spin_parity_idx]*div0(
                    numerator = B_pixel_sum[Ex_idx, :, spin_parity_idx],
                    denominator = B_pixel_count[Ex_idx, :, spin_parity_idx]
                )

    gsf_time = time.perf_counter() - gsf_time
    avg_gsf_time = time.perf_counter()

    if include_only_nonzero_in_average:
        """
        Update 20171009 (Midtbø): Took proper care to only average over
        the non-zero f(Eg,Ex,J,parity_initial) pixels.

        NOTE: Probably not necessary to set an upper limit on gSF
        due to the input adjustment of Ex_max.
        """
        gSF_currentExrange = gSF[Ex_min_idx:Ex_max_idx + 1, :, :]
        gSF_ExJpiavg = div0(
            numerator = gSF_currentExrange.sum(axis = (0, 2)),
            denominator = (gSF_currentExrange != 0).sum(axis = (0, 2))
        )
    else:
        """
        NOTE: Probably not necessary to set an upper limit on gSF
        due to the input adjustment of Ex_max.
        """
        gSF_ExJpiavg = gSF[Ex_min_idx:Ex_max_idx + 1, :, :].mean(axis=(0, 2))
        msg = "Including non-zero values when averaging the gamma strength"
        msg += " function is not correct and should be used with care!"
        warnings.warn(msg, RuntimeWarning)
    
    avg_gsf_time = time.perf_counter() - avg_gsf_time

    bins = np.linspace(0, Ex_max, n_bins + 1)
    bins = (bins[:-1] + bins[1:])/2   # Middle point of the bins.
    bins = bins[:len(gSF_ExJpiavg)]

    total_gsf_time = time.perf_counter() - total_gsf_time
    if flags["debug"]:
        transit_total_skips = \
            sum([skip_counter[key] for key in skip_counter if key.startswith("Transit")])
        level_density_total_skips = \
            sum([skip_counter[key] for key in skip_counter if key.startswith("Level density")])
        n_transitions_included = n_transitions - transit_total_skips
        n_levels_included = n_levels - level_density_total_skips
        print("--------------------------------")
        print(f"{transit_gsf_time = } s")
        print(f"{level_density_gsf_time = } s")
        print(f"{gsf_time = } s")
        print(f"{avg_gsf_time = } s")
        print(f"{total_gsf_time = } s")
        print(f"{multipole_type = }")
        for elem in skip_counter:
            print(f"Skips: {elem}: {skip_counter[elem]}")
        # print(f"{skip_counter = }")
        print(f"{transit_total_skips = }")
        print(f"{n_transitions = }")
        print(f"{n_transitions_included = }")
        print(f"{level_density_total_skips = }")
        print(f"{n_levels = }")
        print(f"{n_levels_included = }")
        print("--------------------------------")

    if return_n_transitions:
        return bins, gSF_ExJpiavg, n_transitions_array
    else:
        return bins, gSF_ExJpiavg

def level_plot(
    levels: np.ndarray,
    include_n_levels: int = 1_000,
    filter_spins: Union[None, list] = None,
    filter_parity: Union[None, str] = None,
    ax: Union[None, plt.Axes] = None,
    color: Union[None, str] = None,
    line_width: float = 0.4,
    x_offset_scale: float = 1.0,
    ):
    """
    Generate a level plot for a single isotope. Spin on the x axis,
    energy on the y axis.

    Parameters
    ----------
    levels : np.ndarray
        NxM array of [[energy, spin, parity, index], ...]. This is the
        instance attribute 'levels' of ReadKshellOutput.
    
    include_n_levels : int
        The maximum amount of states to plot for each spin. Default set
        to a large number to indicate ≈ no limit.

    filter_spins : Union[None, list]
        Which spins to include in the plot. If None, all spins are
        plotted.

    filter_parity : Union[None, str]
        A filter for parity. If None (default) then the parity of the
        ground state will be used. `+` is positive, `-` is negative,
        while `both` gives both parities.

    ax : Union[None, plt.Axes]
        matplotlib Axes to plot on. If None, plt.Figure and plt.Axes is
        generated in this function.

    color : Union[None, str]
        Color to use for the levels. If None, the next color in the
        matplotlib color_cycle iterator is used.
    
    line_width : float
        The width of the level lines. Not really supposed to be changed
        by the user. Set to 0.2 for comparison plots when both integer
        and half integer angular momenta are included, 0.4 else.

    x_offset_scale : float
        To scale the x offset for the hlines. This is used to fit
        columns for both integer and half integer angular momenta, as
        well as both parities.
    """
    ax_input = False if (ax is None) else True

    if levels[0, 0] != 0:
        """
        Adjust energies relative to the ground state energy.
        """
        energies = levels[:, 0] - levels[0, 0]
    else:
        energies = levels[:, 0]

    spins = levels[:, 1]/2  # levels[:, 1] is 2*spin.
    parities = levels[:, 2]

    allowed_filter_parity = [None, "+", "-", "both"]
    if filter_parity not in allowed_filter_parity:
        msg = f"Allowed parity filters are: {allowed_filter_parity}."
        raise ValueError(msg)

    if filter_parity is None:
        """
        Default to the ground state parity.
        """
        parity_integer: int = [levels[0, 2]]
        parity_symbol: str = "+" if (levels[0, 2] == 1) else "-"
        x_offset = 0    # No offset needed for single parity plot.

    elif filter_parity == "+":
        parity_symbol: str = filter_parity
        parity_integer: list = [1]
        x_offset = 0

    elif filter_parity == "-":
        parity_symbol: str = filter_parity
        parity_integer: list = [-1]
        x_offset = 0

    elif filter_parity == "both":
        line_width /= 2 # Make room for both parities.
        parity_symbol: str = r"-+"
        parity_integer: list = [-1, 1]
        x_offset = 1/4*x_offset_scale  # Offset for plots containing both parities.
    
    if filter_spins is not None:
        spin_scope = np.unique(filter_spins)    # x values for the plot.
    else:
        spin_scope = np.unique(spins)
    
    counts = {} # Dict to keep tabs on how many levels of each angular momentum have been plotted.

    if not ax_input:
        fig, ax = plt.subplots()

    if color is None:
        color = next(ax._get_lines.color_cycle)

    for i in range(len(energies)):
        if filter_spins is not None:
            if spins[i] not in filter_spins:
                """
                Skip spins which are not in the filter.
                """
                continue

        if parities[i] not in parity_integer:
            continue

        key: str = f"{spins[i]} + {parities[i]}"
        try:
            counts[key] += 1
        except KeyError:
            counts[key] = 1
        
        if counts[key] > include_n_levels:
            """
            Include only the first `include_n_levels` amount of states
            for any of the spins.
            """
            continue

        ax.hlines(
            y = energies[i],
            xmin = spins[i] - line_width + x_offset*parities[i]*0.9,
            xmax = spins[i] + line_width + x_offset*parities[i]*0.9,
            color = color,
            alpha = 0.5,
        )

    ax.set_xticks(spin_scope)
    ax.set_xticklabels([f"{Fraction(i)}" + r"$^{" + f"{parity_symbol}" + r"}$" for i in spin_scope])
    ax.set_xlabel(r"$j^{\pi}$")
    ax.set_ylabel(r"$E$ [MeV]")

    if not ax_input:
        plt.show()

def level_density(
    levels: np.ndarray,
    bin_width: Union[int, float],
    include_n_levels: Union[None, int] = None,
    filter_spins: Union[None, int, list] = None,
    filter_parity: Union[None, str, int] = None,
    E_min: Union[float, int] = 0,
    E_max: Union[float, int] = np.inf,
    return_counts: bool = False,
    plot: bool = False,
    save_plot: bool = False
    ) -> Tuple[np.ndarray, np.ndarray]:
    """
    Calculate the level density for a given bin size.

    Parameters
    ----------
    levels : Union[np.ndarray, list]
        Nx4 array of [[E, 2*spin, parity, idx], ...] or 1D array / list
        of only energies.

    bin_width : Union[int, float]
        Energy interval of which to calculate the density.

    include_n_levels : Union[None, int]
        The number of states per spin to include. Example:
        include_n_levels = 100 will include only the 100 lowest laying
        states for each spin.

    filter_spins : Union[None, int, list]
        Keep only the levels which have angular momenta in the filter.
        If None, all angular momenta are kept. Input must be the actual
        angular momenta values and not 2*j.

    filter_parity : Union[None, str, int]
        Keep only levels of parity 'filter_parity'. +1, -1, '+', '-'
        allowed inputs.

    E_min : Union[None, float, int]
        Minimum energy to include in the calculation.

    E_max : Union[None, float, int]
        Maximum energy to include in the calculation. If input E_max is
        larger than the largest E in the data set, E_max is set to that
        value.

    return_counts : bool
        Return the counts per bin instead of the density
        (density = counts/bin_width).

    plot : bool
        For toggling plotting on / off.

    save_plot : bool    
        Toogle saving of plot (as .png with dpi=300) on / off.

    Returns
    -------
    bins : np.ndarray
        The corresponding bins (x value for plotting).

    density : np.ndarray
        The level density.

    Raises
    ------
    ValueError:
        If any filter is given when energy_levels is a list of only
        energy levels.

    TypeError:
        If input parameters are of the wrong type.
    """
    if not isinstance(levels, np.ndarray):
        levels = np.array(levels)

    if not isinstance(filter_spins, (int, float, list, type(None), np.ndarray)):
        msg = f"'filter_spins' must be of type: int, float, list, None. Got {type(filter_spins)}."
        raise TypeError(msg)

    if not isinstance(include_n_levels, (int, type(None))):
        msg = f"'include_n_levels' must be of type: int, None. Got {type(include_n_levels)}."
        raise TypeError(msg)

    if not isinstance(filter_parity, (type(None), int, str)):
        msg = f"'filter_parity' must be of type: None, int, str. Got {type(filter_parity)}."
        raise TypeError(msg)

    if not isinstance(E_min, (int, float)):
        msg = f"'E_min' must be of type: int, float. Got {type(E_min)}."
        raise TypeError(msg)

    if not isinstance(E_max, (int, float)):
        msg = f"'E_max' must be of type: int, float. Got {type(E_max)}."
        raise TypeError(msg)

    if isinstance(filter_parity, str):
        valid_filter_parity = ["+", "-"]
        if filter_parity not in valid_filter_parity:
            msg = f"Valid parity filters are: {valid_filter_parity}."
            raise ValueError(msg)
        
        filter_parity = 1 if (filter_parity == "+") else -1

    if isinstance(filter_spins, (int, float)):
        filter_spins = [filter_spins]

    if (levels.ndim == 1) and (filter_spins is not None):
        msg = "Spin filter cannot be applied to a list of only energy levels!"
        raise ValueError(msg)

    if (levels.ndim == 1) and (include_n_levels is not None):
        msg = "Cannot choose the number of levels per spin if 'levels' is only a list of energies!"
        raise ValueError(msg)
    
    if (levels.ndim == 1) and (filter_parity is not None):
        msg = "Parity filter cannot be applied to a list of only energy levels!"
        raise ValueError(msg)

    energy_levels = np.copy(levels) # Just in case.
    if energy_levels.ndim == 1:
        """
        'levels' only contains energy values.
        """
        if energy_levels[0] < 0:
            msg = "Please scale energies relative to the ground state"
            msg += " energy before calculating the NLD!"
            raise ValueError(msg)

    elif energy_levels.ndim == 2:
        """
        'levels' is a multidimensional array on the form
        [[E, 2*spin, parity, idx], ...]. Subtract ground state energy
        from all energies.
        """
        energy_levels[:, 0] -= energy_levels[0, 0]
        
        if include_n_levels is not None:
            """
            Include ony 'include_n_levels' of levels per angular
            momentum and parity pair.
            """
            indices = energy_levels[:, 3]  # Counter for the number of levels per spin.
            energy_levels = energy_levels[indices <= include_n_levels]

        if filter_spins is not None:
            """
            filter_spins is a list of angular momenta. Inside this if
            statement we know that 'levels' is a multidimensional array due
            to the check inside the previous except.
            levels: 
            """
            filter_spins = [2*j for j in filter_spins]  # energy_levels has 2*j to avoid fractions.

            mask_list = []
            for j in filter_spins:
                """
                Create a [bool1, bool2, ...] mask for each j.
                """
                mask_list.append(energy_levels[:, 1] == j)

            energy_levels = energy_levels[np.logical_or.reduce(mask_list)]  # Contains only levels of j in the filter.

        if filter_parity is not None:
            energy_levels = energy_levels[energy_levels[:, 2] == filter_parity]
        
        energy_levels = energy_levels[:, 0]

    E_max = min(E_max, energy_levels[-1] + 0.1) # E_max cant be larger than the largest energy in the data set. Add small number to include the final level(s) in the counting.
    if E_max <= E_min:
        """
        This behaviour is OK for angular momentum distribution heatmaps
        where the NLDs for different angular momenta are compared using
        the same bin range.

        In this situation, the density at E_min is of course zero
        because there are no levels of this energy for the given angular
        momentum.
        """
        bins = np.array([E_min])
        density = np.array([0])
        return bins, density

    bins = np.arange(E_min, E_max + bin_width, bin_width)
    bins[-1] = E_max    # arange will mess up the final bin if it does not match the bin width.

    n_bins = len(bins)
    counts = np.zeros(n_bins - 1)

    for i in range(n_bins - 1):
        mask_1 = energy_levels >= bins[i]
        mask_2 = energy_levels < bins[i+1]
        mask_3 = np.logical_and(mask_1, mask_2)
        counts[i] = sum(mask_3)
        # counts[i] = np.sum(bins[i] <= energy_levels[energy_levels < bins[i + 1]])
    
    density = (counts/bin_width)
    # bins = bins[1:]
    bins = bins[:-1]    # Maybe just a matter of preference...?

    if plot:
        fig, ax = plt.subplots()
        if return_counts:
            ax.step(bins, counts, color="black")
            ax.set_ylabel(r"Counts")
        else:
            ax.step(bins, density, color="black")
            ax.set_ylabel(r"NLD [MeV$^{-1}$]")
        ax.set_xlabel("E [MeV]")
        ax.legend([f"{bin_width=} MeV"])
        ax.grid()
        if save_plot:
            fname = "nld.png"
            print(f"NLD saved as '{fname}'")
            fig.savefig(fname=fname, dpi=300)
        plt.show()

    if return_counts:
        return bins, density, counts
    else:
        return bins, density

def porter_thomas(
    transitions: np.ndarray,
    Ei: Union[int, float, list],
    BXL_bin_width: Union[int, float],
    j_list: Union[list, None] = None,
    Ei_bin_width: Union[int, float] = 0.1,
    return_chi2: bool = False,
) -> tuple[np.ndarray, np.ndarray, Optional[np.ndarray]]:
    """
    Calculate the distribution of B(XL)/mean(B(XL)) values scaled to
    a chi-squared distribution of 1 degree of freedom.

    Parameters
    ----------
    transitions : np.ndarray
        Array containing transition data for the specified
        multipolarity.

        [2*spin_initial, parity_initial, idx_initial, Ex_initial,
        2*spin_final, parity_final, idx_final, Ex_final, E_gamma,
        B(XL, i->f), B(XL, f<-i)]

    Ei : int, float, list
        The initial excitation energy of the transitions where the
        distribution will be calculated. If Ei is only a number, then a
        bin of size 'Ei_bin_width' around Ei will be used. If Ei is a
        list, tuple, or array with both a lower and an upper limit, then
        all excitation energies in that interval will be used.

    BXL_bin_width : int, float
        The bin size of the BXL values for the distribution (not the Ei
        bin size!).

    Ei_bin_width : int, float
        The size of the initial energy bin if 'Ei' is only one number.
        Will not be used if 'Ei' is both a lower and an upper limit.

    return_chi2 : bool
        If True, the chi-squared distribution y values will be returned
        as a third return value.

    Returns
    -------
    BXL_bins : np.ndarray
        The BXL bins (x values).

    BXL_counts : np.ndarray
        The number of counts in each BXL_bins bin (y values).

    rv.pdf(BXL_bins) : np.ndarray
        The chi-squared distribution y values.
    """
    pt_prepare_data_time = time.perf_counter()
    if isinstance(Ei, (list, tuple, np.ndarray)):
        """
        If Ei defines a lower and an upper limit.
        """
        Ei_mask = np.logical_and(
            transitions[:, 3] >= Ei[0],
            transitions[:, 3] < Ei[-1]
        )
        BXL = transitions[Ei_mask]
    else:
        BXL = transitions[np.abs(transitions[:, 3] - Ei) < Ei_bin_width] # Consider only levels around Ei.

    if j_list is not None:
        """
        Create a mask of j values for the transitions array. Allow only
        entries with initial angular momenta in j_list.
        """
        if not isinstance(j_list, list):
            msg = f"j_list must be of type list! Got {type(j_list)}."
            raise TypeError(msg)

        j_list = [2*j for j in j_list]  # Angular momenta are stored as 2*j to avoid fractions.

        mask_list = []
        for j in j_list:
            """
            Create a [bool1, bool2, ...] mask for each j.
            """
            mask_list.append(BXL[:, 0] == j)

        BXL = BXL[np.logical_or.reduce(mask_list)]  # Contains only transitions of j in the filter.

    # BXL = np.copy(BXL[:, 9]) # The 9th col. is the reduced decay transition probabilities.
    n_BXL_before = len(BXL)
    idxi_masks = []
    pii_masks = []
    ji_masks = []
    BXL_tmp = []

    initial_indices = np.unique(BXL[:, 2]).astype(int)
    initial_parities = np.unique(BXL[:, 1]).astype(int)
    initial_j = np.unique(BXL[:, 0])

    for idxi in initial_indices:
        idxi_masks.append(BXL[:, 2] == idxi)

    for pii in initial_parities:
        pii_masks.append(BXL[:, 1] == pii)

    for ji in initial_j:
        ji_masks.append(BXL[:, 0] == ji)

    n_B_skips = 0
    for pii in pii_masks:
        for idxi in idxi_masks:
            for ji in ji_masks:
                mask = np.logical_and(ji, np.logical_and(pii, idxi))
                tmp = BXL[mask][:, 9]   # 9 is B decay.
                n_B_skips += sum(tmp == 0)  # Count the number of zero values.
                tmp = tmp[tmp != 0] # Remove all zero values.
                if not tmp.size:
                    """
                    Some combinations of masks might not match any
                    levels.
                    """
                    continue
                
                BXL_tmp.extend(tmp/tmp.mean())

    BXL = np.asarray(BXL_tmp)
    BXL.sort()
    # BXL = BXL/np.mean(BXL)
    n_BXL_after = len(BXL)

    if (n_BXL_before - n_B_skips) != n_BXL_after:
        msg = "The number of BXL values has changed too much during the Porter-Thomas analysis!"
        msg += f" This should not happen! n_BXL_after should be: {n_BXL_before - n_B_skips}, got: {n_BXL_after}."
        msg += f"\n{n_BXL_before = }"
        msg += f"\n{n_BXL_after = }"
        msg += f"\n{n_B_skips = }"
        raise RuntimeError(msg)

    BXL_bins = np.arange(0, BXL[-1] + BXL_bin_width, BXL_bin_width)
    n_BXL_bins = len(BXL_bins)
    BXL_counts = np.zeros(n_BXL_bins)
    pt_prepare_data_time = time.perf_counter() - pt_prepare_data_time
    
    pt_count_time = time.perf_counter()
    for i in range(n_BXL_bins - 1):
        """
        Calculate the number of transitions with BXL values between
        BXL_bins[i] and BXL_bins[i + 1].
        """
        BXL_counts[i] = np.sum(BXL_bins[i] <= BXL[BXL < BXL_bins[i + 1]])
    
    pt_count_time = time.perf_counter() - pt_count_time

    pt_post_process_time = time.perf_counter()
    rv = chi2(1)
    BXL_counts = BXL_counts[1:] # Exclude the first data point because chi2(1) goes to infinity and is hard to work with there.
    BXL_bins = BXL_bins[1:]
    n_BXL_bins -= 1
    # BXL_counts_normalised = BXL_counts/np.trapz(BXL_counts)  # Normalize counts.
    # popt, _ = curve_fit(
    #     f = lambda x, scale: scale*rv.pdf(x),
    #     xdata = BXL_bins,
    #     ydata = BXL_counts,
    #     p0 = [rv.pdf(BXL_bins)[1]/BXL_counts[1]],
    #     method = "lm",
    # )
    # BXL_counts *= popt[0]   # Scale counts to match chi2.
    # BXL_counts_normalised *= np.mean(rv.pdf(BXL_bins)[1:20]/BXL_counts_normalised[1:20])
    """
    Normalise BXL_counts to the chi2(1) distribution, ie. find a
    coefficient which makes BXL_counts become chi2(1) by
    BXL_counts*chi2(1)/BXL_counts = chi2(1). Since any single point in
    the BXL distribution might over or undershoot chi2(1), I have chosen
    to use the mean of 19 ([1:20] slice, pretty arbitrary chosen) of
    these values to make a more stable normalisation coefficient.
    """
    BXL_counts_normalised = BXL_counts*np.mean(rv.pdf(BXL_bins)[1:20]/BXL_counts[1:20])
    pt_post_process_time = time.perf_counter() - pt_post_process_time

    if flags["debug"]:
        print("--------------------------------")
        print(f"Porter-Thomas: Prepare data time: {pt_prepare_data_time:.3f} s")
        print(f"Porter-Thomas: Count time: {pt_count_time:.3f} s")
        print(f"Porter-Thomas: Post process time: {pt_post_process_time:.3f} s")
        print(f"{sum(BXL_counts) = }")
        print(f"{sum(BXL_counts_normalised) = }")
        print(f"{n_B_skips = }")
        print(f"{Ei = }")
        print(f"{Ei_bin_width = }")
        print("--------------------------------")

    if return_chi2:
        # return BXL_bins, BXL_counts, rv.pdf(BXL_bins)
        return BXL_bins, BXL_counts_normalised, rv.pdf(BXL_bins)
    else:
        # return BXL_bins, BXL_counts
        return BXL_bins, BXL_counts_normalised

def nuclear_shell_model():
    """
    Generate a diagram of the nuclear shell model shell structure.
    """
    plt.rcParams.update({
        "backend": "pgf",
        "text.usetex": True,
        "font.family": "serif",
        "font.serif": ["roman"],
        "legend.fontsize": 14,
        "xtick.labelsize": 15,
        "ytick.labelsize": 15,
        "axes.labelsize": 14,
        "axes.titlesize": 15,
    })
    fig, ax = plt.subplots(figsize=(6.4, 8))
    ax.axis(False)
    fontsize = 15
    x_offset = 0.6
    x_text_offset = x_offset - 0.5

    first_layer_labels = [
        r"$1s$", r"$1p$", r"$1d$", r"$2s$", r"$1f$", r"$2p$", r"$1g$",
        r"$2d$", r"$3s$"
    ]
    first_layer_y = [1, 2.4, 4.2, 4.45, 6.3, 6.8, 9, 10.0, 10.5]
    second_layer_labels = [
        r"$1s_{1/2}$", r"$1p_{3/2}$", r"$1p_{1/2}$", r"$1d_{5/2}$",
        r"$2s_{1/2}$", r"$1d_{3/2}$", r"$1f_{7/2}$", r"$2p_{3/2}$",
        r"$1f_{5/2}$", r"$2p_{1/2}$", r"$1g_{9/2}$", r"$2d_{5/2}$",
        r"$1g_{7/2}$", r"$3s_{1/2}$", r"$2d_{3/2}$"
    ]
    second_layer_y = [
        first_layer_y[0], first_layer_y[1] - 0.15, first_layer_y[1] + 0.15,
        first_layer_y[2] - 0.3, first_layer_y[3], first_layer_y[2] + 0.51,
        first_layer_y[4] - 0.6, first_layer_y[5] - 0.10, first_layer_y[4] + 0.7,
        first_layer_y[5] + 0.5, first_layer_y[6] - 1.0, first_layer_y[7] - 0.4,
        first_layer_y[6] + 0.9, first_layer_y[7] + 0.8, first_layer_y[8]
    ]
    dash_layer = [
        [2 + x_offset, first_layer_y[0], 2.5 + x_offset, second_layer_y[0]],
        [2 + x_offset, first_layer_y[1], 2.5 + x_offset, second_layer_y[1]],
        [2 + x_offset, first_layer_y[1], 2.5 + x_offset, second_layer_y[2]],
        [2 + x_offset, first_layer_y[2], 2.5 + x_offset, second_layer_y[3]],
        [2 + x_offset, first_layer_y[2], 2.5 + x_offset, second_layer_y[5]],
        [2 + x_offset, first_layer_y[3], 2.5 + x_offset, second_layer_y[4]],
        [2 + x_offset, first_layer_y[4], 2.5 + x_offset, second_layer_y[6]],
        [2 + x_offset, first_layer_y[4], 2.5 + x_offset, second_layer_y[8]],
        [2 + x_offset, first_layer_y[5], 2.5 + x_offset, second_layer_y[7]],
        [2 + x_offset, first_layer_y[5], 2.5 + x_offset, second_layer_y[9]],
        [2 + x_offset, first_layer_y[6], 2.5 + x_offset, second_layer_y[10]],
        [2 + x_offset, first_layer_y[7], 2.5 + x_offset, second_layer_y[11]],
        [2 + x_offset, first_layer_y[6], 2.5 + x_offset, second_layer_y[12]],
        [2 + x_offset, first_layer_y[7], 2.5 + x_offset, second_layer_y[13]],
        [2 + x_offset, first_layer_y[8], 2.5 + x_offset, second_layer_y[14]],
    ]
    core_layer_labels = [
        r"$^{16}$O", r"$^{40}$Ca", r"$^{56}$Ni"
    ]
    core_layer_y = [
        second_layer_y[2] + 0.5, second_layer_y[5] + 0.5, second_layer_y[6] + 0.5
    ]

    occupations = [
        2, 4, 2, 6, 2, 4, 8, 4, 6, 2, 10, 6, 8, 2, 4
    ]
    occupations_y = second_layer_y
    cum_occupations = [
        2, 8, 20, 28, 50
    ]
    cum_occupations_y = [
        second_layer_y[0], second_layer_y[2], second_layer_y[5],
        second_layer_y[6], second_layer_y[10]
    ]
    ax.hlines(  # To force the width of the figure.
        y = 1,
        xmin = 3.5 + x_offset,
        xmax = 4.5 + x_offset,
        color = "white"
    )
    ax.hlines(  # To force the width of the figure.
        y = 1,
        xmin = 1,
        xmax = 2,
        color = "white"
    )
    for y, label in zip(first_layer_y, first_layer_labels):
        ax.hlines(
            y = y,
            xmin = 1 + x_offset,
            xmax = 2 + x_offset,
            color = "black",
        )
        fig.text(
            x = 0.12 + x_text_offset,
            y = y/13.95 + 0.067,
            s = label,
            fontsize = fontsize
        )

    for y, label in zip(second_layer_y, second_layer_labels):
        ax.hlines(
            y = y,
            xmin = 2.5 + x_offset,
            xmax = 3.5 + x_offset,
            color = "black",
        )
        fig.text(
            x = 0.6 + x_text_offset,
            y = y/14.2 + 0.067,
            s = label,
            fontsize = fontsize
        )

    for x1, y1, x2, y2 in dash_layer:
        ax.plot([x1, x2], [y1, y2], linestyle="dashed", color="black")

    for occupation, y in zip(occupations, occupations_y):
        fig.text(
            x = 0.69 + x_text_offset,
            y = y/14.2 + 0.067,
            s = occupation,
            fontsize = fontsize - 1
        )

    for occupation, y in zip(cum_occupations, cum_occupations_y):
        fig.text(
            x = 0.73 + x_text_offset,
            y = y/14.2 + 0.067,
            s = occupation,
            fontsize = fontsize - 1
        )

    for y, label in zip(core_layer_y, core_layer_labels):
        fig.text(
            x = 0.77 + x_text_offset,
            y = y/14 + 0.067,
            s = label,
            fontsize = fontsize - 1
        )
        fig.text(
            x = 0.73 + x_text_offset,
            y = y/14 + 0.064,
            s = "---------",
            fontsize = fontsize - 1
        )

    # USD
    x1 = 1.35
    x2 = 1.25
    y1 = 4.9
    y2 = 3.83
    ax.vlines(
        x = x2,
        ymin = y2,
        ymax = y1,
        color = "darkorange",
    )
    fig.text(
        x = 0.15,
        y = 0.37,
        s = "USD",
        fontsize = 12,
        rotation = "vertical",
        color = "darkorange"
    )
    # GXPF
    y3 = 7.5
    y4 = 5.6
    ax.vlines(
        x = x2,
        ymin = y4,
        ymax = y3,
        color = "firebrick",
    )
    fig.text(
        x = 0.15,
        y = 0.52,
        s = "GXPF",
        fontsize = 12,
        rotation = "vertical",
        color = "firebrick"
    )
    # SDPF-MU
    x4 = x2 - 0.04
    ax.vlines(
        x = x4,
        ymin = y2,
        ymax = y3,
        color = "royalblue",
    )
    fig.text(
        x = 0.14,
        y = 0.42,
        s = "SDPF-MU",
        fontsize = 12,
        rotation = "vertical",
        color = "royalblue"
    )
    #JUN45
    y7 = 6.5
    y8 = 8.2
    x6 = x4 - 0.04
    ax.vlines(
        x = x6,
        ymin = y8,
        ymax = y7,
        color = "green",
    )
    fig.text(
        x = 0.14,
        y = 0.59,
        s = "JUN45",
        fontsize = 12,
        rotation = "vertical",
        color = "green"
    )
    # SDPF-SDG
    ax.vlines(
        x = x6 - 0.04,
        ymin = y2,
        ymax = 11,
        color = "mediumorchid",
    )
    fig.text(
        x = 0.15,
        y = 0.66,
        s = "SDPF-SDG",
        fontsize = 12,
        rotation = "vertical",
        color = "mediumorchid"
    )
    # Spectroscopic notation
    fig.text(
        x = 0.45,
        y = 0.93,
        s = r"$s \;\; p \;\; d \;\; f \;\; g \;\; h$",
        fontsize = fontsize - 1,
        color = "black",
    )
    fig.text(
        x = 0.45,
        y = 0.92,
        s = "------------------",
        fontsize = fontsize - 1,
        color = "black",
    )
    fig.text(
        x = 0.415,
        y = 0.90,
        s = r"$l: 0 \;\; 1 \;\; 2 \;\; 3 \;\; 4 \;\; 5$",
        fontsize = fontsize - 1,
        color = "black",
    )
    fig.text(
        x = 0.405,
        y = 0.88,
        s = r"$\pi:+  -  +  - \, + \, -$",
        fontsize = fontsize - 1,
        color = "black",
    )

    fig.savefig(fname="nuclear_shell_model.png", dpi=500)#, format="eps")
    plt.show()
<file_sep>__author__ = "<NAME>, <NAME>"
__version__ = "1.5.1.0"
__credits__ = "<NAME>, <NAME>"

"""
Version legend:
a.b.c.d

a: major release
b: new functionality added
c: new feature to existing functionality
d: bug fixes
"""

from .kshell_utilities import *
from .general_utilities import *
from .kshell_exceptions import *
from .count_dim import *
from .parameters import *
from .compare import Compare
from .collect_logs import collect_logs, check_multipolarities
from .script_editing import edit_and_queue_executables
from .partition_editor import partition_editor, test_partition_editor, test_partition_editor_2
from . import loaders
from . import data_structures
<file_sep>from __future__ import annotations
from typing import Union, Tuple
from itertools import cycle, islice
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from .kshell_utilities import ReadKshellOutput
from .general_utilities import (
    level_plot, level_density, gamma_strength_function_average
)

class Compare:
    """
    Plot levels, level density and gamma strength function for easy
    comparison between multiple kshell outputs.
    """
    def __init__(self,
            kshell_outputs: list[ReadKshellOutput],
            legend_labels: Union[None, list[str]] = None,
        ):
        """
        Initialize instance with the given kshell outputs and a default
        color palette.

        Parameters
        ----------
        kshell_outputs : list[ReadKshellOutput]
            list of instances of the `ReadKshellOutput` class to be
            plotted together.

        legend_labels : Union[None, list[str]]
            A list of labels for the legends of the plots. The number of
            labels must equal the number of elements in
            `kshell_outputs`.
        """
        type_error_msg = 'kshell_outputs must be a list of ReadKshellOutput'
        type_error_msg += ' instances (the return value of ksutil.loadtxt).'
        type_error_msg += ' Eg: ksutil.Compare(kshell_outputs=[V50, V51]).'
        if not isinstance(kshell_outputs, (list, tuple)):
            raise TypeError(type_error_msg)
        
        else:
            if not all(isinstance(i, ReadKshellOutput) for i in kshell_outputs):
                raise TypeError(type_error_msg)
            
        if isinstance(legend_labels, list):
            if len(legend_labels) != len(kshell_outputs):
                msg = (
                    "The number of labels must equal the number of"
                    " ReadKshellOutput instances!"
                )
                raise RuntimeError(msg)

        self._legend_labels = legend_labels
        self._kshell_outputs = kshell_outputs
        self._color_palette = sns.color_palette(
            palette = "tab10",
            n_colors = len(self._kshell_outputs))

    def set_color_palette(
            self, 
            color_palette: Union[str, list[str], list[tuple]]):
        """
        Set the color palette to use for the plots.

        Parameters
        ----------
        color_palette : Union[str, list[str], list[tuple]]
            name of the color palette to use or a list of colors.
        """
        required_number_of_colors = len(self._kshell_outputs)
        if isinstance(color_palette, str):
            self._color_palette = sns.color_palette(
                palette = color_palette,
                n_colors = required_number_of_colors)
        elif len(color_palette) >= required_number_of_colors:
            self._color_palette = color_palette
        else:
            warning_message = (
                "The supplied color palette is too short. "
                "Some colors will be repeated.")
            print(warning_message)
            self._color_palette = list(islice(
                cycle(color_palette), 
                required_number_of_colors))

    def plot_levels(
        self,
        ax: Union[None, plt.Axes] = None,
        include_n_levels: int = 1_000,
        filter_spins: Union[None, list] = None,
        filter_parity: Union[None, str] = None,
        ):
        """
        Draw level plots for all kshell outputs. 

        Parameters
        ----------
        ax : Union[None, plt.Axes]
            matplotlib Axes on which to plot. If None, plt.Figure and
            plt.Axes is generated in this function.

        See level_plot in general_utilities.py for details on the other
        parameters.
        """
        ax_input, fig, ax = self._get_fig_and_ax(ax)
        xticks = {}

        is_half_integer: bool = any(self._kshell_outputs[0].levels[:, 1]%2 == 1)
        for kshell_output in self._kshell_outputs:
            if is_half_integer != any(kshell_output.levels[:, 1]%2 == 1):
                """
                If both integer and half integer angular momentum
                exsists in two or more kshell data sets, then the line
                width must be halved to make room for both.
                """
                line_width = 0.2
                x_offset_scale = 0.5
                break

        else:
            """
            In this case, the kshell outputs contain either integer or
            half integer angular momenta, not both.
            """
            line_width = 0.4
            x_offset_scale = 1

        if self._legend_labels is None:
            labels: list = [i.nucleus for i in self._kshell_outputs]
        else:
            labels: list = self._legend_labels

        for color, kshell_output, label in zip(
                self._color_palette, self._kshell_outputs, labels
            ):

            level_plot(
                levels = kshell_output.levels,
                include_n_levels = include_n_levels,
                filter_spins = filter_spins,
                filter_parity = filter_parity,
                ax = ax,
                color = color,
                line_width = line_width,
                x_offset_scale = x_offset_scale,
            )

            for tick_position, tick_label in zip(
                    ax.get_xticks(),
                    [l.get_text() for l in ax.get_xticklabels()]
                    ):

                xticks[tick_position] = tick_label

            ax.plot([], [], label=label, color=color)

        ax.set_xticks(ticks=list(xticks.keys()), labels=list(xticks.values()))

        ax.legend(loc="lower right")

        if not ax_input:
            plt.show()

    def plot_level_densities(
        self,
        ax: Union[None, plt.Axes] = None,
        bin_width: Union[int, float] = 0.2,
        include_n_levels: Union[None, int] = None,
        filter_spins: Union[None, int, list] = None,
        filter_parity: Union[None, str, int] = None,
        E_min: Union[float, int] = 0,
        E_max: Union[float, int] = np.inf
        ):
        """
        Draw level density plots for all kshell outputs.

        Parameters
        ----------
        ax : Union[None, plt.Axes]
            matplotlib Axes on which to plot. If None, plt.Figure and plt.Axes 
            is generated in this function.

        See level_density in general_utilities.py for details on the other
        parameters.
        """
        ax_input, fig, ax = self._get_fig_and_ax(ax)

        ax.set_ylabel(r"NLD [MeV$^{-1}$]")
        ax.set_xlabel("E [MeV]")
        ax.legend([f"{bin_width=} MeV"])

        for color, kshell_output in zip(self._color_palette,
                                        self._kshell_outputs):
            bins, density = level_density(
                levels = kshell_output.levels,
                bin_width = bin_width,
                include_n_levels = include_n_levels,
                filter_spins = filter_spins,
                filter_parity = filter_parity,
                E_min = E_min,
                E_max = E_max,
                return_counts = False,
                plot = False,
                save_plot = False
                )

            ax.step(bins, density, color=color)

        if not ax_input:
            plt.show()

    def plot_gamma_strength_functions(
        self,
        ax: Union[None, plt.Axes] = None,
        bin_width: Union[float, int] = 0.2,
        Ex_min: Union[float, int] = 5,
        Ex_max: Union[float, int] = 50,
        multipole_type: str = "M1",
        prefactor_E1: Union[None, float] = None,
        prefactor_M1: Union[None, float] = None,
        prefactor_E2: Union[None, float] = None,
        initial_or_final: str = "initial",
        partial_or_total: str = "partial",
        include_only_nonzero_in_average: bool = True,
        include_n_levels: Union[None, int] = None,
        filter_spins: Union[None, list] = None,
        filter_parities: str = "both",
        return_n_transitions: bool = False,
        ):
        """
        Draw gamma strength function plots for all kshell outputs.

        Parameters
        ----------
        ax : Union[None, plt.Axes]
            matplotlib Axes on which to plot. If None, plt.Figure and plt.Axes 
            is generated in this function.

        See level_density in general_utilities.py for details on the other
        parameters.
         """
        ax_input, fig, ax = self._get_fig_and_ax(ax)

        for color, kshell_output in zip(self._color_palette,
                                        self._kshell_outputs):
            transitions_dict = {
                "M1": kshell_output.transitions_BM1,
                "E2": kshell_output.transitions_BE2,
                "E1": kshell_output.transitions_BE1
            }

            bins, gsf = gamma_strength_function_average(
                levels = kshell_output.levels,
                transitions = transitions_dict[multipole_type],
                bin_width = bin_width,
                Ex_min = Ex_min,
                Ex_max = Ex_max,
                multipole_type = multipole_type,
                prefactor_E1 = prefactor_E1,
                prefactor_M1 = prefactor_M1,
                prefactor_E2 = prefactor_E2,
                initial_or_final = initial_or_final,
                partial_or_total = partial_or_total,
                include_only_nonzero_in_average = include_only_nonzero_in_average,
                include_n_levels = include_n_levels,
                filter_spins = filter_spins,
                filter_parities = filter_parities,
                return_n_transitions = return_n_transitions
                )
            ax.plot(bins, gsf, color=color)

        if not ax_input:
            plt.show()

    @staticmethod
    def _get_fig_and_ax(
        ax: Union[None, plt.Axes]
        ) -> Tuple[bool, plt.Figure, plt.Axes]:
        """
        Return a matplotlib Figure and Axes on which to plot, and whether these
        were generated in this function or existed previously.

        Parameter:
        ----------
        ax : Union[None, plt.Axes]
            matplotlib Axes on which to plot. If None, plt.Figure and plt.Axes
            is generated by this function.

        Returns:
        --------
        ax_input : bool
            whether a matplotlib Axes was passed.
        fig : plt.Figure
            matplotlib Figure on which to plot.
        ax : plt.Axes
            matplotlib Axes on which to plot. If a matplotlib Axes object was
            passed, it is returned unchanged.
        """
        ax_input = False if (ax is None) else True
        if not ax_input:
            fig, ax = plt.subplots()
        else:
            fig = ax.get_figure()
        return (ax_input, fig, ax)
<file_sep>from typing import Union
import numpy as np

def generate_states(
    start: int = 0,
    stop: int = 14,
    n_states: int = 100,
    parity: Union[str, int] = "both"
    ):
    """
    Generate correct string for input to `kshell_ui.py` when asked for
    which states to calculate. Copy the string generated by this
    function and paste it into `kshell_ui.py` when it prompts for
    states.

    DEPRECATED: RANGE FUNCTIONALITY WAS ADDED IN kshell_ui.py MAKING
    THIS FUNCTION OBSOLETE. WILL BE REMOVED.

    Parameters
    ----------
    start : int
        The lowest spin value.

    stop : int
        The largest spin value.

    n_states : int
        The number of states per spin value.

    parity : Union[str, int]
        The parity of the states. Allowed values are: 1, -1, 'both',
        'positive', 'negative', 'pos', 'neg', '+', '-'.

    Examples
    --------
    ``` python
    >>> import kshell_utilities as ksutil
    >>> ksutil.generate_states(start=0, stop=3, n_states=100, parity="both")
    0+100, 0.5+100, 1+100, 1.5+100, 2+100, 2.5+100, 3+100, 0-100, 0.5-100, 1-100, 1.5-100, 2-100, 2.5-100, 3-100,
    ```
    """
    allowed_positive_parity_inputs = ["positive", "pos", "+", "1", "+1", 1, "both"]
    allowed_negative_parity_inputs = ["negative", "neg", "-", "-1", -1, "both"]
    
    def correct_syntax(lst):
        for elem in lst:
            print(elem, end=", ")
    
    if parity in allowed_positive_parity_inputs:
        positive = [f"{i:g}{'+'}{n_states}" for i in np.arange(start, stop+0.5, 0.5)]
        correct_syntax(positive)
    
    if parity in allowed_negative_parity_inputs:
        negative = [f"{i:g}{'-'}{n_states}" for i in np.arange(start, stop+0.5, 0.5)]
        correct_syntax(negative)<file_sep>from __future__ import annotations
import os, warnings
from .parameters import GS_FREE_PROTON, GS_FREE_NEUTRON

def _ask_for_time_input(unit: str, lower_lim: int, upper_lim: int) -> int:
    """
    Input prompting generalized for seconds, minutes, hours, and days.

    Parameters
    ----------
    unit : str
        'seconds', 'minutes', 'hours', or 'days'.
    
    lower_lim : int
        Lower allowed limit for the unit input.

    upper_lim : int
        Upper allowed limit for the unit input.

    Returns
    -------
    ans : int
        User input for the unit within the limits.
    """
    while True:
        try:
            ans = int(input(f"Number of {unit}: "))
        except ValueError:
            print("Only integers are allowed!")
            continue
        
        if lower_lim <= ans <= upper_lim:
            return ans
        else:
            print(f"The number of {unit} must be in the interval [{lower_lim}, {upper_lim}]")
            continue

def edit_and_queue_executables():
    """
    Simple wrapper for quietly quitting on KeyboardInterrupt without
    having to indent the entire _edit_and_queue_executables function.
    """
    try:
        _edit_and_queue_executables()
    except KeyboardInterrupt:
        print("\nExiting...")
        return
    
def _prompt_user_for_values(
    available_parameters: dict[str, bool],
    available_parameters_values: dict[str, str | int | float]
):
    """
    Ask the user for the actual values for each of the parameters which
    are to be changed.
    """
    if available_parameters["job_name"]:
        available_parameters_values["job_name"] = input("Job name: ")

    if available_parameters["nodes"]:
        while True:
            """
            Ask for new number of nodes.
            """
            try:
                available_parameters_values["nodes"] = int(input("Number of nodes: "))
                break
            except ValueError:
                print("Only integers are allowed!")
                continue

    if available_parameters["seconds"]:
        available_parameters_values["seconds"] = _ask_for_time_input(
            unit = "seconds",
            lower_lim = 0,
            upper_lim = 59
        )
        tmp = available_parameters_values["time_parameter_new"].split(":")
        available_parameters_values["time_parameter_new"] = f"{tmp[0]}:{tmp[1]}:{available_parameters_values['seconds']:02d}"

    if available_parameters["minutes"]:
        available_parameters_values["minutes"] = _ask_for_time_input(
            unit = "minutes",
            lower_lim = 0,
            upper_lim = 59
        )
        tmp = available_parameters_values["time_parameter_new"].split(":")
        available_parameters_values["time_parameter_new"] = f"{tmp[0]}:{available_parameters_values['minutes']:02d}:{tmp[2]}"

    if available_parameters["hours"]:
        available_parameters_values["hours"] = _ask_for_time_input(
            unit = "hours",
            lower_lim = 0,
            upper_lim = 23
        )
        tmp = available_parameters_values["time_parameter_new"].split(":")
        tmp_days, _ = tmp[0].split("-")
        available_parameters_values["time_parameter_new"] = f"{tmp_days}-{available_parameters_values['hours']:02d}:{tmp[1]}:{tmp[2]}"

    if available_parameters["days"]:
        available_parameters_values["days"] = _ask_for_time_input(
            unit = "days",
            lower_lim = 0,
            upper_lim = 4
        )
        tmp = available_parameters_values["time_parameter_new"].split(":")
        _, tmp_hours = tmp[0].split("-")
        available_parameters_values["time_parameter_new"] = f"{available_parameters_values['days']}-{tmp_hours}:{tmp[1]}:{tmp[2]}"

    if available_parameters["quench"]:
        while True:
            try:
                available_parameters_values["quench"] = float(input("quench: "))
                break
            except ValueError:
                print("Only integers and floats are allowed!")
                continue

def _edit_and_queue_executables():
    """
    Loop over .sh files generated by 'kshell_ui' and adjust queue system
    parameters and if possible, queue the .sh file.

    Currently hard-coded to the slurm system.
    """
    available_parameters: dict[str, bool] = {
        "nodes": False,
        "seconds": False,
        "minutes": False,
        "hours": False,
        "days": False,
        "job_name": False,
        "quench": False,
        "queue": False,
        "apply same values to several or all .sh files": False,
        # "ask for confirmation per file": False,
    }
    available_parameters_values: dict[str, str | int | float] = {key: "" for key in available_parameters.items()}
    available_parameters_values["time_parameter_new"] = ""  # For storing the entire new formatted time parameter string.
    prompt_user_input: bool = True

    print("Please choose what parameters you want to alter (y/n) (default n):")
    for parameter in available_parameters:
        """
        Decide what to do with all parameters.
        """
        while True:
            ans = input(f"{parameter}: ").lower()
            if ans == "y":
                available_parameters[parameter] = True
                break
            elif ans == "n":
                break
            elif ans == "":
                break
            else:
                continue

    content = [i for i in os.listdir() if i.endswith(".sh")]
    content.sort()

    for elem in content:
        """
        Loop over all .sh files in directory.
        """
        edit_ans = input(f"\nEdit {elem}? (y/n) (default y): ").lower()
        if edit_ans == "n":
            print(f"Skipping {elem}")
            continue

        print(f"Loading {elem}")
        time_parameter = ""
        nodes_parameter = ""
        job_name_parameter = ""
        gs_parameter = ""
        content = ""

        with open(elem, "r") as infile:
            """
            Extract all file content and other specific lines.
            """
            for line in infile:
                content += line
                if "#SBATCH --time=" in line:
                    time_parameter += line
                elif "#SBATCH --nodes=" in line:
                    nodes_parameter += line
                elif "#SBATCH --job-name=" in line:
                    job_name_parameter += line
                elif "gs =" in line:
                    gs_parameter += line

        if (not content) or ((not time_parameter) and (not nodes_parameter) and (not job_name_parameter)):
            print(f"Could not extract info from {elem}. Skipping ...")
            continue
        
        # time_parameter_new = time_parameter.split("=")[-1].strip()    # '0-00:00:00' d-hh:mm:ss

        if prompt_user_input:
            if available_parameters["apply same values to several or all .sh files"]:
                """
                Ask the user only once and start with the same time
                template for all .sh files.
                """
                prompt_user_input = False
                available_parameters_values["time_parameter_new"] = '#SBATCH --time=0-00:00:00'.split("=")[-1].strip()    # '0-00:00:00' d-hh:mm:ss
            else:
                available_parameters_values["time_parameter_new"] = time_parameter.split("=")[-1].strip()    # '0-00:00:00' d-hh:mm:ss

            _prompt_user_for_values(
                available_parameters = available_parameters,
                available_parameters_values = available_parameters_values
            )
  
        if time_parameter:
            """
            Insert new time parameter.
            """
            content_tmp = content.replace(time_parameter, f"#SBATCH --time={available_parameters_values['time_parameter_new']}\n")

            if (content_tmp == content) and (available_parameters['seconds'] or available_parameters['minutes'] or available_parameters['hours'] or available_parameters['days']):
                msg = "Time parameter is unchanged. Either str.replace could"
                msg += " not find a match or new time parameter is identical"
                msg += " to old time parameter."
                warnings.warn(msg, RuntimeWarning)
            else:
                content = content_tmp

        if job_name_parameter and available_parameters['job_name']:
            """
            Insert new job name parameter.
            """
            content_tmp = content.replace(job_name_parameter, f"#SBATCH --job-name={available_parameters_values['job_name']}\n")

            if content_tmp == content:
                msg = "Job name parameter is unchanged. Either str.replace"
                msg += " could not find a match or new job name parameter is"
                msg += " identical to old job name parameter."
                warnings.warn(msg, RuntimeWarning)
            else:
                content = content_tmp

        if nodes_parameter and available_parameters["nodes"]:
            """
            Insert new nodes parameter.
            """
            nodes_parameter_new = nodes_parameter.split("=")[0]
            nodes_parameter_new += f"={available_parameters_values['nodes']}\n"
            content_tmp = content.replace(nodes_parameter, nodes_parameter_new)
            
            if content_tmp == content:
                msg = "Nodes parameter is unchanged. Either str.replace could"
                msg += " not find a match or new nodes parameter is identical"
                msg += " to old nodes parameter."
                warnings.warn(msg, RuntimeWarning)
            else:
                content = content_tmp

        if available_parameters["quench"]:
            quenching_factor = available_parameters_values["quench"]
            gs_parameter_new = gs_parameter.split("=")[0]
            gsp = round(quenching_factor*GS_FREE_PROTON, 2)
            gsn = round(quenching_factor*GS_FREE_NEUTRON, 2)
            gs_parameter_new += f"= {gsp}, {gsn}\n"
            content_tmp = content.replace(gs_parameter, gs_parameter_new)

            if content_tmp == content:
                msg = "gs parameter is unchanged. Either str.replace could"
                msg += " not find a match or new nodes parameter is identical"
                msg += " to old nodes parameter."
                warnings.warn(msg, RuntimeWarning)
            else:
                content = content_tmp

        with open(elem, "w") as outfile:
            outfile.write(content)

        if available_parameters["queue"]:
            os.system(f"sbatch {elem}")
<file_sep>from __future__ import annotations
import time, os, curses, warnings
from math import inf
from typing import Callable
from vum import Vum
from .count_dim import count_dim
from .kshell_exceptions import KshellDataStructureError
from .parameters import (
    spectroscopic_conversion, shell_model_order, major_shell_order
)
from .data_structures import (
    OrbitalParameters, Configuration, ModelSpace, Interaction, Partition, Skips
)
from .partition_tools import (
    _prompt_user_for_interaction_and_partition, _calculate_configuration_parity,
    _sanity_checks, configuration_energy,
)
from .loaders import load_partition, load_interaction
from .partition_compare import _partition_compare

warnings.filterwarnings("ignore", category=UserWarning)
DELAY: int = 2  # Delay time for time.sleep(DELAY) in seconds
PARITY_CURRENT_Y_COORD = 5
# is_duplicate_warning = True
y_offset: int = 0

class ScreenDummy:
    def clear(self):
        return

class VumDummy:
    def __init__(self) -> None:
        self.screen = ScreenDummy()
        self.command_log_length = 0
        self.command_log_length = 0
        self.n_rows = 0
        self.n_cols = 0
        self.blank_line = ""

    def addstr(self, y, x, string, is_blank_line=None):
        return

    def input(self, _) -> str:
        return "n"
    
class VumDummy2(VumDummy):
    def __init__(self) -> None:
        super().__init__()
        # self.answers: list[str] = [
        #     "n",    # Analyse?
        #     "y",    # Add config?
        #     "p",    # Proton or neutron?
        #     "r",    # Single or range?
        #     "2",    # N particle N hole?
        #     "2",    # How many excitations to add? NOTE: Currently not in use.
        #     "pf",   # Initial major shell?
        #     "sdg",  # Final major shell?
        #     "n",    # Add config?
        # ]
        self.answers: list[str] = [
            "n",
            "y",
            "p",
            "s",
            "f",
            "f",
            "2",
            "2",
            "2",
            "2",
            "2",
            "2",
        ]
    
    def input(self, _) -> str:
        ans = self.answers.pop(0)
        print(ans)
        return ans

def _analyse_existing_configuration(
    vum: Vum,
    proton_configurations: list[Configuration],
    neutron_configurations: list[Configuration],
    input_wrapper: Callable,
    model_space: list[OrbitalParameters],
    y_offset: int,
    n_proton_orbitals: int,
) -> None:
    """
    Prompt the user for an index of an existing configuration and show
    an analysis of the configuration in question. This function is only
    used once and exists simply to make the amount of code lines in
    `_partition_editor` smaller.

    Parameters
    ----------
    vum : Vum
        The Vum object to draw the map with.
    
    proton_configurations : list[Configuration]
        The proton configurations to choose from.

    neutron_configurations : list[Configuration]
        The neutron configurations to choose from.

    input_wrapper : Callable
        The input wrapper to use for user input.

    model_space : list[OrbitalParameters]
        The model space orbitals to use for the calculation.

    y_offset : int
        The y offset to use for drawing the map.
    """
    pn_configuration_dict: dict[str, list[Configuration]] = {
        "p": proton_configurations,
        "n": neutron_configurations,
    }
    while True:
        if input_wrapper("Analyse existing configuration? (y/n)") == "y":
            while True:
                p_or_n = input_wrapper("Proton or neutron configuration? (p/n)")
                if (p_or_n == "p"):
                    is_proton = True
                    is_neutron = False
                    orbital_idx_offset: int = 0
                    break
                if (p_or_n == "n"):
                    is_proton = False
                    is_neutron = True
                    orbital_idx_offset: int = n_proton_orbitals # This is needed to index the neutron orbitals correctly when drawing the map.
                    break
            
            draw_shell_map(vum=vum, model_space=model_space, is_proton=is_proton, is_neutron=is_neutron)
            
            while True:
                configuration_idx = input_wrapper(f"Choose a {p_or_n} configuration index in 1, 2, ..., {len(pn_configuration_dict[p_or_n])} (q to quit)")
                if configuration_idx == "q": break
                
                try:
                    configuration_idx = int(configuration_idx)
                except ValueError:
                    draw_shell_map(vum=vum, model_space=model_space, is_proton=is_proton, is_neutron=is_neutron)
                    continue
                
                configuration_idx -= 1    # List indices are from 0 while indices in .ptn are from 1.
                if configuration_idx in range(len(pn_configuration_dict[p_or_n])):

                    current_configuration = pn_configuration_dict[p_or_n][configuration_idx].configuration
                    configuration_parity = 1
                    for i in range(len(current_configuration)):
                        configuration_parity *= model_space[i].parity**current_configuration[i] if current_configuration[i] else 1  # Avoid multiplication by 0.
                        draw_shell_map(
                            vum = vum,
                            model_space = model_space,
                            is_proton = is_proton,
                            is_neutron = is_neutron,
                            occupation = (orbital_idx_offset + i, current_configuration[i]),
                        )

                    vum.addstr(y_offset + PARITY_CURRENT_Y_COORD, 0, f"parity current configuration = {configuration_parity}")
                
                else:
                    draw_shell_map(vum=vum, model_space=model_space, is_proton=is_proton, is_neutron=is_neutron)

        else: break # If answer from user is not 'y'.

    vum.addstr(y_offset + PARITY_CURRENT_Y_COORD, 0, "parity current configuration = None")
    draw_shell_map(vum=vum, model_space=model_space, is_proton=True, is_neutron=True)

def draw_shell_map(
    vum: Vum,
    model_space: list[OrbitalParameters],
    is_proton: bool,
    is_neutron: bool,
    occupation: tuple[int, int] | None = None,
    n_holes: int = 0,
) -> None:
    """
    Draw a simple map of the model space orbitals of the current
    interaction file. Sort the orbitals based on the shell_model_order
    dict.

    Parameters
    ----------
    vum : Vum
        The Vum object to draw the map with.

    model_space : list[OrbitalParameters]
        The model space orbitals to draw.

    is_proton : bool
        Whether to draw the proton map.

    is_neutron : bool
        Whether to draw the neutron map.

    occupation : tuple[int, int] | None, optional
        The occupation of the model space orbitals. The tuple contains
        (orbital index, occupation). If None, draw the entire map with
        no occupation, by default None.

    n_holes : int
        For indicating the original position of excited particles.

    Returns
    -------
    None
    """
    y_offset: int = 14
    model_space_copy = sorted(  # Sort the orbitals based on the shell_model_order dict.
        model_space,
        key = lambda orbital: shell_model_order[f"{orbital.n}{spectroscopic_conversion[orbital.l]}{orbital.j}"].idx,
        reverse = True
    )
    model_space_proton = [orbital for orbital in model_space_copy if orbital.tz == -1]
    model_space_neutron = [orbital for orbital in model_space_copy if orbital.tz == 1]

    max_neutron_j: int = max([orbital.j for orbital in model_space_neutron], default=0)
    max_proton_j: int = max([orbital.j for orbital in model_space_proton], default=0)  # Use the max j value to center the drawn orbitals.

    if is_proton:
        proton_offset: int = 17 + (max_proton_j + 1)*2  # The offset for the neutron map.
        is_blank_line: bool = False
        if occupation is None:
            """
            Draw the entire map with no occupation.
            """
            for i in range(len(model_space_proton)):
                string = (
                    f"{model_space_proton[i].idx + 1:2d}"
                    f" {model_space_proton[i].name}"
                    f" {model_space_proton[i].parity:2d} " +
                    " "*(max_proton_j - model_space_proton[i].j) + "-" + " -"*(model_space_proton[i].j + 1)
                )
                vum.addstr(i + y_offset, 0, string)
        else:
            """
            Re-draw a single orbital with the user inputted occupation.
            """
            location: int = 0
            for orbital in model_space_proton:
                if orbital.idx == occupation[0]: break
                location += 1
            else:
                msg = (
                    f"Orbital index {occupation[0]} not found in the"
                    " current model space!"
                )
                raise RuntimeError(msg)
            
            string = (
                f"{orbital.idx + 1:2d}"
                f" {orbital.name}"
                f" {orbital.parity:2d} " +
                " "*(max_proton_j - orbital.j) + "-"
            )
            if not n_holes:
                string += "o-"*occupation[1] + " -"*(orbital.j + 1 - occupation[1])
            else:
                string += "o-"*(occupation[1] - n_holes) + "x-"*n_holes + " -"*(orbital.j + 1 - occupation[1])
            vum.addstr(location + y_offset, 0, string)

    else:
        proton_offset: int = 0
        is_blank_line: bool = True

    if is_neutron:
        if occupation is None:
            for i in range(len(model_space_neutron)):        
                string = (
                    f"{model_space_neutron[i].idx + 1:2d}"
                    f" {model_space_neutron[i].name}"
                    f" {model_space_neutron[i].parity:2d} " +
                    " "*(max_neutron_j - model_space_neutron[i].j) + "-" + " -"*(model_space_neutron[i].j + 1)
                )
                vum.addstr(i + y_offset, proton_offset, string, is_blank_line=is_blank_line)
        
        else:
            """
            Re-draw a single orbital with the user inputted occupation.
            """
            location: int = 0
            for orbital in model_space_neutron:
                if orbital.idx == occupation[0]: break
                location += 1
            else:
                msg = (
                    f"Orbital index {occupation[0]} not found in the"
                    " current model space!"
                )
                raise RuntimeError(msg)
            
            string = (
                f"{orbital.idx + 1:2d}"
                f" {orbital.name}"
                f" {orbital.parity:2d} " +
                " "*(max_neutron_j - orbital.j) + "-"
            )
            string += "o-"*occupation[1] + " -"*(orbital.j + 1 - occupation[1])
            vum.addstr(location + y_offset, proton_offset, string)

def _summary_information(
    vum: Vum,
    skips: Skips,
    filename_interaction: str,
    filename_partition: str,
    partition_proton: Partition,
    partition_neutron: Partition,
    partition_combined: Partition,
    interaction: Interaction,
    M: list[int],
    mdim: list[int],
    # n_proton_skips: int,
    # n_neutron_skips: int,
    # n_parity_skips: int,
    # n_ho_skips: int,
    y_offset: int,
    mdim_original: int | None = None,
):
    vum.addstr(
        y = y_offset,
        x = 0,
        string = f"{filename_interaction} -- {filename_partition} -- parity: {partition_combined.parity} -- core pn: ({interaction.n_core_protons}, {interaction.n_core_neutrons}) -- valence pn: ({interaction.model_space_proton.n_valence_nucleons}, {interaction.model_space_neutron.n_valence_nucleons})"
    )
    
    if mdim_original is None:
        vum.addstr(
            y = y_offset + 1,
            x = 0,
            string = f"M-scheme dim (M={M[-1]}): {mdim[-1]:d} ({mdim[-1]:.2e})"
        )
    else:
        vum.addstr(
            y = y_offset + 1,
            x = 0,
            string = f"M-scheme dim (M={M[-1]}): {mdim[-1]:d} ({mdim[-1]:.2e}) (original {mdim_original:d} ({mdim_original:.2e}))"
        )
    # vum.addstr(
    #     y = y_offset + 2,
    #     x = 0,
    #     string = f"n valence protons, neutrons: {interaction.model_space_proton.n_valence_nucleons}, {interaction.model_space_neutron.n_valence_nucleons}"
    # )
    # vum.addstr(
    #     y = y_offset + 3,
    #     x = 0,
    #     string = f"n core protons, neutrons: {interaction.n_core_protons}, {interaction.n_core_neutrons}"
    # )
    vum.addstr(
        y = y_offset + 2,
        x = 0,
        string = f"n proton +, -, sum : ({partition_proton.n_existing_positive_configurations} + {partition_proton.n_new_positive_configurations}), ({partition_proton.n_existing_negative_configurations} + {partition_proton.n_new_negative_configurations}), ({partition_proton.n_existing_configurations} + {partition_proton.n_new_configurations})",
    )
    vum.addstr(
        y = y_offset + 3,
        x = 0,
        string = f"n neutron +, -, sum : ({partition_neutron.n_existing_positive_configurations} + {partition_neutron.n_new_positive_configurations}), ({partition_neutron.n_existing_negative_configurations} + {partition_neutron.n_new_negative_configurations}), ({partition_neutron.n_existing_configurations} + {partition_neutron.n_new_configurations})"
    )
    vum.addstr(
        y = y_offset + 4,
        x = 0,
        string = f"n combined: {partition_combined.n_configurations}"
    )
    vum.addstr(
        y = y_offset + PARITY_CURRENT_Y_COORD,
        x = 0,
        string = "parity current configuration = None"
    )
    vum.addstr(
        y = y_offset + 6,
        x = 0,
        string = f"n proton, neutron configurations will be deleted because of parity or H.O. mismatch: {skips.n_proton_skips, skips.n_neutron_skips}"
    )
    vum.addstr(
        y = y_offset + 7,
        x = 0,
        string = f"n parity, H.O. skips: {skips.n_parity_skips, skips.n_ho_skips}"
    )
    vum.addstr(
        y = y_offset + 8,
        x = 0,
        string = f"Monopole trunc skips: {skips.n_monopole_skips}"
    )
    if partition_combined.max_configuration_energy_original == partition_combined.max_configuration_energy:
        vum.addstr(
            y = y_offset + 9,
            x = 0,
            string = f"Min, max, diff configuration energy: {partition_combined.min_configuration_energy:.2f}, {partition_combined.max_configuration_energy:.2f}, {abs(partition_combined.min_configuration_energy - partition_combined.max_configuration_energy):.2f}"
        )
    else:
        vum.addstr(
            y = y_offset + 9,
            x = 0,
            string = f"Min, max, diff configuration energy: {partition_combined.min_configuration_energy:.2f}, {partition_combined.max_configuration_energy:.2f} (original {partition_combined.max_configuration_energy_original:.2f}), {abs(partition_combined.min_configuration_energy - partition_combined.max_configuration_energy):.2f}"
        )
    vum.addstr(
        y = y_offset + 10,
        x = 0,
        string = f"Min H.O.: {partition_combined.ho_quanta_min}, max H.O.: {partition_combined.ho_quanta_max}"
    )

def _generate_total_configurations(
    interaction: Interaction,
    partition_proton: Partition,
    partition_neutron: Partition,
    partition_combined: Partition,
    partition_file_parity: int,
    skips: Skips,
    threshold_energy: float,
    allow_invalid: bool = True,
    is_recursive: bool = False,
):# -> tuple[int, int]:
    """
    Generate all the possible combinations of proton and neutron
    configurations. The parities of the proton and neutron
    configurations must combine multiplicatively to the parity of
    the partition file. The `combined_configurations` list will be
    cleared before the new configurations are added.

    Parameters
    ----------
    partition_file_parity : int
        The parity of the partition file.

    Returns
    -------
    None

    Raises
    ------
    KshellDataStructureError
        If any proton or neutron configuration is unused. This happens
        if the parity of the configuration does not multiplicatively
        combine with any of the configurations of the opposite nucleon
        to the parity of the partition file.
    """
    ho_quanta_min_before = partition_combined.ho_quanta_min_this_parity
    ho_quanta_max_before = partition_combined.ho_quanta_max_this_parity
    assert partition_combined.ho_quanta_min != +1000    # Check that they're not at the default values.
    assert partition_combined.ho_quanta_max != -1000

    partition_combined.clear()
    neutron_configurations_count: list[int] = [0]*partition_neutron.n_configurations
    neutron_configurations_parity_skips: list[int] = [0]*partition_neutron.n_configurations
    neutron_configurations_ho_skips: list[int] = [0]*partition_neutron.n_configurations
    proton_configurations_count: list[int] = [0]*partition_proton.n_configurations
    proton_configurations_parity_skips: list[int] = [0]*partition_proton.n_configurations
    proton_configurations_ho_skips: list[int] = [0]*partition_proton.n_configurations
    n_monopole_skips = 0

    partition_proton.configurations.sort(key=lambda elem: elem.configuration)   # Sort configs in lexicographic order.
    partition_neutron.configurations.sort(key=lambda elem: elem.configuration)
    
    for p_idx in range(partition_proton.n_configurations):
        for n_idx in range(partition_neutron.n_configurations):
            parity_tmp: int = partition_proton.configurations[p_idx].parity*partition_neutron.configurations[n_idx].parity
            if parity_tmp != partition_file_parity:
                """
                Only combinations of proton and neutron orbitals with
                the same parity as the parity of the partition file are
                accepted.
                """
                neutron_configurations_parity_skips[n_idx] += 1
                proton_configurations_parity_skips[p_idx] += 1
                continue

            ho_quanta_tmp: int = (
                partition_proton.configurations[p_idx].ho_quanta + 
                partition_neutron.configurations[n_idx].ho_quanta
            )
            if not (partition_combined.ho_quanta_min <= ho_quanta_tmp <= partition_combined.ho_quanta_max):
                """
                The combined harmonic oscillator quanta of the combined
                proton and neutron orbitals must be within the initial
                limits.
                """
                neutron_configurations_ho_skips[n_idx] += 1
                proton_configurations_ho_skips[p_idx] += 1
                continue
            
            is_original = partition_proton.configurations[p_idx].is_original and partition_neutron.configurations[n_idx].is_original
            energy = configuration_energy(
                interaction = interaction,
                proton_configuration = partition_proton.configurations[p_idx],
                neutron_configuration = partition_neutron.configurations[n_idx],
            )
            if (energy > threshold_energy):# and not is_original:
                """
                Skip configurations with energies over the threshold
                energy only if they are newly generated configurations
                and not originally existing configurations.
                """
                n_monopole_skips += 1
                continue
            
            if energy < partition_combined.min_configuration_energy:
                """
                I am not yet sure if this is OK or not.
                """
                msg = (
                    "A newly calculated pn configuration energy is lower than"
                    " the initial lowest energy pn configuration!"
                )
                raise RuntimeError(msg)

            neutron_configurations_count[n_idx] += 1
            proton_configurations_count[p_idx] += 1

            if partition_file_parity == -1: partition_combined.n_new_negative_configurations += 1
            if partition_file_parity == +1: partition_combined.n_new_positive_configurations += 1

            partition_combined.ho_quanta_min_this_parity = min(partition_combined.ho_quanta_min_this_parity, ho_quanta_tmp)
            partition_combined.ho_quanta_max_this_parity = max(partition_combined.ho_quanta_max_this_parity, ho_quanta_tmp)

            partition_combined.configurations.append(
                Configuration(
                    configuration = [p_idx, n_idx],
                    parity = parity_tmp,
                    ho_quanta = ho_quanta_tmp,
                    energy = energy,
                    is_original = is_original,
                )
            )
    if is_recursive:
        assert len(proton_configurations_count) == len(partition_proton.configurations)
        assert len(neutron_configurations_count) == len(partition_neutron.configurations)
        assert 0 not in proton_configurations_count, proton_configurations_count # Every config should be used now.
        assert 0 not in neutron_configurations_count, neutron_configurations_count
        assert (
            partition_proton.n_existing_positive_configurations +
            partition_proton.n_existing_negative_configurations +
            partition_proton.n_new_positive_configurations +
            partition_proton.n_new_negative_configurations
        ) == partition_proton.n_configurations
        assert (
            partition_neutron.n_existing_positive_configurations +
            partition_neutron.n_existing_negative_configurations +
            partition_neutron.n_new_positive_configurations +
            partition_neutron.n_new_negative_configurations
        ) == partition_neutron.n_configurations
        assert (
            partition_combined.n_existing_positive_configurations +
            partition_combined.n_existing_negative_configurations +
            partition_combined.n_new_positive_configurations +
            partition_combined.n_new_negative_configurations
        ) == partition_combined.n_configurations
        
        return # The recursive call should only re-generate combined configs.
    
    if not allow_invalid:
        """
        Remove proton and neutron configurations which have not been
        counted.
        """
        for p_idx in reversed(range(partition_proton.n_configurations)):
            """
            Traverse configurations in reverse to preserve the order of
            deletion. Consider the list

                ['0', '1', '2', '3', '4'].

            If we are to delete elements 1 and 3, starting from index 0,
            we would first delete '1' at position 1 and then '4' at
            position 3. In reverse however, we would delete '3' at
            position 3 and then '1' at position 1.
            """
            if proton_configurations_count[p_idx] == 0:
                if partition_proton.configurations[p_idx].parity == -1:
                    partition_proton.n_new_negative_configurations -= 1

                elif partition_proton.configurations[p_idx].parity == +1:
                    partition_proton.n_new_positive_configurations -= 1

                else:
                    raise KshellDataStructureError
                
                partition_proton.configurations.pop(p_idx)

        for n_idx in reversed(range(partition_neutron.n_configurations)):
            if neutron_configurations_count[n_idx] == 0:
                if partition_neutron.configurations[n_idx].parity == -1:
                    partition_neutron.n_new_negative_configurations -= 1

                elif partition_neutron.configurations[n_idx].parity == +1:
                    partition_neutron.n_new_positive_configurations -= 1

                else:
                    raise KshellDataStructureError
                
                partition_neutron.configurations.pop(n_idx)

        _generate_total_configurations( # Re-generate combined partition after removing invalid p and n configs.
            interaction = interaction,
            partition_proton = partition_proton,
            partition_neutron = partition_neutron,
            partition_combined = partition_combined,
            partition_file_parity = partition_file_parity,
            skips = skips,
            threshold_energy = threshold_energy,
            is_recursive = True,
        )

    # assert ho_quanta_min_before == partition_combined.ho_quanta_min_this_parity, f"{ho_quanta_min_before} != {partition_combined.ho_quanta_min_this_parity}"
    # assert ho_quanta_max_before == partition_combined.ho_quanta_max_this_parity, f"{ho_quanta_max_before} != {partition_combined.ho_quanta_max_this_parity}"
    
    partition_combined.max_configuration_energy = max([configuration.energy for configuration in partition_combined.configurations])
    skips.n_proton_skips = len([i for i in proton_configurations_count if i == 0])    # The number of proton configurations which will be removed at the end of this program.
    skips.n_neutron_skips = len([i for i in neutron_configurations_count if i == 0])  # The number of neutron configurations which will be removed at the end of this program.
    skips.n_parity_skips = sum(proton_configurations_parity_skips) + sum(neutron_configurations_parity_skips)
    skips.n_ho_skips = sum(proton_configurations_ho_skips) + sum(neutron_configurations_ho_skips)
    skips.n_monopole_skips = n_monopole_skips

def _check_configuration_duplicate(
    new_configuration: list[int],
    existing_configurations: list[Configuration],
    vum: Vum,
    interaction: Interaction,
    is_proton: bool,
    is_neutron: bool,
    init_orb_idx: int,
    n_particles_choice: int,
    is_duplicate_warning: bool,
) -> bool:
    """
    allow_duplicate_warning : bool
        The duplicate warnings should just be displayed if the user
        requested configurations have duplicates in the existing
        configurations. This program might have some overlap in the
        configuration generation algorithms and in these cases no
        duplicate warning should be displayed.
    """
    
    for i, configuration in enumerate(existing_configurations):
        if new_configuration == configuration.configuration:
            duplicate_configuration =  [str(i), configuration]
            break
    else:
        return False
    
    if not is_duplicate_warning: return True
    
    vum.addstr(
        vum.n_rows - 3 - vum.command_log_length, 0,
        f"DUPLICATE: {new_configuration = }, {duplicate_configuration = }"
    )
    draw_shell_map(vum=vum, model_space=interaction.model_space.orbitals, is_proton=is_proton, is_neutron=is_neutron)

    for i in range(len(new_configuration)):
        draw_shell_map(
            vum = vum,
            model_space = interaction.model_space.orbitals,
            is_proton = is_proton,
            is_neutron = is_neutron,
            occupation = (i, new_configuration[i]),
        )
    draw_shell_map(
        vum = vum,
        model_space = interaction.model_space.orbitals,
        is_proton = is_proton,
        is_neutron = is_neutron,
        occupation = (init_orb_idx, new_configuration[init_orb_idx] + n_particles_choice),
        n_holes = n_particles_choice,
    )
        
    return True

def _add_npnh_excitations(
    vum: Vum,
    input_wrapper: Callable,
    model_space_slice: ModelSpace,
    interaction: Interaction,
    partition: Partition,
    partition_proton: Partition,
    partition_neutron: Partition,
    partition_combined: Partition,
    nucleon_choice: str,
    is_proton: bool,
    is_neutron: bool,
) -> bool:
    while True:
        """
        Prompt user for the number of particles to excite.
        """
        n_particles_choice = input_wrapper("N-particle N-hole (N)")
        if n_particles_choice == "q": return False
        
        try:
            n_particles_choice = int(n_particles_choice)
        except ValueError:
            continue

        if n_particles_choice > 2:
            vum.addstr(
                vum.n_rows - 3 - vum.command_log_length, 0,
                "INVALID: Only supports 1 and 2 particle excitations!"
            )
            continue

        if n_particles_choice < 1:
            vum.addstr(
                vum.n_rows - 3 - vum.command_log_length, 0,
                "INVALID: The number of particles must be larger than 0."
            )
            continue

        break

    while True:
        """
        Prompt the user for which major shells to include
        in the N-particle N-hole excitation.
        """
        initial_major_shell_choice = input_wrapper(f"Initial major shell? ({model_space_slice.major_shell_names})")
        if initial_major_shell_choice == "q": return False

        if initial_major_shell_choice in model_space_slice.major_shell_names: break

    while True:
        """
        Prompt the user for which major shells to include
        in the N-particle N-hole excitation.
        """
        final_major_shell_choice = input_wrapper(f"Final major shell? ({model_space_slice.major_shell_names})")
        if final_major_shell_choice == "q": return False

        if final_major_shell_choice in model_space_slice.major_shell_names: break
    
    if initial_major_shell_choice == final_major_shell_choice:
        vum.addstr(
            vum.n_rows - 3 - vum.command_log_length, 0,
            "INVALID: Initial and final major shell cannot be the same!"
        )
        return False

    if major_shell_order[initial_major_shell_choice] > major_shell_order[final_major_shell_choice]:
        vum.addstr(
            vum.n_rows - 3 - vum.command_log_length, 0,
            "INVALID: Initial major shell cannot be higher energy than final major shell!"
        )
        return False

    vum.addstr( # Remove any error messages.
        vum.n_rows - 3 - vum.command_log_length, 0, " "
    )
    initial_orbital_indices: list[int] = [] # Store the indices of the orbitals which are in the initial major shell.
    final_orbital_indices: list[int] = [] # Store the indices of the orbitals which are in the final major shell.
    final_orbital_degeneracy: dict[int, int] = {}    # Accompanying degeneracy of the orbital.
    new_configurations: list[Configuration] = []    # Will be merged with partition.configurations at the end of this function.
    is_duplicate_warning: bool = True
    n_duplicate_skips: int = 0

    for orb in model_space_slice.orbitals:
        """
        Extract indices and degeneracies of the initial and
        final orbitals.
        """
        if orb.order.major_shell_name == initial_major_shell_choice:
            if is_proton:
                init_orb_idx = orb.idx
            elif is_neutron:
                init_orb_idx = orb.idx - interaction.model_space_proton.n_orbitals
            else:
                msg = "'is_proton' and 'is_neutron' should never both be True or False at the same time!"
                raise ValueError(msg)
            
            initial_orbital_indices.append(init_orb_idx)
        
        elif orb.order.major_shell_name == final_major_shell_choice:
            if is_proton:
                final_orb_idx = orb.idx
            elif is_neutron:
                final_orb_idx = orb.idx - interaction.model_space_proton.n_orbitals

            final_orbital_indices.append(final_orb_idx)
            final_orbital_degeneracy[final_orb_idx] = orb.j + 1    # j is stored as j*2.

    while True:
        remove_initial_orbital_choice = input_wrapper(f"Remove any of the initial orbitals? ({[idx + 1 for idx in initial_orbital_indices]}/n)")
        if remove_initial_orbital_choice == "n": break
        if remove_initial_orbital_choice == "q": return

        try:
            remove_initial_orbital_choice = int(remove_initial_orbital_choice)
        except ValueError:
            continue

        remove_initial_orbital_choice -= 1

        try:
            initial_orbital_indices.remove(remove_initial_orbital_choice)
        except ValueError:
            continue

        if not initial_orbital_indices: return

    while True:
        remove_final_orbital_choice = input_wrapper(f"Remove any of the final orbitals? ({[idx + 1 for idx in final_orbital_indices]}/n)")
        if remove_final_orbital_choice == "n": break
        if remove_final_orbital_choice == "q": return

        try:
            remove_final_orbital_choice = int(remove_final_orbital_choice)
        except ValueError:
            continue

        remove_final_orbital_choice -= 1

        try:
            final_orbital_indices.remove(remove_final_orbital_choice)
        except ValueError:
            continue

        if not final_orbital_indices: return

    if is_proton:
        new_ho_quanta_max = partition_proton.ho_quanta_max_this_parity + partition_neutron.ho_quanta_max_this_parity
        new_ho_quanta_max -= n_particles_choice*interaction.model_space_proton.orbitals[initial_orbital_indices[-1]].ho_quanta
        new_ho_quanta_max += n_particles_choice*interaction.model_space_proton.orbitals[final_orbital_indices[-1]].ho_quanta

    if is_neutron:
        new_ho_quanta_max = partition_proton.ho_quanta_max_this_parity + partition_neutron.ho_quanta_max_this_parity
        new_ho_quanta_max -= n_particles_choice*interaction.model_space_neutron.orbitals[initial_orbital_indices[-1]].ho_quanta
        new_ho_quanta_max += n_particles_choice*interaction.model_space_neutron.orbitals[final_orbital_indices[-1]].ho_quanta

    if new_ho_quanta_max > partition_combined.ho_quanta_max:
        while True:
            ho_quanta_max_choice = input_wrapper(f"This NpNh will exceed the max ho quanta ({partition_combined.ho_quanta_max}) with up to {new_ho_quanta_max - partition_combined.ho_quanta_max}. New max?")
            if ho_quanta_max_choice == "q": return False

            try:
                ho_quanta_max_choice = int(ho_quanta_max_choice)
            except ValueError:
                continue

            if ho_quanta_max_choice < partition_combined.ho_quanta_max:
                vum.addstr(
                    vum.n_rows - 3 - vum.command_log_length, 0,
                    "INVALID: New max H.O. quanta must be larger than the old value!"
                )
                continue

            partition_combined.ho_quanta_max = ho_quanta_max_choice
            break
            
    """
    NOTE: Gjøre slik at jeg kan velge å legge til n NpNh-eksitasjoner og
    at de er de med n lavest energi. `init_orb_idx` og `final_orb_idx`
    inneholder nok det jeg trenger for å gjøre dette valget. Et
    alternativ er å begrense hele orbitaler i stedet for å velge antall
    eksitasjoner. F.eks. å bare helt droppe den orbitalen med høyest
    energi.
    """
    for configuration in partition.configurations:
        """
        Loop over every existing configuration.
        """
        for init_orb_idx in initial_orbital_indices:
            """
            Case: N particles from from the same initial
            orbital are excited to the same final orbital.
            """
            if configuration.configuration[init_orb_idx] < n_particles_choice: continue   # Cannot excite enough particles.
            
            for final_orb_idx in final_orbital_indices:
                """
                Both particles to the same final orbital.
                """
                max_additional_occupation = final_orbital_degeneracy[final_orb_idx] - configuration.configuration[final_orb_idx]
                assert max_additional_occupation >= 0, "'max_additional_occupation' should never be negative!"  # Development sanity test.
                
                if max_additional_occupation < n_particles_choice: continue # Cannot excite enough particles.
                
                new_configuration = configuration.configuration.copy()
                new_configuration[init_orb_idx] -= n_particles_choice
                new_configuration[final_orb_idx] += n_particles_choice

                is_duplicate_configuration = _check_configuration_duplicate(
                    new_configuration = new_configuration,
                    existing_configurations = new_configurations,
                    vum = vum,
                    interaction = interaction,
                    is_proton = is_proton,
                    is_neutron = is_neutron,
                    init_orb_idx = init_orb_idx,
                    n_particles_choice = n_particles_choice,
                    is_duplicate_warning = False
                )

                if is_duplicate_configuration:
                    """
                    Checking against the newly generated configuraions.
                    """
                    continue
                
                is_duplicate_configuration = _check_configuration_duplicate(
                    new_configuration = new_configuration,
                    existing_configurations = partition.configurations,
                    vum = vum,
                    interaction = interaction,
                    is_proton = is_proton,
                    is_neutron = is_neutron,
                    init_orb_idx = init_orb_idx,
                    n_particles_choice = n_particles_choice,
                    is_duplicate_warning = is_duplicate_warning,
                )
                if is_duplicate_configuration:
                    """
                    Check that the newly generated configuration does
                    not already exist.
                    """
                    n_duplicate_skips += 1
                    if not is_duplicate_warning: continue
                    duplicate_choice = vum.input("Enter any char to continue or 'i' to ignore duplicate warnings (they will still be deleted)")
                    if duplicate_choice == "i": is_duplicate_warning = False
                    continue

                parity_tmp = _calculate_configuration_parity(
                    configuration = new_configuration,
                    model_space = model_space_slice.orbitals,
                )
                if   parity_tmp == -1: partition.n_new_negative_configurations += 1
                elif parity_tmp == +1: partition.n_new_positive_configurations += 1

                ho_quanta_tmp = sum([   # The number of harmonic oscillator quanta for each configuration.
                    n*orb.ho_quanta for n, orb in zip(new_configuration, model_space_slice.orbitals)
                ])
                new_configurations.append(
                    Configuration(
                        configuration = new_configuration.copy(),
                        parity = parity_tmp,
                        ho_quanta = ho_quanta_tmp,
                        energy = None,
                        is_original = False,
                    )
                )

    if n_particles_choice == 1:
        """
        The next loop of this function only applies for the 2 particle
        excitation case.
        """
        partition.configurations.extend(new_configurations)
        if n_duplicate_skips:
            vum.addstr(
                vum.n_rows - 3 - vum.command_log_length, 0,
                f"DUPLICATE SKIPS: {n_duplicate_skips}"
            )
        return True
    
    for configuration in partition.configurations:
        """
        NOTE: Currently hard-coded for 2 particle excitation. This loop
        deals with the case where 1 particle from 2 different initial
        orbitals are excited into the final major shell.
        """
        # for init_orb_idx in initial_orbital_indices:
        for i1 in range(len(initial_orbital_indices)):
            init_orb_idx_1 = initial_orbital_indices[i1]

            if configuration.configuration[init_orb_idx_1] < 1: continue   # Cannot excite enough particles.

            for i2 in range(i1+1, len(initial_orbital_indices)):
                init_orb_idx_2 = initial_orbital_indices[i2]

                if configuration.configuration[init_orb_idx_2] < 1: continue   # Cannot excite enough particles.


                for f1 in range(len(final_orbital_indices)):
                    final_orb_idx_1 = final_orbital_indices[f1]

                    max_additional_occupation_1 = final_orbital_degeneracy[final_orb_idx_1] - configuration.configuration[final_orb_idx_1]
                    assert max_additional_occupation_1 >= 0, "'max_additional_occupation' should never be negative!"  # Development sanity test.
                    
                    if max_additional_occupation_1 < 1: continue # Cannot excite enough particles.

                    for f2 in range(f1+1, len(final_orbital_indices)):
                        final_orb_idx_2 = final_orbital_indices[f2]

                        max_additional_occupation_2 = final_orbital_degeneracy[final_orb_idx_2] - configuration.configuration[final_orb_idx_2]
                        assert max_additional_occupation_2 >= 0, "'max_additional_occupation' should never be negative!"  # Development sanity test.
                        
                        if max_additional_occupation_2 < 1: continue # Cannot excite enough particles.

                        new_configuration = configuration.configuration.copy()
                        new_configuration[init_orb_idx_1]  -= 1
                        new_configuration[init_orb_idx_2]  -= 1
                        new_configuration[final_orb_idx_1] += 1
                        new_configuration[final_orb_idx_2] += 1
                        
                        is_duplicate_configuration = _check_configuration_duplicate(
                            new_configuration = new_configuration,
                            existing_configurations = new_configurations,
                            vum = vum,
                            interaction = interaction,
                            is_proton = is_proton,
                            is_neutron = is_neutron,
                            init_orb_idx = init_orb_idx,
                            n_particles_choice = n_particles_choice,
                            is_duplicate_warning = False
                        )

                        if is_duplicate_configuration:
                            """
                            Checking against the newly generated configuraions.
                            """
                            continue
                        
                        is_duplicate_configuration = _check_configuration_duplicate(
                            new_configuration = new_configuration,
                            existing_configurations = partition.configurations,
                            vum = vum,
                            interaction = interaction,
                            is_proton = is_proton,
                            is_neutron = is_neutron,
                            init_orb_idx = init_orb_idx,
                            n_particles_choice = n_particles_choice,
                            is_duplicate_warning = is_duplicate_warning,
                        )
                        if is_duplicate_configuration:
                            """
                            Check that the newly generated configuration does
                            not already exist.
                            """
                            n_duplicate_skips += 1
                            if not is_duplicate_warning: continue
                            duplicate_choice = vum.input("Enter any char to continue or 'i' to ignore duplicate warnings (they will still be deleted)")
                            if duplicate_choice == "i": is_duplicate_warning = False
                            continue

                        parity_tmp = _calculate_configuration_parity(
                            configuration = new_configuration,
                            model_space = model_space_slice.orbitals
                        )
                        if   parity_tmp == -1: partition.n_new_negative_configurations += 1
                        elif parity_tmp == +1: partition.n_new_positive_configurations += 1

                        ho_quanta_tmp = sum([   # The number of harmonic oscillator quanta for each configuration.
                            n*orb.ho_quanta for n, orb in zip(new_configuration, model_space_slice.orbitals)
                        ])
                        new_configurations.append(
                            Configuration(
                                configuration = new_configuration.copy(),
                                parity = parity_tmp,
                                ho_quanta = ho_quanta_tmp,
                                energy = None,
                                is_original = False,
                            )
                        )

                for f1 in range(len(final_orbital_indices)):
                    """
                    Case: Both particles are excited to the same final
                    orbital.
                    """
                    final_orb_idx_1 = final_orbital_indices[f1]

                    max_additional_occupation_1 = final_orbital_degeneracy[final_orb_idx_1] - configuration.configuration[final_orb_idx_1]
                    assert max_additional_occupation_1 >= 0, "'max_additional_occupation' should never be negative!"  # Development sanity test.
                    
                    if max_additional_occupation_1 < n_particles_choice: continue # Cannot excite enough particles.

                    new_configuration = configuration.configuration.copy()
                    new_configuration[init_orb_idx_1]  -= 1
                    new_configuration[init_orb_idx_2]  -= 1
                    new_configuration[final_orb_idx_1] += n_particles_choice

                    is_duplicate_configuration = _check_configuration_duplicate(
                        new_configuration = new_configuration,
                        existing_configurations = new_configurations,
                        vum = vum,
                        interaction = interaction,
                        is_proton = is_proton,
                        is_neutron = is_neutron,
                        init_orb_idx = init_orb_idx,
                        n_particles_choice = n_particles_choice,
                        is_duplicate_warning = False
                    )

                    if is_duplicate_configuration:
                        """
                        Checking against the newly generated configuraions.
                        """
                        continue
                    
                    is_duplicate_configuration = _check_configuration_duplicate(
                        new_configuration = new_configuration,
                        existing_configurations = partition.configurations,
                        vum = vum,
                        interaction = interaction,
                        is_proton = is_proton,
                        is_neutron = is_neutron,
                        init_orb_idx = init_orb_idx,
                        n_particles_choice = n_particles_choice,
                        is_duplicate_warning = is_duplicate_warning,
                    )
                    if is_duplicate_configuration:
                        """
                        Check that the newly generated configuration does
                        not already exist.
                        """
                        n_duplicate_skips += 1
                        if not is_duplicate_warning: continue
                        duplicate_choice = vum.input("Enter any char to continue or 'i' to ignore duplicate warnings (they will still be deleted)")
                        if duplicate_choice == "i": is_duplicate_warning = False
                        continue

                    parity_tmp = _calculate_configuration_parity(
                        configuration = new_configuration,
                        model_space = model_space_slice.orbitals
                    )
                    if   parity_tmp == -1: partition.n_new_negative_configurations += 1
                    elif parity_tmp == +1: partition.n_new_positive_configurations += 1

                    ho_quanta_tmp = sum([   # The number of harmonic oscillator quanta for each configuration.
                        n*orb.ho_quanta for n, orb in zip(new_configuration, model_space_slice.orbitals)
                    ])
                    new_configurations.append(
                        Configuration(
                            configuration = new_configuration.copy(),
                            parity = parity_tmp,
                            ho_quanta = ho_quanta_tmp,
                            energy = None,
                            is_original = False,
                        )
                    )

    partition.configurations.extend(new_configurations)
    if n_duplicate_skips:
        vum.addstr(
            vum.n_rows - 3 - vum.command_log_length, 0,
            f"DUPLICATE SKIPS: {n_duplicate_skips}"
        )
    return True

def partition_editor(
    filename_partition_edited: str | None = None,
):  
    """
    Wrapper for error handling.
    """
    try:
        vum: Vum = Vum()
        vum.screen.clear()  # Clear the screen between interactive sessions.
        program_choice = vum.input("Edit or compare? (e/c)")
        if program_choice == "e":
            msg = _partition_editor(
                filename_partition_edited = filename_partition_edited,
                vum_wrapper = vum,
            )
        elif program_choice == "c":
            msg = _partition_compare(vum=vum)
        
        else: return

    except KeyboardInterrupt:
        curses.endwin()
        print("Exiting without saving changes...")

    except curses.error as e:
        curses.endwin()
        raise(e)
    
    except Exception as e:
        curses.endwin()
        raise(e)

    else:
        curses.endwin()
        print(msg)

def test_partition_editor(
    filename_partition: str,
    filename_interaction: str,
    filename_partition_edited: str = "test_partition.ptn",
):
    try:
        os.remove(filename_partition_edited)
    except FileNotFoundError:
        pass
    
    vum: VumDummy = VumDummy()
    _partition_editor(
        filename_interaction = filename_interaction,
        filename_partition = filename_partition,
        filename_partition_edited = filename_partition_edited,
        vum_wrapper = vum,
    )
    n_lines: int = 0
    with open(filename_partition_edited, "r") as infile_edited, open(filename_partition, "r") as infile:
        for line_edited, line in zip(infile_edited, infile):
            n_lines += 1
            if line_edited != line:
                line_edited = line_edited.rstrip('\n')
                line = line.rstrip('\n')
                msg = (
                    f"{line_edited} != {line}"
                )
                raise KshellDataStructureError(msg)
        # else:
        #     print(f"All {n_lines} lines are identical in the files {filename_partition} and {filename_partition_edited}!")

    try:
        os.remove(filename_partition_edited)
    except FileNotFoundError:
        pass

def test_partition_editor_2(
    filename_partition_edited = "test_partition.ptn",
    filename_partition = "Ni67_gs8_n.ptn",
):
    try:
        os.remove(filename_partition_edited)
    except FileNotFoundError:
        pass
    
    vum: VumDummy2 = VumDummy2()
    ans = _partition_editor(
        filename_interaction = "gs8.snt",
        filename_partition = filename_partition,
        filename_partition_edited = filename_partition_edited,
        vum_wrapper = vum,
    )

    try:
        os.remove(filename_partition_edited)
    except FileNotFoundError:
        pass

    print(ans)

def _partition_editor(
    vum_wrapper: Vum,
    filename_interaction: str | None = None,
    filename_partition: str | None = None,
    filename_partition_edited: str | None = None,
    is_recursive = False,
) -> tuple[int, int] | None:
    """
    Extract the model space orbitals from an interaction file. Extract
    proton and neutron partitions from an accompanying partition file.
    Prompt the user for new proton and neutron configurations and 
    regenerate the proton-neutron partition based on the existing and
    new proton and neutron configurations.

    Here, 'configuration' refers to a specific occupation of protons or
    neutrons in the model space orbitals while 'partition' refers to the
    set of all the configurations. Hence, a 'proton configuration' is
    one specific configuration, like
    ```
    1     6  4  2  0  2  4  2  0  0  0  0  0
    ```
    while the 'proton partition' is the set of all the proton
    configurations, like
    ```
    1     6  4  2  0  2  4  2  0  0  0  0  0
    2     6  4  2  0  3  3  2  0  0  0  0  0
    3     6  4  2  0  3  4  1  0  0  0  0  0
    ...
    ```

    Parameters
    ----------
    filename_partition_edited : str | None
        The name of the edited partition file. If None, a name will be
        generated automatically.
    """
    global return_string
    return_string = ""  # For returning a message to be printed after the screen closes.
    
    vum: Vum = vum_wrapper
    screen = vum.screen
    input_wrapper = vum.input

    # screen.clear()  # Clear the screen between interactive sessions.

    if (filename_interaction is None) or (filename_partition is None):
        tmp = _prompt_user_for_interaction_and_partition(vum=vum, is_compare_mode=False)
        if not tmp: return "Exiting without saving changes..."
        filename_interaction, filename_partition = tmp
    
    filename_partition_opposite_parity: str = \
        filename_partition[0:-5] + ("n" if (filename_partition[-5] == "p") else "p") + filename_partition[-4:]

    partition_proton: Partition = Partition()
    partition_neutron: Partition = Partition()
    partition_combined: Partition = Partition()
    interaction: Interaction = Interaction()
    skips: Skips = Skips()
    threshold_energy: float = inf    # Threshold energy for monopole truncation. inf means no truncation.
    
    if os.path.isfile(filename_partition_opposite_parity) and not is_recursive:
        """
        Perform a recursive call to this function to determine the min
        and max ho quanta of the opposite parity partition.
        """
        partition_combined.ho_quanta_min_opposite_parity, partition_combined.ho_quanta_max_opposite_parity = _partition_editor(
            filename_interaction = filename_interaction,
            filename_partition = filename_partition_opposite_parity,
            vum_wrapper = VumDummy(),
            is_recursive = True,
        )
    elif not os.path.isfile(filename_partition_opposite_parity):
        """
        This is not a problem if the calculations only have a single
        parity or if there is no hw truncation.
        """
        vum.addstr(vum.n_rows - 2 - vum.command_log_length, 0, (
            f"WARNING! Could not find opposite parity partition file,"
            " thus cannot properly calculate min and max H.O. quanta "
            " for proper treatment of hw truncation."
        ))
        time.sleep(DELAY)

    if filename_partition_edited is None:
        filename_partition_edited = f"{filename_partition.split('.')[0]}_edited.ptn"
    
    load_interaction(filename_interaction=filename_interaction, interaction=interaction)
    draw_shell_map(vum=vum, model_space=interaction.model_space.orbitals, is_proton=True, is_neutron=True)
    header = load_partition(
        filename_partition = filename_partition,
        interaction = interaction,
        partition_proton = partition_proton,
        partition_neutron = partition_neutron,
        partition_combined = partition_combined,
    )

    if is_recursive:
        """
        This happens only inside the recursive function call.
        """
        return partition_combined.ho_quanta_min_this_parity, partition_combined.ho_quanta_max_this_parity

    combined_existing_configurations_expected: list[Configuration] = partition_combined.configurations.copy()    # For sanity check.

    while True:
        _generate_total_configurations(
            interaction = interaction,
            threshold_energy = threshold_energy,
            partition_proton = partition_proton,
            partition_neutron = partition_neutron,
            partition_combined = partition_combined,
            partition_file_parity = partition_combined.parity,
            skips = skips,
        )
        if len(combined_existing_configurations_expected) != partition_combined.n_configurations:
            msg = (
                "The number of combined configurations is not correct!"
                f" Expected: {len(combined_existing_configurations_expected)}, calculated: {partition_combined.n_configurations}"
            )
            vum.addstr(vum.n_rows - 1 - vum.command_log_length - 2, 0, msg)
            msg = (
                "This can happen if the .ptn file was generated with monopole"
                " truncation."
            )
            vum.addstr(vum.n_rows - 1 - vum.command_log_length - 1, 0, msg)
            while True:
                check_choice = input_wrapper(".ptn monopole threshold energy or skip the check? (energy/s)")
                if check_choice == "s": break

                try:
                    threshold_energy = float(check_choice)
                except ValueError:
                    continue
                else:
                    threshold_energy = partition_combined.min_configuration_energy + threshold_energy
                    break

            if check_choice == "s": break

        else: break

    vum.addstr(vum.n_rows - 1 - vum.command_log_length - 2, 0, vum.blank_line)
    vum.addstr(vum.n_rows - 1 - vum.command_log_length - 1, 0, vum.blank_line)

        # assert len(combined_existing_configurations_expected) == partition_combined.n_configurations, msg

    for i in range(partition_combined.n_configurations):
        if partition_combined.configurations[i].configuration != combined_existing_configurations_expected[i].configuration:
            p_idx_expected = combined_existing_configurations_expected[i].configuration[0]
            n_idx_expected = combined_existing_configurations_expected[i].configuration[1]
            p_idx_calculated = partition_combined.configurations[i].configuration[0]
            n_idx_calculated = partition_combined.configurations[i].configuration[1]
            msg = (
                f"Combined config with index {i} is incorrectly calculated!"
                f" Expected: {p_idx_expected, n_idx_expected}, calculated: {p_idx_calculated, n_idx_calculated}"
                f"\nproton_configurations_formatted [{p_idx_expected:5d}] = {partition_proton.configurations[p_idx_expected].configuration}, ho = {partition_proton.configurations[p_idx_expected].ho_quanta}, parity = {partition_proton.configurations[p_idx_expected].parity}"
                f"\nneutron_configurations_formatted[{n_idx_expected:5d}] = {partition_neutron.configurations[n_idx_expected].configuration}, ho = {partition_neutron.configurations[n_idx_expected].ho_quanta}, parity = {partition_neutron.configurations[n_idx_expected].parity}"
                f"\nproton_configurations_formatted [{p_idx_calculated:5d}] = {partition_proton.configurations[p_idx_calculated].configuration}, ho = {partition_proton.configurations[p_idx_calculated].ho_quanta}, parity = {partition_proton.configurations[p_idx_calculated].parity}"
                f"\nneutron_configurations_formatted[{n_idx_calculated:5d}] = {partition_neutron.configurations[n_idx_calculated].configuration}, ho = {partition_neutron.configurations[n_idx_calculated].ho_quanta}, parity = {partition_neutron.configurations[n_idx_calculated].parity}"
                f"\nho expected = {partition_proton.configurations[p_idx_expected].ho_quanta + partition_neutron.configurations[n_idx_expected].ho_quanta}"
                f"\nho calculated = {partition_proton.configurations[p_idx_calculated].ho_quanta + partition_neutron.configurations[n_idx_calculated].ho_quanta}"
                f"\nho min: {partition_combined.ho_quanta_min_this_parity}"
                f"\nho max: {partition_combined.ho_quanta_max_this_parity}"
            )
            if partition_combined.configurations[i].configuration not in [c.configuration for c in combined_existing_configurations_expected]:
                msg += f"\n{partition_combined.configurations[i].configuration} should never appear!"
            
            raise KshellDataStructureError(msg)

    if isinstance(vum, VumDummy):
        """
        Skip dim calculation for testing.
        """
        M = [0]
        mdim = [0]
    else:
        M, mdim, jdim = count_dim(
            model_space_filename = filename_interaction,
            partition_filename = None,
            print_dimensions = False,
            debug = False,
            parity = partition_combined.parity,
            proton_partition = [configuration.configuration for configuration in partition_proton.configurations],
            neutron_partition = [configuration.configuration for configuration in partition_neutron.configurations],
            total_partition = [configuration.configuration for configuration in partition_combined.configurations],
        )
    mdim_original: int = mdim[-1]
    _summary_information(
        vum = vum,
        skips = skips,
        filename_interaction = filename_interaction,
        filename_partition = filename_partition,
        partition_proton = partition_proton,
        partition_neutron = partition_neutron,
        partition_combined = partition_combined,
        interaction = interaction,
        M = M,
        mdim = mdim,
        y_offset = y_offset,
        mdim_original = None,
        # n_proton_skips = 0,
        # n_neutron_skips = 0,
        # n_parity_skips = 0,
        # n_ho_skips = 0,
    )
    _analyse_existing_configuration(
        vum = vum,
        proton_configurations = partition_proton.configurations,
        neutron_configurations = partition_neutron.configurations,
        input_wrapper = input_wrapper,
        model_space = interaction.model_space.orbitals,
        y_offset = y_offset,
        n_proton_orbitals = interaction.model_space_proton.n_orbitals,
    )
    while True:
        configuration_choice = input_wrapper("Add configuration? (y/n)")
        if configuration_choice == "y":
            pass
        elif configuration_choice in ["n", "q"]:
            break
        else: continue

        while True:
            nucleon_choice = input_wrapper("Proton or neutron? (p/n)")
            if nucleon_choice == "p":
                is_proton = True
                is_neutron = False
                nucleon = "proton"
                model_space_slice = interaction.model_space_proton
                partition = partition_proton
                n_valence_nucleons = interaction.model_space_proton.n_valence_nucleons
                break
            
            elif nucleon_choice == "n":
                is_proton = False
                is_neutron = True
                nucleon = "neutron"
                model_space_slice = interaction.model_space_neutron
                partition = partition_neutron
                n_valence_nucleons = interaction.model_space_neutron.n_valence_nucleons
                break

            elif nucleon_choice == "q":
                break

            else: continue

        while True:
            draw_shell_map(vum=vum, model_space=interaction.model_space.orbitals, is_proton=is_proton, is_neutron=is_neutron)
            configuration_type_choice = input_wrapper("Single or range of configurations? (s/r)")
            if configuration_type_choice == "s":
                """
                Prompt the user for single specific configurations.
                """
                while True:
                    new_configuration = _prompt_user_for_configuration(
                        vum = vum,
                        nucleon = nucleon,
                        model_space = model_space_slice.orbitals,
                        n_valence_nucleons = n_valence_nucleons,
                        input_wrapper = input_wrapper,
                        y_offset = y_offset,
                    )
                    if new_configuration:
                        is_duplicate_configuration = _check_configuration_duplicate(
                            new_configuration = new_configuration,
                            existing_configurations = partition.configurations,
                            vum = vum,
                            interaction = interaction,
                            is_proton = is_proton,
                            is_neutron = is_neutron,
                            init_orb_idx = 0,
                            n_particles_choice = 0,
                            is_duplicate_warning = False,
                        )
                        if is_duplicate_configuration:
                            msg = (
                                "This configuration already exists! Skipping..."
                            )
                            vum.addstr(vum.n_rows - 1 - vum.command_log_length - 2, 0, "DUPLICATE")
                            vum.addstr(vum.n_rows - 1 - vum.command_log_length - 1, 0, msg)
                            time.sleep(DELAY)
                            draw_shell_map(vum=vum, model_space=interaction.model_space.orbitals, is_proton=is_proton, is_neutron=is_neutron)
                            continue

                        parity_tmp = _calculate_configuration_parity(
                            configuration = new_configuration,
                            model_space = model_space_slice.orbitals
                        )
                        if   parity_tmp == -1: partition.n_new_negative_configurations += 1
                        elif parity_tmp == +1: partition.n_new_positive_configurations += 1

                        ho_quanta_tmp = sum([   # The number of harmonic oscillator quanta for each configuration.
                            n*orb.ho_quanta for n, orb in zip(new_configuration, interaction.model_space_proton.orbitals)
                        ])
                        return_string += f"\n{len(partition_proton.configurations) = }, {len(partition_neutron.configurations) = }, {len(partition_combined.configurations) = }"
                        partition.configurations.append(
                            Configuration(
                                configuration = new_configuration,
                                parity = parity_tmp,
                                ho_quanta = ho_quanta_tmp,
                                energy = None,
                                is_original = False,
                            )
                        )
                        return_string += f"\n{len(partition_proton.configurations) = }, {len(partition_neutron.configurations) = }, {len(partition_combined.configurations) = }"
                        _generate_total_configurations(
                            interaction = interaction,
                            threshold_energy = threshold_energy,
                            partition_proton = partition_proton,
                            partition_neutron = partition_neutron,
                            partition_combined = partition_combined,
                            partition_file_parity = partition_combined.parity,
                            allow_invalid = True,
                            skips = skips,
                        )
                        # try:
                        #     n_proton_skips, n_neutron_skips = _generate_total_configurations(
                        #         partition_proton = partition_proton,
                        #         partition_neutron = partition_neutron,
                        #         partition_combined = partition_combined,
                        #         partition_file_parity = partition_combined.parity,
                        #         allow_invalid = True,
                        #     )
                        # except KshellDataStructureError:
                        #     """
                        #     Accept invalid configuration parity for now
                        #     since the user might specify more
                        #     configurations which may multiplicatively
                        #     combine parities to match the .ptn parity.
                        #     """
                        #     pass
                        
                        # print(f"n proton, neutron configurations will be skipped because of parity or H.O. mismatch: {n_proton_skips, n_neutron_skips}")
                        return_string += f"\n{len(partition_proton.configurations) = }, {len(partition_neutron.configurations) = }, {len(partition_combined.configurations) = }"
                        # return return_string
                        M, mdim, jdim = count_dim(
                            model_space_filename = filename_interaction,
                            partition_filename = None,
                            print_dimensions = False,
                            debug = False,
                            parity = partition_combined.parity,
                            proton_partition = [configuration.configuration for configuration in partition_proton.configurations],
                            neutron_partition = [configuration.configuration for configuration in partition_neutron.configurations],
                            total_partition = [configuration.configuration for configuration in partition_combined.configurations],
                        )
                        # vum.addstr(y_offset + 1, 0, f"M-scheme dim (M={M[-1]}): {mdim[-1]:d} ({mdim[-1]:.2e}) (original {mdim_original:d} ({mdim_original:.2e}))")
                        # vum.addstr(y_offset + 2, 0, f"n proton, neutron configurations: {partition_proton.n_existing_configurations} + {partition_proton.n_new_configurations}, {partition_neutron.n_existing_configurations} + {partition_neutron.n_new_configurations}")
                        # vum.addstr(y_offset + 6, 0, f"n proton +, - : {partition_proton.n_existing_positive_configurations} + {partition_proton.n_new_positive_configurations}, {partition_proton.n_existing_negative_configurations} + {partition_proton.n_new_negative_configurations}")
                        # vum.addstr(y_offset + 7, 0, f"n neutron +, - : {partition_neutron.n_existing_positive_configurations} + {partition_neutron.n_new_positive_configurations}, {partition_neutron.n_existing_negative_configurations} + {partition_neutron.n_new_negative_configurations}")
                        # vum.addstr(y_offset + 9, 0, f"n pn configuration combinations: {partition_combined.n_configurations}")
                        _summary_information(
                            vum = vum,
                            filename_interaction = filename_interaction,
                            filename_partition = filename_partition,
                            partition_proton = partition_proton,
                            partition_neutron = partition_neutron,
                            partition_combined = partition_combined,
                            interaction = interaction,
                            M = M,
                            mdim = mdim,
                            y_offset = y_offset,
                            mdim_original = mdim_original,
                            skips = skips,
                            # n_proton_skips = n_proton_skips,
                            # n_neutron_skips = n_neutron_skips,
                            # n_parity_skips = n_parity_skips,
                            # n_ho_skips = n_ho_skips,
                        )
                        vum.addstr(vum.n_rows - 1 - vum.command_log_length - 2, 0, "Configuration added!")
                        time.sleep(DELAY)

                    elif new_configuration is None:
                        """
                        Quit signal. Do not keep the current configuration,
                        but keep earlier defined new configurations and quit
                        the prompt.
                        """
                        draw_shell_map(vum=vum, model_space=interaction.model_space.orbitals, is_proton=is_proton, is_neutron=is_neutron)
                        vum.addstr(vum.n_rows - 3 - vum.command_log_length, 0, " ")
                        # vum.addstr(vum.n_rows - 2 - vum.command_log_length, 0, "Current configuration discarded")
                        vum.addstr(vum.n_rows - 2 - vum.command_log_length, 0, " ")
                        break
                break

            elif configuration_type_choice == "r":
                """
                Prompt the user for a range of configurations. Currently
                only supports 2p2h excitations.
                """
                if input_wrapper("Monopole truncation? (y/n)") == "y":
                    while True:
                        try:
                            threshold_energy = input_wrapper("Monopole trunc threshold energy (relative to min)")
                            threshold_energy = float(threshold_energy)

                        except ValueError:
                            continue

                        if threshold_energy > 0:
                            threshold_energy = partition_combined.min_configuration_energy + threshold_energy
                            break

                if not _add_npnh_excitations(
                    vum = vum,
                    input_wrapper = input_wrapper,
                    model_space_slice = model_space_slice,
                    interaction = interaction,
                    partition = partition,
                    partition_proton = partition_proton,
                    partition_neutron = partition_neutron,
                    partition_combined = partition_combined,
                    nucleon_choice = nucleon_choice,
                    is_proton = is_proton,
                    is_neutron = is_neutron,
                ):
                    continue
                try:
                    _generate_total_configurations(
                        interaction = interaction,
                        threshold_energy = threshold_energy,
                        partition_proton = partition_proton,
                        partition_neutron = partition_neutron,
                        partition_combined = partition_combined,
                        partition_file_parity = partition_combined.parity,
                        skips = skips,
                    )
                except KshellDataStructureError:
                    """
                    Allow parity mismatch for now.
                    """
                    pass

                M, mdim, jdim = count_dim(
                    model_space_filename = filename_interaction,
                    partition_filename = None,
                    print_dimensions = False,
                    debug = False,
                    parity = partition_combined.parity,
                    proton_partition = [configuration.configuration for configuration in partition_proton.configurations],
                    neutron_partition = [configuration.configuration for configuration in partition_neutron.configurations],
                    total_partition = [configuration.configuration for configuration in partition_combined.configurations],
                )
                # vum.addstr(y_offset + 1, 0, f"M-schemelol dim (M={M[-1]}): {mdim[-1]:d} ({mdim[-1]:.2e}) (original {mdim_original:d} ({mdim_original:.2e}))")
                # vum.addstr(y_offset + 2, 0, f"n proton, neutron configurations: {partition_proton.n_existing_configurations} + {partition_proton.n_new_configurations}, {partition_neutron.n_existing_configurations} + {partition_neutron.n_new_configurations}")
                # vum.addstr(y_offset + 6, 0, f"n proton +, - : {partition_proton.n_existing_positive_configurations} + {partition_proton.n_new_positive_configurations}, {partition_proton.n_existing_negative_configurations} + {partition_proton.n_new_negative_configurations}")
                # vum.addstr(y_offset + 7, 0, f"n neutron +, - : {partition_neutron.n_existing_positive_configurations} + {partition_neutron.n_new_positive_configurations}, {partition_neutron.n_existing_negative_configurations} + {partition_neutron.n_new_negative_configurations}")
                # vum.addstr(y_offset + 9, 0, f"n pn configuration combinations: {partition_combined.n_configurations}")
                _summary_information(
                    vum = vum,
                    filename_interaction = filename_interaction,
                    filename_partition = filename_partition,
                    partition_proton = partition_proton,
                    partition_neutron = partition_neutron,
                    partition_combined = partition_combined,
                    interaction = interaction,
                    M = M,
                    mdim = mdim,
                    y_offset = y_offset,
                    mdim_original = mdim_original,
                    skips = skips,
                    # n_proton_skips = n_proton_skips,
                    # n_neutron_skips = n_neutron_skips,
                    # n_parity_skips = n_parity_skips,
                    # n_ho_skips = n_ho_skips,
                )
                # vum.addstr(vum.n_rows - 1 - vum.command_log_length - 2, 0, "Configurations added!")
                break

            elif configuration_type_choice == "q":
                break

            else: continue

    _generate_total_configurations(
        interaction = interaction,
        partition_proton = partition_proton,
        partition_neutron = partition_neutron,
        partition_combined = partition_combined,
        partition_file_parity = partition_combined.parity,
        threshold_energy = threshold_energy,    # NOTE: Maybe this should be inf to avoid removing existing configurations?
        allow_invalid = False,
        skips = skips,
    )
    n_new_proton_configurations = partition_proton.n_new_negative_configurations + partition_proton.n_new_positive_configurations
    n_new_neutron_configurations = partition_neutron.n_new_negative_configurations + partition_neutron.n_new_positive_configurations

    if not isinstance(vum, VumDummy):
        """
        Dont create a .ptn file if no changes are made, but allow it for
        unit tests using VumDummy.
        """
        if (not n_new_neutron_configurations) and (not n_new_proton_configurations):
            return_string += "No new configurations. Skipping creation of new .ptn file."
            return return_string
    
    _sanity_checks(
        partition_proton = partition_proton,
        partition_neutron = partition_neutron,
        partition_combined = partition_combined,
        interaction = interaction,
    )
    while True:
        save_changes_choice = input_wrapper("Save changes to new partition file? (y/n)")
        
        if save_changes_choice == "n":
            return_string += "Exiting without saving changes..."
            return return_string
        
        elif save_changes_choice == "y": break

    with open(filename_partition_edited, "w") as outfile:
        """
        Write edited data to new partition file.
        """
        outfile.write(header)
        outfile.write(f" {partition_proton.n_configurations} {partition_neutron.n_configurations}\n")
        outfile.write("# proton partition\n")

        for i in range(partition_proton.n_configurations):
            outfile.write(
                f"{i + 1:6d}     "    # +1 because .ptn indices start at 1.
                f"{str(partition_proton.configurations[i].configuration).strip('[]').replace(',', ' ')}"
                "\n"
            )
        outfile.write("# neutron partition\n")

        for i in range(partition_neutron.n_configurations):
            outfile.write(
                f"{i + 1:6d}     "    # +1 because .ptn indices start at 1.
                f"{str(partition_neutron.configurations[i].configuration).strip('[]').replace(',', ' ')}"
                "\n"
            )
        outfile.write("# partition of proton and neutron\n")
        outfile.write(f"{partition_combined.n_configurations}\n")

        for configuration in partition_combined.configurations:
            outfile.write(f"{configuration.configuration[0] + 1:5d} {configuration.configuration[1] + 1:5d}\n") # +1 because .ptn indices start at 1.

    return_string += (
        f"New configuration saved to {filename_partition_edited}"
        f" with {n_new_proton_configurations} new proton configurations"
        f" and {n_new_neutron_configurations} new neutron configurations"
    )
    return return_string

def _prompt_user_for_configuration(
    vum: Vum,
    nucleon: str,
    model_space: list[OrbitalParameters],
    n_valence_nucleons: int,
    input_wrapper: Callable,
    y_offset: int,
) -> list | None:
        """
        Prompt the user for a new proton or neutron configuration,
        meaning that this function prompts the user for the occupation
        of all the orbitals in the model space and returns the
        configuration as a list of occupation numbers. Note that the
        order of the orbitals that the user is prompted for is sorted
        by the `shell_model_order` dictionary but the occupation number
        list returned by this function is sorted by the order of the
        orbitals as defined in the interaction file.

        Parameters
        ----------
        vum : Vum
            The Vum object to use for displaying information.

        nucleon : str
            The nucleon type to prompt the user for. Must be either
            "proton" or "neutron".
        
        model_space : list[OrbitalParameters]
            The model space to use for the prompt.

        n_valence_nucleons : int
            The number of valence protons or neutrons. Used for making
            sure that the user does not enter more than this number of
            nucleons in the configuration.

        input_wrapper : Callable
            The input wrapper to use for getting user input.

        y_offset : int
            The y-offset to use for displaying information.

        Returns
        -------
        list | None
            The occupation of each model space orbitals as a list of
            occupation numbers. If the user enters "q" to quit, None is
            returned.
        """
        is_blank_map: bool = True
        if nucleon == "proton":
            is_proton = True
            is_neutron = False

        elif nucleon == "neutron":
            is_proton = False
            is_neutron = True

        else:
            msg = f"Invalid nucleon parameter '{nucleon}'!"
            raise ValueError(msg)

        model_space_copy = sorted(  # Sort the orbitals based on the shell_model_order dict.
            model_space,
            key = lambda orbital: shell_model_order[f"{orbital.n}{spectroscopic_conversion[orbital.l]}{orbital.j}"].idx,
        )
        n_remaining_nucleons: int = n_valence_nucleons
        vum.addstr(vum.n_rows - 3 - vum.command_log_length, 0, " ")
        vum.addstr(vum.n_rows - 2 - vum.command_log_length, 0, f"Please enter {nucleon} orbital occupation (f to fill, q to quit):")
        occupation: list[tuple[int, int]] = []
        occupation_parity: int = 1

        for orbital in model_space_copy:
            if n_remaining_nucleons == 0:
                """
                If there are no more valence nucleons to use, set the
                remaining occupations to 0.
                """
                occupation.append((orbital.idx, 0))
                vum.addstr(vum.n_rows - 3 - vum.command_log_length, 0, "Occupation of remaining orbitals set to 0.")
                continue

            while True:
                ans = input_wrapper(f"{orbital.idx + 1:2d} {orbital} (remaining: {n_remaining_nucleons})")
                if is_blank_map:
                    is_blank_map = False
                    draw_shell_map(
                        vum = vum,
                        model_space = model_space,
                        is_proton = is_proton,
                        is_neutron = is_neutron,
                    )

                if (ans == "q") or (ans == "quit") or (ans == "exit"): return None
                
                if ans == "f":
                    """
                    Fill the orbital completely or with as many
                    remaining nucleons as possible.
                    """
                    ans = min(orbital.j + 1, n_remaining_nucleons)
                
                try:
                    ans = int(ans)
                except ValueError:
                    continue

                if (ans > (orbital.j + 1)) or (ans < 0):
                    vum.addstr(
                        vum.n_rows - 3 - vum.command_log_length, 0,
                        f"Allowed occupation for this orbital is [0, 1, ..., {orbital.j + 1}]"
                    )
                    continue
                
                occupation_parity *= orbital.parity**ans
                vum.addstr(y_offset + PARITY_CURRENT_Y_COORD, 0, f"parity current configuration = {occupation_parity}")
                
                n_remaining_nucleons -= ans
                if ans > 0:
                    """
                    No point in re-drawing the line if the occupation is
                    unchanged.
                    """
                    draw_shell_map(
                        vum = vum,
                        model_space = model_space,
                        is_proton = is_proton,
                        is_neutron = is_neutron,
                        occupation = (orbital.idx, ans)
                    )
                break

            occupation.append((orbital.idx, ans))

            cum_occupation = sum(tup[1] for tup in occupation)
            if cum_occupation > n_valence_nucleons:
                vum.addstr(
                    vum.n_rows - 3 - vum.command_log_length, 0,
                    f"INVALID: Total occupation ({cum_occupation}) exceeds the number of valence {nucleon}s ({n_valence_nucleons})"
                )
                draw_shell_map(
                    vum = vum,
                    model_space = model_space,
                    is_proton = is_proton,
                    is_neutron = is_neutron,
                )
                return []
        
        cum_occupation = sum(tup[1] for tup in occupation)
        if cum_occupation < n_valence_nucleons:
            vum.addstr(
                vum.n_rows - 3 - vum.command_log_length, 0,
                f"INVALID: Total occupation ({cum_occupation}) does not use the total number of valence {nucleon}s ({n_valence_nucleons})"
            )
            draw_shell_map(
                vum = vum,
                model_space = model_space,
                is_proton = is_proton,
                is_neutron = is_neutron,
            )
            return []
        
        vum.addstr(vum.n_rows - 2 - vum.command_log_length, 0, vum.blank_line)
        occupation.sort(key=lambda tup: tup[0]) # Sort list of tuples based on the orbital.idx.
        return [tup[1] for tup in occupation]   # Return only the occupation numbers now sorted based on the orbital order of the .snt file.
<file_sep>from __future__ import annotations
import multiprocessing, curses
import kshell_utilities as ksutil
from vum import Vum
import numpy as np
from kshell_utilities.data_structures import (
    Partition, Interaction
)
from .partition_tools import _prompt_user_for_interaction_and_partition

def _duplicate_checker(
    partition_a: Partition,
    partition_b: Partition,
) -> np.ndarray[bool]:
    """
    Check if the conffigurations in `partition_a` also occur in
    `partition_b`.
    """
    is_in_both_partitions: np.ndarray[bool] = np.zeros(partition_a.n_configurations, dtype=bool)

    for idx_a in range(partition_a.n_configurations):
        config_a = partition_a.configurations[idx_a].configuration
        
        for idx_b in range(partition_b.n_configurations):
            config_b = partition_b.configurations[idx_b].configuration
            
            if config_a == config_b:
                is_in_both_partitions[idx_a] = True
                break

        else:
            is_in_both_partitions[idx_a] = False

    return is_in_both_partitions

def _sanity_checks(
    filename_partition_a: str,
    filename_partition_b: str,
    filename_interaction: str,
):
    ksutil.test_partition_editor(   # Proof that partition_editor perfectly replicates the combined configs.
        filename_partition = filename_partition_a,
        filename_partition_edited = filename_partition_a + "testfile.ptn",
        filename_interaction = filename_interaction,
    )
    ksutil.test_partition_editor(   # Proof that partition_editor perfectly replicates the combined configs.
        filename_partition = filename_partition_b,
        filename_partition_edited = filename_partition_b + "testfile.ptn",
        filename_interaction = filename_interaction,
    )

def partition_compare():
    """
    Wrapper function for error handling.
    """
    try:
        vum: Vum = Vum()
        msg = _partition_compare(vum)

    except KeyboardInterrupt:
        print("Exiting without saving changes...")

    except curses.error as e:
        raise(e)
    
    except Exception as e:
        raise(e)

    else:
        print(msg)

def _partition_compare(vum: Vum):
    tmp = _prompt_user_for_interaction_and_partition(
        vum = vum,
        is_compare_mode = True,
    )
    filename_interaction, (filename_partition_a, filename_partition_b) = tmp
    # filename_interaction = "gs8.snt"
    interaction: Interaction = Interaction()
    ksutil.load_interaction(
        filename_interaction = filename_interaction,
        interaction = interaction
    )
    
    # filename_partition_a = "Ni67_gs8_n_edited_2.ptn"
    partition_proton_a: Partition = Partition()
    partition_neutron_a: Partition = Partition()
    partition_combined_a: Partition = Partition()
    ksutil.load_partition(
        filename_partition = filename_partition_a,
        interaction = interaction,
        partition_proton = partition_proton_a,
        partition_neutron = partition_neutron_a,
        partition_combined = partition_combined_a,
    )

    # filename_partition_b = "Ni67_gs8_n.ptn"
    partition_proton_b: Partition = Partition()
    partition_neutron_b: Partition = Partition()
    partition_combined_b: Partition = Partition()
    ksutil.load_partition(
        filename_partition = filename_partition_b,
        interaction = interaction,
        partition_proton = partition_proton_b,
        partition_neutron = partition_neutron_b,
        partition_combined = partition_combined_b,
    )
    
    sanity_check_process = multiprocessing.Process( # Sanity checks in a separate process to make the program more snappy.
        target = _sanity_checks,
        args = (filename_partition_a, filename_partition_b, filename_interaction),
    )
    sanity_check_process.start()
    
    _, mdim_a, _ = ksutil.count_dim(
        model_space_filename = filename_interaction,
        partition_filename = None,
        print_dimensions = False,
        debug = False,
        parity = partition_combined_a.parity,
        proton_partition = [configuration.configuration for configuration in partition_proton_a.configurations],
        neutron_partition = [configuration.configuration for configuration in partition_neutron_a.configurations],
        total_partition = [configuration.configuration for configuration in partition_combined_a.configurations],
    )
    _, mdim_b, _ = ksutil.count_dim(
        model_space_filename = filename_interaction,
        partition_filename = None,
        print_dimensions = False,
        debug = False,
        parity = partition_combined_b.parity,
        proton_partition = [configuration.configuration for configuration in partition_proton_b.configurations],
        neutron_partition = [configuration.configuration for configuration in partition_neutron_b.configurations],
        total_partition = [configuration.configuration for configuration in partition_combined_b.configurations],
    )

    vum.addstr(0, 0,
        string = f"Compare Partitions :: {filename_interaction} :: (A) {filename_partition_a} :: (B) {filename_partition_b}"
    )
    vum.addstr(2, 33,
        f"A"
    )
    vum.addstr(2, 40 + 6,
        f"B",
        is_blank_line = False,
    )
    vum.addstr(3, 0,
        f"dim                   {mdim_a[-1]:14.2e} {mdim_b[-1]:14.2e}"
    )
    vum.addstr(4, 0,
        f"n proton configs      {partition_proton_a.n_configurations:14d} {partition_proton_b.n_configurations:14d}"
    )
    vum.addstr(5, 0,
        f"n neutron configs     {partition_neutron_a.n_configurations:14d} {partition_neutron_b.n_configurations:14d}"
    )
    vum.addstr(6, 0,
        f"n combined configs    {partition_combined_a.n_configurations:14d} {partition_combined_b.n_configurations:14d}"
    )
    
    min_max_a = str((partition_proton_a.ho_quanta_min_this_parity, partition_proton_a.ho_quanta_max_this_parity))[1:-1]#.replace(" ", "")
    min_max_b = str((partition_proton_b.ho_quanta_min_this_parity, partition_proton_b.ho_quanta_max_this_parity))[1:-1]#.replace(" ", "")
    vum.addstr(7, 0,
        f"H.O. min,max proton   {min_max_a:>14s} {min_max_b:>14s}"
    )
    min_max_a = str((partition_neutron_a.ho_quanta_min_this_parity, partition_neutron_a.ho_quanta_max_this_parity))[1:-1]#.replace(" ", "")
    min_max_b = str((partition_neutron_b.ho_quanta_min_this_parity, partition_neutron_b.ho_quanta_max_this_parity))[1:-1]#.replace(" ", "")
    vum.addstr(8, 0,
        f"H.O. min,max neutron  {min_max_a:>14s} {min_max_b:>14s}"
    )
    min_max_a = str((partition_combined_a.ho_quanta_min_this_parity, partition_combined_a.ho_quanta_max_this_parity))[1:-1]#.replace(" ", "")
    min_max_b = str((partition_combined_b.ho_quanta_min_this_parity, partition_combined_b.ho_quanta_max_this_parity))[1:-1]#.replace(" ", "")
    vum.addstr(9, 0,
        f"H.O. min,max combined  {min_max_a:>14s} {min_max_b:>14s}"
    )

    vum.addstr(11, 29,
        f"A in B"
    )
    vum.addstr(11, 36 + 6,
        f"B in A",
        is_blank_line = False,
    )

    a_in_b = _duplicate_checker(
        partition_a = partition_proton_a,
        partition_b = partition_proton_b,
    )
    a_in_b = str(f"{a_in_b.sum()}/{partition_proton_a.n_configurations}")
    b_in_a = _duplicate_checker(
        partition_a = partition_proton_b,
        partition_b = partition_proton_a,
    )
    b_in_a = str(f"{b_in_a.sum()}/{partition_proton_b.n_configurations}")
    vum.addstr(12, 0,
        f"n proton configs    {a_in_b:>14s} {b_in_a:>14s}"
    )
    
    a_in_b = _duplicate_checker(
        partition_a = partition_neutron_a,
        partition_b = partition_neutron_b,
    )
    a_in_b = str(f"{a_in_b.sum()}/{partition_neutron_a.n_configurations}")
    b_in_a = _duplicate_checker(
        partition_a = partition_neutron_b,
        partition_b = partition_neutron_a,
    )
    b_in_a = str(f"{b_in_a.sum()}/{partition_neutron_b.n_configurations}")
    vum.addstr(13, 0,
        f"n neutron configs   {a_in_b:>14s} {b_in_a:>14s}"
    )

    vum.input("Enter any char to exit")

if __name__ == "__main__":
    # main()
    partition_compare()
<file_sep>from setuptools import setup

# try:
#     import pypandoc
#     long_description = pypandoc.convert('README.md', 'rst')
# except (IOError, ImportError, OSError):
#     long_description = ""

setup(
    name = 'kshell-utilities',
    version = '1.5.1.0',    
    description = 'Handy utilities for handling nuclear shell model calculations from KSHELL',
    # long_description = long_description,
    url = 'https://github.com/GaffaSnobb/kshell-utilities',
    author = ['<NAME>', '<NAME>'],
    author_email = '<EMAIL>',
    packages = ['kshell_utilities', 'tests'],
    install_requires = ['numpy', 'matplotlib', 'seaborn', 'scipy', 'numba', 'vum'],

    classifiers = [
        'Development Status :: 5 - Production/Stable',
        'Intended Audience :: Science/Research',
        'Operating System :: MacOS',
        'Operating System :: Unix',
        'Programming Language :: Python :: 3.11',
        'Topic :: Scientific/Engineering :: Physics',
    ],
)
<file_sep>from __future__ import annotations
import os
from vum import Vum
from .data_structures import (
    OrbitalParameters, Partition, Interaction, Configuration
)

def _sanity_checks(
    partition_proton: Partition,
    partition_neutron: Partition,
    partition_combined: Partition,
    interaction: Interaction,
):
    """
    A few different sanity checks to make sure that the new
    configurations are physical.
    """
    for i, configuration in enumerate(partition_proton.configurations):
        assert len(configuration.configuration) == interaction.model_space_proton.n_orbitals
        assert sum(configuration.configuration) == interaction.model_space_proton.n_valence_nucleons
        assert sum([n*orb.ho_quanta for n, orb in zip(configuration.configuration, interaction.model_space_proton.orbitals)]) == configuration.ho_quanta

        for j in range(i+1, partition_proton.n_configurations):
            assert partition_proton.configurations[i].configuration != partition_proton.configurations[j].configuration, f"Duplicate proton configs {i} and {j}!"

        for orbital, occupation in zip(interaction.model_space_proton.orbitals, configuration.configuration):
            """
            Check that max degeneracy is respected.
            """
            assert occupation <= (orbital.j + 1), "Occupation should never be lager than the max degeneracy!"

    assert (i + 1) == partition_proton.n_configurations

    for i, configuration in enumerate(partition_neutron.configurations):
        assert len(configuration.configuration) == interaction.model_space_neutron.n_orbitals
        assert sum(configuration.configuration) == interaction.model_space_neutron.n_valence_nucleons
        assert sum([n*orb.ho_quanta for n, orb in zip(configuration.configuration, interaction.model_space_neutron.orbitals)]) == configuration.ho_quanta

        for j in range(i+1, partition_neutron.n_configurations):
            assert partition_neutron.configurations[i].configuration != partition_neutron.configurations[j].configuration, f"Duplicate neutron configs {i} and {j}!"

        for orbital, occupation in zip(interaction.model_space_neutron.orbitals, configuration.configuration):
            """
            Check that max degeneracy is respected.
            """
            assert occupation <= (orbital.j + 1), "Occupation should never be lager than the max degeneracy!"

    assert (i + 1) == partition_neutron.n_configurations

    for configuration in partition_combined.configurations:
        p_idx, n_idx = configuration.configuration
        
        assert partition_combined.parity == partition_proton.configurations[p_idx].parity*partition_neutron.configurations[n_idx].parity
        assert configuration.ho_quanta == partition_proton.configurations[p_idx].ho_quanta + partition_neutron.configurations[n_idx].ho_quanta

def _calculate_configuration_parity(
    configuration: list[int],
    model_space: list[OrbitalParameters]
) -> int:
    """
    Calculate the parity of a configuration.

    Parameters
    ----------
    configuration : list[int]
        The configuration to calculate the parity of.
    
    model_space : list[OrbitalParameters]
        The model space orbitals to use for the calculation.
    """
    if not configuration:
        msg = "Configuration is empty! Undefined behaviour."
        raise ValueError(msg)

    parity: int = 1
    for i in range(len(configuration)):
        if not configuration[i]: continue   # Empty orbitals do not count towards the total parity.
        parity *= model_space[i].parity**configuration[i]

    return parity

def _prompt_user_for_interaction_and_partition(
    vum: Vum,
    is_compare_mode: bool,
):  
    """
    Parameterss
    -----------
    is_compare_mode : bool
        User will be prompted for two partition files and they will be
        returned as a list of the two names.
    """
    filenames_interaction = sorted([i for i in os.listdir() if i.endswith(".snt")])
    filenames_partition = sorted([i for i in os.listdir() if i.endswith(".ptn")])
    
    if not filenames_interaction:
        return f"No interaction file present in {os.getcwd()}. Exiting..."
    if not filenames_partition:
        return f"No partition file present in {os.getcwd()}. Exiting..."

    if len(filenames_interaction) == 1:
        filename_interaction = filenames_interaction[0]
        vum.screen.addstr(0, 0, f"{filename_interaction} chosen")
        vum.screen.refresh()

    elif len(filenames_interaction) > 1:
        interaction_choices: str = ""
        for i in range(len(filenames_interaction)):
            interaction_choices += f"{filenames_interaction[i]} ({i}), "
        
        vum.screen.addstr(vum.n_rows - 1 - vum.command_log_length - 1, 0, "Several interaction files detected.")
        vum.screen.addstr(vum.n_rows - 1 - vum.command_log_length, 0, interaction_choices)
        vum.screen.refresh()
        
        while True:
            ans = vum.input("Several interaction files detected. Please make a choice")
            if ans == "q": return False
            try:
                ans = int(ans)
            except ValueError:
                continue
            
            try:
                filename_interaction = filenames_interaction[ans]
            except IndexError:
                continue

            break

    vum.screen.addstr(0, 0, vum.blank_line)
    vum.screen.addstr(1, 0, vum.blank_line)
    vum.screen.refresh()
    
    if len(filenames_partition) == 1:
        if is_compare_mode: return "Cannot compare, only one partition file found!"
        filename_partition = filenames_partition[0]
        vum.screen.addstr(0, 0, f"{filename_partition} chosen")
        vum.screen.refresh()

    elif len(filenames_partition) > 1:
        if is_compare_mode: filename_partition: list[str] = []
        partition_choices: str = ""
        for i in range(len(filenames_partition)):
            partition_choices += f"{filenames_partition[i]} ({i}), "
        
        while True:
            vum.screen.addstr(vum.n_rows - 1 - vum.command_log_length - 1, 0, "Several partition files detected.")
            vum.screen.addstr(vum.n_rows - 1 - vum.command_log_length, 0, partition_choices)
            vum.screen.refresh()
            ans = vum.input("Several partition files detected. Please make a choice")
            if ans == "q": return False
            try:
                ans = int(ans)
            except ValueError:
                continue
            
            try:
                if is_compare_mode: filename_partition.append(filenames_partition[ans])
                else: filename_partition = filenames_partition[ans]
            except IndexError:
                continue
            
            if is_compare_mode and (len(filename_partition) < 2): continue
            break
    
    vum.screen.addstr(vum.n_rows - 1 - vum.command_log_length - 1, 0, vum.blank_line)
    vum.screen.addstr(vum.n_rows - 1 - vum.command_log_length, 0, vum.blank_line)
    vum.screen.refresh()

    return filename_interaction, filename_partition

# def is_configuration_below_energy_threshold(
#     interaction: Interaction,
#     threshold_energy: float,
#     proton_configuration: Configuration,
#     neutron_configuration: Configuration,
# ) -> bool:
#     """
#     Calculate the energy of a pn configuration. Return True if the
#     energy is lower than the threshold energy and False otherwise. Used
#     for monopole truncation.
#     """
#     energy = configuration_energy(
#         interaction = interaction,
#         proton_configuration = proton_configuration,
#         neutron_configuration = neutron_configuration,
#     )

#     return energy < threshold_energy

def configuration_energy(
    interaction: Interaction,
    proton_configuration: Configuration,
    neutron_configuration: Configuration,
) -> float:
    """
    Code from https://github.com/GaffaSnobb/kshell/blob/6e6edd6ac7ae70513b4bdaa94099cd3a3e32351d/bin/espe.py#L164
    """
    combined_configuration = proton_configuration.configuration + neutron_configuration.configuration
    if interaction.fmd_mass == 0:
        fmd = 1
    else:
        mass = interaction.n_core_neutrons + interaction.n_core_protons
        mass += interaction.model_space.n_valence_nucleons
        fmd = (mass/interaction.fmd_mass)**interaction.fmd_power

    assert len(interaction.spe) == len(combined_configuration)

    energy: float = 0.0
    
    for spe, occupation in zip(interaction.spe, combined_configuration, strict=True):
        energy += occupation*spe

    for i0 in range(len(combined_configuration)):
        for i1 in range(len(combined_configuration)):
            if i0 < i1: continue    # Why?
            
            elif i0 == i1:
                energy += combined_configuration[i0]*(combined_configuration[i0] - 1)*0.5*interaction.vm[i0, i0]*fmd
            else:
                energy += combined_configuration[i0]*combined_configuration[i1]*interaction.vm[i0, i1]*fmd

    return energy
<file_sep>from itertools import zip_longest
import numpy as np
import kshell_utilities.kshell_utilities

def test_file_read_levels():
    """
    Test that kshell_utilities.loadtxt successfully reads excitation
    energy output from KSHELL. Note that -1 spin states are supposed to
    be skipped.

    Raises
    ------
    AssertionError
        If the read values are not exactly equal to the expected values.
    """
    res = kshell_utilities.loadtxt(
        path = "summary_test_text_file.txt",
        load_and_save_to_file = False,
        old_or_new = "old"
    )

    E_expected = [
        -114.552, -114.044, -113.972, -113.898, -113.762, -113.602, -41.394,
        -33.378, -9.896, -9.052
    ]
    spin_expected = [
        2*2, 2*5, 2*4, 2*7, 2*2, 2*3, 2*0, 2*3, 2*3/2, 2*7/2
    ]
    parity_expected = [
        1, 1, 1, 1, 1, 1, 1, 1, -1, -1
    ]

    for calculated, expected in zip_longest(res.levels[:, 0], E_expected):
        msg = f"Error in Ex. Expected: {expected}, got: {calculated}."
        assert calculated == expected, msg

    for calculated, expected in zip_longest(res.levels[:, 1], spin_expected):
        msg = f"Error in spin. Expected: {expected}, got: {calculated}."
        assert calculated == expected, msg

    for calculated, expected in zip_longest(res.levels[:, 2], parity_expected):
        msg = f"Error in parity. Expected: {expected}, got: {calculated}."
        assert calculated == expected, msg

def test_int_vs_floor():
    """
    Check that floor yields the same result as casting to int. This does
    not work for negative values.
    """
    res = kshell_utilities.loadtxt(
        path = "summary_test_text_file.txt",
        load_and_save_to_file = False,
        old_or_new = "old"
    )

    res_1 = np.floor(res.transitions_BM1[:, 5])
    res_2 = res.transitions_BM1[:, 5].astype(int)

    for elem_1, elem_2 in zip(res_1, res_2):
        assert elem_1 == elem_2

def test_file_read_transitions():
    """
    Test that all values in res.transitions are as expected.
    """
    res = kshell_utilities.loadtxt(
        path = "summary_test_text_file.txt",
        load_and_save_to_file = False,
        old_or_new = "old"
    )

    # BE2
    # ---
    spin_initial_expected = np.array([5, 4, 4, 3, 9/2, 9/2])*2
    for i in range(len(res.transitions_BE2[:, 0])):
        msg = "BE2: Error in spin_initial."
        msg += f" Expected: {spin_initial_expected[i]}, got {res.transitions_BE2[i, 0]}."
        assert res.transitions_BE2[i, 0] == spin_initial_expected[i], msg

    parity_initial_expected = np.array([1, 1, 1, 1, -1, -1])
    for i in range(len(res.transitions_BE2[:, 1])):
        msg = "BE2: Error in parity_initial."
        msg += f" Expected: {parity_initial_expected[i]}, got {res.transitions_BE2[i, 1]}."
        assert res.transitions_BE2[i, 1] == parity_initial_expected[i], msg

    Ex_initial_expected = np.array([0.176, 0.182, 0.182, 0.221, 27.281, 27.281])
    for i in range(len(res.transitions_BE2[:, 3])):
        msg = "BE2: Error in Ex_initial."
        msg += f" Expected: {Ex_initial_expected[i]}, got {res.transitions_BE2[i, 3]}."
        assert res.transitions_BE2[i, 3] == Ex_initial_expected[i], msg

    spin_final_expected = np.array([6, 6, 5, 5, 13/2, 5/2])*2
    for i in range(len(res.transitions_BE2[:, 4])):
        msg = "BE2: Error in spin_final."
        msg += f" Expected: {spin_final_expected[i]}, got {res.transitions_BE2[i, 4]}."
        assert res.transitions_BE2[i, 4] == spin_final_expected[i], msg

    parity_final_expected = np.array([1, 1, 1, 1, -1, -1])
    for i in range(len(res.transitions_BE2[:, 5])):
        msg = "BE2: Error in parity_final."
        msg += f" Expected: {parity_final_expected[i]}, got {res.transitions_BE2[i, 5]}."
        assert res.transitions_BE2[i, 5] == parity_final_expected[i], msg

    Ex_final_expected = np.array([0, 0, 0.176, 0.176, 17.937, 18.349])
    for i in range(len(res.transitions_BE2[:, 7])):
        msg = "BE2: Error in Ex_final."
        msg += f" Expected: {Ex_final_expected[i]}, got {res.transitions_BE2[i, 7]}."
        assert res.transitions_BE2[i, 7] == Ex_final_expected[i], msg

    E_gamma_expected = np.array([0.176, 0.182, 0.006, 0.045, 9.344, 8.932])
    for i in range(len(res.transitions_BE2[:, 8])):
        msg = "BE2: Error in E_gamma."
        msg += f" Expected: {E_gamma_expected[i]}, got {res.transitions_BE2[i, 8]}."
        assert res.transitions_BE2[i, 8] == E_gamma_expected[i], msg

    B_decay_expected = np.array([157, 44.8, 3.1, 35, 0, 1.1])
    for i in range(len(res.transitions_BE2[:, 9])):
        msg = "BE2: Error in B_decay."
        msg += f" Expected: {B_decay_expected[i]}, got {res.transitions_BE2[i, 9]}."
        assert res.transitions_BE2[i, 9] == B_decay_expected[i], msg

    B_excite_expected = np.array([132.9, 31, 2.5, 22.3, 0.0, 1.9])
    for i in range(len(res.transitions_BE2[:, 10])):
        msg = "BE2: Error in B_excite."
        msg += f" Expected: {B_excite_expected[i]}, got {res.transitions_BE2[i, 10]}."
        assert res.transitions_BE2[i, 10] == B_excite_expected[i], msg

    # BM1
    # ---
    spin_initial_expected = np.array([2, 2, 1, 2, 7/2, 7/2])*2
    for i in range(len(res.transitions_BM1[:, 0])):
        msg = "BM1: Error in spin_initial."
        msg += f" Expected: {spin_initial_expected[i]}, got {res.transitions_BM1[i, 0]}."
        assert res.transitions_BM1[i, 0] == spin_initial_expected[i], msg

    parity_initial_expected = np.array([1, 1, 1, 1, -1, -1])
    for i in range(len(res.transitions_BM1[:, 1])):
        msg = "BM1: Error in parity_initial."
        msg += f" Expected: {parity_initial_expected[i]}, got {res.transitions_BM1[i, 1]}."
        assert res.transitions_BM1[i, 1] == parity_initial_expected[i], msg

    Ex_initial_expected = np.array([5.172, 17.791, 19.408, 18.393, 24.787, 24.787])
    for i in range(len(res.transitions_BM1[:, 3])):
        msg = "BM1: Error in Ex_initial."
        msg += f" Expected: {Ex_initial_expected[i]}, got {res.transitions_BM1[i, 3]}."
        assert res.transitions_BM1[i, 3] == Ex_initial_expected[i], msg

    spin_final_expected = np.array([0, 0, 2, 2, 5/2, 5/2])*2
    for i in range(len(res.transitions_BM1[:, 4])):
        msg = "BM1: Error in spin_final."
        msg += f" Expected: {spin_final_expected[i]}, got {res.transitions_BM1[i, 4]}."
        assert res.transitions_BM1[i, 4] == spin_final_expected[i], msg

    parity_final_expected = np.array([1, 1, 1, 1, -1, -1])
    for i in range(len(res.transitions_BM1[:, 5])):
        msg = "BM1: Error in parity_final."
        msg += f" Expected: {parity_final_expected[i]}, got {res.transitions_BM1[i, 5]}."
        assert res.transitions_BM1[i, 5] == parity_final_expected[i], msg

    Ex_final_expected = np.array([0, 0, 17.791, 17.791, 21.486, 21.564])
    for i in range(len(res.transitions_BM1[:, 7])):
        msg = "BM1: Error in Ex_final."
        msg += f" Expected: {Ex_final_expected[i]}, got {res.transitions_BM1[i, 7]}."
        assert res.transitions_BM1[i, 7] == Ex_final_expected[i], msg

    E_gamma_expected = np.array([5.172, 17.791, 1.617, 0.602, 3.301, 3.222])
    for i in range(len(res.transitions_BM1[:, 8])):
        msg = "BM1: Error in E_gamma."
        msg += f" Expected: {E_gamma_expected[i]}, got {res.transitions_BM1[i, 8]}."
        assert res.transitions_BM1[i, 8] == E_gamma_expected[i], msg

    B_decay_expected = np.array([20.5, 0.0, 5.7, 0.1, 0.069, 0.463])
    for i in range(len(res.transitions_BM1[:, 9])):
        msg = "BM1: Error in B_decay."
        msg += f" Expected: {B_decay_expected[i]}, got {res.transitions_BM1[i, 9]}."
        assert res.transitions_BM1[i, 9] == B_decay_expected[i], msg

    B_excite_expected = np.array([102.3, 0.0, 3.4, 0.1, 0.092, 0.617])
    for i in range(len(res.transitions_BM1[:, 10])):
        msg = "BM1: Error in B_excite."
        msg += f" Expected: {B_excite_expected[i]}, got {res.transitions_BM1[i, 10]}."
        assert res.transitions_BM1[i, 10] == B_excite_expected[i], msg

def _test_file_read_levels_jem():
    """
    For JEM syntax.

    NOTE: Currently not working because of https://github.com/GaffaSnobb/kshell-utilities/blob/d0bc1eae743e353671c562040c544b7df4c28d05/kshell_utilities/kshell_utilities.py#L437

    Raises
    ------
    AssertionError
        If the read values are not exactly equal to the expected values.
    """
    res = kshell_utilities.loadtxt(
        path = "summary_Zn60_jun45_jem_syntax.txt",
        load_and_save_to_file = False,
        old_or_new = "jem"
    )[0]

    E_expected = [
        -50.42584, -49.42999, -47.78510, -46.30908, -46.24604,
        -45.97338, -45.85816, -45.69709, -45.25258, -45.22663
    ]
    spin_expected = [
        0, 4, 8, 4, 0, 4, 12, 2, 8, 6
    ]
    parity_expected = [
        1, 1, 1, 1, 1, 1, 1, 1, 1, 1
    ]

    for calculated, expected in zip_longest(res.levels[:, 0], E_expected):
        msg = f"Error in Ex. Expected: {expected}, got: {calculated}."
        assert calculated == expected, msg

    for calculated, expected in zip_longest(res.levels[:, 1], spin_expected):
        msg = f"Error in spin. Expected: {expected}, got: {calculated}."
        assert calculated == expected, msg

    for calculated, expected in zip_longest(res.levels[:, 2], parity_expected):
        msg = f"Error in parity. Expected: {expected}, got: {calculated}."
        assert calculated == expected, msg

def _test_file_read_transitions_jem():
    """
    Test that all values in res.transitions are as expected.

    NOTE: Currently not working because of https://github.com/GaffaSnobb/kshell-utilities/blob/d0bc1eae743e353671c562040c544b7df4c28d05/kshell_utilities/kshell_utilities.py#L437
    """
    res = kshell_utilities.loadtxt(
        path = "summary_Zn60_jun45_jem_syntax.txt",
        load_and_save_to_file = False,
        old_or_new = "jem"
    )[0]

    E_gs = 50.42584
    
    # BE2
    # ---
    spin_initial_expected = np.array([
        4, 4, 4, 4, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 10, 10, 10, 10, 10
    ])
    for i in range(len(res.transitions_BE2[:, 0])):
        msg = "BE2: Error in spin_initial."
        msg += f" Expected: {spin_initial_expected[i]}, got {res.transitions_BE2[i, 0]}."
        assert res.transitions_BE2[i, 0] == spin_initial_expected[i], msg

    parity_initial_expected = np.array([
        -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1
    ])
    for i in range(len(res.transitions_BE2[:, 1])):
        msg = "BE2: Error in parity_initial."
        msg += f" Expected: {parity_initial_expected[i]}, got {res.transitions_BE2[i, 1]}."
        assert res.transitions_BE2[i, 1] == parity_initial_expected[i], msg
    
    Ex_initial_expected = E_gs - np.abs(np.array([
        -42.834, -42.061, -41.992, -41.210, -41.117, -40.802, -40.802, -40.802,
        -40.802, -40.802, -40.802, -40.802, -40.802, -50.426, -50.426, -50.426,
        -50.426, -50.426, -50.426, -50.426, -50.426, -43.879, -43.879, -43.879,
        -43.879, -43.879
    ]))
    for i in range(len(res.transitions_BE2[:, 3])):
        msg = "BE2: Error in Ex_initial."
        msg += f" Expected: {Ex_initial_expected[i]}, got {res.transitions_BE2[i, 3]}."
        assert res.transitions_BE2[i, 3] == Ex_initial_expected[i], msg

    spin_final_expected = np.array([
        0, 0, 0, 0, 0, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 10, 10, 10, 10, 10
    ])
    for i in range(len(res.transitions_BE2[:, 4])):
        msg = "BE2: Error in spin_final."
        msg += f" Expected: {spin_final_expected[i]}, got {res.transitions_BE2[i, 4]}."
        assert res.transitions_BE2[i, 4] == spin_final_expected[i], msg

    parity_final_expected = np.array([
        -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, +1, +1, +1, +1, +1,
        +1, +1, +1, +1, +1, +1, +1, +1
    ])
    for i in range(len(res.transitions_BE2[:, 5])):
        msg = "BE2: Error in parity_final."
        msg += f" Expected: {parity_final_expected[i]}, got {res.transitions_BE2[i, 5]}."
        assert res.transitions_BE2[i, 5] == parity_final_expected[i], msg

    Ex_final_expected = E_gs - np.abs(np.array([
        -40.802, -40.802, -40.802, -40.802, -40.802, -40.754, -40.666, -40.412,
        -40.375, -40.085, -40.060, -39.988, -39.802, -49.430, -46.309, -45.973,
        -44.988, -44.810, -44.615, -44.470, -43.619, -43.483, -42.618, -42.610,
        -42.173, -42.140
    ]))
    for i in range(len(res.transitions_BE2[:, 7])):
        msg = "BE2: Error in Ex_final."
        msg += f" Expected: {Ex_final_expected[i]}, got {res.transitions_BE2[i, 7]}."
        assert res.transitions_BE2[i, 7] == Ex_final_expected[i], msg

    E_gamma_expected = np.array([
        2.032, 1.259, 1.190, 0.408, 0.315, 0.048, 0.136, 0.390, 0.427, 0.717,
        0.742, 0.814, 1.000, 0.996, 4.117, 4.452, 5.437, 5.616, 5.811, 5.956,
        6.806, 0.396, 1.261, 1.269, 1.706, 1.739
    ])
    for i in range(len(res.transitions_BE2[:, 8])):
        msg = "BE2: Error in E_gamma."
        msg += f" Expected: {E_gamma_expected[i]}, got {res.transitions_BE2[i, 8]}."
        assert res.transitions_BE2[i, 8] == E_gamma_expected[i], msg

    B_decay_expected = np.array([
        0.51104070, 16.07765200, 4.57882840, 8.57092480, 0.46591730,
        0.13763550, 13.32346160, 2.80866440, 266.60020590, 2.26964570,
        1.78483210, 0.08375330, 23.60063140, 675.71309460, 0.32412310,
        0.97507780, 19.98328650, 0.28167560, 5.18973520, 0.05644800,
        0.90882360, 0.07640840, 1.66964180, 0.84907290, 0.28217940, 5.27947770
    ])
    for i in range(len(res.transitions_BE2[:, 9])):
        msg = "BE2: Error in B_decay."
        msg += f" Expected: {B_decay_expected[i]}, got {res.transitions_BE2[i, 9]}."
        assert res.transitions_BE2[i, 9] == B_decay_expected[i], msg

    B_excite_expected = np.array([
        2.55520380, 80.38826020, 22.89414210, 42.85462440, 2.32958690,
        0.02752710, 2.66469230, 0.56173280, 53.32004120, 0.45392910,
        0.35696640, 0.01675060, 4.72012620, 135.14261890, 0.06482460,
        0.19501550, 3.99665730, 0.05633510, 1.03794700, 0.01128960, 0.18176470,
        0.07640840, 1.66964180, 0.84907290, 0.28217940, 5.27947770
    ])
    for i in range(len(res.transitions_BE2[:, 10])):
        msg = "BE2: Error in B_excite."
        msg += f" Expected: {B_excite_expected[i]}, got {res.transitions_BE2[i, 10]}."
        assert res.transitions_BE2[i, 10] == B_excite_expected[i], msg

    # BM1
    # ---
    spin_initial_expected = np.array([
        2, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 2, 2, 2, 2, 2, 2
    ])
    for i in range(len(res.transitions_BM1[:, 0])):
        msg = "BM1: Error in spin_initial."
        msg += f" Expected: {spin_initial_expected[i]}, got {res.transitions_BM1[i, 0]}."
        assert res.transitions_BM1[i, 0] == spin_initial_expected[i], msg

    parity_initial_expected = np.array([
        -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, +1, +1, +1, +1, +1, +1, +1,
        +1, +1, +1
    ])
    for i in range(len(res.transitions_BM1[:, 1])):
        msg = "BM1: Error in parity_initial."
        msg += f" Expected: {parity_initial_expected[i]}, got {res.transitions_BM1[i, 1]}."
        assert res.transitions_BM1[i, 1] == parity_initial_expected[i], msg

    Ex_initial_expected = E_gs -  np.abs(np.array([
        -43.539, -42.609, -40.802, -40.802, -40.802, -40.802, -40.802, -40.802,
        -40.802, -40.802, -40.802, -42.467, -42.467, -42.467, -45.697, -45.001,
        -44.512, -44.439, -44.017, -43.294, -43.063
    ]))
    for i in range(len(res.transitions_BM1[:, 3])):
        msg = "BM1: Error in Ex_initial."
        msg += f" Expected: {Ex_initial_expected[i]}, got {res.transitions_BM1[i, 3]}."
        assert res.transitions_BM1[i, 3] == Ex_initial_expected[i], msg

    spin_final_expected = np.array([
        0, 0, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 0, 0, 0, 0, 0, 0, 0
    ])
    for i in range(len(res.transitions_BM1[:, 4])):
        msg = "BM1: Error in spin_final."
        msg += f" Expected: {spin_final_expected[i]}, got {res.transitions_BM1[i, 4]}."
        assert res.transitions_BM1[i, 4] == spin_final_expected[i], msg

    parity_final_expected = np.array([
        -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, +1, +1, +1, +1, +1, +1, +1,
        +1, +1, +1
    ])
    for i in range(len(res.transitions_BM1[:, 5])):
        msg = "BM1: Error in parity_final."
        msg += f" Expected: {parity_final_expected[i]}, got {res.transitions_BM1[i, 5]}."
        assert res.transitions_BM1[i, 5] == parity_final_expected[i], msg

    Ex_final_expected = E_gs -  np.abs(np.array([
        -40.802, -40.802, -40.786, -40.700, -40.178, -40.068, -39.878, -39.615,
        -39.376, -39.254, -39.167, -35.222, -35.217, -35.174, -42.351, -42.351,
        -42.351, -42.351, -42.351, -42.351, -42.351
    ]))
    for i in range(len(res.transitions_BM1[:, 7])):
        msg = "BM1: Error in Ex_final."
        msg += f" Expected: {Ex_final_expected[i]}, got {res.transitions_BM1[i, 7]}."
        assert res.transitions_BM1[i, 7] == Ex_final_expected[i], msg

    E_gamma_expected = np.array([
        2.737, 1.808, 0.016, 0.101, 0.624, 0.734, 0.924, 1.187, 1.425, 1.548,
        1.635, 7.246, 7.250, 7.293, 3.346, 2.649, 2.161, 2.087, 1.665, 0.943,
        0.711
    ])
    for i in range(len(res.transitions_BM1[:, 8])):
        msg = "BM1: Error in E_gamma."
        msg += f" Expected: {E_gamma_expected[i]}, got {res.transitions_BM1[i, 8]}."
        assert res.transitions_BM1[i, 8] == E_gamma_expected[i], msg

    B_decay_expected = np.array([
        0.00000640, 0.00013110, 0.00023580, 0.00232120, 6.16214940, 0.00008630,
        0.00388720, 0.00055320, 0.00001990, 0.72079000, 0.00458170, 0.00000170,
        0.00142630, 0.00016040, 0.00543100, 0.44305110, 0.00001740, 0.00056840,
        0.00007460, 0.10609890, 0.02403140
    ])
    for i in range(len(res.transitions_BM1[:, 9])):
        msg = "BM1: Error in B_decay."
        msg += f" Expected: {B_decay_expected[i]}, got {res.transitions_BM1[i, 9]}."
        assert res.transitions_BM1[i, 9] == B_decay_expected[i], msg

    B_excite_expected = np.array([
        0.00001940, 0.00039350, 0.00007860, 0.00077370, 2.05404980, 0.00002870,
        0.00129570, 0.00018440, 0.00000660, 0.24026330, 0.00152720, 0.00000050,
        0.00047540, 0.00005340, 0.01629320, 1.32915340, 0.00005220, 0.00170520,
        0.00022400, 0.31829690, 0.07209440
    ])
    for i in range(len(res.transitions_BM1[:, 10])):
        msg = "BM1: Error in B_excite."
        msg += f" Expected: {B_excite_expected[i]}, got {res.transitions_BM1[i, 10]}."
        assert res.transitions_BM1[i, 10] == B_excite_expected[i], msg

if __name__ == "__main__":
    test_file_read_levels()
    test_int_vs_floor()
    test_file_read_transitions()
    # test_file_read_levels_jem()
    # test_file_read_transitions_jem()