text: """ Handles reading and interpolation of atmopshere (grid) data """
  import logging
  
  import numpy as np
  from flex.extensions.bindata import MultipleDataExtension
  
  from ..abund import Abund
  from ..data_structure import (
      Collection,
      CollectionFactory,
      absolute,
      array,
      asfloat,
      asstr,
      lowercase,
      oneof,
      this,
      uppercase,
  )
  
  logger = logging.getLogger(__name__)
  
  
  class AtmosphereError(RuntimeError):
      """Something went wrong with the atmosphere interpolation"""
  
  
  @CollectionFactory
  class Atmosphere(Collection):
      """
      Atmosphere structure
      contains all information to describe the solar atmosphere
      i.e. temperature etc in the different layers
      as well as stellar parameters and abundances
      """
  
      # fmt: off
      _fields = [
          ("teff", 5770, asfloat, this, "float: effective temperature in Kelvin"),
          ("logg", 4.0, asfloat, this, "float: surface gravity in log10(cgs)"),
          ("abund", Abund.solar(), this, this, "Abund: elemental abundances"),
          ("vturb", 0, absolute, this, "float: turbulence velocity in km/s"),
          ("lonh", 0, asfloat, this, "float: ?"),
          ("source", "marcs2014.sav", this, this, "str: datafile name of this data, or atmosphere grid/file"),
          ("method", "grid", lowercase(oneof("grid", "embedded")), this,
              "str: whether the data source is a grid or a fixed atmosphere"),
          ("geom", "PP", uppercase(oneof("PP", "SPH", None)), this,
              "str: the geometry of the atmopshere model"),
          ("radius", 0, asfloat, this, "float: radius of the spherical model"),
          ("height", None, array(None, "f8"), this, "array: height of the spherical model"),
          ("opflag", [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], array(20, int), this,
              "array of size (20,): opacity flags"),
          ("wlstd", 5000, asfloat, this, "float: wavelength standard deviation"),
          ("depth", None, uppercase(oneof(None, "RHOX", "TAU")), this,
              "str: flag that determines whether to use RHOX or TAU for calculations"),
          ("interp", None, uppercase(oneof(None, "RHOX", "TAU")), this,
              "str: flag that determines whether to use RHOX or TAU for interpolation"),
          ("rhox", None, array(None, "f8"), this,
              "array: mass column density"),
          ("tau", None, array(None, "f8"), this,
              "array: continuum optical depth"),
          ("temp", None, array(None, "f8"), this,
              "array: temperature profile in Kelvin"),
          ("rho", None, array(None, "f8"), this,
              "array: density profile"),
          ("xna", None, array(None, "f8"), this,
              "array: number density of atoms in 1/cm**3"),
          ("xne", None, array(None, "f8"), this,
              "array: number density of electrons in 1/cm**3"),
          ("citation_info", "", asstr, this, "str: citation for this atmosphere"),
      ]
      # fmt: on
  
      # TODO: pick the right geometry for the grid, based on whether it has height/radius values or not
  
      def __init__(self, **kwargs):
          monh = kwargs.pop("monh", kwargs.pop("feh", 0))
          abund = kwargs.pop("abund", "empty")
          abund_format = kwargs.pop("abund_format", "sme")
          if "geom" in kwargs.keys() and kwargs["geom"] == "":
              kwargs["geom"] = None
          super().__init__(**kwargs)
          self.abund = Abund(monh=monh, pattern=abund, type=abund_format)
  
      @property
      def monh(self):
          """float: metallicity"""
          return self.abund.monh
  
      @monh.setter
      def monh(self, value):
          self.abund.monh = value
  
      @property
      def names(self):
          return self._names + ["monh"]
  
      @property
      def dtype(self):
          obj = lambda: None
          obj.names = [n.lower() for n in self.names]
          return obj
  
      @property
      def ndep(self):
          scale = self.temp
          if scale is not None:
              return scale.shape[0]
          return None
  
      @ndep.setter
      def ndep(self, value):
          pass
  
      def _save(self):
          data = {}
          ext2 = self.abund._save()
          header = ext2.header
          data["abund"] = ext2.data
          header["abund_format"] = header["type"]
          del header["type"]
  
          for name in self._names:
              value = self[name]
  
              if isinstance(value, np.ndarray):
                  data[name] = value
              elif value is None or isinstance(value, Abund):
                  pass
              else:
                  # if isinstance(value, (np.floating, np.integer, np.str, float, int, str)):
                  header[name] = value
  
          ext = MultipleDataExtension(header, data)
  
          return ext
  
      @classmethod
      def _load(cls, ext):
          header = ext.header
          header.update(ext.data)
          obj = cls(**header)
          return obj
  
  
  @CollectionFactory
  class AtmosphereGrid(np.recarray):
      """
      A grid of atmospheres, used for the interpolation
      of model atmospheres. Each entry represents one
      atmosphere model.
      """
  
      # fmt: off
      _fields = [
          ("source", None, asstr, this, "str: datafile name of this data"),
          ("method", "grid", lowercase(oneof("grid", "embedded")), this,
              "str: whether the data source is a grid or a fixed atmosphere"),
          ("geom", None, uppercase(oneof(None, "PP", "SPH")), this,
              "str: the geometry of the atmopshere model"),
          ("depth", None, uppercase(oneof(None, "RHOX", "TAU")), this,
              "str: flag that determines whether to use RHOX or TAU for calculations"),
          ("interp", None, uppercase(oneof(None, "RHOX", "TAU")), this,
              "str: flag that determines whether to use RHOX or TAU for interpolation"),
          ("abund_format", "sme", oneof(*Abund._formats), this,
              "str: format of the Abundance field, as defined by the Abund class"),
          ("citation_info", "", asstr, this, "str: Citation text to cite in your papers"),
      ]
      # fmt: on
  
      def __new__(cls, natmo, npoints, **kwargs):
          dtype = [
              ("teff", "f4"),
              ("logg", "f4"),
              ("monh", "f4"),
              ("vturb", "f4"),
              ("lonh", "f4"),
              ("radius", "f4"),
              ("height", f"({npoints},)f4"),
              ("wlstd", "f4"),
              ("opflag", "(20,)i4"),
              ("abund", "(99,)f4"),
              ("temp", f"({npoints},)f4"),
              ("rhox", f"({npoints},)f4"),
              ("tau", f"({npoints},)f4"),
              ("rho", f"({npoints},)f4"),
              ("xna", f"({npoints},)f4"),
              ("xne", f"({npoints},)f4"),
          ]
  
          names = [s[0].lower() for s in dtype]
          titles = [s[0].upper() for s in dtype]
  
          data = super().__new__(cls, (natmo,), dtype=dtype, names=names, titles=titles)
  
          data.interp = "TAU"
          data.depth = "RHOX"
          data.method = "grid"
          data.geom = "PP"
          data.source = ""
          data.citation_info = ""
          data.abund_format = "sme"
          data.info = ""
          return data
  
      def __array_finalize__(self, obj):
          if obj is None:
              return
          self.interp = getattr(obj, "interp", "TAU")
          self.depth = getattr(obj, "depth", "RHOX")
          self.method = getattr(obj, "method", "grid")
          self.geom = getattr(obj, "geom", "PP")
          self.source = getattr(obj, "source", "")
          self.citation_info = getattr(obj, "citation_info", "")
          self.abund_format = getattr(obj, "abund_format", "sme")
          self.info = getattr(obj, "info", "")
  
      def __reduce__(self):
          # Get the parent's __reduce__ tuple
          pickled_state = super(AtmosphereGrid, self).__reduce__()
          # Create our own tuple to pass to __setstate__
          new_state = pickled_state[2] + (
              self.interp,
              self.depth,
              self.source,
              self.geom,
              self.citation_info,
              self.method,
              self.abund_format,
              self.info,
          )
          # Return a tuple that replaces the parent's __setstate__ tuple with our own
          return (pickled_state[0], pickled_state[1], new_state)
  
      def __setstate__(self, state):
          self.interp = state[-8]
          self.depth = state[-7]
          self.source = state[-6]
          self.geom = state[-5]
          self.citation_info = state[-4]
          self.method = state[-3]
          self.abund_format = state[-2]
          self.info = state[-1]
  
          # Call the parent's __setstate__ with the other tuple elements.
          super(AtmosphereGrid, self).__setstate__(state[0:-8])
  
      def __getitem__(self, key):
          """Overwrite the getitem routine, so we keep additional
          properties and/or return an atmosphere object, when only
          one record is returned"""
          cls = self.__class__
          value = super().__getitem__(key)
          if isinstance(value, cls) and value.size == 1:
              return value[0]
  
          if isinstance(value, np.record):
              kwargs = {s: value[s] for s in value.dtype.names}
              value = Atmosphere(**kwargs)
          if isinstance(value, (Atmosphere, cls)):
              for name in self._names:
                  setattr(value, name, getattr(self, name))
          return value
  
      def __str__(self):
          return self.source
  
      def __repr__(self):
          return str(self)
  
      def get(self, teff, logg, monh):
          mask = self.teff == teff
          mask &= self.logg == logg
          mask &= self.monh == monh
          return self[mask]
  
      @property
      def ndep(self):
          return self.shape[1]
  
      def save(self, filename):
          """Save the Atmopshere grid to a file using a numpy save format"""
          header = {
              "interp": self.interp,
              "depth": self.depth,
              "source": self.source,
              "geom": self.geom,
              "citation_info": self.citation_info,
              "method": self.method,
              "abund_format": self.abund_format,
              "info": self.info,
          }
          names = list(header.keys())
          values = [header[n] for n in names]
          header = np.rec.array(values, names=names)
          np.savez(filename, data=self, header=header)
          return
  
      @classmethod
      def load(cls, filename):
          """Load the atmosphere grid data from disk"""
          data = np.load(filename)
          self = data["data"].view(cls)
          header = data["header"]
          for k in header.dtype.names:
              v = header[k][()]
              setattr(self, k, v)
          return self
  