"""Define and solve models for LGS systems"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/api/lgssystem.ipynb.

# %% auto 0
__all__ = ['optree', 'ScipyGCROTMKSpILUSolver', 'LGSSystem']

# %% ../nbs/api/lgssystem.ipynb
import importlib
from IPython.display import SVG, Markdown

from pymor.vectorarrays.interface import VectorArray
from pymor.operators.interface import Operator
from pymor.basic import IdentityOperator, VectorOperator

from .imports import *
from .utilities.nbdev import DictTbl, AttributeTbl
from .utilities.testing import test_array
from .utilities.sparse import sparse_kronecker_matrix, sparse_toeplitz, sparse_identity, sparse_diag, sparse
from .utilities.numpy import sym_range
from .pymor.parameters import *
from .pymor.vectorarrays import *
from .pymor.operators import *
from .pymor.timestepping import *
from .pymor.grids import *
from .pymor.models import *
from .velocitygroups import *
from .dmelements import *

# %% ../nbs/api/lgssystem.ipynb
def optree(op):
    match op:
        case LincombOperator(): 
            return [op.name + "(Lincomb)", [optree(o) for o in op.operators]]
        case ProductOperator(): 
            return [op.name + "(Product)", [optree(o) for o in op.operators]]
        # case IdentityOperator():
        #     return "Identity"
        case _:
            return f'{op.name}({op.__class__.__name__.removesuffix("Operator")})'

# %% ../nbs/api/lgssystem.ipynb
from scipy.sparse.linalg import LinearOperator, gcrotmk, spilu
from pymor.solvers.interface import Solver

# %% ../nbs/api/lgssystem.ipynb
class ScipyGCROTMKSpILUSolver(Solver):

    # @defaults('tol', 'maxiter', 'spilu_drop_tol', 'spilu_fill_factor', 'spilu_drop_rule', 'spilu_permc_spec')
    # @defaults('check_finite')
    def __init__(self, check_finite=None, atol=0., rtol=1e-5, maxiter=1000,
                 spilu_drop_tol=1e-4, spilu_fill_factor=20, spilu_drop_rule=None, spilu_permc_spec='COLAMD', preconditioner=None):
        super().__init__()
        self.__auto_init(locals())

    def _solve(self, operator, V, mu, initial_guess):
        operator = operator.assemble(mu)
        matrix = to_matrix(operator).tocsc()
        precond = to_matrix(self.preconditioner.assemble(mu)).tocsc() if self.preconditioner else matrix
            
        V = V.to_numpy()
        initial_guess = initial_guess.to_numpy() if initial_guess is not None else None
        promoted_type = np.promote_types(matrix.dtype, V.dtype)

        R = np.empty((matrix.shape[1], V.shape[1]), dtype=promoted_type, order='F')
        ilu = spilu(precond, drop_tol=self.spilu_drop_tol, fill_factor=self.spilu_fill_factor,
                    drop_rule=self.spilu_drop_rule, permc_spec=self.spilu_permc_spec)
        precond = LinearOperator(matrix.shape, ilu.solve)
        for i in range(V.shape[1]):
            R[:, i], info = gcrotmk(matrix, V[:, i], initial_guess[:, i] if initial_guess is not None else None,
                                     atol=self.atol, rtol=self.rtol, maxiter=self.maxiter, M=precond)
            if info != 0:
                if info > 0:
                    raise InversionError(f'gcrotmk failed to converge after {info} iterations')
                else:
                    raise InversionError(f'gcrotmk failed with error code {info} (illegal input or breakdown)')

        if self.check_finite:
            if not np.isfinite(np.sum(R)):
                raise InversionError('Result contains non-finite values')

        return operator.source.from_numpy(R), {}

# %% ../nbs/api/lgssystem.ipynb
class LGSSystem(ParametricObject):
    """`LGSSystem` creates numerical models for the chosen laser guide star atomic system."""
    
    block_names = [
        "A_ind", # Part of density-matrix evolution matrix that is independent of atomic velocity
        "A_dop", # Part of density-matrix evolution matrix describing Doppler shift
        "A_rec", # Part of density-matrix evolution matrix describing atomic recoil
        "A_vcc", # Part of density-matrix evolution matrix describing velocity-changing collisions
        "b", # Constant vector in the density-matrix evolution equation describing atomic transit
        "Flux" # Operator giving return flux from each transition
    ]
    
    def __init__(
        self, 
        system:str, # Name of the atomic system to be loaded from file
        fixed_params:dict={} # Values for system parameters that will be held fixed
    ):
        self.path = LGSSystem._get_path(system)
        fixed_params = Mu(fixed_params)
        self.__auto_init(locals())
        for name in self.block_names:
            setattr(
                self, 
                name,
                LincombOperator.from_file(self.path/name).partial_assemble(self.fixed_params).with_(name=name)
            )
        self.Flux = XarrayMatrixOperator(self.Flux.matrix.rename({"Transition": "Transition (range)"}))
        self.parameters = Parameters.of(*(getattr(self, name) for name in self.block_names))
        self.initial = LincombOperator.from_file(self.path/'b').partial_assemble(Mu(BeamTransitRatePerS=1)).as_range_array()
        # self.initial = self.Flux.source.from_xarray(
        #     LincombOperator.from_file(self.path/'b').partial_assemble(Mu(BeamTransitRatePerS=1)).matrix.rename({'Density matrix (range)': 'Density matrix (source)'})
        # )
        self.n_variables = self.Flux.source.dim 
        self.dm_elements = self.Flux.assemble().matrix['Density matrix' + Lbl.SRC].rename({"Density matrix (source)": "Density matrix"})
        self.level_population = level_population(self.dm_elements)
        self.population = population(self.dm_elements) 
        self.total_population = total_population(self.dm_elements)

# %% ../nbs/api/lgssystem.ipynb
@patch(cls_method=True)
def _is_valid_dir(cls:LGSSystem, dir):
    """Check if the supplied directory has the required files to define an LGSSystem."""
    dir = Path(dir)
    if not dir.is_dir(): return False
    files = {f.name for f in dir.glob("*.mtxn")}
    required_files = {'A_dop.mtxn', 'A_ind.mtxn', 'A_rec.mtxn', 'A_vcc.mtxn', 'Flux.mtxn', 'b.mtxn'}
    return required_files <= files

# %% ../nbs/api/lgssystem.ipynb
@patch(cls_method=True)
def _builtin_paths(cls:LGSSystem):
    return [dir for dir in importlib.resources.files("pylgs.systems").iterdir() if LGSSystem._is_valid_dir(dir)]

# %% ../nbs/api/lgssystem.ipynb
@patch(cls_method=True)
def builtins(cls:LGSSystem):
    return [dir.name for dir in LGSSystem._builtin_paths()]

# %% ../nbs/api/lgssystem.ipynb
@patch(cls_method=True)
def _get_path(cls:LGSSystem, name):
    if LGSSystem._is_valid_dir(name): return Path(name)
    for dir in LGSSystem._builtin_paths(): 
        if dir.name == name: return dir
    raise FileNotFoundError(f"{name} does not correspond to a built-in or supplied LGSSystem.")

# %% ../nbs/api/lgssystem.ipynb
@patch(cls_method=True)
def diagram(
    cls:LGSSystem, 
    name:str | Path, # Name of built-in LGS system definition, or directory of supplied system
    kind="Toy" # "Toy" to show level diagram and pumped transitions, "ToScale" to show hyperfine structure to scale, "NotToScale" to show unscaled hf structure
):
    """Draw a level diagram for the system of the specified kind."""
    return SVG(LGSSystem._get_path(name)/f"LevelDiagram{kind}.svg")

# %% ../nbs/api/lgssystem.ipynb
@patch(cls_method=True)
def info(
    cls:LGSSystem, 
    name:str | Path, # Name of built-in LGS system definition, or directory of supplied system
):
    """Display information about LGS system `name`."""
    return Markdown(filename=LGSSystem._get_path(name)/"info.md")

# %% ../nbs/api/lgssystem.ipynb
@patch
def _reaction(self:LGSSystem, vg:VelocityGroups):
    return (
        (vg.identity().with_(name="VG") * self.A_ind).with_(name="VG ind.") 
        + (vg.velocity_diagonal().with_(name="Velocity") * self.A_dop).with_(name="VG dop.")
    ).with_(name="Reaction")

# %% ../nbs/api/lgssystem.ipynb
@patch
def _vcc(self:LGSSystem, vg:VelocityGroups):
    return vg.n_times_1().with_(name="MaxBoltz") * self.A_vcc

# %% ../nbs/api/lgssystem.ipynb
@patch
def _recoil(self:LGSSystem, vg:VelocityGroups):
    return vg.drho_dv().with_(name='drho_dv') * self.A_rec

# %% ../nbs/api/lgssystem.ipynb
@patch
def operator(
    self:LGSSystem, 
    vg:VelocityGroups # `VelocityGroups` specification
)->Operator:
    """The evolution operatior for the system, given the `VelocityGroups` `vg`."""
    return (-1 * (self._reaction(vg) + self._vcc(vg) + self._recoil(vg))).with_(
        name='DM evolution',
        # solver_options={'inverse': {'type': 'scipy_lgmres_spilu', 'preconditioner_bandwidth': self.n_variables}}
    )

# %% ../nbs/api/lgssystem.ipynb
@patch
def rhs(
    self:LGSSystem, 
    vg:VelocityGroups # `VelocityGroups` specification
)->VectorArray:
    """The right-hand side vector for this system given the `VelocityGroups` `vg`."""
    return velocity_density_vector(vg) * self.b.as_range_array()

# %% ../nbs/api/lgssystem.ipynb
@patch
def _initial_dm(self:LGSSystem, vg):
    return velocity_density_vector(vg) * self.initial

# %% ../nbs/api/lgssystem.ipynb
@patch
def _products(self:LGSSystem, vg): 
    return {
        'population': self.total_population, # Fraction of total atomic population in each velocity group
        'total_population': vg.sum() * self.total_population, # Total atomic population (should equal 1)
        'population_distribution': vg.normalize() * self.total_population, # Atomic population per unit velocity interval
        'level_population': self.level_population, # Population of each atomic level in each velocity group
        'total_level_population': vg.sum() * self.level_population, # Total atomic population in each level
        'level_population_distribution': vg.normalize() * self.level_population, # Level population per unit velocity interval
        'flux': self.Flux, # Return flux on each transition from each velocity group
        'total_flux': vg.sum() * self.Flux, # Total return flux
        'flux_distribution': vg.normalize() * self.Flux, # Return flux per unit velocity interval
        'velocity_sum': vg.sum(), # Sum over velocity groups
        'velocity_normalize': vg.normalize(), # Normalize by velocity-group width
    }

# %% ../nbs/api/lgssystem.ipynb
@patch
def stationary_model(
    self:LGSSystem, 
    vg:int|Iterable|VelocityGroups=6, # Specification of velocity groups 
    # output='total_flux' # Operator to be applied to solution to obtain the output
) -> StationaryModel: # Model for the steady-state LGS system
    """Return a pyMOR `StationaryModel` for the steady state cw LGS system."""
    vg = VelocityGroups(vg)
    return StationaryModel(
        self.operator(vg).with_(solver=ScipyGCROTMKSpILUSolver(spilu_fill_factor=40, preconditioner=-self._reaction(vg))),
        self.rhs(vg), 
        products=self._products(vg),
        data={
            'velocity_groups': vg,
            'grid': OnedUnstructuredGrid(vg.edges)
        }
    )

# %% ../nbs/api/lgssystem.ipynb
@patch
def adaptive_stationary_model(
    self:LGSSystem, 
    mu, # Parameter values to use during velocity-group refinement
    vg=6, # Initial velocity groups
    max_weight=0.01 # Maximum fraction of the return flux from any one velocity group
)->StationaryModel: # Steady-state LGS model with more velocity groups in regions that produce more return flux
    """Create a `StationaryModel` with narrower velocity bins in velocity regions with higher flux."""
    vg = VelocityGroups(vg)
    while True:
        m = self.stationary_model(vg)
        u = m.solve(mu)
        o = m.flux(u).array.sum('Transition').squeeze().data
        o /= np.sum(o)
        mask = o > max_weight
        if not np.sum(mask):
            break
        vg = vg.subdivide(mask)
    return self.stationary_model(vg)

# %% ../nbs/api/lgssystem.ipynb
@patch
def instationary_model(
    self:LGSSystem, 
    vg:int|Iterable|VelocityGroups=6, # Specification of velocity groups 
    # output:str='total_flux',  # Operator to be applied to solution to obtain the output
    T=1.e-6, # Final time (seconds)
    num_values=100, # Number of time step values to return
    initial_rho=None, # Initial density matrix. `None` means to use the thermal distribution
    time_stepper=BDFTimeStepper() # Solver to use
) -> InstationaryModel: # Model for the LGS system dynamics
    """Create an `InstationaryModel` for the LGS system dynamics."""
    vg = VelocityGroups(vg)
    return InstationaryModel(
        T,
        ifnone(initial_rho, self._initial_dm(vg)),
        self.operator(vg), 
        self.rhs(vg), 
        mass=None,
        time_stepper=time_stepper,
        num_values=num_values,
        # output_functional=self._products(vg)[output],
        products=self._products(vg),
        data={
            'velocity_groups': vg,
            'grid': OnedUnstructuredGrid(vg.edges)
        }
    )

# %% ../nbs/api/lgssystem.ipynb
def _to_nmo(op, format='dense'):
    return NumpyMatrixOperator(to_matrix(op, format=format))

# %% ../nbs/api/lgssystem.ipynb
@patch
@delegates(LGSSystem.instationary_model)
def numpy_instationary_model(self:LGSSystem, vg, **kwargs):
    model = self.instationary_model(vg, **kwargs)
    operator = expand(model.operator)
    operator = contract(LincombOperator([_to_nmo(o) for o in operator.operators], operator.coefficients))
    rhs = _to_nmo(model.rhs)
    initial_data = _to_nmo(model.initial_data)
    vg = model.data['velocity_groups']
    # vg_sum = XarrayMatrixOperator(np.ones([1, vg['VGNumber']]), source=vg.sum().source, range=vg.sum().range)
    output_functional = _to_nmo(IdentityOperator(vg.sum().source) * model.products['flux'])
    return InstationaryModel(
        T=model.T,
        initial_data=initial_data,
        operator=operator,
        rhs=rhs,
        output_functional=output_functional,
        mass=None,
        time_stepper=model.time_stepper,
        num_values=model.num_values,
        products=None,
        error_estimator=None,
        visualizer=None,
        name=None,
    )

# %% ../nbs/api/lgssystem.ipynb
@patch
def _floquet_products(self:LGSSystem, vg, n_vector, T): 
    return {
        # 'reconstruct': floquet_reconstruction_operator(n_vector, T),
        # 'total_population': velocity_sum(vg) * self.total_population, 
        # 'population_distribution': velocity_normalize(vg) * self.total_population, 
        'reconstructed_level_population': floquet_reconstruction_operator(n_vector, T) * vg.sum() * self.level_population, 
        # 'level_population_distribution': velocity_normalize(vg) * self.level_population, 
        # 'flux': self.flux,
        'flux_distribution': vg.normalize() * self.Flux,
        'reconstructed_flux_distribution': floquet_reconstruction_operator(n_vector, T) * vg.normalize() * self.Flux,
        'reconstructed_total_flux': floquet_reconstruction_operator(n_vector, T) * (vg.sum() * self.Flux),
        # 'velocity_sum': velocity_sum(vg),
        # 'velocity_normalize': velocity_normalize(vg),
    }

# %% ../nbs/api/lgssystem.ipynb
@patch
def stationary_floquet_model(
    self:LGSSystem, 
    mu, # Parameter specification for parameters modulated at a frequency `nu` 
    time, # List of time values to use for reconstruction
    n_vector, # Number of terms in the system Fourier series
    nu:str='nu', # Name of the frequency variable
    n_operator=None, # Number of terms in the operator Fourier series. `None` defaults to `n_vector` - 1
    vg:int|Iterable|VelocityGroups=6, # Specification of velocity groups
    solver=None,
) -> StationaryFloquetModel: # Model for the Fourier coefficients of the periodic LGS system
    """Return a `StationaryFloquetModel` for the Fourier coefficients of the periodic state of a modulated LGS system."""
    vg = VelocityGroups(vg)
    return StationaryFloquetModel(
        operator=expand(self.operator(vg)),
        rhs=self.rhs(vg), 
        mu=mu,
        nu=nu,
        time=time,
        n_vector=n_vector,
        n_operator=n_operator,
        products=self._products(vg) | self._floquet_products(vg, n_vector, time),
        data={
            'velocity_groups': vg,
            'grid': OnedUnstructuredGrid(vg.edges)
        },
        solver=solver or ScipyGCROTMKSpILUSolver(spilu_fill_factor=40),
        preconditioner=-self._reaction(vg)
    )

# %% ../nbs/api/lgssystem.ipynb
@patch
def adaptive_stationary_floquet_model(
    self:LGSSystem, 
    mu, # Parameter specification for parameters modulated at a frequency `nu`
    mu_refine,
    time, # List of time values to use for reconstruction
    n_vector, # Number of terms in the system Fourier series
    nu:str='nu', # Name of the frequency variable
    n_operator=None, # Number of terms in the operator Fourier series. `None` defaults to `n_vector` - 1
    vg=6, # Initial velocity groups
    solver=None,
    max_weight=0.01, # Maximum fraction of the return flux from any one velocity group
    subdivisions=4,
)->StationaryFloquetModel: # Steady-state LGS model with more velocity groups in regions that produce more return flux
    """Create a `StationaryModel` with narrower velocity bins in velocity regions with higher flux."""
    vg = VelocityGroups(vg)
    while True:
        print(f'Creating model with {vg['VGNumber']} velocity groups')
        m = self.stationary_floquet_model(mu, time, n_vector, nu, n_operator, vg, solver)
        print('Solving model')
        u = m.solve(mu_refine)
        o = m.reconstructed_flux_distribution(u, mu_refine).real.array.sum(['Time', 'Transition']).squeeze().data
        o /= np.sum(o)
        mask = o > max_weight
        if not np.sum(mask):
            break
        vg = vg.subdivide(mask, subdivisions)
    return self.stationary_floquet_model(mu, time, n_vector, nu, n_operator, vg, solver), u
