"""
File Purpose: solving tfbi
"""
import os

from ...tools import ImportFailed
try:
    import tfbi_theory as tt
except ImportError as err:
    tt = ImportFailed("tfbi_theory", err=err, locals=locals(), abbrv='tt')

from ...defaults import DEFAULTS
from ...errors import InputError, InputMissingError
from ...tools import (
    format_docstring, alias, simple_property,
    UNSET,
    ProgressUpdater,
    xarray_mergeload,
)

_paramdocs_tfbi_solving = {
    'solving_pattern': '''import tfbi_theory as tt
        import PlasmaCalcs as pc
        cc = ... # any PlasmaCalculator object from PlasmaCalcs.
        ds0 = cc.tfbi_ds()
        kp = tt.kPickerLowres(ds0)
        dsk = kp.get_ds()   # copy of ds0, but with ds['k'] = k from kPicker.
        drel = tt.TfbiDisprelC.from_ds(dsk)
        dsR = drel.solve()  # copy of dsk, but with ds['omega'] = solution to TFBI theory!''',
    'solving_pattern_notes': '''if cc has more than ~4 ions, you will want to drop some or group them somehow.
                e.g. for MhdMultifluidCalculator, before calling tfbi_ds():
                    cc.use_mix_heavy_ions(0.3, m_mean_mode='density')
            You can also pick ions directly during cc.tfbi_ds().
                e.g. for Bifrost chromosphere analysis, where n[He_II] < 1e-6 * ne, I use:
                    cc.tfbi_ds(ions=[i for i in cc.fluids if i.q==1 and i!='He_II'])
            After finishing solving, you might want to save the result,
                e.g. dsR.pc.save('filename') saves result to 'filename.pcxarr';
                can load it later via pc.xarray_load('filename.pcxarr').
            If you are solving across "many" points (e.g., more than ~5000),
                you might want to use chunking instead. see TfbiChunkSolver;
                also note cc.tfbi_solver(...) automatically uses TfbiChunkSolver
                if appropriate, i.e. if cc.chunks is set (where cc is a TfbiLoader).''',
    'cc': '''PlasmaCalculator
        PlasmaCalculator object used to load the data.
        Should be a TfbiLoader subclass. (PlasmaCalculator satisfies this by default,
            assuming successful import SymSolver and import tfbi_theory.)''',
    'ions': f'''None or specifier of multiple fluids (e.g. slice, or list of strs)
        None --> use cc.fluids.ions()
        ions are determined when called, not during __init__.
        print warning if this specifies more than DEFAULTS.ADDONS.TFBI_MAX_NUM_IONS ions
            (default: {DEFAULTS.ADDONS.TFBI_MAX_NUM_IONS}), because then solving will be slow and may be inaccurate.''',
    'kres': ''''low', 'mid', or 'high'
        resolution in k-space. Tells which self.kPicker_cls to use.
        'low' --> tfbi_theory.kPickerLowres. Recommended if solving across many (e.g. >1000) points.
        'mid' --> tfbi_theory.kPickerMidres. Recommended if solving across a few (e.g. 10 to 100) points.
        'high' --> tfbi_theory.kPickerHighres. Recommended if solving at only 1 point.''',
    'mod_lmod_ang': '''UNSET or dict
        passed directly to kPicker if provided. Can specify k values other than the defaults.
        see help(self.kPicker_cls) for more details.''',
    'tfbi_all': '''bool
        whether to compute all relevant tfbi vars, ds0 = cc('tfbi_all').
        False --> compute only the necessary vars, ds0 = cc('tfbi_inputs').''',
    'drel_cls': '''None, str, or class
        tfbi_theory class to use for solving TFBI theory.
        None --> use self.drel_cls default: tt.TfbiDisprelC
        str --> use getattr(tt, drel_cls) to get the class.''',
}


### --------------------- TfbiSolver --------------------- ###

@format_docstring(**_paramdocs_tfbi_solving)
class TfbiSolver():
    '''high-level interface for solving TFBI theory across many physical parameters.

    Call TfbiSolver to solve TFBI.
    Example:
        import PlasmaCalcs as pc
        cc = pc.PlasmaCalculator(...)   # <-- your PlasmaCalculator of choice
        solver = pc.TfbiSolver(cc)
        solution = solver()             # alias: solver.solve()
    solution is an xarray.Dataset with all relevant quantities (see TfbiLoader.get_tfbi_all),
        and 'omega' telling roots with largest imaginary part.

    If you need more precise control over the solving process, use the pattern:
        {solving_pattern}

        Notes:
            {solving_pattern_notes}

    TfbiSolver internally stores self.cc, ds0, kp, dsk, drel, and dsR,
        as defined by the pattern above.

    cc: {cc}
    ions: {ions}
    kres: {kres}
    mod, lmod, ang: {mod_lmod_ang}
    tfbi_all: {tfbi_all}
    drel_cls: {drel_cls}
    '''
    kPickerLowres_cls_name = 'kPickerLowres'  # used to make self.kPicker_cls if kres='low'
    kPickerMidres_cls_name = 'kPickerMidres'  # used to make self.kPicker_cls if kres='mid'
    kPickerHighres_cls_name = 'kPickerHighres'  # used to make self.kPicker_cls if kres='high'

    @property
    def kPicker_cls(self):
        '''kPicker class to use for choosing wavevectors to consider. Depends on self.kres:
        'low' --> tfbi_theory.kPickerLowres
        'mid' --> tfbi_theory.kPickerMidres
        'high' --> tfbi_theory.kPickerHighres
        '''
        if self.kres == 'low':
            return getattr(tt, self.kPickerLowres_cls_name)
        elif self.kres == 'mid':
            return getattr(tt, self.kPickerMidres_cls_name)
        elif self.kres == 'high':
            return getattr(tt, self.kPickerHighres_cls_name)
        else:
            raise InputError(f"Invalid kres. Expected 'low' or 'high'; got {self.kres!r}")

    drel_cls_name = 'TfbiDisprelC'  # used to make self.drel_cls.
    # (drel_cls = tt.TfbiDisprelC here would crash immediately if tt import failed.
    # By using drel_cls property, we avoid this crash until it's actually relevant.)
    drel_cls = simple_property('_drel_cls', setdefaultvia='_default_drel_cls',
        doc='''class to use for solving TFBI theory. Default: tt.TfbiDisprelC.''')
    def _default_drel_cls(self):
        '''default drel_cls. getattr(tt, self.drel_cls_name)'''
        return getattr(tt, self.drel_cls_name)

    def __init__(self, cc, ions=None, *, kres='low', mod=UNSET, lmod=UNSET, ang=UNSET, tfbi_all=True,
                 drel_cls=None):
        self.cc = cc
        self.ions = ions
        self.kres = kres
        self.mod = mod
        self.lmod = lmod
        self.ang = ang
        self.tfbi_all = tfbi_all
        if drel_cls is not None:
            if isinstance(drel_cls, str):
                self.drel_cls = getattr(tt, drel_cls)
            elif isinstance(drel_cls, type):
                self.drel_cls = drel_cls
            else:
                raise InputError(f"Invalid drel_cls. Expected None, str, or class; got {drel_cls!r}")

    @property
    def ions_explicit(self):
        '''list of ions from self.cc which would be used during self.cc.tfbi_ds()'''
        if self.ions is None:
            result = self.cc.fluids.ions()
        else:
            result = self.cc.fluids.get(self.ions)
        return result

    def _warn_if_too_many_ions(self):
        '''print warning if self.ions specifies too many ions.'''
        default_nmax = DEFAULTS.ADDONS.TFBI_MAX_NUM_IONS
        ions = self.ions_explicit
        if len(ions) > default_nmax:
            warnmsg = (f">>> Warning: TfbiSolver has {len(ions)} ions, which may be slow and inaccurate. <<<\n"
                       f"    (To avoid this warning, either use fewer ions, or "
                            f"increase DEFAULTS.ADDONS.TFBI_MAX_NUM_IONS (={default_nmax}, currently).)")
            print(warnmsg)

    __call__ = alias('solve')

    @format_docstring(**_paramdocs_tfbi_solving, sub_ntab=1)
    def solve(self, *, verbose=True, **kw_growth_root):
        '''solve TFBI theory. Assigns self.ds0, kp, dsk, drel, and dsR.
        Does not do any caching at this level; if the code crashes you will need to restart it completely.
        For more precise control including possibility for caching, use the pattern directly:
            {solving_pattern}

        verbose: bool
            whether to print progress updates (highly recommended).
        additional kwargs get passed directly to drel.solve().
            (options include: ncpu, ncoarse, careful)
        '''
        updater = ProgressUpdater(print_freq=0 if verbose else -1)  # print all statements if verbose, else nothing.
        self._warn_if_too_many_ions()
        # get ds0 -- manage progress updates here.
        ds0_updater = ProgressUpdater(print_freq=0 if verbose else -1)
        ds0_updater.print(f'getting ds0 via {type(self.cc).__name__}.tfbi_ds() ...')
        self.ds0 = self.cc.tfbi_ds(ions=self.ions, all=self.tfbi_all)
        ds0_updater.finalize(f'solver.ds0 = {type(self.cc).__name__}.tfbi_ds()', end='\n')
        # picking k is usually fast; no need for progress updates.
        kw_kPicker = dict(mod=self.mod, lmod=self.lmod, ang=self.ang)
        kw_kPicker = {k: v for k, v in kw_kPicker.items() if v is not UNSET}
        self.kp = self.kPicker_cls(self.ds0, **kw_kPicker)
        # follow the rest of the pattern. tt manages its own progress updates.
        self.dsk = self.kp.get_ds()
        self.drel = self.drel_cls.from_ds(self.dsk)
        self.dsR = self.drel.solve(verbose=verbose, **kw_growth_root)
        updater.print_clear(force=verbose)  # clear print statements on current line, if we're printing anything.
        updater.finalize(f'{type(self).__name__}.solve() (all steps from start to finish)', always=True)
        return self.dsR

    solution = alias('dsR', doc='''result of self.solve(); alias to self.dsR.''')
    solved = property(lambda self: hasattr(self, 'dsR'), doc='''tells whether self.solution exists''')

    def _repr_contents(self):
        '''return list of strings for contents inside of __repr__'''
        contents = []
        cc_str = f'title={self.cc.title!r}' if hasattr(self.cc, 'title') else '...'
        contents.append(f'cc={type(self.cc).__name__}({cc_str})')
        if self.ions is not None:
            contents.append(f'ions={self.ions!r}')
        if isinstance(self.mod, dict) and len(self.mod) > 0:
            contents.append(f'mod={self.mod!r}')
        if isinstance(self.lmod, dict) and len(self.lmod) > 0:
            contents.append(f'lmod={self.lmod!r}')
        if isinstance(self.ang, dict) and len(self.ang) > 0:
            contents.append(f'ang={self.ang!r}')
        if self.tfbi_all != True:
            contents.append(f'tfbi_all={self.tfbi_all}')
        contents.append(f'solved={self.solved!r}')
        return contents

    def __repr__(self):
        contents = self._repr_contents()
        return f'{type(self).__name__}({", ".join(contents)})'


@format_docstring(**_paramdocs_tfbi_solving, sub_ntab=1)
class TfbiChunkSolver(TfbiSolver):
    '''high-level interface for solving TFBI across many physical parameters, using chunks.
    Chunking details are determined by cc.chunks; see help(cc.chunker) for more info.
    Aside from the chunking, this works very similarly to TfbiSolver.

    Call TfbiChunkSolver to solve TFBI.
    Example:
        import PlasmaCalcs as pc
        cc = pc.PlasmaCalculator(...)   # <-- your PlasmaCalculator of choice
        cc.chunks = dict(...)  # <-- your chunking parameters of choice
        solver = pc.TfbiChunkSolver(cc)
        solution = solver()             # alias: solver.solve()
    solution is an xarray.Dataset with all relevant quantities (see TfbiLoader.get_tfbi_all),
        and 'omega' telling roots with largest imaginary part.

    One good choice of chunks is to subsample spatial dimensions, like:
        XSUBSAMPLE, XBIGSTEP = 64, 256
        XSLICES = [slice(ix0, None, XBIGSTEP) for ix0 in range(0, XBIGSTEP, XSUBSAMPLE)]
        cc.chunks = dict(x_slices=XSLICES)
        # this will do 4 chunks total, each chunk handling every 256th point along x.
        # so, every chunk gives the "full" picture, though very low resolution.
        # but, the chunks are offset from each other by 64, so altogether
        # they actual combine to give every 64th point along x.
        # then, one can rerun later with smaller XSUBSAMPLE (e.g. 32)
        # to improve resolution while re-using existing results!!!

    For more precise control of solving process, consider:
        (copied from TfbiSolver docs): the following pattern & notes: 
            {solving_pattern}

            Notes (copied from TfbiSolver docs):
                {solving_pattern_notes}

        BUT, the pattern above doesn't include chunking.
        The FULL pattern, with chunking included, looks roughly like:
            slicers = cc.chunker().slicers()
            locs = [os.path.join(dst, cc.title_with_slices(s)+'.pcxarr') for s in slicers]
            todo = [(s, loc) for (s, loc) in zip(slicers, locs) if not os.path.exists(loc)]
            for s, p in todo:
                cc.slices = s
                cc.chunks = None
                # (*1) (point of interest, referenced below)
                dsR = ... # <-- copy-paste most of the pattern from above!
                dsR.pc.save(loc)
                # (*2) (point of interest, referenced below)

        [TODO][EFF] would be slightly more efficient to avoid recomputing casc for each chunk.
        However, for "reasonable chunk sizes", this is only a ~10% speedup.

    INPUTS -- used here, not just passed directly to super().
        cc: {cc}
            cc.chunks MUST be nonempty; otherwise, use TfbiSolver instead.
        dst: str
            directory where chunk results should be saved. Internally stored as abspath.
        pre_hook: None or callable of cc
            if provided, called as cc = pre_hook(cc) at point (*1) in the pattern above,
                i.e. before solving each chunk.
            E.g. def pre_hook(cc): cc.tfbi_mask(); return cc
                would apply tfbi_mask() to each chunk before solving.
                (Might want this because slicing after masking is hard,
                but masking after slicing is easy, in present implementation.)
            E.g., to artificially set E_un0_perpmod_B to test constant E/B speed across box:
                def pre_hook(cc):
                    result = cc.copy()
                    EBspeed_grid = pc.xr1d([1000, 3000, 5000, 7000], 'EBspeed')
                    result.set('E_un0_perpmod_B', cc('mod_B') * EBspeed_grid)
                    return result
        post_hook: None or callable of (cc, dsR, loc)
            if provided, called as post_hook(cc, dsR, loc) at point (*2) in the pattern above,
                i.e. after solving each chunk and saving the result.
            E.g., to save a growth rate plot in each chunk's loc directory:
                def post_hook(cc, dsR, loc):
                    dsR.pc.unmask_var('omega').it.growth_kmax().it.growthplot()
                    plt.savefig(os.path.join(loc, 'growthplot'))
                    plt.close()

    INPUTS -- passed directly to super(), i.e. works the same way as TfbiSolver.
        ions: {ions}
        kres: {kres}
        mod, lmod, ang: {mod_lmod_ang}
        tfbi_all: {tfbi_all}
        drel_cls: {drel_cls}
    '''
    def __init__(self, cc, dst, ions=None, *, kres='low', mod=UNSET, lmod=UNSET, ang=UNSET, tfbi_all=True,
                 drel_cls=None, pre_hook=None, post_hook=None):
        if not cc.chunks:
            raise InputMissingError('empty cc.chunks. Set chunks, or use TfbiSolver instead.')
        self.dst = os.path.abspath(dst)
        super().__init__(cc, ions=ions, kres=kres, mod=mod, lmod=lmod, ang=ang,
                         tfbi_all=tfbi_all, drel_cls=drel_cls)
        self.pre_hook = pre_hook
        self.post_hook = post_hook

    def chunk_filepath(self, slicer):
        '''return filepath to use for results from chunk defined by slicer.
        Slicer should be a dict of slices, appropriate as cc.slices = slicer.
        '''
        return os.path.join(self.dst, self.cc.title_with_slices(slicer)+'.pcxarr')

    def slicers_and_filepaths(self):
        '''return list of (slicer, filepath) tuples for all chunks implied by self.cc.chunks.'''
        slicers = self.cc.chunker().slicers()
        filepaths = [self.chunk_filepath(s) for s in slicers]
        return list(zip(slicers, filepaths))

    def done(self):
        '''list of (slicer, filepath) tuples for all chunks which have already been solved.
        (To check if solved, just checks whether filepath exists!)
        '''
        return [(s, p) for (s, p) in self.slicers_and_filepaths() if os.path.exists(p)]

    def todo(self):
        '''list of (slicer, filepath) tuples for all chunks which have NOT yet been solved.
        (To check if solved, just checks whether filepath exists!)
        '''
        return [(s, p) for (s, p) in self.slicers_and_filepaths() if not os.path.exists(p)]

    def n_chunks(self):
        '''total number of chunks implied by self.cc.chunks.'''
        return len(self.cc.chunker())

    @property
    def solved(self):
        '''tells number of chunks solved and total, as "Nsolved_of_Ntotal" string.
        E.g., "3_of_10" if 3 out of 10 chunks have existing solutions.
        '''
        nsolved = len(self.done())
        return f'{nsolved}_of_{self.n_chunks()}'

    def solve(self, *, verbose=2, mergeload=True, **kw_growth_root):
        '''solve TFBI theory using chunking.
        See help(type(self)) for more details.

        Writes results to files, as determined by self.dst.
        Never overwrite any files; assumes pre-existing files are correct solutions.

        verbose: bool or int
            whether to print progress updates (highly recommended).
            0: print nothing
            >=1: print updates about which chunk is being done.
            >=2: print tfbi_solver updates within each chunk.
                (chunks might each take a few minutes)
        additional kwargs get passed directly to drel.solve().
            (options include: ncpu, ncoarse, careful)

        mergeload: bool
            whether to return full solution, mergeloaded across all chunks.
            True --> return full solution.
            False --> return filepath to where chunks are stored (i.e., self.dst).
            Either way, the full solution can later be loaded via pc.xarray_mergeload(dst).
        '''
        updater = ProgressUpdater(print_freq=0 if verbose else -1)
        self._warn_if_too_many_ions()
        # think about chunks:
        todo = self.todo()
        if len(todo) > 0:
            n_done = self.n_chunks() - len(todo)
            updater.print(f'already did {n_done}; need to do {len(todo)}.', end='\n', print_time=False)
        with self.cc.maintaining('chunks', 'slices'):  # restore cc.chunks and cc.slices afterward.
            _chunk_dims = self.cc.chunker().chunk_dims  # (just for prettier print below)
            self.cc.chunks = None  # disable internal chunking within each chunk.
            for i, (slicer, loc) in enumerate(todo):
                _chunk_slices = {d: slicer[d] for d in _chunk_dims if d in slicer}  # (just for pretty print)
                updater.print(f'--- doing {i+1} of {len(todo)}: {_chunk_slices} ---', end='\n')
                self.cc.slices = slicer
                cc = self.cc
                if self.pre_hook is not None:
                    cc = self.pre_hook(self.cc)
                # <-- could just use super(). BUT, might want to customize more later,
                #   e.g. might avoid recomputing casc for each chunk.
                #   (Also, don't want to save ds0, kp, dsk, drel, dsR to self here.)
                #   So, copy the code, instead of calling super().
                ds0_updater = ProgressUpdater(print_freq=0 if verbose>=2 else -1)
                ds0_updater.print(f'getting ds0 via {type(cc).__name__}.tfbi_ds() ...')
                ds0 = cc.tfbi_ds(ions=self.ions, all=self.tfbi_all)
                ds0_updater.finalize(f'solver.ds0 = {type(cc).__name__}.tfbi_ds()', end='\n')
                # picking k is usually fast; no need for progress updates.
                kw_kPicker = dict(mod=self.mod, lmod=self.lmod, ang=self.ang)
                kw_kPicker = {k: v for k, v in kw_kPicker.items() if v is not UNSET}
                kp = self.kPicker_cls(ds0, **kw_kPicker)
                # follow the rest of the pattern. tt manages its own progress updates.
                dsk = kp.get_ds()
                drel = self.drel_cls.from_ds(dsk)
                dsR = drel.solve(verbose=(verbose>=2), **kw_growth_root)
                dsR.pc.save(loc)
                if self.post_hook is not None:
                    self.post_hook(cc, dsR, loc)
        updater.finalize(f'--- {type(self).__name__}.solve() (all chunks) ---',
                         end='\n' if mergeload else '')
        if mergeload:
            return xarray_mergeload(self.dst, verbose=verbose)
        else:
            return self.dst

    def _repr_contents(self):
        '''return list of strings for contents inside of __repr__'''
        contents = super()._repr_contents()
        contents.append(f'dst={self.dst!r}')
        if self.pre_hook is not None:
            contents.append(f'pre_hook={self.pre_hook!r}')
        if self.post_hook is not None:
            contents.append(f'post_hook={self.post_hook!r}')
        return contents
