"""
File Purpose: electric field in Bifrost
"""
import numpy as np
import xarray as xr

from ...errors import FormulaMissingError, LoadingNotImplementedError
from ...quantities import PlasmaDriftsLoader
from ...tools import alias, simple_property


### --------------------- BifrostEfieldLoader --------------------- ###

class BifrostEfieldLoader(PlasmaDriftsLoader):
    '''quantities related to electric field in Bifrost'''

    # # # Q & ETA # # #

    # Q & ETA FOR OHMIC TERM (E_etaJ) #
    @known_var(dims=['snap'])
    def get_qjoule(self):
        '''joule heating (?). eta * |J|**2. Directly from Bifrost aux.
        qjoule in aux is stored in raw units.
            for more aux units see: https://ita-solar.github.io/Bifrost/aux_variables/
        '''
        if 'qjoule' not in self.directly_loadable_vars():
            raise FormulaMissingError('qjoule when not saved to aux')
        ufactor = self.u('energy_density time-1', convert_from='raw')
        return self.load_maindims_var_across_dims('qjoule', u=ufactor, dims=['snap'])

    @known_var(deps=['qjoule', 'mod2_J'])
    def get_eta(self):
        '''eta (scalar), such that E = -u x B + eta J + ....'''
        return self('qjoule') / self('mod2_J')

    # ETA FOR HALL TERM (E_hall) #
    cls_behavior_attrs.register('eta_hall_mode', default='best')
    ETA_HALL_MODE_OPTIONS = {
        'best': '''use 'aux' if possible, else |B|/(|qe| * self('ne', ne_mode='best'))''',
        'aux': '''use 'eta_hall' from aux''',
        'ne_best': '''|B|/(|qe| * self('ne', ne_mode='best'))''',
        'neq': '''|B| / (|qe| * self('ne_neq'))''',
        'table': '''|B| / (|qe| * self('ne_fromtable'))''',
        'ne': '''|B| / (|qe| * self('ne')) using present value of self.ne_mode.''',
    }
    eta_hall_mode = simple_property('_eta_hall_mode', default='best',
        validate_from='ETA_HALL_MODE_OPTIONS',
        doc='''tells where to get eta_hall from. See self.ETA_HALL_MODE_OPTIONS for details.''')

    eta_hall_mode_explicit = alias('eta_hall_mode',
        doc='''self.eta_hall_mode, but if 'best' convert to the explicit mode to use.''')
    @eta_hall_mode_explicit.getter
    def eta_hall_mode_explicit(self):
        mode = self.eta_hall_mode
        if mode == 'best':
            if 'eta_hall' in self.directly_loadable_vars():
                mode = 'aux'
            else:
                mode = 'ne_best'
        return mode

    _ETA_HALL_MODE_EXPLICIT_TO_NE_DEPS = {
        'aux': [],
        'ne_best': ['mod_B', 'abs_qe', ('ne', {'ne_mode': 'best'})],
        'neq': ['mod_B', 'abs_qe', 'ne_neq'],
        'table': ['mod_B', 'abs_qe', 'ne_fromtable'],
        'ne': ['mod_B', 'abs_qe', 'ne'],
    }
    @known_var(dims=['snap'],
               attr_deps=[('eta_hall_mode_explicit', '_ETA_HALL_MODE_EXPLICIT_TO_NE_DEPS')])
    def get_eta_hall(self):
        '''eta_hall (scalar), such that E = -u x B + eta_hall J x Bhat + ....
        eta_hall in aux is stored in raw units.
        eta_hall is equivalent to |B| / (ne * |qe|).

        self.eta_hall_mode controls whether to load from aux or use a formula;
            see self.ETA_HALL_MODE_OPTIONS for options & descriptions.

        result will have 'eta_hall_mode' attr telling the mode used.
            (if 'best', will convert to the explicit mode used, for clarity.)
        '''
        mode = self.eta_hall_mode_explicit
        if mode == 'aux':
            if 'eta_hall' not in self.directly_loadable_vars():
                raise FormulaMissingError('eta_hall when not saved to aux')
            ufactor = self.u('e_field current_density-1', convert_from='raw')
            result = self.load_maindims_var_across_dims('eta_hall', u=ufactor, dims=['snap'])
        else:
            mode_to_ne_mode = {'ne_best': 'best', 'neq': 'neq', 'table': 'table', 'ne': self.ne_mode}
            result = self('(mod_B)/(abs_qe*ne)', ne_mode=mode_to_ne_mode[mode])
        return result.assign_attrs(eta_hall_mode=mode)

    # Q & ETA FOR AMBIPOLAR TERM (E_amb) #
    @known_var(dims=['snap'])
    def get_qjamb(self):
        '''ambipolar joule heating. E_amb dot J.
        qjamb in aux is stored in raw units.
        '''
        if 'qjamb' not in self.directly_loadable_vars():
            raise FormulaMissingError('qjamb when not saved to aux')
        ufactor = self.u('energy_density time-1', convert_from='raw')
        return self.load_maindims_var_across_dims('qjamb', u=ufactor, dims=['snap'])

    cls_behavior_attrs.register('eta_amb_mode', default='best')
    ETA_AMB_MODE_OPTIONS = {
        'best': '''use 'aux' if possible, else 'qJ'.''',
        'aux': '''use 'eta_amb' directly from aux''',
        'qJ': '''compute via J, B, and aux 'qjamb': qjamb = |E_amb dot J|
                --> eta_amb = qjamb / |(((J x Bhat) x Bhat) dot J)|''',
        'qJ2': '''compute estimate via J, B, and aux 'qjamb':
                --> eta_amb = qjamb / |J|**2
                (reasonable if J dot B is small).'''
    }
    eta_amb_mode = simple_property('_eta_amb_mode', default='best',
        validate_from='ETA_AMB_MODE_OPTIONS',
        doc='''tells where to get eta_amb from. See self.ETA_AMB_MODE_OPTIONS for details.''')

    eta_amb_mode_explicit = alias('eta_amb_mode',
        doc='''self.eta_amb_mode, but if 'best' convert to the explicit mode to use.''')
    @eta_amb_mode_explicit.getter
    def eta_amb_mode_explicit(self):
        mode = self.eta_amb_mode
        if mode == 'best':
            if 'eta_amb' in self.directly_loadable_vars():
                mode = 'aux'
            else:
                mode = 'qJ'
        return mode

    _ETA_AMB_MODE_EXPLICIT_TO_NE_DEPS = {
        'aux': [],
        'qJ': ['qjamb', 'J_crosscrossdotself_hat_B'],
        'qJ2': ['qjamb', 'mod2_J'],
    }
    @known_var(dims=['snap'],
               attr_deps=[('eta_amb_mode_explicit', '_ETA_AMB_MODE_EXPLICIT_TO_NE_DEPS')])
    def get_eta_amb(self):
        '''eta_amb (scalar), such that E = -u x B - eta_amb (J x Bhat) x Bhat + ....
        eta_amb in aux is stored in raw units.
        eta_amb physical formula is a bit complicated. See e.g. eqn3 in Martinez-Sykora 2017
            (https://iopscience.iop.org/article/10.3847/1538-4357/aa8866).

        self.eta_amb_mode controls whether to load from aux or use a formula;
            see self.ETA_AMB_MODE_OPTIONS for options & descriptions.

        result will have 'eta_amb_mode' attr telling the mode used.
            (if 'best', will convert to the explicit mode used, for clarity.)
        '''
        mode = self.eta_amb_mode_explicit
        if mode == 'aux':
            if 'eta_amb' not in self.directly_loadable_vars():
                raise FormulaMissingError('eta_amb when not saved to aux')
            ufactor = self.u('e_field current_density-1', convert_from='raw')
            result = self.load_maindims_var_across_dims('eta_amb', u=ufactor, dims=['snap'])
        elif mode == 'qJ':
            result = self('qjamb') / np.abs(self('J_crosscrossdotself_hat_B'))
        elif mode == 'qJ2':
            result = self('qjamb') / self('mod2_J')
        else:
            _known = list(self.ETA_AMB_MODE_OPTIONS.keys())
            raise LoadingNotImplementedError(f'eta_amb_mode: {mode!r}; expected one of {_known}')
        return result.assign_attrs(eta_amb_mode=mode)


    # # # E TERMS # # #

    @known_var(deps=['(SF_u)_cross_B'], aliases=['E_motional'])
    def get_E_uxB(self, *, _B=None):
        '''E_uxB = -u x B, the motional electric field. E = E_uxB + ...
        [EFF] for efficiency, can provide B if already known.
            CAUTION: if providing B, any missing components assumed to be 0.
        '''
        result = -self('(SF_u)_cross_B', _B=_B)
        # don't include fluid=SINGLE_FLUID in result.
        result = result.drop_vars('fluid')  # [TODO] handle 'fluid not in result coords' case?
        return result

    @known_var(deps=['eta', 'J'])
    def get_E_etaJ(self, *, _J=None):
        '''E_etaJ = eta * J. Goes into E = E_uxB + E_etaJ + ...
        [EFF] for efficiency, can provide J if already known.
        '''
        J = self('J') if _J is None else _J
        return self('eta') * J

    @known_var(deps=['eta_hall', 'J_cross_B'])
    def get_E_hall(self, *, _J=None, _B=None, _JxB=None, _mod_B=None):
        '''E_hall = eta_hall * J x B / |B|. Goes into E = E_uxB + E_hall + ...
        [EFF] for efficiency, can provide J, B, JxB, and/or _mod_B if already known.
            CAUTION: if providing vector, any missing components assumed to be 0.
        '''
        eta_hall = self('eta_hall')
        JxB = self('J_cross_B', _J=_J, _B=_B) if _JxB is None else _JxB
        mod_B = self('mod_B', _B=_B) if _mod_B is None else _mod_B
        return eta_hall * JxB / mod_B

    @known_var(deps=['eta_amb', 'J', 'B'])
    def get_E_amb(self, *, _J=None, _B=None, _JxB=None, _mod_B=None):
        '''E_amb = -eta_amb (J x B) x B / |B|**2. Goes into E = E_uxB + E_amb + ...
        [EFF] for efficiency, can provide J, B, JxB, and/or _mod_B if already known.
            CAUTION: if providing vector, any missing components assumed to be 0.
        '''
        eta_amb = self('eta_amb')
        JxB = self('J_cross_B', _J=_J, _B=_B) if _JxB is None else _JxB
        JxBxB = self('(J_cross_B)_cross_B', _val0=JxB, _B=_B)
        mod_B = self('mod_B', _B=_B) if _mod_B is None else _mod_B
        return -eta_amb * JxBxB / mod_B**2


    # # # TOTAL E # # #

    cls_behavior_attrs.register('E_from_aux', default=None)
    E_from_aux = simple_property('_E_from_aux', default=None,
        doc='''whether to load bifrost E from aux (ex ey ez).
        None --> True if ex ey ez all in aux, else False.
        True --> always load from ex ey ez aux (crash if not available)
        False --> never load from ex ey ez aux.''')

    E_from_aux_explicit = alias('E_from_aux',
        doc='''self.E_from_aux, but if None convert to the explicit mode to use.''')
    @E_from_aux_explicit.getter
    def E_from_aux_explicit(self):
        if self.E_from_aux is None:
            return all(x in self.directly_loadable_vars() for x in ['ex', 'ey', 'ez'])
        else:
            return self.E_from_aux

    @known_var(attr_deps=[('E_from_aux_explicit',
                           {True: ['E_uxB'], False: ['E_etaJ', 'E_hall', 'E_amb']})])
    def get_E_u0(self, *, _J=None, _B=None, _JxB=None, _mod_B=None):
        '''E without motional electric field contribution. E_u0 = E_etaJ + E_hall + E_amb.
        if self.E_from_aux_explicit, instead computed via: E_u0 = E (from aux ex ey ez) - E_uxB.

        [EFF] for efficiency, can provide J, B, JxB, and/or _mod_B if already known.
            CAUTION: if providing vector, any missing components assumed to be 0.
        '''
        from_aux = self.E_from_aux_explicit
        if from_aux:
            result = self('E-E_uxB')
        else:
            with self.using(component=None):  # all 3 vector components here
                J = self('J') if _J is None else _J
                B = self('B') if _B is None else _B
                JxB = self('J_cross_B', _J=_J, _B=_B) if _JxB is None else _JxB
                mod_B = self('mod_B', _B=_B) if _mod_B is None else _mod_B
            E_etaJ = self('E_etaJ', _J=J)
            E_hall = self('E_hall', _J=J, _B=B, _JxB=JxB, _mod_B=mod_B)
            E_amb  = self('E_amb',  _J=J, _B=B, _JxB=JxB, _mod_B=mod_B)
            result = E_etaJ + E_hall + E_amb
        return result.assign_attrs(E_from_aux=from_aux)

    @known_var(deps=['E_uxB', 'E_u0'])
    def get_E(self):
        '''electric field from Bifrost. E = E_uxB + E_u0.
        loaded from aux ex ey ez if self.E_from_aux_explicit, else from E_uxB + E_u0.
        '''
        from_aux = self.E_from_aux_explicit
        if from_aux:
            result = self.load_maindims_var_across_dims('e', u='e_field', dims=['snap', 'component'])
        else:
            with self.using(component=None):  # all 3 vector components here
                B = self('B')
            E_uxB = self('E_uxB', _B=B)
            E_u0  = self('E_u0',  _B=B)
            result = E_uxB + E_u0
        return result.assign_attrs(E_from_aux=from_aux)

    E_UN0_MODE_OPTIONS = {**PlasmaDriftsLoader.E_UN0_MODE_OPTIONS,
        'E0_EetaJ_perpB': '''E_un0 = E0_un0_perpB + E_etaJ_perp_B.
            The E0 is E from zeroth-order multifluid physics.
            The E_etaJ is the ohmic term. Physically, ohmic may be included in E0,
            however, in MHD simulations ohmic is primarily from numerical effects,
            and much larger than the physical contribution.''',
        'E0_Eauxu0_perpB': '''E_un0 = E0_un0_perpB + E_u0_perp_B, with E_u0 from aux.
            The E0 is E from zeroth-order multifluid physics.
            The Eaux_u0 is E from aux (ex, ey, ez), minus the motional term, E_uxB.
            This makes some physical sense to do if E from aux includes only E_etaJ,
            which occurs e.g. if Hall & ambipolar physics are disabled,
            in which case the result is conceptually similar to when mode='E0_EetaJ_perpB'.
            Will crash if E not included in aux.''',
    }

    _E_UN0_MODE_TO_TYPE = {**PlasmaDriftsLoader._E_UN0_MODE_TO_TYPE,
        'E0_EetaJ_perpB': 'E0_EetaJ_perpB',
        'E0_Eauxu0_perpB': 'E0_Eauxu0_perpB',
    }

    _E_UN0_TYPE_TO_DEPS = {**PlasmaDriftsLoader._E_UN0_TYPE_TO_DEPS,
        'E0_EetaJ_perpB': ['E0_un0_perpB', 'E_etaJ_perp_B'],
        'E0_Eauxu0_perpB': ['E0_un0_perpB', ('E_u0_perp_B', {'E_from_aux': True})],
    }

    @known_var(value_deps=[('E_un0_type', '_E_UN0_TYPE_TO_DEPS')])
    def get_E_un0(self):
        '''electric field in u_n=0 frame.
        Result depends on self.E_un0_mode; see help(type(self).E_un0_mode) for details.
        Note: using E_un0_mode=None (the default) and self.assume_un='u',
            is equivalent to using E_un0_mode = 'un=u', which will just give self('E_u0').
        Else, E_un0 = E + u_n x B, if can get u_n (crash if can't get u_n).
        '''
        kind = self('E_un0_type').item()
        if kind == 'E+unxB':  # slightly more helpful error than super(); u_neutral will probably crash.
            try:
                un = self('u_neutral', component=None)  # check if can get u_n... probably cant...
            except NotImplementedError:
                errmsg = 'E_un0, when E_un0_mode=None, u_neutral unknown, and self.assume_un != "u".'
                raise FormulaMissingError(errmsg)
            else:
                result = self('E')
                if np.any(un != 0):
                    result = result + self('u_neutral_cross_B', _u_neutral=un)
        elif kind == 'E0_EetaJ_perpB':
            E0_perpB = self('E0_un0_perpB')
            E_etaJ_perpB = self('E_etaJ_perp_B')
            result = E0_perpB + E_etaJ_perpB
        elif kind == 'E0_Eauxu0_perpB':
            E0_perpB = self('E0_un0_perpB')
            E_u0_perpB = self('E_u0_perp_B', E_from_aux=True)
            result = E0_perpB + E_u0_perpB
        else:
            result = super().get_E_un0()
        return result

    @known_var
    def get_E_un0_type(self):
        '''string telling method that will be used to get E_un0. Based on self.E_un0_mode and assume_un.
        possible results include (but may be adjusted further by subclasses):
            'nan' <--> will crash. (e.g. this occurs if E_un0_mode = 'E0_perpmodB_min')
            'E+unxB' <--> self('E') + self('u_neutral_cross_B')
            'E_u0' <--> self('E_u0')
            'E+uxB' <--> self('E') + self('u_cross_B')
            'E' <--> self('E')
            'E0_perpB' <--> self('E0_un0_perpB')
            'E0_EetaJ_perpB' <--> self('E0_un0_perpB') + self('E_etaJ_perpmod_B')
            'E0_Eauxu0_perpB' <--> self('E0_un0_perpB') + self('E_u0_perpmod_B', E_from_aux=True)
        '''
        if self.E_un0_mode is None and self.assume_un == 'u':
            return xr.DataArray('E_u0')
        else:  # super includes lookup in _E_UN0_MODE_TO_TYPE so it includes new types defined above.
            return super().get_E_un0_type()
