"""
File Purpose: base quantities for MuramCalculator
"""

from ...mhd import MhdBasesLoader
import xarray as xr

### --------------------- MuramBasesLoader --------------------- ###

class MuramBasesLoader(MhdBasesLoader):
    '''base quantities based on Muram output.'''
    @known_var(dims=['snap'])
    def get_r(self):
        '''mass density. (directly from Muram)
        assumes single fluid mode, i.e. result corresponds to the single fluid from Muram.
        '''
        self.assert_single_fluid_mode('r')
        return self.load_maindims_var_across_dims('result_prim_r', u='mass_density', dims=['snap'])

    @known_var(dims=['snap', 'component'])
    def get_u(self):
        '''velocity. (directly from Muram)
        assumes single fluid mode, i.e. result corresponds to the single fluid from Muram.
        '''
        self.assert_single_fluid_mode('u')
        return self.load_maindims_var_across_dims('result_prim_u', u='speed', dims=['snap', 'component'])

    @known_var(dims=['snap'])
    def get_e(self):
        '''energy density. (directly from Muram)
        Per unit volume, e.g. the SI units would be Joules / meter^3.
        assumes single fluid mode, i.e. result corresponds to the single fluid from Muram.
        '''
        self.assert_single_fluid_mode('e')
        return self.load_maindims_var_across_dims('result_prim_e', u='energy_density', dims=['snap'])

    @known_var(dims=['snap', 'component'])
    def get_B(self):
        '''magnetic field. (directly from Muram)'''
        return self.load_maindims_var_across_dims('result_prim_b', u='b', dims=['snap', 'component'])

    @known_var
    def get_gamma(self):
        '''adiabatic index. Assumed to be 5/3.'''
        return 5/3.

    @known_pattern(r'd([xyz])')
    def get_ds(self, var, *, _match=None):
        '''dx (or dy, or dz) --> grid cell size along the x (or y, or z) axis.
        result is probably a 0D xr.DataArray.

        Caution: this is the "underlying grid cell spacing", not numerically altered by self.slices.
        See also: self.get_maindims_dim_dslice(), which is more suitable if integrating.
        '''
        x, = _match.groups()   
        dx = self.params[f'd{x}'] * self.u('length')
        return xr.DataArray(dx)

    def get_maindims_dgrid(self, dim):
        '''return grid spacing associated with dim values (dim in self.maindims).
        Result is a 1D xarray.DataArray, with coords according to self.get_maindims_coords()[dim].
        Result's values use self.units; coords use self.coords_units (default==self.units).
            (e.g., units='cgs', coord_units='si' --> values are in cgs, but coords are in SI.)

        Caution: this is the "underlying grid cell size", regardless of self.slices.
            E.g., slicing with slice(None,None,4) will NOT lead to 4 times larger result.

        For Muram, this value comes from self.params.
        '''
        self._assert_is_maindim(dim)
        ds = self(f'd{dim}')
        coord = self.get_maindims_coord_xarrays(dim)
        result = xr.ones_like(coord) * ds
        return result
