try:
    # framework is running
    from .startup_choice import *
except ImportError as _excp:
    # class is imported by itself
    if (
        'attempted relative import with no known parent package' in str(_excp)
        or 'No module named \'omfit_classes\'' in str(_excp)
        or "No module named '__main__.startup_choice'" in str(_excp)
    ):
        from startup_choice import *
    else:
        raise

from omfit_classes.utils_math import *
from omfit_classes.omfit_base import OMFITtree
from omfit_classes.omfit_osborne import OMFITpFile
from omfit_classes.omfit_eqdsk import OMFITgeqdsk
from omfit_classes.omfit_data import OMFITncDataset, OMFITncDynamicDataset, importDataset
from omas import ODS, omas_environment, cocos_transform, define_cocos
from omfit_classes.omfit_omas_utils import add_generic_OMFIT_info_to_ods

import inspect
import numpy as np
from uncertainties import unumpy, ufloat
from uncertainties.unumpy import std_devs, nominal_values
from scipy import constants

np.seterr(invalid='ignore', divide='ignore')

__all__ = ['OMFITprofiles', 'OMFITprofilesDynamic']

model_tree_species = ['e', '2H1', '4He2', '6Li3', '10B5', '12C6', '14N7', '20Ne10']
# fmt: off
model_tree_quantities = ['angular_momentum', 'angular_momentum_density', 'angular_momentum_density_{species}', 'angular_momentum_{species}',
'dn_{species}_dpsi', 'dT_{species}_dpsi', 'ELM_phase', 'ELM_since_last', 'ELM_until_next', 'fpol', 'jboot_sauter', 'lnLambda',
'n_{species}', 'nclass_sigma', 'nu_star_{species}', 'nu_{species}', 'omega_gyro_{species}_midplane', 'omega_LX_{species}_midplane',
'omega_N_{species}', 'omega_NTV0_{species}', 'omega_P_{species}', 'omega_plasma_{species}', 'omega_RX_{species}_midplane',
'omega_T_{species}', 'omega_tor_{species}', 'omega_tor_{species}_KDG', 'P_brem', 'P_rad', 'P_rad_cNi', 'P_rad_cW', 'P_rad_nNi', 'P_rad_nW',
'P_rad_ZnNi', 'P_rad_ZnW', 'p_tot', 'p_{species}', 'pres', 'psi', 'psi_n', 'q', 'R_midplane', 'resistivity', 'rho', 'T_{species}', 'time',
'Total_Zeff', 'V_pol_{species}_KDG', 'Zavg_Ni', 'Zavg_W', 'Zeff']
# fmt: on


def ugrad1(a2d):
    """
    Gradient along second axis with uncertainty propagation.
    :param a2d: 2D array or uarray
    :return:
    """
    if isinstance(a2d, DataArray):
        a2d = a2d.values
    if is_uncertain(a2d):
        dy = np.gradient(nominal_values(a2d), axis=1)
        ye = std_devs(a2d)
        sigma = np.zeros_like(ye)
        sigma[:, 1:-1] = 0.5 * np.sqrt(ye[:, :-2] ** 2 + ye[:, 2:] ** 2)
        sigma[:, 0] = 0.5 * np.sqrt(ye[:, 0] ** 2 + ye[:, 1] ** 2)
        sigma[:, -1] = 0.5 * np.sqrt(ye[:, -2] ** 2 + ye[:, -1] ** 2)
        result = unumpy.uarray(dy, sigma)
    else:
        result = np.gradient(a2d, axis=1)
    return result


def mZ(species):
    """
    Parse subscript strings and return ion mass and charge

    :param species: subscript strings such as `e`, `12C6`, `2H1`, 'fast_2H1`, ...

    :return: m and Z
    """
    species = str(species).replace('fast_', '')
    if species == 'e':
        Z = -1
        m = constants.m_e
    else:
        m = int(re.sub('([0-9]+)([a-zA-Z]+)([0-9]+)', r'\1', species))
        name = re.sub('([0-9]+)([a-zA-Z]+)([0-9]+)', r'\2', species)
        Z = int(re.sub('([0-9]+)([a-zA-Z]+)([0-9]+)', r'\3', species))
        m *= constants.m_u
    return m, Z


def get_species(derived):
    """
    Identify species and ions that have density information
    """
    species = []
    for key in list(derived.data_vars.keys()):
        if not re.match('^[nT](_fast)?_([0-9]+[a-zA-Z]+[0-9]|e)$', key):
            continue
        s = key.split('_')[-1]
        if '_fast_' in key:
            s = 'fast_' + s
        species.append(s)
    species = tolist(np.unique(species))
    ions = [s for s in species if s not in ['e']]
    ions_with_dens = [i for i in ions if 'n_' + i in derived]
    ions_with_fast = [i.replace('fast_', '') for i in ions if 'fast_' in i]
    return species, ions, ions_with_dens, ions_with_fast


class OMFITprofiles(OMFITncDataset):
    """
    Data class used by OMFITprofiles, CAKE and other
    OMFIT modules for storing experimental profiles data
    """

    def __init__(self, filename, data_vars=None, coords=None, attrs=None, comment=''):
        """
        :param filename: filename of the NetCDF file where data will be saved

        :param data_vars: see xarray.Dataset

        :param coords: see xarray.Dataset

        :param attrs: see xarray.Dataset

        :param comment: String that if set will show in the OMFIT tree GUI
        """
        self.dynaLoad = False
        super().__init__(filename, data_vars=data_vars, coords=coords, attrs=attrs)
        self.OMFITproperties['comment'] = comment

    @property
    def comment(self):
        return self.OMFITproperties['comment']

    @comment.setter
    def comment(self, comment):
        self.OMFITproperties['comment'] = comment

    def __tree_repr__(self):
        if self.comment:
            return self.__class__.__name__ + ': ' + self.comment, []
        else:
            return super().__tree_repr__()

    def to_omas(self, ods=None, times=None):
        """
        :param ods: ODS to which data will be appended

        :return: ods
        """

        if ods is None:
            ods = ODS()
            eq_times = None
        else:
            # Determine if equilibrium is avaliable and for what times.
            eq_times = ods.time('equilibrium') * 1e3  # ODS times are in s, and omfit_profiles are in ms.

        if 'device' in self.attrs:
            ods['dataset_description.data_entry.machine'] = self.attrs['device']
        if 'shot' in self.attrs:
            ods['dataset_description.data_entry.pulse'] = self.attrs['shot']

        # identify fitting coordinate
        for fit_coordinate in ['rho', 'psi_n', None]:
            if fit_coordinate in self.dims:
                break
        if fit_coordinate is None:
            raise ValueError("Fit coordinate should be 'rho' or 'psi_n'")

        # general info
        species, ions, ions_with_dens, ions_with_fast = get_species(self)
        nion = len(ions)
        if times is None:
            times = self['time'].values
        else:
            times = np.atleast_1d(times)

        # figure out time match between eq (if available) and profiles
        if eq_times is None:
            printw("No equilibrium data is avaliable to to_omas(). Some info will not be stored.")
        elif np.all([time in eq_times for time in times]):  # aka is all times have avaliable eq
            printd("Matching equilibrium times found")
        elif np.any([time in eq_times for time in times]):
            printw("Some time slices don't have corresponding equilibria!")
            printw("These time slices will be missing some information.")
        else:
            printw("No equilibrium data is avaliable to to_omas(). Some info will not be stored.")

        # assign both core_profies and edge_profiles but with data from different spatial ranges
        for coredgestring in ['core_profiles', 'edge_profiles']:
            if coredgestring.startswith('core'):
                derived = self.where(self[fit_coordinate] <= 1.0, drop=True)
                core = True
            else:
                derived = self.where(self[fit_coordinate] > 0.8, drop=True)
                core = False

            coredge = ods[coredgestring]
            prop = coredge['ids_properties']
            prop['comment'] = 'Data from OMFITprofiles.to_omas()'
            prop['homogeneous_time'] = True

            coredge['time'] = times / 1e3

            for ti, time in enumerate(times):
                profs = coredge[f'profiles_1d.{ti}']
                profs['time'] = time / 1e3

                # get corresponding eq and extract needed info
                geq = None
                R = None
                Bt = None
                Bp = None  # Make sure we don't accidentally use values from last timeslice
                if eq_times is not None and len(eq_times) > 0 and time in eq_times:
                    i_eq = np.where(time == eq_times)[0][0]  # just the first index
                    geq = OMFITgeqdsk('g0.0').from_omas(ods, time_index=i_eq)
                    # Not the most efficient but can access exiting functions better.
                    if fit_coordinate == 'psi_n':
                        psin = derived['psi_n'].values
                    else:
                        psin = derived['psi_n'].sel(time=time).values
                    psin_eq = np.linspace(0.0, 1.0, len(geq['PRES']))

                    # Re-interpolate
                    R = np.interp(psin, psin_eq, geq['fluxSurfaces']['midplane']['R'])
                    Bt = np.interp(psin, psin_eq, geq['fluxSurfaces']['midplane']['Bt'])
                    Bp = np.interp(psin, psin_eq, geq['fluxSurfaces']['midplane']['Bp'])

                for q in derived.variables:
                    fit = derived[q]
                    if q in ['time']:
                        continue
                    if q == 'rho':
                        if fit_coordinate == 'rho':
                            profs['grid.rho_tor_norm'] = derived['rho'].values
                        else:
                            profs['grid.rho_tor_norm'] = fit.sel(time=time).values
                    elif q == 'psi_n':
                        if fit_coordinate == 'psi_n':
                            profs['grid.rho_pol_norm'] = derived['psi_n'].values
                        else:
                            profs['grid.rho_pol_norm'] = fit.sel(time=time).values
                    elif q == 'n_e':
                        profs['electrons.density_thermal'] = fit.sel(time=time).values
                    elif q == 'T_e':
                        profs['electrons.temperature'] = fit.sel(time=time).values
                    elif q == 'omega_P_e' and core:  # this location is not valid OMAS location for edge profiles
                        profs['electrons.rotation.diamagnetic'] = fit.sel(time=time).values
                    elif '_' in q and q.split('_', 1)[1] in ions:
                        continue

                # thermal ions
                ni = 0
                for ion in ions[::-1]:
                    if ion == 'b' or ion.startswith('fast_'):
                        continue
                    profi = profs[f'ion.{ni}']
                    profi['density_thermal'] = derived[f'n_{ion}'].sel(time=time).values
                    if f'T_{ion}' in derived:
                        profi['temperature'] = derived[f'T_{ion}'].sel(time=time).values
                    ion_details = list(atomic_element(symbol=ion).values())[0]
                    profi['label'] = ion_details['symbol']
                    profi['z_ion'] = float(ion_details['Z_ion'])
                    profi['multiple_states_flag'] = 0
                    profi['element[0].atoms_n'] = 1
                    profi['element[0].z_n'] = float(ion_details['Z'])
                    profi['element[0].a'] = float(ion_details['A'])
                    profi['multiple_states_flag'] = 0
                    if f'V_tor_{ion}' in derived and not (f'omega_tor_{ion}' in derived and 'R_midplane' in derived):
                        profi['velocity.toroidal'] = derived[f'V_tor_{ion}'].sel(time=time).values
                    elif f'omega_tor_{ion}' in derived and 'R_midplane' in derived:
                        profi['velocity.toroidal'] = (
                            derived[f'omega_tor_{ion}'].sel(time=time) * derived['R_midplane'].sel(time=time)
                        ).values
                    if f'V_pol_{ion}' in derived:
                        profi['velocity.poloidal'] = derived[f'V_pol_{ion}'].sel(time=time).values
                    if core:  # extra rotation info for the core profiles. (Not valid nodes for edge)
                        if f'omega_P_{ion}' in derived:
                            profi['rotation.diamagnetic'] = derived[f'omega_P_{ion}'].sel(time=time).values
                        if f'omega_tor_{ion}' in derived:
                            profi['rotation_frequency_tor'] = derived[f'omega_tor_{ion}'].sel(time=time).values
                        if f'V_pol_{ion}' in derived and Bp is not None:
                            # Save to parallel streaming function, this will allow omegp to be calculated from ods
                            profi['rotation.parallel_stream_function'] = derived[f'V_pol_{ion}'].sel(time=time).values / Bp

                    # Advance ni; its important that if it is fast ion and the loop-iteration is skipped, then ni do not advance
                    ni += 1

                # fast ions
                for ion in ions:
                    if ion != 'b' and not ion.startswith('fast_'):
                        continue
                    # Get the 'base' ion for the fast population
                    if ion.startswith('fast_'):
                        base_ion = ion.replace('fast_', '')
                    elif ion == 'b':
                        base_ion == '2H1'  # back compat for '_b' notations
                    ion_details = list(atomic_element(symbol=base_ion).values())[0]

                    # Determin the corresponding ion index
                    ni = len(profs['ion'])
                    for nii in profs['ion']:
                        profi = profs[f'ion.{nii}']
                        if profi['label'] == ion_details['symbol']:
                            ni = nii
                            break
                    profi = profs[f'ion.{ni}']

                    # Add fast_ion data.
                    profi['density_fast'] = derived[f'n_{ion}'].sel(time=time).values
                    if f'p_{ion}' in derived:
                        pfast = derived[f'p_{ion}'].sel(time=time)  #'ion' here would have the form 'fast_2H1' for example
                    else:
                        pfast = derived[f'T_{ion}'].sel(time=time).values * constants.e * derived[f'n_{ion}'].sel(time=time).values
                    profi['pressure_fast_perpendicular'] = (
                        1.0 / 3.0 * pfast
                    )  # Assume isotropic fast ions. Also OMAS treats p_xxx_perp as pressure in one of the perp directions, I think.
                    profi['pressure_fast_parallel'] = 1.0 / 3.0 * pfast

                    # Attach atomic data (from base_ion)
                    profi['label'] = ion_details['symbol']
                    profi['z_ion'] = float(ion_details['Z_ion'])
                    profi['multiple_states_flag'] = 0
                    profi['element[0].atoms_n'] = 1
                    profi['element[0].z_n'] = float(ion_details['Z'])
                    profi['element[0].a'] = float(ion_details['A'])
                    profi['multiple_states_flag'] = 0

                if 'Total_Zeff' in derived:
                    profs['zeff'] = derived['Total_Zeff'].sel(time=time).values

        # Populate total pressure nodes under 'profiles_1d`
        ods.physics_core_profiles_pressures()
        # ods.physics_edge_profiles_pressures() # This function does not exist, but really should.
        return ods

    def model_tree_quantities(self):
        """
        Returns list of MDS+ model_tree_quantities for all species.
        The function will warn if some of the `model_tree_quantities` are missing in OMFIT-source/omfit/omfit_classes/omfit_profiles.py
        and the model tree should be updated

        :return: list of strings
        """
        new_model_tree_quantities = set(model_tree_quantities)
        for item in self.variables:
            match = False
            for s in model_tree_species:
                tmp = item.replace(f'_{s}', '_{species}')
                if tmp != item:
                    match = True
                    break
            if match:
                new_model_tree_quantities.add(tmp)
            else:
                new_model_tree_quantities.add(item)
        new_model_tree_quantities = sorted(list(new_model_tree_quantities), key=lambda x: x.lower())
        if new_model_tree_quantities != model_tree_quantities:
            import textwrap

            printe('WARNING!: Update model_tree_quantities in OMFIT-source/omfit/omfit_classes/omfit_profiles.py')
            printe('WARNING!: and update the OMFIT_PROFS MDS+ model tree')
            printe('-' * 140)
            printe('# fmt: off')
            printe(textwrap.fill(f'model_tree_quantities = {repr(new_model_tree_quantities)}', width=140))
            printe('# fmt: on')
            printe('-' * 140)

        quantities = []
        for item in new_model_tree_quantities:
            if '{' in item:
                for s in model_tree_species:
                    quantities.append(item.format(species=s))
            else:
                quantities.append(item)
        return quantities

    def create_model_tree(self, server, treename='OMFIT_PROFS'):
        """
        Generate MDS+ model tree

        :param server: MDS+ server

        :param treename: MDS+ treename
        """
        from omfit_classes.omfit_mds import OMFITmdsConnection

        conn = OMFITmdsConnection(server)

        quantities = {self.mds_translator(k): None for k in self.model_tree_quantities()}
        quantities['__content__'] = ''
        conn.create_model_tree(treename, '', quantities, clear_subtree=True)

    def to_mds(self, server, shot, treename='OMFIT_PROFS'):
        """
        This script writes the OMFITproflies datase to DIII-D MDS+ and updates d3drdb accordingly

        :param server: MDS+ server

        :param shot: shot to store the data to

        :param treename: MDS+ treename

        :return: runid, treename
        """
        from omfit_classes.omfit_mds import OMFITmdsConnection

        conn = OMFITmdsConnection(server)

        quantities = self.model_tree_quantities()

        # find next available runid in d3drdb for this shot
        from omfit_classes.omfit_rdb import OMFITrdb

        rdb = OMFITrdb(db='code_rundb', server='d3drdb', by_column=True)
        runs = rdb.select(f"select run_id from plasmas where shot={shot} and code_name='OMFITprofiles'", by_column=True)
        for run in range(1, 9999):
            runid = int('%06d%04d' % (shot, run))
            if not len(runs) or runid not in runs['run_id']:
                break
        print(f'Writing OMFITprofiles to MDS+ {runid}')

        # write to MDS+
        quantities = conn.write_dataset(
            treename=treename,
            shot=runid,
            subtree='',
            xarray_dataset=self,
            quantities=quantities,
            translator=lambda x: self.mds_translator(x),
        )
        conn.write(treename=treename, shot=runid, node='__content__', data=';'.join(quantities))

        # add data to d3drdb (after MDS+ so that this is done only if MDS+ write is successful)
        data = {
            'code_name': 'OMFITprofiles',
            'shot': shot,
            'experiment': 'DIII-D',
            'run_type': 'user',
            'tree': treename,
            'start_time': np.min(self['time'].values),
            'stop_time': np.max(self['time'].values),
            'mdsserver': server,
            'run_id': runid,
        }
        rdb.insert('plasmas', data, verbose=False)
        pprint(data)

        return runid, treename

    def mds_translator(self, inv=None):
        """
        Converts strings OMFITprofiles dataset keys to MDS+ nodes less than 12 chars long

        :param inv: tring to which to apply the transformation
                    if `None` the transformation is applied to all of the OMFITprofiles.model_tree_quantities for sanity check

        :return: transformed sting or if inv is None the `mapped_model_2_mds` and `mapped_mds_2_model` dictionaries
        """
        mapper = SortedDict()
        mapper['angular_momentum_density'] = 'mom_dens'
        mapper['angular_momentum'] = 'mom'
        mapper['midplane'] = 'mid'
        mapper['omega_gyro_'] = 'gyrof_'
        mapper['omega_'] = 'w_'
        mapper['2H1'] = 'D'
        mapper['4He2'] = 'He'
        mapper['6Li3'] = 'Li'
        mapper['10B5'] = 'B'
        mapper['12C6'] = 'C'
        mapper['14N7'] = 'N'
        mapper['20Ne10'] = 'Ne'
        mapper['since_last'] = '_last'
        mapper['until_next'] = '_next'

        if inv is not None:
            for match, sub in mapper.items():
                inv = inv.replace(match, sub)
            if len(inv) > 12:
                raise Exception(
                    f'MDS+ OMFITprofiles quantity is longer than 12 chars: {inv}\nUpdate the mds_translator function accordingly'
                )
            return inv
        else:
            model_tree_quantities = self.model_tree_quantities()
            mapped_model_2_mds = SortedDict()
            mapped_mds_2_model = SortedDict()
            for item0 in model_tree_quantities:
                item = item0
                for match, sub in mapper.items():
                    item = item.replace(match, sub)
                if len(item) > 12:
                    raise Exception(f'MDS+ string is longer than 12 chars: {item}')
                if item0 != item and item in model_tree_quantities:
                    raise Exception(f'MDS+ string shadows something else: {item}')
                if item in mapped_mds_2_model:
                    raise Exception(f'Multiple items map to the same quantity: {item0} {mapped_mds_2_model[item]}')
                mapped_model_2_mds[item0] = item
                mapped_mds_2_model[item] = item0
            return mapped_model_2_mds, mapped_mds_2_model

    def to_pFiles(self, eq, times=None, shot=None):
        """
        :param eq: ODS() or dict. (OMFITtree() is a dict)  Needs to contain equilibria information, either in the form
            of the ODS with needed eq already loaded, or as OMFITgeqdsk() objects in the Dict with the time[ms] as
            keys. Times for the eq need to be strict matches to profiles times coord.

        :param times: array like. time for which you would like p files to be generated.

        :param shot: int. shot number, only relevant in generating p file names.

        :return: OMFITtree() containing a series of OMFITpfile objs.
        """

        # Generate times if needed
        if times is None:
            try:
                times = self['time'].values
            except KeyError:
                # Just here to add a helpful hint
                printw("Looking like your OMFITprofiles obj is missing 'time'. Is it properly initialized?")
                raise

        # Get shot if it can be found
        if shot is None:
            if 'shot' in self.attrs:
                shot = self.attrs['shot']
            else:
                shot = 0  # Dummy number for p file names.

        # get good_times
        good_times = []
        if isinstance(eq, ODS):
            # Check times
            good_times = [t for t in times if t in eq.time('equilibrium')]  # ODS().time('...') will throw error if time is inconsistent

        elif isinstance(eq, dict):
            for i, time in enumerate(times):
                d3_time = int(time)
                if d3_time in eq:
                    good_times.append(time)
                else:
                    printw(f"Missing eq data for {time}, it will be skipped!")
        else:
            printw("Input arg 'eq' is in unrecognized format. This will fail!")

        good_times = [t for t in good_times if t in self['time'].values]
        good_times = array(good_times)

        if len(good_times) == 0:
            printw("No valid time found! Each timesilce needs profiles and equilibrium.")
            return
        else:
            printd(f"The following time was found to be good: {good_times}")
            printd("pFiles will be produced for these times only.")

        ods = ODS()
        # Now inject eq info, but only for good_times
        if isinstance(eq, ODS):
            prof_ods.equilibrium.time = np.array(good_times)
            for i, time in enumerate(good_times):
                j = where(eq.times('equilibrium') == time)
                ods[f'equilibrium.time_slice.{i}'] = eq[f'equilibrium.time_slice.{j}']

            prof_ods.physics_consistent_times()  # Generate missing time array
        elif isinstance(eq, dict):
            for i, time in enumerate(good_times):
                d3_time = int(time)
                ods = eq[d3_time].to_omas(ods=ods, time_index=i)
                # This is not very efficient, but does make to_omas() more standardized. Open to change as needs change.
        ods = self.to_omas(times=good_times, ods=ods)

        out_tree = OMFITtree()
        for i, time in enumerate(good_times):
            d3_time = int(time)
            out_tree[d3_time] = OMFITpFile(f'p{shot:06}.{d3_time:05}').from_omas(ods, time_index=i)

        return out_tree

    def get_xkey(self):
        """
        Get the key of the x-coord associated with this data array.

        :returns: str. key of the x-coord associated with the profiles, like 'psi_n'.
        """
        dims = list(self.dims.keys())
        return dims[dims.index('time') - 1]

    def diamagnetic_frequencies(self, spc, update=True):
        """
        Calculate the diamagnetic frequency, and its density / temperature components.

        Formula: \omega_P = -\frac{T}{nZe}\frac{dn}{d\psi} - \frac{1}{Ze}\frac{dT}{d\psi}

        :param spc: Species for which temperature is fit

        :param update: bool. Set to true to update self, if False, only returns and does not update self.
            Gradients, if missing, will always update though.

        """
        freqs = xarray.Dataset()

        m, Z = mZ(spc)

        # calculate gradients if not avaliable
        if f'dn_{spc}_dpsi' not in self and f'n_{spc}' in self:
            dn_dpsi = self.xderiv('n_' + spc, coord='psi')

        if f'dT_{spc}_dpsi' not in self and f'T_{spc}' in self:
            dT_dpsi = self.xderiv('T_' + spc, coord='psi')

        # density part
        if self.check_keys(keys=[f'n_{spc}', f'T_{spc}'] + [f'dn_{spc}_dpsi'], name=f'omega_N_{spc}'):
            omega_N = -self[f'dn_{spc}_dpsi'] * self[f'T_{spc}'] * constants.eV / (self[f'n_{spc}'] * Z * constants.e)
            omega_N.attrs['long_name'] = r"$\omega_N = -\frac{T}{nZe}\frac{dn}{d\psi}$"
            freqs.update({f'omega_N_{spc}': omega_N})
        # temperature part
        if self.check_keys(keys=[f'T_{spc}'] + [f'dT_{spc}_dpsi'], name=f'omega_T_{spc}'):
            omega_T = -self[f'dT_{spc}_dpsi'] * constants.eV / (Z * constants.e)
            omega_T.attrs['long_name'] = r"$\omega_T = -\frac{1}{Ze}\frac{dT}{d\psi}$"
            freqs.update({f'omega_T_{spc}': omega_T})
        # total
        if len(freqs.data_vars) == 2:
            omega_P = omega_N + omega_T
            omega_P.attrs['long_name'] = r"$\omega_{p," + spc + r"} = -\frac{T}{nZe}\frac{dn}{d\psi} -\frac{1}{Ze}\frac{dT}{d\psi}$"
            freqs.update({f'omega_P_{spc}': omega_P})

        if update:
            self.update(freqs)
        return freqs

    def xderiv(self, key, coord='psi', update=True):
        """
        Returns the derivative of the value corresponding to key on the spatial coordinate coord.

        :param key: str. The variable

        :param coord: str. The radial coordinate with respect to which the derivative is taken

        :param update: bool. Set to true to update self, if False, only returns and does not update self.

        :return: Dataset
        """
        result = xarray.Dataset()
        xkey = self.get_xkey()
        if not self.check_keys([key, coord], name=f'd{key}_d{coord}'):
            return result
        dc = np.gradient(np.atleast_2d(self[coord].values), axis=-1)
        dkey = 'd' + key + '_d' + xkey
        if dkey in self and any(isfinite(v) for v in nominal_values(self[dkey].values).flat) and coord != xkey:
            # if we have derivative UQ from the fit method, keep it and just swap coordinates
            dx = np.gradient(np.atleast_2d(self[xkey].values), axis=-1)
            dkdc = DataArray(self[dkey] * (dx / dc), coords=self[key].coords)
        else:
            # otherwise we have to calculate the derivative numerically - huge uncertainties when propagated
            try:
                dkdc = DataArray(ugrad1(self[key].values) / dc, coords=self[key].coords)
            except ZeroDivisionError:  # This happens if dc array have 0s and self[key].values is a uarray/has uncertainty.
                i_bad = dc == 0
                dc[i_bad] = 1
                dkdc_arr = ugrad1(self[key].values) / dc
                dkdc_arr[i_bad] = ufloat(np.nan, np.nan)
                dkdc = DataArray(ugrad1(self[key].values) / dc, coords=self[key].coords)
        if coord in ['rho', 'psi_n']:
            dkdc.attrs['units'] = self[key].attrs.get('units', '')
        dkdc.attrs['long_name'] = r'd{:} / d$\{:}$'.format(key, coord)
        result['d' + key + '_d' + coord] = dkdc

        if update:
            self.update(result)
        return result

    def check_keys(self, keys=[], name='', print_error=True):
        """
        Check to make sure required data is available
        """
        missing = []
        for k in keys:
            if not k in self:
                missing.append(k)
        if missing:
            if not name:
                name = 'value'
            if print_error:
                printw('  WARNING: Could not form {:}. Missing {:}'.format(name, ', '.join(['`%s`' % x for x in missing])))
            return False
        return True


class OMFITprofilesDynamic(OMFITncDynamicDataset):
    """
    Class for dynamic calculation of derived quantities

    :Examples:

    Initialize the class with a filename and FIT Dataset.
    >> tmp=OMFITprofiles('test.nc', fits=root['OUTPUTS']['FIT'], equilibrium=root['OUTPUTS']['SLICE']['EQ'], root['SETTINGS']['EXPERIMENT']['gas'])

    Accessing a quantity will dynamically calculate it.
    >> print tmp['Zeff']

    Quantities are then stored (they are not calculated twice).
    >> tmp=OMFITprofiles('test.nc',
                          fits=root['OUTPUTS']['FIT'],
                          equilibrium=root['OUTPUTS']['SLICE']['EQ'],
                          main_ion='2H1')
    >> uband(tmp['rho'],tmp['n_2H1'])
    """

    def __init__(self, filename, fits=None, equilibrium=None, main_ion='2H1', **kw):
        OMFITncDataset.__init__(self, filename, **kw)

        if fits:
            self.update(fits)

            # profile fit dimension
            dims = list(self.dims.keys())
            xkey = dims[dims.index('time') - 1]

            # collect some meta data about which particle species have what info available
            self.attrs['main_ion'] = str(main_ion)
            species, ions, ions_with_dens, ions_with_fast = get_species(self)
            self['species'] = DataArray(species, dims=['species'])
            ions += [self.attrs['main_ion']]
            self['ions'] = DataArray(ions, dims=['ions'])
            ions_with_dens += [self.attrs['main_ion']]
            self['ions_with_dens'] = DataArray(ions_with_dens, dims=['ions_with_dens'])
            ions_with_rot = [key.replace('omega_tor_', '') for key in self if key.startswith('omega_tor') and key != 'omega_tor_e']
            self['ions_with_rot'] = DataArray(ions_with_rot, dims=['ions_with_rot'])

            # interpolate other radial coordinates from equilibrium
            printd('- Interpolating equilibrium quantities')
            needed = {
                'avg_R',
                'avg_R**2',
                'avg_1/R',
                'avg_1/R**2',
                'avg_Btot',
                'avg_Btot**2',
                'avg_vp',
                'avg_q',
                'avg_fc',
                'avg_F',
                'avg_P',
                'geo_psi',
                'geo_R',
                'geo_Z',
                'geo_a',
                'geo_eps',
                'geo_vol',
            }
            eqcoords = needed.intersection(list(equilibrium.keys()))
            for meas in eqcoords:
                if 'profile_' + xkey != meas and 'profile_psi_n' in equilibrium[meas]:
                    yy = []
                    for t in self['time'].values:
                        eq_t = equilibrium.sel(time=t, method='nearest')
                        x = np.squeeze(eq_t['profile_' + xkey])
                        y = np.squeeze(nominal_values(eq_t[meas].values))
                        yy.append(interp1e(x, y)(self[xkey].values))

                        # Ensure that 'q' is not extrapolated outside the separatrix
                        if meas == 'profile_q':
                            mask = self[xkey].values > 1.0
                            yy[-1][mask] = np.nan
                    yy = np.array(yy)
                    key = meas.replace('profile_', '')
                    self[key] = DataArray(
                        yy, coords=[fits['time'], fits[xkey]], dims=['time', xkey], attrs=copy.deepcopy(equilibrium[meas].attrs)
                    )
            self.set_coords(eqcoords)
            # reassign global attrs clobbered when assigning eq DataArrays with attrs
            self.attrs['main_ion'] = str(main_ion)
            self.save()

        self.update_dynamic_keys(self.__class__)

    def __getitem__(self, key):
        # map specific quantities to class functions
        mapper = {}
        mapper['n_' + self.attrs['main_ion']] = 'calc_n_main_ion'
        mapper['T_' + self.attrs['main_ion']] = 'calc_T_main_ion'

        # resolve mappings
        if key not in self:
            if key in mapper:
                getattr(self, mapper[key])()
                if mapper[key] in self._dynamic_keys:
                    self._dynamic_keys.pop(self._dynamic_keys.index(mapper[key]))

        # return value
        return OMFITncDynamicDataset.__getitem__(self, key)

    def calc_n_main_ion(self):
        """
        Density of the main ion species.
        Assumes quasi-neutrality.

        :return: None. Updates the instance's Dataset in place.

        """
        main_ion = str(self.attrs['main_ion'])
        mg, zg = mZ(main_ion)
        nz = self['n_e']
        for key in self['ions_with_dens'].values:
            if key != main_ion:
                nz -= self['n_' + key].values * mZ(key)[1]
        self['n_' + main_ion] = nz / zg
        invalid = np.where(self['n_' + main_ion].values <= 0)[0]
        if len(invalid) > 0:
            printe('  Had to force main ion density to be always positive!')
            printe('  This will likely present a problem when running transport codes!')
            valid = np.where(self['n_' + main_ion].values > 0)[0]
            self['n_' + main_ion].values[invalid] = np.nanmin(self['n_' + main_ion].values[valid])

    def calc_T_main_ion(self):
        """
        Temperature of the main ion species.
        Assumes it is equal to the measured ion species temperature.
        If there are multiple impurity temperatures measured, it uses the first one.

        :return: None. Updates the instance's Dataset in place.

        """
        main_ion = str(self.attrs['main_ion'])
        impurities_with_temp = [k for k in self['ions'].values if k != 'b' and 'T_' + k in list(self.keys())]
        nwith = len(impurities_with_temp)
        if nwith == 0:
            raise OMFITexception("No main or impurity ion temperatures measured")
        if nwith > 1:
            printw(
                "WARNING: Multiple impurities temperatures measured, setting main ion temperature based on {:}".format(
                    impurities_with_temp[0]
                )
            )
        for ion in impurities_with_temp:
            self['T_' + main_ion] = self[f'T_{ion}'] * 1
            break

    def calc_Zeff(self):
        r"""
        Effective charge of plasma.

        Formula: Z_{eff} = \sum{n_s Z_s^2} / \sum{n_s Z_s}

        :return: None. Updates the instance's Dataset in place.

        """
        # calculate Zeff (not assuming quasi-neutrality)
        nz_sum = np.sum([self['n_' + i].values * mZ(i)[1] for i in self['ions_with_dens'].values], axis=0)
        nz2_sum = np.sum([self['n_' + i].values * mZ(i)[1] ** 2 for i in self['ions_with_dens'].values], axis=0)
        z_eff = nz2_sum / nz_sum + 0 * self['n_e'].rename('Zeff')
        z_eff.attrs['long_name'] = r'$Z_{eff}$'
        self['Zeff'] = z_eff

    def calc_Total_Zeff(self):
        r"""
        Effective charge of plasma.

        Formula: Z_{eff} = \sum{n_s Z_s^2} / \sum{n_s Z_s}

        :return: None. Updates the instance's Dataset in place.

        """
        main_ion = str(self.attrs['main_ion'])
        mg, zg = mZ(main_ion)
        nz = self['n_e']
        for key in self['ions_with_dens'].values:
            if key != main_ion:
                nz -= self['n_' + key].values * mZ(key)[1]
        self['n_' + main_ion] = nz / zg
        invalid = np.where(self['n_' + main_ion].values <= 0)[0]
        if len(invalid) > 0:
            printe('  Had to force main ion density to be always positive!')
            printe('  This will likely present a problem when running transport codes!')
            valid = np.where(self['n_' + main_ion].values > 0)[0]
            self['n_' + main_ion].values[invalid] = np.nanmin(self['n_' + main_ion].values[valid])


if __name__ == '__main__':

    # ensure that all specified model_tree_quantities can be translated to have <=12 chars
    for s in model_tree_species:
        for q in model_tree_quantities:
            item0 = q.format(species=s)
            item = OMFITprofiles.mds_translator(None, item0)
