from dataclasses import dataclass, field, asdict
from copy import copy
from enum import Enum, auto
from typing import Sequence, TYPE_CHECKING, Optional, Union, TypeVar, Type
import numpy as np
from .witt import witt
import lightweaver.constants as Const
from numpy.polynomial.legendre import leggauss
from .utils import ConvergenceError, view_flatten, get_data_path
from .atomic_table import PeriodicTable, AtomicAbundance, DefaultAtomicAbundance
import astropy.units as u
import pickle

if TYPE_CHECKING:
    from .LwCompiled import LwSpectrum

class ScaleType(Enum):
    '''
    Atmospheric scales used in the definition of 1D atmospheres to allow the
    correct conversion to a height based system.
    Options:

        - `Geometric`

        - `ColumnMass`

        - `Tau500`

    '''
    Geometric = 0
    ColumnMass = auto()
    Tau500 = auto()

class BoundaryCondition:
    '''
    Base class for boundary conditions.

    Defines the interface; do not use directly.

    Attributes
    ----------
    These attributes are only available after set_required_angles has been called.

    mux : np.ndarray
        The mu_x to return from compute_bc (in order).
    muy : np.ndarray
        The mu_y to return from compute_bc (in order).
    muz : np.ndarray
        The mu_z to return from compute_bc (in order).
    indexVector : np.ndarray
        A 2D array of integer shape (mu, toObs) - where mu is the mu index on
        the associated atmosphere - relating each index of the second (Nrays)
        axis of a pair of (mu, toObs). Used to construct and destructure this
        array.

    '''
    def compute_bc(self, atmos: 'Atmosphere', spect: 'LwSpectrum') -> np.ndarray:
        '''
        Called when the radiation boundary condition is needed by the backend.

        Parameters
        ----------
        atmos : Atmosphere
            The atmospheric object in which to compute the radiation.
        spect : LwSpectrum
            The computational spectrum object provided by the Context.

        Returns
        -------
        result : np.ndarray
            This function needs to return a contiguous array of shape [Nwave,
            Nrays, Nbc], where Nwave is the number of wavelengths in the
            wavelength grid, Nrays is the number of rays in the angular
            quadrature (also including up/down directions) ordered as
            specified by the mux/y/z and indexVector variables on this
            object, Nbc is the number of spatial positions the boundary
            condition needs to be defined at ordered in a flattened [Nz, Ny,
            Nx] fashion. (dtype: <f8)

        '''
        raise NotImplementedError

    def set_required_angles(self, mux, muy, muz, indexVector):
        '''
        The angles (and their ordering) to be used for this boundary
        condition (in the case of a callable)
        '''
        self.mux = mux
        self.muy = muy
        self.muz = muz
        self.indexVector = indexVector

class NoBc(BoundaryCondition):
    '''
    Indicates no boundary condition on the axis because it is invalid for the
    current simulation.
    Used only by the backend.
    '''
    pass

class ZeroRadiation(BoundaryCondition):
    '''
    Zero radiation boundary condition.
    Commonly used for coronal situations.
    '''
    pass

class ThermalisedRadiation(BoundaryCondition):
    '''
    Thermalised radiation (blackbody) boundary condition.
    Commonly used for photospheric situations.
    '''
    pass

class PeriodicRadiation(BoundaryCondition):
    '''
    Periodic boundary condition.
    Commonly used on the x-axis in 2D simulations.
    '''
    pass

def get_top_pressure(eos: witt, temp, ne=None, rho=None):
    '''
    Return a pressure for the top of atmosphere.
    For internal use.

    In order this is deduced from:
        - the electron density `ne`, if provided
        - the mass density `rho`, if provided
        - the electron pressure present in FALC

    Returns
    -------
    pressure : float
        pressure IN CGS [dyn cm-2]

    '''
    if ne is not None:
        pe = ne * Const.CM_TO_M**3 * eos.BK * temp
        return eos.pg_from_pe(temp, pe)
    elif rho is not None:
        return eos.pg_from_rho(temp, rho)

    pgasCgs = np.array([0.70575286, 0.59018545, 0.51286639, 0.43719268, 0.37731009,
        0.33516886, 0.31342915, 0.30604891, 0.30059491, 0.29207645,
        0.2859011 , 0.28119224, 0.27893046, 0.27949676, 0.28299726,
        0.28644693, 0.28825946, 0.29061192, 0.29340255, 0.29563072,
        0.29864548, 0.30776456, 0.31825915, 0.32137574, 0.3239401 ,
        0.32622212, 0.32792196, 0.3292243 , 0.33025437, 0.33146736,
        0.3319676 , 0.33217821, 0.3322355 , 0.33217166, 0.33210297,
        0.33203833, 0.33198508])
    tempCoord = np.array([  7600.,   7780.,   7970.,   8273.,   8635.,   8988.,   9228.,
        9358.,   9458.,   9587.,   9735.,   9983.,  10340.,  10850.,
        11440.,  12190.,  13080.,  14520.,  16280.,  17930.,  20420.,
        24060.,  27970.,  32150.,  36590.,  41180.,  45420.,  49390.,
        53280.,  60170.,  66150.,  71340.,  75930.,  83890.,  90820.,
        95600., 100000.])

    ptop = np.interp(temp, tempCoord, pgasCgs)
    return ptop

@dataclass
class Stratifications:
    '''
    Stores the optional derived z-stratifications of an atmospheric model.

    Attributes
    ----------
    cmass : np.ndarray
        Column mass [kg m-2].
    tauRef : np.ndarray
        Reference optical depth at 500 nm.
    '''
    cmass: np.ndarray
    tauRef: np.ndarray

    def dimensioned_view(self, shape) -> 'Stratifications':
        '''
        Makes an instance of `Stratifications` reshaped to the provided
        shape for multi-dimensional atmospheres.
        For internal use.

        Parameters
        ----------
        shape : tuple
            Shape to reform the stratifications, provided by
            `Layout.dimensioned_shape`.

        Returns
        -------
        stratifications : Stratifications
            Reshaped stratifications.
        '''
        strat = copy(self)
        strat.cmass = self.cmass.reshape(shape)
        strat.tauRef = self.tauRef.reshape(shape)
        return strat

    def unit_view(self) -> 'Stratifications':
        '''
        Makes an instance of `Stratifications`  with the correct `astropy.units`
        For internal use.

        Returns
        -------
        stratifications : Stratifications
            The same data with units applied.
        '''
        strat = copy(self)
        strat.cmass = self.cmass << u.kg / u.m**2
        strat.tauRef = self.tauRef << u.dimensionless_unscaled
        return strat

    def dimensioned_unit_view(self, shape) -> 'Stratifications':
        '''
        Makes an instance of `Stratifications` reshaped to the provided shape
        with the correct `astropy.units` for multi-dimensional atmospheres.
        For internal use.

        Parameters
        ----------
        shape : tuple
            Shape to reform the stratifications, provided by
            `Layout.dimensioned_shape`.

        Returns
        -------
        stratifications : Stratifications
            Reshaped stratifications with units.
        '''
        strat = self.dimensioned_view(shape)
        return strat.unit_view()

@dataclass
class Layout:
    '''
    Storage for basic atmospheric parameters whose presence is determined by problem dimensionality, boundary conditions and optional stratifications.

    Attributes
    ----------
    Ndim : int
        Number of dimensions in model.

    x : np.ndarray
        Ordinates of grid points along the x-axis (present for Ndim >= 2) [m].
    y : np.ndarray
        Ordinates of grid points along the y-axis (present for Ndim == 3) [m].
    z : np.ndarray
        Ordinates of grid points along the z-axis (present for all Ndim) [m].
    vx : np.ndarray
        x component of plasma velocity (present for Ndim >= 2) [m/s].
    vy : np.ndarray
        y component of plasma velocity (present for Ndim == 3) [m/s].
    vz : np.ndarray
        z component of plasma velocity (present for all Ndim) [m/s]. Aliased to `vlos` when `Ndim==1`
    xLowerBc : BoundaryCondition
        Boundary condition for the plane of minimal x-coordinate.
    xUpperBc : BoundaryCondition
        Boundary condition for the plane of maximal x-coordinate.
    yLowerBc : BoundaryCondition
        Boundary condition for the plane of minimal y-coordinate.
    yUpperBc : BoundaryCondition
        Boundary condition for the plane of maximal y-coordinate.
    zLowerBc : BoundaryCondition
        Boundary condition for the plane of minimal z-coordinate.
    zUpperBc : BoundaryCondition
        Boundary condition for the plane of maximal z-coordinate.
    '''

    Ndim: int
    x: np.ndarray
    y: np.ndarray
    z: np.ndarray
    vx: np.ndarray
    vy: np.ndarray
    vz: np.ndarray
    xLowerBc: BoundaryCondition
    xUpperBc: BoundaryCondition
    yLowerBc: BoundaryCondition
    yUpperBc: BoundaryCondition
    zLowerBc: BoundaryCondition
    zUpperBc: BoundaryCondition
    stratifications: Optional[Stratifications] = None

    @classmethod
    def make_1d(cls, z: np.ndarray, vz: np.ndarray,
                lowerBc: BoundaryCondition, upperBc: BoundaryCondition,
                stratifications: Optional[Stratifications]=None) -> 'Layout':
        '''
        Construct 1D Layout.
        '''

        return cls(Ndim=1, x=np.array(()), y=np.array(()),
                   z=z, vx=np.array(()), vy=np.array(()),
                   vz=vz, xLowerBc=NoBc(), xUpperBc=NoBc(),
                   yLowerBc=NoBc(), yUpperBc=NoBc(),
                   zLowerBc=lowerBc, zUpperBc=upperBc,
                   stratifications=stratifications)

    @classmethod
    def make_2d(cls, x: np.ndarray, z: np.ndarray,
                vx: np.ndarray, vz: np.ndarray,
                xLowerBc: BoundaryCondition, xUpperBc: BoundaryCondition,
                zLowerBc: BoundaryCondition, zUpperBc: BoundaryCondition,
                stratifications: Optional[Stratifications]=None) -> 'Layout':
        '''
        Construct 2D Layout.
        '''

        Bc = BoundaryCondition
        return cls(Ndim=2, x=x, y=np.array(()), z=z,
                   vx=vx, vy=np.array(()), vz=vz,
                   xLowerBc=xLowerBc, xUpperBc=xUpperBc,
                   yLowerBc=NoBc(), yUpperBc=NoBc(),
                   zLowerBc=zLowerBc, zUpperBc=zUpperBc,
                   stratifications=stratifications)

    @classmethod
    def make_3d(cls, x: np.ndarray, y: np.ndarray, z: np.ndarray,
                vx: np.ndarray, vy: np.ndarray, vz: np.ndarray,
                xLowerBc: BoundaryCondition, xUpperBc: BoundaryCondition,
                yLowerBc: BoundaryCondition, yUpperBc: BoundaryCondition,
                zLowerBc: BoundaryCondition, zUpperBc: BoundaryCondition,
                stratifications: Optional[Stratifications]=None) -> 'Layout':
        '''
        Construct 3D Layout.
        '''

        return cls(Ndim=3, x=x, y=y, z=z,
                   vx=vx, vy=vy, vz=vz,
                   xLowerBc=xLowerBc, xUpperBc=xUpperBc,
                   yLowerBc=yLowerBc, yUpperBc=yUpperBc,
                   zLowerBc=zLowerBc, zUpperBc=zUpperBc,
                   stratifications=stratifications)

    @property
    def Nx(self) -> int:
        '''
        Number of grid points along the x-axis.
        '''
        return self.x.shape[0]

    @property
    def Ny(self) -> int:
        '''
        Number of grid points along the y-axis.
        '''
        return self.y.shape[0]

    @property
    def Nz(self) -> int:
        '''
        Number of grid points along the z-axis.
        '''
        return self.z.shape[0]

    @property
    def Noutgoing(self) -> int:
        '''
        Number of grid points at which the outgoing radiation is computed.
        '''
        return max(1, self.Nx, self.Nx * self.Ny)

    @property
    def vlos(self) -> np.ndarray:
        if self.Ndim > 1:
            raise ValueError('vlos is ambiguous when Ndim > 1, use vx, vy, or vz instead.')
        return self.vz

    @property
    def Nspace(self) -> int:
        '''
        Number of spatial points present in the grid.
        '''
        if self.Ndim == 1:
            return self.Nz
        elif self.Ndim == 2:
            return self.Nx * self.Nz
        elif self.Ndim == 3:
            return self.Nx * self.Ny * self.Nz
        else:
            raise ValueError('Invalid Ndim: %d, check geometry initialisation' % self.Ndim)

    @property
    def tauRef(self):
        '''
        Alias to `self.stratifications.tauRef`, if computed.
        '''
        if self.stratifications is not None:
            return self.stratifications.tauRef
        else:
            raise ValueError('tauRef not computed for this Atmosphere')

    @property
    def cmass(self):
        '''
        Alias to `self.stratifications.cmass`, if computed.
        '''
        if self.stratifications is not None:
            return self.stratifications.cmass
        else:
            raise ValueError('tauRef not computed for this Atmosphere')

    @property
    def dimensioned_shape(self):
        '''
        Tuple defining the shape to which the arrays of atmospheric paramters
        can be reshaped to be indexed in a 1/2/3D fashion.
        '''
        if self.Ndim == 1:
            shape = (self.Nz,)
        elif self.Ndim == 2:
            shape = (self.Nz, self.Nx)
        elif self.Ndim == 3:
            shape = (self.Nz, self.Ny, self.Nx)
        else:
            raise ValueError('Unreasonable Ndim (%d)' % self.Ndim)
        return shape

    def dimensioned_view(self) ->  'Layout':
        '''
        Returns a view over the contents of Layout reshaped so all data has
        the correct (1/2/3D) dimensionality for the atmospheric model, as
        these are all stored under a flat scheme.
        '''
        layout = copy(self)
        shape = self.dimensioned_shape
        if self.stratifications is not None:
            layout.stratifications = self.stratifications.dimensioned_view(shape)
        if self.vx.size > 0:
            layout.vx = self.vx.reshape(shape)
        if self.vy.size > 0:
            layout.vy = self.vy.reshape(shape)
        if self.vz.size > 0:
            layout.vz = self.vz.reshape(shape)
        return layout

    def unit_view(self) -> 'Layout':
        '''
        Returns a view over the contents of the Layout with the correct
        `astropy.units`.
        '''
        layout = copy(self)
        layout.x = self.x << u.m
        layout.y = self.y << u.m
        layout.z = self.z << u.m
        layout.vx = self.vx << u.m / u.s
        layout.vy = self.vy << u.m / u.s
        layout.vz = self.vz << u.m / u.s
        if self.stratifications is not None:
            layout.stratifications = self.stratifications.unit_view()
        return layout

    def dimensioned_unit_view(self) -> 'Layout':
        '''
        Returns a view over the contents of Layout reshaped so all data has
        the correct (1/2/3D) dimensionality for the atmospheric model, and
        the correct `astropy.units`.
        '''
        layout = self.dimensioned_view()
        return layout.unit_view()

@dataclass
class Atmosphere:
    '''
    Storage for all atmospheric data. These arrays will be shared directly
    with the backend, so a modification here also modifies the data seen by
    the backend. Be careful to modify these arrays *in place*, as their data
    is shared by direct memory reference. Use the class methods to construct
    atmospheres of different dimensionality.

    Attributes
    ----------
    structure : Layout
        A layout structure holding the atmospheric stratification, and
        velocity description.
    temperature : np.ndarray
        The atmospheric temperature structure.
    vturb : np.ndarray
        The atmospheric microturbulent velocity structure.
    ne : np.ndarray
        The electron density structure in the atmosphere.
    nHTot : np.ndarray
        The total hydrogen number density distribution throughout the
        atmosphere.
    B : np.ndarray, optional
        The magnitude of the stratified magnetic field throughout the
        atmosphere (Tesla).
    gammaB : np.ndarray, optional
        Co-altitude of magnetic field vector (radians) throughout the
        atmosphere from the local vertical.
    chiB : np.ndarray, optional
        Azimuth of magnetic field vector (radians) in the x-y plane, measured
        from the x-axis.
    '''

    structure: Layout
    temperature: np.ndarray
    vturb: np.ndarray
    ne: np.ndarray
    nHTot: np.ndarray
    B: Optional[np.ndarray] = None
    gammaB: Optional[np.ndarray] = None
    chiB: Optional[np.ndarray] = None

    @property
    def Ndim(self) -> int:
        '''
        Ndim : int
            The dimensionality (1, 2, or 3) of the atmospheric model.
        '''
        return self.structure.Ndim

    @property
    def Nx(self) -> int:
        '''
        Nx : int
            The number of points in the x-direction discretisation.
        '''
        return self.structure.Nx

    @property
    def Ny(self) -> int:
        '''
        Ny : int
            The number of points in the y-direction discretisation.
        '''
        return self.structure.Ny

    @property
    def Nz(self) -> int:
        '''
        Nz : int
            The number of points in the y-direction discretisation.
        '''
        return self.structure.Nz

    @property
    def Noutgoing(self) -> int:
        '''
        Noutgoing : int
            The number of cells at the top of the atmosphere (that each produce a
            spectrum).
        '''
        return self.structure.Noutgoing

    @property
    def vx(self) -> np.ndarray:
        '''
        vx : np.ndarray
            x component of plasma velocity (present for Ndim >= 2) [m/s].
        '''
        return self.structure.vx

    @property
    def vy(self) -> np.ndarray:
        '''
        vy : np.ndarray
            y component of plasma velocity (present for Ndim == 3) [m/s].
        '''
        return self.structure.vy

    @property
    def vz(self) -> np.ndarray:
        '''
        vz : np.ndarray
            z component of plasma velocity (present for all Ndim) [m/s]. Aliased
            to `vlos` when `Ndim==1`
        '''
        return self.structure.vz

    @property
    def vlos(self) -> np.ndarray:
        '''
        vz : np.ndarray
            z component of plasma velocity (present for all Ndim) [m/s]. Only
            available when Ndim==1`.
        '''
        return self.structure.vlos

    @property
    def cmass(self) -> np.ndarray:
        '''
        cmass : np.ndarray
            Column mass [kg m-2].
        '''
        return self.structure.cmass

    @property
    def tauRef(self) -> np.ndarray:
        '''
        tauRef : np.ndarray
            Reference optical depth at 500 nm.
        '''
        return self.structure.tauRef

    @property
    def height(self) -> np.ndarray:
        return self.structure.z

    @property
    def x(self) -> np.ndarray:
        '''
        x : np.ndarray
            Ordinates of grid points along the x-axis (present for Ndim >= 2) [m].
        '''
        return self.structure.x

    @property
    def y(self) -> np.ndarray:
        '''
        y : np.ndarray
            Ordinates of grid points along the y-axis (present for Ndim == 3) [m].
        '''
        return self.structure.y

    @property
    def z(self) -> np.ndarray:
        '''
        z : np.ndarray
            Ordinates of grid points along the z-axis (present for all Ndim) [m].
        '''
        return self.structure.z

    @property
    def zLowerBc(self) -> BoundaryCondition:
        '''
        zLowerBc : BoundaryCondition
            Boundary condition for the plane of minimal z-coordinate.
        '''
        return self.structure.zLowerBc

    @property
    def zUpperBc(self) -> BoundaryCondition:
        '''
        zUpperBc : BoundaryCondition
            Boundary condition for the plane of maximal z-coordinate.
        '''
        return self.structure.zUpperBc

    @property
    def yLowerBc(self) -> BoundaryCondition:
        '''
        yLowerBc : BoundaryCondition
            Boundary condition for the plane of minimal y-coordinate.
        '''
        return self.structure.yLowerBc

    @property
    def yUpperBc(self) -> BoundaryCondition:
        '''
        yUpperBc : BoundaryCondition
            Boundary condition for the plane of maximal y-coordinate.
        '''
        return self.structure.yUpperBc

    @property
    def xLowerBc(self) -> BoundaryCondition:
        '''
        xLowerBc : BoundaryCondition
            Boundary condition for the plane of minimal x-coordinate.
        '''
        return self.structure.xLowerBc

    @property
    def xUpperBc(self) -> BoundaryCondition:
        '''
        xUpperBc : BoundaryCondition
            Boundary condition for the plane of maximal x-coordinate.
        '''
        return self.structure.xUpperBc

    @property
    def Nspace(self):
        '''
        Nspace : int
            Total number of points in the atmospheric spatial discretistaion.
        '''
        return self.structure.Nspace

    @property
    def Nrays(self):
        '''
        Nrays : int
            Number of rays in angular discretisation used.
        '''
        if self.muz is None:
            raise AttributeError('Nrays not set, call atmos.rays or .quadrature first')

        return self.muz.shape[0]

    def dimensioned_view(self):
        '''
        Returns a view over the contents of Layout reshaped so all data has
        the correct (1/2/3D) dimensionality for the atmospheric model, as
        these are all stored under a flat scheme.
        '''
        shape = self.structure.dimensioned_shape
        atmos = copy(self)
        atmos.structure = self.structure.dimensioned_view()
        atmos.temperature = self.temperature.reshape(shape)
        atmos.vturb = self.vturb.reshape(shape)
        atmos.ne = self.ne.reshape(shape)
        atmos.nHTot = self.nHTot.reshape(shape)
        if self.B is not None:
            atmos.B = self.B.reshape(shape)
            atmos.chiB = self.chiB.reshape(shape)
            atmos.gammaB = self.gammaB.reshape(shape)
        return atmos

    def unit_view(self):
        '''
        Returns a view over the contents of the Layout with the correct
        `astropy.units`.
        '''
        atmos = copy(self)
        atmos.structure = self.structure.unit_view()
        atmos.temperature = self.temperature << u.K
        atmos.vturb = self.vturb << u.m / u.s
        atmos.ne = self.ne << u.m**(-3)
        atmos.nHTot = self.nHTot << u.m**(-3)
        if self.B is not None:
            atmos.B = self.B << u.T
            atmos.chiB = self.chiB << u.rad
            atmos.gammaB = self.gammaB << u.rad
        return atmos

    def dimensioned_unit_view(self):
        '''
        Returns a view over the contents of Layout reshaped so all data has
        the correct (1/2/3D) dimensionality for the atmospheric model, and
        the correct `astropy.units`.
        '''
        atmos = self.dimensioned_view()
        return atmos.unit_view()

    @classmethod
    def make_1d(cls, scale: ScaleType, depthScale: np.ndarray,
                temperature: np.ndarray, vlos: np.ndarray,
                vturb: np.ndarray, ne: Optional[np.ndarray]=None,
                hydrogenPops: Optional[np.ndarray]=None,
                nHTot: Optional[np.ndarray]=None,
                B: Optional[np.ndarray]=None,
                gammaB: Optional[np.ndarray]=None,
                chiB: Optional[np.ndarray]=None,
                lowerBc: Optional[BoundaryCondition]=None,
                upperBc: Optional[BoundaryCondition]=None,
                convertScales: bool=True,
                abundance: Optional[AtomicAbundance]=None,
                logG: float=2.44,
                Pgas: Optional[np.ndarray]=None,
                Pe: Optional[np.ndarray]=None,
                Ptop: Optional[float]=None,
                PeTop: Optional[float]=None,
                verbose: bool=False):
        '''
        Constructor for 1D Atmosphere objects. Optionally will use an
        equation of state (EOS) to estimate missing parameters.

        If sufficient information is provided (i.e. all required parameters
        and ne and (hydrogenPops or nHTot)) then the EOS is not invoked to
        estimate any thermodynamic properties. If both of nHTot and
        hydrogenPops are omitted, then the electron pressure will be used
        with the Wittmann equation of state to estimate the mass density, and
        the hydrogen number density will be inferred from this and the
        abundances. If, instead, ne is omitted, then the mass density will be
        used with the Wittmann EOS to estimate the electron pressure.
        If both of these are omitted then the EOS will be used to estimate
        both. If:

            - Pgas is provided, then this gas pressure will define the
              atmospheric stratification and will be used with the EOS.

            - Pe is provided, then this electron pressure will define the
              atmospheric stratification and will be used with the EOS.

            - Ptop is provided, then this gas pressure at the top of the
              atmosphere will be used with the log gravitational acceleration
              logG, and the EOS to estimate the missing parameters assuming
              hydrostatic equilibrium.

            - PeTop is provided, then this electron pressure at the top of
              the atmosphere will be used with the log gravitational
              acceleration logG, and the EOS to estimate the missing parameters
              assuming hydrostatic equilibrium.

            - If all of Pgas, Pe, Ptop, PeTop are omitted then Ptop will be
              estimated from the gas pressure in the FALC model at the
              temperature at the top boundary. The hydrostatic reconstruction
              will then continue as usual.

        convertScales will substantially slow down this function due to the
        slow calculation of background opacities used to compute tauRef. If
        an atmosphere is constructed with a Geometric stratification, and an
        estimate of tauRef is not required before running the main RT module,
        then this can be set to False.
        All of these parameters can be provided as astropy Quantities, and
        will be converted in the constructor.

        Parameters
        ----------
        scale : ScaleType
            The type of stratification used along the z-axis.
        depthScale : np.ndarray
            The z-coordinates used along the chosen stratification. The
            stratification is expected to start at the top of the atmosphere
            (closest to the observer), and descend along the observer's line
            of sight.
        temperature : np.ndarray
            Temperature structure of the atmosphere [K].
        vlos : np.ndarray
            Velocity structure of the atmosphere along z [m/s].
        vturb : np.ndarray
            Microturbulent velocity structure of the atmosphere [m/s].
        ne : np.ndarray
            Electron density structure of the atmosphere [m-3].
        hydrogenPops : np.ndarray, optional
            Detailed (per level) hydrogen number density structure of the
            atmosphere [m-3], 2D array [Nlevel, Nspace].
        nHTot : np.ndarray, optional
            Total hydrogen number density structure of the atmosphere [m-3]
        B : np.ndarray, optional.
            Magnetic field strength [T].
        gammaB : np.ndarray, optional
            Co-altitude of magnetic field vector [radians].
        chiB : np.ndarray, optional
            Azimuth of magnetic field vector (in x-y plane, from x) [radians].
        lowerBc : BoundaryCondition, optional
            Boundary condition for incoming radiation at the minimal z
            coordinate (default: ThermalisedRadiation).
        upperBc : BoundaryCondition, optional
            Boundary condition for incoming radiation at the maximal z
            coordinate (default: ZeroRadiation).
        convertScales : bool, optional
            Whether to automatically compute tauRef and cmass for an
            atmosphere given in a stratification of m (default: True).
        abundance: AtomicAbundance, optional
            An instance of AtomicAbundance giving the abundances of each
            atomic species in the given atmosphere, only used if the EOS is
            invoked. (default: DefaultAtomicAbundance)
        logG: float, optional
            The log10 of the magnitude of gravitational acceleration [m/s2]
            (default: 2.44).
        Pgas: np.ndarray, optional
            The gas pressure stratification of the atmosphere [Pa],
            optionally used by the EOS.
        Pe: np.ndarray, optional
            The electron pressure stratification of the atmosphere [Pa],
            optionally used by the EOS.
        Ptop: np.ndarray, optional
            The gas pressure at the top of the atmosphere [Pa], optionally
            used by the EOS for a hydrostatic reconstruction.
        Petop: np.ndarray, optional
            The electron pressure at the top of the atmosphere [Pa],
            optionally used by the EOS for a hydrostatic reconstruction.
        verbose: bool, optional
            Explain decisions made with the EOS to estimate missing
            parameters (if invoked) through print calls (default: False).

        Raises
        ------
        ValueError
            if incorrect arguments or unable to construct estimate missing
            parameters.
        '''
        if scale == ScaleType.Geometric:
            depthScale = (depthScale << u.m).value
        elif scale == ScaleType.ColumnMass:
            depthScale = (depthScale << u.kg / u.m**2).value
        temperature = (temperature << u.K).value
        vlos = (vlos << u.m / u.s).value
        vturb = (vturb << u.m / u.s).value
        if ne is not None:
            ne = (ne << u.m**(-3)).value
        if hydrogenPops is not None:
            hydrogenPops = (hydrogenPops << u.m**(-3)).value
        if nHTot is not None:
            nHTot = (nHTot << u.m**(-3)).value
        if B is not None:
            B = (B << u.T).value
        if gammaB is not None:
            gammaB = (gammaB << u.rad).value
        if chiB is not None:
            chiB = (chiB << u.rad).value

        if lowerBc is None:
            lowerBc = ThermalisedRadiation()
        elif isinstance(lowerBc, PeriodicRadiation):
            raise ValueError('Cannot set periodic boundary conditions for 1D atmosphere')
        if upperBc is None:
            upperBc = ZeroRadiation()
        elif isinstance(upperBc, PeriodicRadiation):
            raise ValueError('Cannot set periodic boundary conditions for 1D atmosphere')

        if scale != ScaleType.Geometric and not convertScales:
            raise ValueError('Height scale must be provided if scale conversion is not applied')

        if nHTot is None and hydrogenPops is not None:
            nHTot = np.sum(hydrogenPops, axis=0)

        if np.any(temperature < 2000):
            # NOTE(cmo): Minimum value was decreased in NICOLE so should be safe
            raise ValueError('Minimum temperature too low for EOS (< 2000 K)')

        if abundance is None:
            abundance = DefaultAtomicAbundance

        wittAbundances = np.array([abundance[e] for e in PeriodicTable.elements])
        eos = witt(abund_init=wittAbundances)

        Nspace = depthScale.shape[0]
        if nHTot is None and ne is not None:
            if verbose:
                print('Setting nHTot from electron pressure.')
            pe = ne * Const.CM_TO_M**3 * eos.BK * temperature
            rho = np.zeros(Nspace)
            for k in range(Nspace):
                rho[k] = eos.rho_from_pe(temperature[k], pe[k])
            nHTot = np.copy(rho / (Const.CM_TO_M**3 / Const.G_TO_KG) / (Const.Amu * abundance.massPerH))
        elif ne is None and nHTot is not None:
            if verbose:
                print('Setting ne from mass density.')
            rho = Const.Amu * abundance.massPerH * nHTot * Const.CM_TO_M**3 / Const.G_TO_KG
            pe = np.zeros(Nspace)
            for k in range(Nspace):
                pe[k] = eos.pe_from_rho(temperature[k], rho[k])
            ne = np.copy(pe / (eos.BK * temperature) / Const.CM_TO_M**3)
        elif ne is None and nHTot is None:
            if Pgas is not None and Pgas.shape[0] != Nspace:
                raise ValueError('Dimensions of Pgas do not match atmospheric depth')
            if Pe is not None and Pe.shape[0] != Nspace:
                raise ValueError('Dimensions of Pe do not match atmospheric depth')

            if Pgas is not None and Pe is None:
                if verbose:
                    print('Setting ne, nHTot from provided gas pressure.')
                # Convert to cgs for eos
                pgas = Pgas * (Const.CM_TO_M**2 / Const.G_TO_KG)
                pe = np.zeros(Nspace)
                for k in range(Nspace):
                    pe[k] = eos.pe_from_pg(temperature[k], pgas[k])
            elif Pe is not None and Pgas is None:
                if verbose:
                    print('Setting ne, nHTot from provided electron pressure.')
                # Convert to cgs for eos
                pe = Pe * (Const.CM_TO_M**2 / Const.G_TO_KG)
                pgas = np.zeros(Nspace)
                for k in range(Nspace):
                    pgas[k] = eos.pg_from_pe(temperature[k], pe[k])
            elif Pgas is None and Pe is None:
                # Doing Hydrostatic Eq. based here on NICOLE implementation
                gravAcc = 10**logG / Const.CM_TO_M
                Avog = 6.022045e23 # Avogadro's Number
                if Ptop is None and PeTop is not None:
                    if verbose:
                        print('Setting ne, nHTot to hydrostatic equilibrium (logG=%f) from provided top electron pressure.' % logG)
                    PeTop *= (Const.CM_TO_M**2 / Const.G_TO_KG)
                    Ptop = eos.pg_from_pe(temperature[0], PeTop)
                elif Ptop is not None and PeTop is None:
                    if verbose:
                        print('Setting ne, nHTot to hydrostatic equilibrium (logG=%f) from provided top gas pressure.' % logG)
                    Ptop *= (Const.CM_TO_M**2 / Const.G_TO_KG)
                    PeTop = eos.pe_from_pg(temperature[0], Ptop)
                elif Ptop is None and PeTop is None:
                    if verbose:
                        print('Setting ne, nHTot to hydrostatic equilibrium (logG=%f) from FALC gas pressure at upper boundary temperature.' % logG)
                    Ptop = get_top_pressure(eos, temperature[0])
                    PeTop = eos.pe_from_pg(temperature[0], Ptop)
                else:
                    raise ValueError("Cannot set both Ptop and PeTop")

                if scale == ScaleType.Tau500:
                    tau = depthScale
                elif scale == ScaleType.Geometric:
                    height = depthScale / Const.CM_TO_M
                else:
                    cmass = depthScale / Const.G_TO_KG * Const.CM_TO_M**2

                # NOTE(cmo): Compute HSE following the NICOLE method.
                rho = np.zeros(Nspace)
                chi_c = np.zeros(Nspace)
                pgas = np.zeros(Nspace)
                pe = np.zeros(Nspace)
                pgas[0] = Ptop
                pe[0] = PeTop
                chi_c[0] = eos.contOpacity(temperature[0], pgas[0], pe[0], np.array([5000.0]))
                avg_mol_weight = lambda k: abundance.massPerH / (abundance.totalAbundance + pe[k] / pgas[k])
                rho[0] = Ptop * avg_mol_weight(0) / Avog / eos.BK / temperature[0]
                chi_c[0] /= rho[0]

                for k in range(1, Nspace):
                    chi_c[k] = chi_c[k-1]
                    rho[k] = rho[k-1]
                    for it in range(200):
                        if scale == ScaleType.Tau500:
                            dtau = tau[k] - tau[k-1]
                            pgas[k] = pgas[k-1] + gravAcc * dtau / (0.5 * (chi_c[k-1] + chi_c[k]))
                        elif scale == ScaleType.Geometric:
                            pgas[k] = pgas[k-1] * np.exp(-gravAcc / Avog / eos.BK * avg_mol_weight(k-1) * 0.5 * (1.0 / temperature[k-1] + 1.0 / temperature[k]) * (height[k] - height[k-1]))
                        else:
                            pgas[k] = gravAcc * cmass[k]

                        pe[k] = eos.pe_from_pg(temperature[k], pgas[k])
                        prevChi = chi_c[k]
                        chi_c[k] = eos.contOpacity(temperature[k], pgas[k], pe[k], np.array([5000.0]))
                        rho[k] = pgas[k] * avg_mol_weight(k) / Avog / eos.BK / temperature[k]
                        chi_c[k] /= rho[k]

                        change = np.abs(prevChi - chi_c[k]) / (prevChi + chi_c[k])
                        if change < 1e-5:
                            break
                    else:
                        raise ConvergenceError('No convergence in HSE at depth point %d, last change %2.4e' % (k, change))
            nHTot = np.copy(rho / (Const.CM_TO_M**3 / Const.G_TO_KG) / (Const.Amu * abundance.massPerH))
            ne = np.copy(pe / (eos.BK * temperature) / Const.CM_TO_M**3)

        # NOTE(cmo): Compute final pgas, pe from EOS that will be used for
        # background opacity.
        rhoSI = Const.Amu * abundance.massPerH * nHTot
        rho = Const.Amu * abundance.massPerH * nHTot * Const.CM_TO_M**3 / Const.G_TO_KG
        pgas = np.zeros_like(depthScale)
        pe = np.zeros_like(depthScale)
        for k in range(Nspace):
            pgas[k] = eos.pg_from_rho(temperature[k], rho[k])
            pe[k] = eos.pe_from_rho(temperature[k], rho[k])

        chi_c = np.zeros_like(depthScale)
        for k in range(depthScale.shape[0]):
            chi_c[k] = eos.contOpacity(temperature[k], pgas[k], pe[k], np.array([5000.0])) / Const.CM_TO_M

        # NOTE(cmo): We should now have a uniform minimum set of data (other
        # than the scale type), allowing us to simply convert between the
        # scales we do have!
        if convertScales:
            if scale == ScaleType.ColumnMass:
                height = np.zeros_like(depthScale)
                tau_ref = np.zeros_like(depthScale)
                cmass = depthScale

                height[0] = 0.0
                tau_ref[0] = chi_c[0] / rhoSI[0] * cmass[0]
                for k in range(1, cmass.shape[0]):
                    height[k] = height[k-1] - 2.0 * (cmass[k] - cmass[k-1]) / (rhoSI[k-1] + rhoSI[k])
                    tau_ref[k] = tau_ref[k-1] + 0.5 * (chi_c[k-1] + chi_c[k]) * (height[k-1] - height[k])

                hTau1 = np.interp(1.0, tau_ref, height)
                height -= hTau1
            elif scale == ScaleType.Geometric:
                cmass = np.zeros(Nspace)
                tau_ref = np.zeros(Nspace)
                height = depthScale

                cmass[0] = (nHTot[0] * abundance.massPerH + ne[0]) * (Const.KBoltzmann * temperature[0] / 10**logG)
                tau_ref[0] = 0.5 * chi_c[0] * (height[0] - height[1])
                if tau_ref[0] > 1.0:
                    tau_ref[0] = 0.0

                for k in range(1, Nspace):
                    cmass[k] = cmass[k-1] + 0.5 * (rhoSI[k-1] + rhoSI[k]) * (height[k-1] - height[k])
                    tau_ref[k] = tau_ref[k-1] + 0.5 * (chi_c[k-1] + chi_c[k]) * (height[k-1] - height[k])
            elif scale == ScaleType.Tau500:
                cmass = np.zeros(Nspace)
                height = np.zeros(Nspace)
                tau_ref = depthScale

                cmass[0] = (tau_ref[0] / chi_c[0]) * rhoSI[0]
                for k in range(1, Nspace):
                    height[k] = height[k-1] - 2.0 * (tau_ref[k] - tau_ref[k-1]) / (chi_c[k-1] + chi_c[k])
                    cmass[k] = cmass[k-1] + 0.5 * (chi_c[k-1] + chi_c[k]) * (height[k-1] - height[k])

                hTau1 = np.interp(1.0, tau_ref, height)
                height -= hTau1
            else:
                raise ValueError("Other scales not handled yet")

            stratifications: Optional[Stratifications] = Stratifications(cmass=cmass, tauRef=tau_ref)

        else:
            stratifications = None

        layout = Layout.make_1d(z=height, vz=vlos,
                                lowerBc=lowerBc, upperBc=upperBc,
                                stratifications=stratifications)
        atmos = cls(structure=layout, temperature=temperature, vturb=vturb,
                    ne=ne, nHTot=nHTot, B=B, gammaB=gammaB, chiB=chiB)

        return atmos

    @classmethod
    def make_2d(cls, height: np.ndarray, x: np.ndarray,
                temperature: np.ndarray, vx: np.ndarray,
                vz: np.ndarray, vturb: np.ndarray,
                ne: Optional[np.ndarray]=None,
                nHTot: Optional[np.ndarray]=None,
                B: Optional[np.ndarray]=None,
                gammaB: Optional[np.ndarray]=None,
                chiB: Optional[np.ndarray]=None,
                xUpperBc: Optional[BoundaryCondition]=None,
                xLowerBc: Optional[BoundaryCondition]=None,
                zUpperBc: Optional[BoundaryCondition]=None,
                zLowerBc: Optional[BoundaryCondition]=None,
                abundance: Optional[AtomicAbundance]=None,
                verbose=False):
        '''
        Constructor for 2D Atmosphere objects.

        No provision for estimating parameters using hydrostatic equilibrium
        is provided, but one of ne, or nHTot can be omitted and inferred by
        use of the Wittmann equation of state.
        The atmosphere must be defined on a geometric stratification.
        All atmospheric parameters are expected in a 2D [z, x] array.

        Parameters
        ----------
        height : np.ndarray
            The z-coordinates of the atmospheric grid. The stratification is
            expected to start at the top of the atmosphere (closest to the
            observer), and descend along the observer's line of sight.
        x : np.ndarray
            The (horizontal) x-coordinates of the atmospheric grid.
        temperature : np.ndarray
            Temperature structure of the atmosphere [K].
        vx : np.ndarray
            x-component of the atmospheric velocity [m/s].
        vz : np.ndarray
            z-component of the atmospheric velocity [m/s].
        vturb : np.ndarray
            Microturbulent velocity structure [m/s].
        ne : np.ndarray
            Electron density structure of the atmosphere [m-3].
        nHTot : np.ndarray, optional
            Total hydrogen number density structure of the atmosphere [m-3].
        B : np.ndarray, optional.
            Magnetic field strength [T].
        gammaB : np.ndarray, optional
            Co-altitude of magnetic field vector [radians].
        chiB : np.ndarray, optional
            Azimuth of magnetic field vector (in x-y plane, from x) [radians].
        xLowerBc : BoundaryCondition, optional
            Boundary condition for incoming radiation at the minimal x
            coordinate (default: PeriodicRadiation).
        xUpperBc : BoundaryCondition, optional
            Boundary condition for incoming radiation at the maximal x
            coordinate (default: PeriodicRadiation).
        zLowerBc : BoundaryCondition, optional
            Boundary condition for incoming radiation at the minimal z
            coordinate (default: ThermalisedRadiation).
        zUpperBc : BoundaryCondition, optional
            Boundary condition for incoming radiation at the maximal z
            coordinate (default: ZeroRadiation).
        convertScales : bool, optional
            Whether to automatically compute tauRef and cmass for an
            atmosphere given in a stratification of m (default: True).
        abundance: AtomicAbundance, optional
            An instance of AtomicAbundance giving the abundances of each
            atomic species in the given atmosphere, only used if the EOS is
            invoked. (default: DefaultAtomicAbundance)
        verbose: bool, optional
            Explain decisions made with the EOS to estimate missing
            parameters (if invoked) through print calls (default: False).

        Raises
        ------
        ValueError
            if incorrect arguments or unable to construct estimate missing
            parameters.
        '''

        x = (x << u.m).value
        height = (height << u.m).value
        temperature = (temperature << u.K).value
        vx = (vx << u.m / u.s).value
        vz = (vz << u.m / u.s).value
        vturb = (vturb << u.m / u.s).value
        if ne is not None:
            ne = (ne << u.m**(-3)).value
        if nHTot is not None:
            nHTot = (nHTot << u.m**(-3)).value
        if B is not None:
            B = (B << u.T).value
            flatB = view_flatten(B)
        else:
            flatB = None

        if gammaB is not None:
            gammaB = (gammaB << u.rad).value
            flatGammaB = view_flatten(gammaB)
        else:
            flatGammaB = None

        if chiB is not None:
            chiB = (chiB << u.rad).value
            flatChiB = view_flatten(chiB)
        else:
            flatChiB = None

        if zLowerBc is None:
            zLowerBc = ThermalisedRadiation()
        elif isinstance(zLowerBc, PeriodicRadiation):
            raise ValueError('Cannot set periodic boundary conditions for z-axis.')
        if zUpperBc is None:
            zUpperBc = ZeroRadiation()
        elif isinstance(zUpperBc, PeriodicRadiation):
            raise ValueError('Cannot set periodic boundary conditions for z-axis.')
        if xUpperBc is None:
            xUpperBc = PeriodicRadiation()
        if xLowerBc is None:
            xLowerBc = PeriodicRadiation()
        if abundance is None:

            abundance = DefaultAtomicAbundance

        wittAbundances = np.array([abundance[e] for e in PeriodicTable.elements])
        eos = witt(abund_init=wittAbundances)

        flatHeight = view_flatten(height)
        flatTemperature = view_flatten(temperature)
        Nspace = flatHeight.shape[0]
        if nHTot is None and ne is not None:
            if verbose:
                print('Setting nHTot from electron pressure.')
            flatNe = view_flatten(ne)
            pe = flatNe * Const.CM_TO_M**3 * eos.BK * flatTemperature
            rho = np.zeros(Nspace)
            for k in range(Nspace):
                rho[k] = eos.rho_from_pe(flatTemperature[k], pe[k])
            nHTot = np.copy(rho / (Const.CM_TO_M**3 / Const.G_TO_KG) / (Const.Amu * abundance.massPerH))
        elif ne is None and nHTot is not None:
            if verbose:
                print('Setting ne from mass density.')
            flatNHTot = view_flatten(nHTot)
            rho = Const.Amu * abundance.massPerH * flatNHTot * Const.CM_TO_M**3 / Const.G_TO_KG
            pe = np.zeros(Nspace)
            for k in range(Nspace):
                pe[k] = eos.pe_from_rho(flatTemperature[k], rho[k])
            ne = np.copy(pe / (eos.BK * flatTemperature) / Const.CM_TO_M**3)
        elif ne is None and nHTot is None:
            raise ValueError('Cannot omit both ne and nHTot (currently).')
        flatX = view_flatten(x)
        flatNHTot = view_flatten(nHTot)
        flatNe = view_flatten(ne)
        flatVx = view_flatten(vx)
        flatVz = view_flatten(vz)
        flatVturb = view_flatten(vturb)

        layout = Layout.make_2d(x=flatX, z=flatHeight, vx=flatVx, vz=flatVz,
                                xLowerBc=xLowerBc, xUpperBc=xUpperBc,
                                zLowerBc=zLowerBc, zUpperBc=zUpperBc,
                                stratifications=None)

        atmos = cls(structure=layout, temperature=flatTemperature,
                    vturb=flatVturb, ne=flatNe, nHTot=flatNHTot, B=flatB,
                    gammaB=flatGammaB, chiB=flatChiB)
        return atmos

    def quadrature(self, Nrays: Optional[int]=None,
                   mu: Optional[Sequence[float]]=None,
                   wmu: Optional[Sequence[float]]=None):
        '''
        Compute the angular quadrature for solving the RTE and Kinetic
        Equilibrium in a given atmosphere.

        Procedure varies with dimensionality.

        By convention muz is always positive, as the direction on this axis
        is determined by the toObs term that is used internally to the formal
        solver.

        1D:
            If a number of rays is given (typically 3 or 5), then the
            Gauss-Legendre quadrature for this set is used.
            If mu and wmu are instead given then these will be validated and
            used.

        2+D:
            If the number of rays selected is in the list of near optimal
            quadratures for unpolarised radiation provided by Stepan et al
            2020 (A&A, 646 A24), then this is used. Otherwise an exception is
            raised.

            The available quadratures are:

            +--------+-------+
            | Points | Order |
            +========+=======+
            |   1    |  3    |
            +--------+-------+
            |   3    |  7    |
            +--------+-------+
            |   6    |  9    |
            +--------+-------+
            |   7    |  11   |
            +--------+-------+
            |   10   |  13   |
            +--------+-------+
            |   11   |  15   |
            +--------+-------+

        Parameters
        ----------
        Nrays : int, optional
            The number of rays to use in the quadrature. See notes above.
        mu : sequence of float, optional
            The cosine of the angle made between the between each of the set
            of rays and the z axis, only used in 1D.
        wmu : sequence of float, optional
            The integration weights for each mu, must be provided if mu is provided.

        Raises
        ------
        ValueError
            on incorrect input.
        '''

        if self.Ndim == 1:
            if Nrays is not None and mu is None:
                if Nrays >= 1:
                    x, w = leggauss(Nrays)
                    mid, halfWidth = 0.5, 0.5
                    x = mid + halfWidth * x
                    w *= halfWidth

                    self.muz = x
                    self.wmu = w
                else:
                    raise ValueError('Unsupported Nrays=%d' % Nrays)
            elif Nrays is not None and mu is not None:
                if wmu is None:
                    raise ValueError('Must provide wmu when providing mu')
                if Nrays != len(mu):
                    raise ValueError('mu must be Nrays long if Nrays is provided')
                if len(mu) != len(wmu):
                    raise ValueError('mu and wmu must be the same shape')

                self.muz = np.array(mu)
                self.wmu = np.array(wmu)

            self.muy = np.zeros_like(self.muz)
            self.mux = np.sqrt(1.0 - self.muz**2)
        else:
            with open(get_data_path() + 'Quadratures.pickle', 'rb') as pkl:
                quads = pickle.load(pkl)

            rays = {int(q.split('n')[1]): q for q in quads}
            if Nrays not in rays:
                raise ValueError('For multidimensional cases Nrays must be in %s' % repr(rays))

            quad = quads[rays[Nrays]]

            if self.Ndim == 2:
                Nrays *= 2
                theta = np.deg2rad(quad[:, 1])
                chi = np.deg2rad(quad[:, 2])
                # polar coords:
                # x = sin theta cos chi
                # y = sin theta sin chi
                # z = cos theta
                self.mux = np.zeros(Nrays)
                self.mux[:Nrays // 2] = np.sin(theta) * np.cos(chi)
                self.mux[Nrays // 2:] = -np.sin(theta) * np.cos(chi)
                self.muz = np.zeros(Nrays)
                self.muz[:Nrays // 2] = np.cos(theta)
                self.muz[Nrays // 2:] = np.cos(theta)
                self.wmu = np.zeros(Nrays)
                self.wmu[:Nrays // 2] = quad[:, 0]
                self.wmu[Nrays // 2:] = quad[:, 0]
                self.wmu /= np.sum(self.wmu)
                self.muy = np.sqrt(1.0 - (self.mux**2 + self.muz**2))

            else:
                raise NotImplementedError()

        self.configure_bcs()

    def rays(self, muz: Union[float, Sequence[float]],
             mux: Optional[Union[float, Sequence[float]]]=None,
             muy: Optional[Union[float, Sequence[float]]]=None,
             wmu: Optional[Union[float, Sequence[float]]]=None):
        '''
        Set up the rays on the Atmosphere for computing the intensity in a
        particular direction (or set of directions).

        If only the z angle is set then the ray is assumed in the x-z plane.
        If either muz or muy is omitted then this angle is inferred by
        normalisation of the projection.

        By convention muz is always positive, as the direction on this axis
        is determined by the toObs term that is used internally to the formal
        solver.

        Parameters
        ----------
        muz : float or sequence of float, optional
            The angular projections along the z axis.
        mux : float or sequence of float, optional
            The angular projections along the x axis.
        muy : float or sequence of float, optional
            The angular projections along the y axis.
        wmu : float or sequence of float, optional
            The integration weights for the given ray if J is to be
            integrated for angle set.

        Raises
        ------
        ValueError
            if the angular projections or integration weights are incorrectly
            normalised.
        '''

        if isinstance(muz, float):
            muz = [muz]
        if isinstance(mux, float):
            mux = [mux]
        if isinstance(muy, float):
            muy = [muy]

        if mux is None and muy is None:
            self.muz = np.array(muz)
            self.wmu = np.zeros_like(self.muz)
            self.muy = np.zeros_like(self.muz)
            self.mux = np.sqrt(1.0 - self.muz**2)
        elif muy is None:
            self.muz = np.array(muz)
            self.wmu = np.zeros_like(self.muz)
            self.mux = np.array(mux)
            self.muy = np.sqrt(1.0 - (self.muz**2 + self.mux**2))
        elif mux is None:
            self.muz = np.array(muz)
            self.wmu = np.zeros_like(self.muz)
            self.muy = np.array(muy)
            self.mux = np.sqrt(1.0 - (self.muz**2 + self.muy**2))
        else:
            self.muz = np.array(muz)
            self.mux = np.array(mux)
            self.muy = np.array(muy)
            self.wmu = np.zeros_like(muz)

            if not np.allclose(self.muz**2 + self.mux**2 + self.muy**2, 1):
                raise ValueError('mux**2 + muy**2 + muz**2 != 1.0')

        if not np.all(self.muz > 0):
            raise ValueError('muz must be > 0')

        if wmu is not None:
            self.wmu = np.array(wmu)

            if not np.isclose(self.wmu.sum(), 1.0):
                raise ValueError('sum of wmus is not 1.0')

        self.configure_bcs()

    def configure_bcs(self):
        '''
        Configure the required angular information for all boundary
        conditions on the model.
        '''

        # NOTE(cmo): We always have z-bcs
        # For zLowerBc, muz is positive, and we have all mux, muz
        mux, muy, muz = self.mux, self.muy, self.muz
        # NOTE(cmo): indexVector is of shape (mu, toObs) to allow the core to
        # easily destructure the blob that will be handed to it from
        # compute_bc.
        indexVector = np.ones((self.mux.shape[0], 2), dtype=np.int32) * -1
        indexVector[:, 1] = np.arange(mux.shape[0])
        self.zLowerBc.set_required_angles(mux, muy, muz, indexVector)

        indexVector = np.ones((mux.shape[0], 2), dtype=np.int32) * -1
        indexVector[:, 0] = np.arange(mux.shape[0])
        self.zUpperBc.set_required_angles(-mux, -muy, -muz, indexVector)

        # NOTE(cmo): If 2+D we have x-bcs too
        # xLowerBc has all muz and all mux > 0
        mux, muy, muz = [], [], []
        indexVector = np.ones((self.mux.shape[0], 2), dtype=np.int32) * -1
        count = 0
        musDone = np.zeros(self.muz.shape[0], dtype=np.bool)
        for mu in range(self.muz.shape[0]):
            for equalMu in np.argwhere(np.abs(self.muz) == self.muz[mu]).reshape(-1)[::-1]:
                if musDone[equalMu]:
                    continue
                musDone[equalMu] = True

                for toObsI in range(2):
                    sign = [-1, 1][toObsI]
                    sMux = sign * self.mux[equalMu]
                    if sMux > 0:
                        mux.append(sMux)
                        muy.append(sign * self.muy[equalMu])
                        muz.append(sign * self.muz[equalMu])
                        indexVector[equalMu, toObsI] = count
                        count += 1
            if np.all(musDone):
                break

        mux = np.array(mux)
        muy = np.array(muy)
        muz = np.array(muz)
        self.xLowerBc.set_required_angles(mux, muy, muz, indexVector)

        mux, muy, muz = [], [], []
        indexVector = np.ones((self.mux.shape[0], 2), dtype=np.int32) * -1
        count = 0
        musDone = np.zeros(self.muz.shape[0], dtype=np.bool)
        for mu in range(self.muz.shape[0]):
            for equalMu in np.argwhere(np.abs(self.muz) == self.muz[mu]).reshape(-1):
                if musDone[equalMu]:
                    continue
                musDone[equalMu] = True

                for toObsI in range(2):
                    sign = [-1, 1][toObsI]
                    sMux = sign * self.mux[equalMu]
                    if sMux < 0:
                        mux.append(sMux)
                        muy.append(sign * self.muy[equalMu])
                        muz.append(sign * self.muz[equalMu])
                        indexVector[equalMu, toObsI] = count
                        count += 1
            if np.all(musDone):
                break

        mux = np.array(mux)
        muy = np.array(muy)
        muz = np.array(muz)
        self.xUpperBc.set_required_angles(mux, muy, muz, indexVector)

        self.yLowerBc.set_required_angles(np.zeros((0)), np.zeros((0)), np.zeros((0)),
                                          np.ones((self.mux.shape[0], 2), dtype=np.int32) * -1)
        self.yUpperBc.set_required_angles(np.zeros((0)), np.zeros((0)), np.zeros((0)),
                                          np.ones((self.mux.shape[0], 2), dtype=np.int32) * -1)

        if self.Ndim > 2:
            raise ValueError('Only <= 2D atmospheres supported currently.')


import atsapi as ats
import numpy as np

# hardware specific constants
HARDWARE_SPEC = {
    # minimum record size
    'MIN_REC_SIZE': {
        'ATS9870': 256,
        'ATS9360': 256
    },
    # minimum record size resolution
    'REC_SIZE_RES': {
        'ATS9870': 64,
        'ATS9360': 128
    },
    # external trigger range
    'EXT_TRIG_RANGE': {
        'ATS9870': 'ETR_5V',
        'ATS9360': 'ETR_2V5'
    },
    # internal sample rate
    'SAMPLE_RATE': {
        'ATS9870': [
            'SAMPLE_RATE_1KSPS',
            'SAMPLE_RATE_2KSPS',
            'SAMPLE_RATE_5KSPS',
            'SAMPLE_RATE_10KSPS',
            'SAMPLE_RATE_20KSPS',
            'SAMPLE_RATE_50KSPS',
            'SAMPLE_RATE_100KSPS',
            'SAMPLE_RATE_200KSPS',
            'SAMPLE_RATE_500KSPS',
            'SAMPLE_RATE_1MSPS',
            'SAMPLE_RATE_2MSPS',
            'SAMPLE_RATE_5MSPS',
            'SAMPLE_RATE_10MSPS',
            'SAMPLE_RATE_20MSPS',
            'SAMPLE_RATE_50MSPS',
            'SAMPLE_RATE_100MSPS',
            'SAMPLE_RATE_200MSPS',
            'SAMPLE_RATE_500MSPS',
            'SAMPLE_RATE_1000MSPS'
        ],
        'ATS9360': [
            'SAMPLE_RATE_1KSPS',
            'SAMPLE_RATE_2KSPS',
            'SAMPLE_RATE_5KSPS',
            'SAMPLE_RATE_10KSPS',
            'SAMPLE_RATE_20KSPS',
            'SAMPLE_RATE_50KSPS',
            'SAMPLE_RATE_100KSPS',
            'SAMPLE_RATE_200KSPS',
            'SAMPLE_RATE_500KSPS',
            'SAMPLE_RATE_1MSPS',
            'SAMPLE_RATE_2MSPS',
            'SAMPLE_RATE_5MSPS',
            'SAMPLE_RATE_10MSPS',
            'SAMPLE_RATE_20MSPS',
            'SAMPLE_RATE_50MSPS',
            'SAMPLE_RATE_100MSPS',
            'SAMPLE_RATE_200MSPS',
            'SAMPLE_RATE_500MSPS',
            'SAMPLE_RATE_800MSPS',
            'SAMPLE_RATE_1000MSPS',
            'SAMPLE_RATE_1200MSPS',
            'SAMPLE_RATE_1500MSPS',
            'SAMPLE_RATE_1800MSPS'
        ]
    }
}

SAMPLE_RATE_VALUE = {
    'SAMPLE_RATE_1KSPS': 1e3,
    'SAMPLE_RATE_2KSPS': 2e3,
    'SAMPLE_RATE_5KSPS': 5e3,
    'SAMPLE_RATE_10KSPS': 10e3,
    'SAMPLE_RATE_20KSPS': 20e3,
    'SAMPLE_RATE_50KSPS': 50e3,
    'SAMPLE_RATE_100KSPS': 100e3,
    'SAMPLE_RATE_200KSPS': 200e3,
    'SAMPLE_RATE_500KSPS': 500e3,
    'SAMPLE_RATE_1MSPS': 1e6,
    'SAMPLE_RATE_2MSPS': 2e6,
    'SAMPLE_RATE_5MSPS': 5e6,
    'SAMPLE_RATE_10MSPS': 10e6,
    'SAMPLE_RATE_20MSPS': 20e6,
    'SAMPLE_RATE_25MSPS': 25e6,
    'SAMPLE_RATE_50MSPS': 50e6,
    'SAMPLE_RATE_100MSPS': 100e6,
    'SAMPLE_RATE_125MSPS': 125e6,
    'SAMPLE_RATE_160MSPS': 160e6,
    'SAMPLE_RATE_180MSPS': 180e6,
    'SAMPLE_RATE_200MSPS': 200e6,
    'SAMPLE_RATE_250MSPS': 250e6,
    'SAMPLE_RATE_400MSPS': 400e6,
    'SAMPLE_RATE_500MSPS': 500e6,
    'SAMPLE_RATE_800MSPS': 800e6,
    'SAMPLE_RATE_1000MSPS': 1000e6,
    'SAMPLE_RATE_1200MSPS': 1200e6,
    'SAMPLE_RATE_1500MSPS': 1500e6,
    'SAMPLE_RATE_1600MSPS': 1600e6,
    'SAMPLE_RATE_1800MSPS': 1800e6,
    'SAMPLE_RATE_2000MSPS': 2000e6,
    'SAMPLE_RATE_2400MSPS': 2400e6,
    'SAMPLE_RATE_3000MSPS': 3000e6,
    'SAMPLE_RATE_3600MSPS': 3600e6,
    'SAMPLE_RATE_4000MSPS': 4000e6
}

INPUT_RANGE_VALUE = {
    'INPUT_RANGE_PM_20_MV': 20e-3,
    'INPUT_RANGE_PM_40_MV': 40e-3,
    'INPUT_RANGE_PM_50_MV': 50e-3,
    'INPUT_RANGE_PM_80_MV': 80e-3,
    'INPUT_RANGE_PM_100_MV': 100e-3,
    'INPUT_RANGE_PM_200_MV': 200e-3,
    'INPUT_RANGE_PM_400_MV': 400e-3,
    'INPUT_RANGE_PM_500_MV': 500e-3,
    'INPUT_RANGE_PM_800_MV': 800e-3,
    'INPUT_RANGE_PM_1_V': 1,
    'INPUT_RANGE_PM_2_V': 2,
    'INPUT_RANGE_PM_4_V': 4,
    'INPUT_RANGE_PM_5_V': 5,
    'INPUT_RANGE_PM_8_V': 8,
    'INPUT_RANGE_PM_10_V': 10,
    'INPUT_RANGE_PM_20_V': 20,
    'INPUT_RANGE_PM_40_V': 40,
    'INPUT_RANGE_PM_16_V': 16,
    'INPUT_RANGE_PM_1_V_25': 1.25,
    'INPUT_RANGE_PM_2_V_5': 2.5,
    'INPUT_RANGE_PM_125_MV': 125e-3,
    'INPUT_RANGE_PM_250_MV': 250e-3
}

IMPEDANCE_VALUE = {
    'IMPEDANCE_1M_OHM': 1e6,
    'IMPEDANCE_50_OHM': 50,
    'IMPEDANCE_75_OHM': 75,
    'IMPEDANCE_300_OHM': 300
}

EXT_TRIG_RANGE_VALUE = {
    'ETR_5V': 5,
    'ETR_1V': 1,
    'ETR_2V5': 2.5
}

def choose_internal_sample_rate(model, sample_rate):
    min_v = 1000
    target = np.log(sample_rate)
    for key in HARDWARE_SPEC['SAMPLE_RATE'][model]:
        if abs(np.log(SAMPLE_RATE_VALUE[key])-target) < min_v:
            min_v = abs(np.log(SAMPLE_RATE_VALUE[key])-target)
            mkey = key
    return mkey

def choose_external_sample_rate(model, target):
    """Return appropriate sample rate and decimation."""
    # ATS9870 only supports 1GS/s when using external 10MHz ref. Choose appropriate
    # decimation.
    if model == 'ATS9870':
        sample_rate = 1000000000
        deci = [1,2,4,10,20,40,100,200,400,1000,2000,4000,10000,20000,40000,100000]
        min_i = np.argmin(abs(np.log(sample_rate)-np.log(deci)-np.log(target)))
        return sample_rate, deci[min_i]
    # ATS9360 supports 300MS/s -- 1800MS/s, 1MS/s step
    if model == 'ATS9360':
        target = np.clip(target, 300e6, 1800e6)
        sample_rate = int(target/1e6)*1000000
        return sample_rate, 1
    raise Exception('Not implemented')


import os
import time
import json
import urllib
import requests

import numpy as np

def makeCentQRCode(output_location,geojson_file):
    """
    Usage:
    Not extremely useful - FYI

    GeoJSON format changes, so the **coords** line might have to be altered (change indices) to work

    Give function an output folder (location where QC codes will be exported to) and a geojson file (formatted same a QGIS export - pretty sure)

    Will create 150x150 black and white QR codes
    """
    baseGoogUrl = r"https://www.google.com/maps/place/"
    qrCode = r"https://api.qrserver.com/v1/create-qr-code/?data={}&size=150x150"

    geoms = json.loads(open(geojson_file,"r").read())

    cnt = 1
    for g in geoms["features"]:
        coords = g["geometry"]["coordinates"][0][0]
        centX = np.array(coords)[:,0].mean()
        centY = np.array(coords)[:,1].mean()

        googUrl = baseGoogUrl + "{},{}".format(centY,centX)
        urlComplete = urllib.parse.quote(googUrl)
        reqObj = requests.get(qrCode.format(urlComplete))

        with open(os.path.join(output_location,f"QC_CODE{cnt}.png"),"wb") as outImg:
            outImg.write(reqObj.content)
        cnt+=1
        time.sleep(0.5)


import math
from pyradioconfig.parts.panther.calculators.calc_agc import CALC_AGC_panther
from py_2_and_3_compatibility import *
from scipy import interpolate

#This file contains calculations related to the configuring the AGC

class CALC_AGC_ocelot(CALC_AGC_panther):

    def calc_agc_misc(self, model):
        self._reg_write(model.vars.AGC_AGCPERIOD0_SETTLETIMEIF,6)
        self._reg_write(model.vars.AGC_AGCPERIOD0_SETTLETIMERF,14)
        self._reg_write(model.vars.AGC_CTRL0_AGCRST,0)
        self._reg_write(model.vars.AGC_CTRL0_DISCFLOOPADJ,1)
        self._reg_write(model.vars.AGC_CTRL0_DISRESETCHPWR,0)
        self._reg_write(model.vars.AGC_CTRL0_DSADISCFLOOP,0)
        self._reg_write(model.vars.AGC_CTRL0_DISPNDWNCOMP,0)
        self._reg_write(model.vars.AGC_CTRL0_DISPNGAINUP,0)
        self._reg_write(model.vars.AGC_CTRL0_ENRSSIRESET,0)
        self._reg_write(model.vars.AGC_CTRL1_PWRPERIOD, 1)
        self._reg_write(model.vars.AGC_CTRL2_DISRFPKD,0)
        self._reg_write(model.vars.AGC_CTRL2_DMASEL,0)
        self._reg_write(model.vars.AGC_CTRL2_PRSDEBUGEN,0)
        self._reg_write(model.vars.AGC_CTRL2_REHICNTTHD,7)
        self._reg_write(model.vars.AGC_CTRL2_RELBYCHPWR,3)
        self._reg_write(model.vars.AGC_CTRL2_RELOTHD,4)
        self._reg_write(model.vars.AGC_CTRL2_RELTARGETPWR,236)
        self._reg_write(model.vars.AGC_CTRL2_SAFEMODE,0)
        self._reg_write(model.vars.AGC_CTRL2_SAFEMODETHD,3)
        self._reg_write(model.vars.AGC_CTRL3_IFPKDDEB,1)
        self._reg_write(model.vars.AGC_CTRL3_IFPKDDEBPRD,40)
        self._reg_write(model.vars.AGC_CTRL3_IFPKDDEBRST,10)
        self._reg_write(model.vars.AGC_CTRL3_IFPKDDEBTHD,1)
        self._reg_write(model.vars.AGC_CTRL3_RFPKDDEB,1)
        self._reg_write(model.vars.AGC_CTRL3_RFPKDDEBPRD,40)
        self._reg_write(model.vars.AGC_CTRL3_RFPKDDEBRST,10)
        self._reg_write(model.vars.AGC_CTRL3_RFPKDDEBTHD,1)
        self._reg_write(model.vars.AGC_CTRL4_FRZPKDEN,0)
        self._reg_write(model.vars.AGC_CTRL4_PERIODRFPKD,4000)
        self._reg_write(model.vars.AGC_CTRL4_RFPKDCNTEN,1)
        self._reg_write(model.vars.AGC_CTRL4_RFPKDPRDGEAR,4)
        self._reg_write(model.vars.AGC_CTRL4_RFPKDSEL,1)
        self._reg_write(model.vars.AGC_CTRL4_RFPKDSYNCSEL,1)
        self._reg_write(model.vars.AGC_CTRL5_PNUPDISTHD,48)
        self._reg_write(model.vars.AGC_CTRL5_PNUPRELTHD,4)
        self._reg_write(model.vars.AGC_CTRL5_SEQPNUPALLOW,0)
        self._reg_write(model.vars.AGC_CTRL5_SEQRFPKDEN,0)
        self._reg_write(model.vars.AGC_GAINRANGE_GAININCSTEP,1)
        self._reg_write(model.vars.AGC_GAINRANGE_HIPWRTHD,3)
        self._reg_write(model.vars.AGC_GAINRANGE_LATCHEDHISTEP,0)
        self._reg_write(model.vars.AGC_GAINSTEPLIM0_MAXPWRVAR,0)
        self._reg_write(model.vars.AGC_GAINSTEPLIM0_TRANRSTAGC,0)
        self._reg_write(model.vars.AGC_LBT_CCARSSIPERIOD,0)
        self._reg_write(model.vars.AGC_LBT_ENCCAGAINREDUCED,0)
        self._reg_write(model.vars.AGC_LBT_ENCCARSSIMAX,0)
        self._reg_write(model.vars.AGC_LBT_ENCCARSSIPERIOD,0)
        self._reg_write(model.vars.AGC_LNAMIXCODE0_LNAMIXSLICE1,61)
        self._reg_write(model.vars.AGC_LNAMIXCODE0_LNAMIXSLICE2,46)
        self._reg_write(model.vars.AGC_LNAMIXCODE0_LNAMIXSLICE3,36)
        self._reg_write(model.vars.AGC_LNAMIXCODE0_LNAMIXSLICE4,28)
        self._reg_write(model.vars.AGC_LNAMIXCODE0_LNAMIXSLICE5,21)
        self._reg_write(model.vars.AGC_LNAMIXCODE1_LNAMIXSLICE6,17)
        self._reg_write(model.vars.AGC_LNAMIXCODE1_LNAMIXSLICE7,12)
        self._reg_write(model.vars.AGC_LNAMIXCODE1_LNAMIXSLICE8,10)
        self._reg_write(model.vars.AGC_LNAMIXCODE1_LNAMIXSLICE9,6)
        self._reg_write(model.vars.AGC_LNAMIXCODE1_LNAMIXSLICE10,5)
        self._reg_write(model.vars.AGC_MANGAIN_MANGAINEN,0)
        self._reg_write(model.vars.AGC_MANGAIN_MANGAINIFPGA,0)
        self._reg_write(model.vars.AGC_MANGAIN_MANGAINLNA,0)
        self._reg_write(model.vars.AGC_MANGAIN_MANGAINPN,0)
        self._reg_write(model.vars.AGC_MANGAIN_MANIFHILATRST,0)
        self._reg_write(model.vars.AGC_MANGAIN_MANIFLOLATRST,0)
        self._reg_write(model.vars.AGC_MANGAIN_MANRFLATRST,0)
        self._reg_write(model.vars.AGC_PGACODE0_PGAGAIN1,0)
        self._reg_write(model.vars.AGC_PGACODE0_PGAGAIN2,1)
        self._reg_write(model.vars.AGC_PGACODE0_PGAGAIN3,2)
        self._reg_write(model.vars.AGC_PGACODE0_PGAGAIN4,3)
        self._reg_write(model.vars.AGC_PGACODE0_PGAGAIN5,4)
        self._reg_write(model.vars.AGC_PGACODE0_PGAGAIN6,5)
        self._reg_write(model.vars.AGC_PGACODE0_PGAGAIN7,6)
        self._reg_write(model.vars.AGC_PGACODE0_PGAGAIN8,7)
        self._reg_write(model.vars.AGC_PGACODE1_PGAGAIN9,8)
        self._reg_write(model.vars.AGC_PGACODE1_PGAGAIN10,9)
        self._reg_write(model.vars.AGC_PGACODE1_PGAGAIN11,10)
        self._reg_write(model.vars.AGC_RSSISTEPTHR_RSSIFAST,0)
        self._reg_write(model.vars.RAC_PGACTRL_PGAENLATCHI,1)
        self._reg_write(model.vars.RAC_PGACTRL_PGAENLATCHQ,1)
        self._reg_write(model.vars.AGC_GAINSTEPLIM0_CFLOOPSTEPMAX, 0)
        self._reg_write(model.vars.AGC_GAINSTEPLIM0_HYST, 0)

        # numbers from Arup's email dated 7/2/19 to Xun

    def calc_lnamixenrfpkdlothresh_reg(self, model):
        self._reg_write(model.vars.RAC_RX_LNAMIXENRFPKDLOTHRESH, 1)

    def calc_lnamixrfatt_reg(self, model):

        #Load in model variables
        rf_band = model.vars.rf_band.value

        #Calculate the LNAMIXRFATT values based on band
        if rf_band == model.vars.rf_band.var_enum.BAND_915 or rf_band == model.vars.rf_band.var_enum.BAND_868:
            lnamixrfatt = [0, 63, 141, 238, 502, 940, 1269, 1942, 2526, 3484, 4547, 6035, 7678, 9973, 12989, 16383,
                           5630, 7160, 9180, 11700, 14800, 16383, 16383]
            pnindexmax = 16
        elif rf_band == model.vars.rf_band.var_enum.BAND_490:
            lnamixrfatt = [0, 20, 45, 78, 125, 175, 246, 477, 726, 1012, 1471, 1983, 2703, 3485, 4543, 5835,
                           7582, 9707, 12672, 16359, 16383, 16383, 16383]
            pnindexmax = 20
        elif rf_band == model.vars.rf_band.var_enum.BAND_434:
            lnamixrfatt = [0, 15, 44, 77, 117, 172, 239, 468, 725, 998, 1469, 1983, 2551, 3491, 4525, 5849,
                           7623, 9910, 12750, 16383, 14800, 16383, 16383]
            pnindexmax = 20
        elif rf_band == model.vars.rf_band.var_enum.BAND_315:
            lnamixrfatt = [0, 12, 29, 52, 78, 111, 156, 207, 420, 502, 750, 1187, 1502, 2012, 2720, 3521,
                           4557, 5870, 7623, 9711, 12719, 16383, 16383]
            pnindexmax = 22
        elif rf_band == model.vars.rf_band.var_enum.BAND_169:
            lnamixrfatt = [0, 13, 36, 63, 99, 141, 199, 412, 502, 757, 1190, 1507, 2029, 2743, 3541, 4601,
                           6055, 7811, 9966, 12964, 16383, 16383, 16383]
            pnindexmax = 21

        else:
            lnamixrfatt = [0, 20, 48, 76, 116, 164, 228, 436, 668, 924, 1210, 1530, 2030, 2720, 3480, 4350,
                           5630, 7160, 9180, 11700, 14800, 16383, 16383]
            pnindexmax = 17

        #Write registers
        self._reg_write(model.vars.AGC_PNRFATT0_LNAMIXRFATT1,lnamixrfatt[0])
        self._reg_write(model.vars.AGC_PNRFATT0_LNAMIXRFATT2, lnamixrfatt[1])
        self._reg_write(model.vars.AGC_PNRFATT1_LNAMIXRFATT3, lnamixrfatt[2])
        self._reg_write(model.vars.AGC_PNRFATT1_LNAMIXRFATT4, lnamixrfatt[3])
        self._reg_write(model.vars.AGC_PNRFATT2_LNAMIXRFATT5, lnamixrfatt[4])
        self._reg_write(model.vars.AGC_PNRFATT2_LNAMIXRFATT6, lnamixrfatt[5])
        self._reg_write(model.vars.AGC_PNRFATT3_LNAMIXRFATT7, lnamixrfatt[6])
        self._reg_write(model.vars.AGC_PNRFATT3_LNAMIXRFATT8, lnamixrfatt[7])
        self._reg_write(model.vars.AGC_PNRFATT4_LNAMIXRFATT9, lnamixrfatt[8])
        self._reg_write(model.vars.AGC_PNRFATT4_LNAMIXRFATT10, lnamixrfatt[9])
        self._reg_write(model.vars.AGC_PNRFATT5_LNAMIXRFATT11, lnamixrfatt[10])
        self._reg_write(model.vars.AGC_PNRFATT5_LNAMIXRFATT12, lnamixrfatt[11])
        self._reg_write(model.vars.AGC_PNRFATT6_LNAMIXRFATT13, lnamixrfatt[12])
        self._reg_write(model.vars.AGC_PNRFATT6_LNAMIXRFATT14, lnamixrfatt[13])
        self._reg_write(model.vars.AGC_PNRFATT7_LNAMIXRFATT15, lnamixrfatt[14])
        self._reg_write(model.vars.AGC_PNRFATT7_LNAMIXRFATT16, lnamixrfatt[15])
        self._reg_write(model.vars.AGC_PNRFATT8_LNAMIXRFATT17, lnamixrfatt[16])
        self._reg_write(model.vars.AGC_PNRFATT8_LNAMIXRFATT18, lnamixrfatt[17])
        self._reg_write(model.vars.AGC_PNRFATT9_LNAMIXRFATT19, lnamixrfatt[18])
        self._reg_write(model.vars.AGC_PNRFATT9_LNAMIXRFATT20, lnamixrfatt[19])
        self._reg_write(model.vars.AGC_PNRFATT10_LNAMIXRFATT21, lnamixrfatt[20])
        self._reg_write(model.vars.AGC_PNRFATT10_LNAMIXRFATT22, lnamixrfatt[21])
        self._reg_write(model.vars.AGC_PNRFATT11_LNAMIXRFATT23, lnamixrfatt[22])
        self._reg_write(model.vars.AGC_GAINSTEPLIM1_PNINDEXMAX, pnindexmax)

    def calc_agc_series2(self, model):
        #We need to override this function because the variable fxo_or_fdec8 will go away (replaced by ADC rate calculations)

        mod_format = model.vars.modulation_type.value
        hardmodem_freq_actual = model.vars.hardmodem_freq_actual.value
        f_if = model.vars.if_frequency_hz_actual.value

        # period over which we count how many times we tripped the HI threshold - xtal PLL freq because AGC runs at this clock
        if f_if > 0:
            periodhi = int(py2round(hardmodem_freq_actual / (2 * f_if)))
        else:
            periodhi = 14 # for zero-IF used on FPGA tests fix periodhi to 14

        # Function of PERIODHI and attack vs decay ratio needed. We currently use 3x
        # The scaler 3 could be an input that tunes attack vs decay ratio.
        baudrate = model.vars.baudrate.value

        if (mod_format == model.vars.modulation_type.var_enum.OOK):
            periodlow = int(py2round((hardmodem_freq_actual/(baudrate * 0.9))))
        else:
            periodlow = 3 * periodhi

        self._reg_write(model.vars.AGC_AGCPERIOD0_PERIODHI, int(round(periodhi)))
        self._reg_write(model.vars.AGC_AGCPERIOD1_PERIODLOW, int(round(periodlow)))

        # % There are many possible ways to handle the table but the simplest would be:
        # % this is based on sine wave tripping N times at different gain settings
        self._reg_write(model.vars.AGC_STEPDWN_STEPDWN0, 0)
        self._reg_write(model.vars.AGC_STEPDWN_STEPDWN1, 1)
        self._reg_write(model.vars.AGC_STEPDWN_STEPDWN2, 2)
        self._reg_write(model.vars.AGC_STEPDWN_STEPDWN3, 3)
        self._reg_write(model.vars.AGC_STEPDWN_STEPDWN4, 4)
        self._reg_write(model.vars.AGC_STEPDWN_STEPDWN5, 5)

        hicntregion0 = int(py2round(0.55 * periodhi))
        hicntregion1 = int(py2round(0.75 * periodhi))
        hicntregion2 = int(py2round(0.85 * periodhi))
        hicntregion3 = int(py2round(0.90 * periodhi))
        hicntregion4 = int(py2round(0.93 * periodhi))

        if (hicntregion0 > 255):
            print("  WARNING: AGC_HICNTREGION_HICNTREGION 0 calculated beyond range: hicntregion0 {}, hicntregion1 {}, AGC_AGCPERIOD0_PERIODHI {}. Saturating value to 255 !".format(hicntregion0, hicntregion1, periodhi))
            hicntregion0 = 255

        if (hicntregion1 > 255):
            print("  WARNING: AGC_HICNTREGION_HICNTREGION 1 calculated beyond range: hicntregion0 {}, hicntregion1 {}, AGC_AGCPERIOD0_PERIODHI {}. Saturating value to 255 !".format(hicntregion0, hicntregion1, periodhi))
            hicntregion1 = 255

        self._reg_write(model.vars.AGC_HICNTREGION0_HICNTREGION0, hicntregion0)
        self._reg_write(model.vars.AGC_HICNTREGION0_HICNTREGION1, hicntregion1)
        self._reg_write(model.vars.AGC_HICNTREGION0_HICNTREGION2, hicntregion2)
        self._reg_write(model.vars.AGC_HICNTREGION0_HICNTREGION3, hicntregion3)
        self._reg_write(model.vars.AGC_HICNTREGION1_HICNTREGION4, hicntregion4)

        # % safe way of setting this.
        self._reg_write(model.vars.AGC_AGCPERIOD0_MAXHICNTTHD, hicntregion4)

    def calc_fastloopdel_reg(self, model):
        #This register does not exist in Ocelot
        pass

    def calc_ifpgadel_reg(self, model):
        #This register does not exist in Ocelot
        pass

    def calc_lnaslicesdel_reg(self, model):
        # This register does not exist in Ocelot
        pass

    def calc_pngainstep_reg(self, model):
        etsi_cat1_compatability = model.vars.etsi_cat1_compatible.value

        # PN gain step size should be 1 for ETSI Cat 1 case and 2 in all other cases
        if (etsi_cat1_compatability != model.vars.etsi_cat1_compatible.var_enum.Normal):
            reg = 1
        else:
            reg = 2

        self._reg_write(model.vars.AGC_GAINRANGE_PNGAINSTEP, reg)

    def calc_agc_settling_delay(self, model):
        pass

    def calc_agcperiod_actual(self, model):

        #Read in the actual AGC period register (different name for Ocelot)
        agc_period_reg = model.vars.AGC_CTRL1_PWRPERIOD.value

        #Calculate the actual period value based on the reg
        val = 2 ** (agc_period_reg * 1.0)

        model.vars.agcperiod_actual.value = val

    def calc_cfloopdel_reg(self, model):
        """calculate AGC settling delay which is basically the group delay of decimation and
        channel filters through the datapath plus processing delays

        calculations are in channel filter clock cycles to directly program into CFLOOPDEL

        Args:
            model (ModelRoot) : Data model to read and write variables from
        """
        agc_delay = model.vars.agc_settling_delay.value

        cfloopdel = agc_delay

        if cfloopdel > 127:
            cfloopdel = 127

        self._reg_write(model.vars.AGC_GAINSTEPLIM0_CFLOOPDEL, int(math.ceil(cfloopdel)))

    def calc_agc_cfloopdel_actual(self, model):
        reg = model.vars.AGC_GAINSTEPLIM0_CFLOOPDEL.value
        adc_freq_actual = model.vars.adc_freq_actual.value
        dec0 = model.vars.dec0_actual.value
        dec1 = model.vars.dec1_actual.value
        src2_actual = model.vars.src2_ratio_actual.value

        fsrc2 = adc_freq_actual / 8.0 / dec0 / dec1 * src2_actual
        model.vars.cfloopdel_us_actual.value = reg / fsrc2 * 1e6

        return

    def calc_agc_sub_reg(self, model):
        demod_sel = model.vars.demod_select.value
        adc_freq_actual = model.vars.adc_freq_actual.value
        dec0_actual = model.vars.dec0_actual.value
        dec1_actual = model.vars.dec1_actual.value
        dec2_actual = model.vars.dec2_actual.value
        baudrate = model.vars.baudrate.value
        src2_actual = model.vars.src2_ratio_actual.value

        osr = adc_freq_actual * src2_actual / (dec0_actual * dec1_actual * 8 * dec2_actual * baudrate)
        # when BCR demod is selected we the OSR can be as high as 127
        # this does not fit into the RXBR register which is used in generating a baudrate clock
        # we switch to using SUB registers instead of RXBR for this purpose when BCR is selected
        if (demod_sel == model.vars.demod_select.var_enum.BCR):
            rawndec = model.vars.MODEM_BCRDEMODOOK_RAWNDEC.value #Checking inside of the BCR condition to allow inheritance for Bobcat
            subperiod = 1
            dec = pow(2, rawndec)
            osr_bcr = osr / dec

            # write half of sampling rate into SUB registers
            subfrac = osr_bcr / 2

            subint = math.floor(subfrac)

            frac = subfrac - subint

            best_error = 999
            for den in range(1,256):
                num = round(frac * den)
                error = abs(frac - num/den)
                if error < best_error:
                    subden = den
                    subnum = num
                    best_error = error

                if best_error < 1e-3:
                    break
        elif (demod_sel == model.vars.demod_select.var_enum.TRECS_VITERBI or
              demod_sel == model.vars.demod_select.var_enum.TRECS_SLICER) and osr >= 8:
            subint = model.vars.rxbrint.value
            subnum = model.vars.rxbrnum.value
            subden = model.vars.rxbrden.value
            subperiod = 1

        else:
            subden = 0
            subint = 0
            subnum = 0
            subperiod = 0

        self._reg_write(model.vars.AGC_CTRL7_SUBDEN,subden)
        self._reg_write(model.vars.AGC_CTRL7_SUBINT,subint)
        self._reg_write(model.vars.AGC_CTRL7_SUBNUM,subnum)
        self._reg_write(model.vars.AGC_CTRL7_SUBPERIOD,subperiod)

        return

    def calc_subfrac_actual(self, model):

        # Load model variables into local variables
        subint_actual = model.vars.AGC_CTRL7_SUBINT.value
        subnum_actual = model.vars.AGC_CTRL7_SUBNUM.value
        subden_actual = model.vars.AGC_CTRL7_SUBDEN.value
        subperiod = model.vars.AGC_CTRL7_SUBPERIOD.value

        # Calculate the sub fraction
        if subperiod and subden_actual > 0:
            subfrac_actual = float(subint_actual + float(subnum_actual) / subden_actual)
        else:
            subfrac_actual = 0.0

        # Load local variables back into model variables
        model.vars.subfrac_actual.value = subfrac_actual

    def calc_agc_mode_reg(self, model):

        demod_sel = model.vars.demod_select.value
        preamsch = model.vars.MODEM_TRECPMDET_PREAMSCH.value

        # we like AGC to freeze once timing is detected (MODE=1) to avoid AGC changes while receiving data
        # the only exception to this is the case where we are directly searching for the sync word when using Viterbi
        # in this case we need to set the MODE=2 to freeze the AGC when frame is detected
        if (demod_sel == model.vars.demod_select.var_enum.TRECS_VITERBI or demod_sel == model.vars.demod_select.var_enum.TRECS_SLICER):
            if preamsch:
                mode = 1
            else:
                mode = 2
        else:
            mode = 1

        self._reg_write(model.vars.AGC_CTRL0_MODE, mode)

        return

    def calc_agc_decision_matrix_reg(self, model):
        dualrfpkddec_val = 240296
        self._reg_write(model.vars.AGC_CTRL6_DUALRFPKDDEC, dualrfpkddec_val)

    #Removing legacy AGC calculations

    def calc_agc_scheme(self, model):
        pass

    def calc_agc_index_min_atten_reg(self, model):
        pass

    def calc_agc_index_min_degen_reg(self, model):
        pass

    def calc_agc_index_min_pga_reg(self, model):
        pass

    def calc_agc_index_min_slices_reg(self, model):
        pass

    def calc_agcperiod_reg(self, model):
        pass

    def calc_agc_reg(self, model):
        pass

    def calc_agc_cfloopstepmax_reg(self, model):
        pass

    def calc_hyst_reg(self, model):
        pass

    def calc_agc_baudrate_calculation_mode(self, model):
        demod_select = model.vars.demod_select.value

        if demod_select == model.vars.demod_select.var_enum.COHERENT:
            # : disable RX baudrate calculation used by AGC and instead assume OSR = 2 * RXBRFRAC
            self._reg_write(model.vars.MODEM_CTRL6_RXBRCALCDIS, 1)
        else:
            self._reg_write(model.vars.MODEM_CTRL6_RXBRCALCDIS, 0)

    def calc_dagc_channel_power_accumulator_reg(self, model):
        demod_select = model.vars.demod_select.value
        modtype = model.vars.modulation_type.value
        preamble_pattern_len = model.vars.preamble_pattern_len.value
        sens_calculated = model.vars.sensitivity.value
        target_osr = model.vars.target_osr.value
        agc_period_actual = model.vars.agcperiod_actual.value
        baudrate = model.vars.baudrate.value
        bitrate = model.vars.bitrate.value
        mod_format = model.vars.modulation_type.value

        if hasattr(model.profiles, 'Long_Range'):
            is_long_range = model.profile == model.profiles.Long_Range
        else:
            is_long_range = False

        if demod_select == model.vars.demod_select.var_enum.COHERENT:
            # : accumulate channel power for preamble pattern duration
            # : actual avgwin period is 2^(avgwin+2)*2^pwrperiod_s*OSR
            # : agc_period_actual = 2^pwrperiod
            # : actual avgwin time is 2^(avgwin+2)*agc_period_actual/baudrate
            preamble_time_us = preamble_pattern_len * (1 / bitrate) * 1e6
            target_avgwin_time_us = 4.0 * preamble_time_us
            target_avgwin_period = (target_avgwin_time_us / 1e6) * baudrate
            target_avgwin = round(math.log2(target_avgwin_period / agc_period_actual) - 2)
            if target_avgwin < 0:
                target_avgwin = 0
            avgwin = int(target_avgwin)
            chpwraccudel = 0
        else:
            avgwin = 0
            chpwraccudel = 0

        if demod_select == model.vars.demod_select.var_enum.COHERENT and \
                        mod_format == model.vars.modulation_type.var_enum.OQPSK:
            # : 0 - Channel power is locked when timing is detected
            # : 1 - Channel power is locked when DSA is detected
            self._reg_write(model.vars.MODEM_COH0_COHCHPWRLOCK, 0)
            # : Set to enable automatic restart of channel power - needed to make sure channel power is not dependent on
            # : power of received frames
            self._reg_write(model.vars.MODEM_COH0_COHCHPWRRESTART, 1)
        else:
            self._reg_write(model.vars.MODEM_COH0_COHCHPWRLOCK, 0)
            self._reg_write(model.vars.MODEM_COH0_COHCHPWRRESTART, 0)

        if demod_select == model.vars.demod_select.var_enum.COHERENT and is_long_range:
            self._reg_write(model.vars.MODEM_LONGRANGE1_AVGWIN, avgwin)
            self._reg_write(model.vars.MODEM_LONGRANGE1_CHPWRACCUDEL, chpwraccudel)
        else:
            self._reg_write(model.vars.MODEM_LONGRANGE1_AVGWIN, 0)
            self._reg_write(model.vars.MODEM_LONGRANGE1_CHPWRACCUDEL, 0)

    def calc_dagc_dynamic_bbss_reg(self, model):
        demod_select = model.vars.demod_select.value
        sens_calculated = model.vars.sensitivity.value
        bitrate = model.vars.bitrate.value
        mod_format = model.vars.modulation_type.value

        if hasattr(model.profiles, 'Long_Range'):
            is_long_range = model.profile == model.profiles.Long_Range
        else:
            is_long_range = False

        """ Model calculation range based on DUT power [dBm] """
        min_dut_power = -140 # : This is arbitrarily low power that should be below sensitivity
        max_dut_power = 10 # : This is maximum supported receive power of DUT
        if sens_calculated < min_dut_power:
            min_dut_power = round(sens_calculated) - 10
        dut_power_list = list(range(min_dut_power,max_dut_power,+1))

        """ Model of Q Sample vs. Dut Power """
        estimated_adc_noise_dBm = -162.8694 #  This is measured value
        model_max_Q_list = []
        for dut_power in  dut_power_list:
            # : Model ADC resolution
            adc_enob = (dut_power - estimated_adc_noise_dBm - 1.76) / 6.02

            # : Model analog frontend noise
            minimum_detectable_signal_dBm = -173.9 + 10*math.log10(bitrate) + 4.0
            if bitrate <= 2.8e3:
                analog_noise_bitwidth = 0.1938 * minimum_detectable_signal_dBm + 33.296
            elif bitrate < 67e3: # : TODO change to ADC mode.
                analog_noise_bitwidth = 0.1565 * minimum_detectable_signal_dBm + 27.928
            else:
                analog_noise_bitwidth = 0.1599 * minimum_detectable_signal_dBm + 28.119

            # : Determine signal + noise level
            signal_max_q = math.pow(2,adc_enob) + math.pow(2,analog_noise_bitwidth)
            signal_max_q_bitwidth = math.log2(signal_max_q)
            if signal_max_q_bitwidth > 17:
                signal_max_q_bitwidth = 17

            model_max_Q_list.append(signal_max_q_bitwidth)

        # : Calculate bbss shift based on the Q sample model
        bbss_model = []
        for model_max_Q in model_max_Q_list:
            bbss_at_max_Q_calc = round(model_max_Q - 5.5)
            bbss_model.append(bbss_at_max_Q_calc)

        """ Calculate chpwraccumux model """
        model_chpwraccumux_list = []
        for dut_power in dut_power_list:
            calc_chpwraccumux = dut_power + 139.0

            # : Based on measurement, chpwraccumux does not go below 14 due to noise
            if calc_chpwraccumux < 14:
                calc_chpwraccumux = 14
            # : By design, maximum chpwraccumux is 80 due signal power after AGC
            elif calc_chpwraccumux > 80:
                calc_chpwraccumux = 80

            model_chpwraccumux_list.append(calc_chpwraccumux)

        """  """
        # : Calculate chpwraccumux at sensitivity. This can be used to calculate noise level
        chpwraccumux_interp_func = interpolate.interp1d(dut_power_list, model_chpwraccumux_list)
        chpwr_accumux_noise = float(chpwraccumux_interp_func(sens_calculated))

        """ """
        bbss_interp_func = interpolate.interp1d(dut_power_list, bbss_model)
        starting_BBSS = math.ceil(bbss_interp_func(sens_calculated))

        """ Calculate BBSS transitions """
        bbss_transition = []
        bbss_transition_threshold = []
        for dut_power_index in range(len(dut_power_list)):
            bbss_val = bbss_model[dut_power_index]
            if bbss_val > starting_BBSS:
                if bbss_model[dut_power_index-1] < bbss_val:
                    bbss_transition.append(int(bbss_val))
                    bbss_transition_threshold.append(int(round(model_chpwraccumux_list[dut_power_index])))

        # : Add two bbss transitions below sensitivity
        bbss_transition.insert(0, starting_BBSS)
        bbss_transition.insert(0, starting_BBSS - 1)
        bbss_transition_threshold.insert(0, int(round(chpwr_accumux_noise-6)))

        # : fill up remaining registers
        max_bbss_shift = 12
        while len(bbss_transition_threshold) < 11:
            bbss_transition_threshold.append(80)
        while len(bbss_transition) < 12:
            bbss_transition.append(max_bbss_shift)

        """ Set calculated value """
        model.vars.chpwraccu_noise.value = chpwr_accumux_noise

        """ Set registers related to dynamic bbss mode"""
        if demod_select == model.vars.demod_select.var_enum.COHERENT and \
                        mod_format == model.vars.modulation_type.var_enum.OQPSK:
            # : Enable dynamic BBSS adjustment
            self._reg_write(model.vars.MODEM_COH0_COHDYNAMICBBSSEN, 1)

            # : BBSS hysteresis
            self._reg_write(model.vars.MODEM_LONGRANGE1_HYSVAL, 3)
        else:
            self._reg_write(model.vars.MODEM_COH0_COHDYNAMICBBSSEN, 0)
            self._reg_write(model.vars.MODEM_LONGRANGE1_HYSVAL, 0)

        """ Set registers related to dynamic thresholds """
        if demod_select == model.vars.demod_select.var_enum.COHERENT and is_long_range:
            # : BBSS thresholds
            self._reg_write(model.vars.MODEM_LONGRANGE2_LRCHPWRTH1, bbss_transition_threshold[0])
            self._reg_write(model.vars.MODEM_LONGRANGE2_LRCHPWRTH2, bbss_transition_threshold[1])
            self._reg_write(model.vars.MODEM_LONGRANGE2_LRCHPWRTH3, bbss_transition_threshold[2])
            self._reg_write(model.vars.MODEM_LONGRANGE2_LRCHPWRTH4, bbss_transition_threshold[3])
            self._reg_write(model.vars.MODEM_LONGRANGE3_LRCHPWRTH5, bbss_transition_threshold[4])
            self._reg_write(model.vars.MODEM_LONGRANGE3_LRCHPWRTH6, bbss_transition_threshold[5])
            self._reg_write(model.vars.MODEM_LONGRANGE3_LRCHPWRTH7, bbss_transition_threshold[6])
            self._reg_write(model.vars.MODEM_LONGRANGE3_LRCHPWRTH8, bbss_transition_threshold[7])
            self._reg_write(model.vars.MODEM_LONGRANGE4_LRCHPWRTH9, bbss_transition_threshold[8])
            self._reg_write(model.vars.MODEM_LONGRANGE4_LRCHPWRTH10, bbss_transition_threshold[9])
            self._reg_write(model.vars.MODEM_LONGRANGE6_LRCHPWRTH11, bbss_transition_threshold[10])

            # : BBSS Shifts
            self._reg_limit_write(model.vars.MODEM_LONGRANGE4_LRCHPWRSH1, bbss_transition[0], max_bbss_shift)
            self._reg_limit_write(model.vars.MODEM_LONGRANGE4_LRCHPWRSH2, bbss_transition[1], max_bbss_shift)
            self._reg_limit_write(model.vars.MODEM_LONGRANGE4_LRCHPWRSH3, bbss_transition[2], max_bbss_shift)
            self._reg_limit_write(model.vars.MODEM_LONGRANGE4_LRCHPWRSH4, bbss_transition[3], max_bbss_shift)
            self._reg_limit_write(model.vars.MODEM_LONGRANGE5_LRCHPWRSH5, bbss_transition[4], max_bbss_shift)
            self._reg_limit_write(model.vars.MODEM_LONGRANGE5_LRCHPWRSH6, bbss_transition[5], max_bbss_shift)
            self._reg_limit_write(model.vars.MODEM_LONGRANGE5_LRCHPWRSH7, bbss_transition[6], max_bbss_shift)
            self._reg_limit_write(model.vars.MODEM_LONGRANGE5_LRCHPWRSH8, bbss_transition[7], max_bbss_shift)
            self._reg_limit_write(model.vars.MODEM_LONGRANGE5_LRCHPWRSH9, bbss_transition[8], max_bbss_shift)
            self._reg_limit_write(model.vars.MODEM_LONGRANGE5_LRCHPWRSH10, bbss_transition[9], max_bbss_shift)
            self._reg_limit_write(model.vars.MODEM_LONGRANGE5_LRCHPWRSH11, bbss_transition[10], max_bbss_shift)
            self._reg_limit_write(model.vars.MODEM_LONGRANGE6_LRCHPWRSH12, bbss_transition[11], max_bbss_shift)
        else:
            self._reg_write(model.vars.MODEM_LONGRANGE2_LRCHPWRTH1, 0)
            self._reg_write(model.vars.MODEM_LONGRANGE2_LRCHPWRTH2, 0)
            self._reg_write(model.vars.MODEM_LONGRANGE2_LRCHPWRTH3, 0)
            self._reg_write(model.vars.MODEM_LONGRANGE2_LRCHPWRTH4, 0)
            self._reg_write(model.vars.MODEM_LONGRANGE3_LRCHPWRTH5, 0)
            self._reg_write(model.vars.MODEM_LONGRANGE3_LRCHPWRTH6, 0)
            self._reg_write(model.vars.MODEM_LONGRANGE3_LRCHPWRTH7, 0)
            self._reg_write(model.vars.MODEM_LONGRANGE3_LRCHPWRTH8, 0)
            self._reg_write(model.vars.MODEM_LONGRANGE4_LRCHPWRTH9, 0)
            self._reg_write(model.vars.MODEM_LONGRANGE4_LRCHPWRTH10, 0)
            self._reg_write(model.vars.MODEM_LONGRANGE6_LRCHPWRTH11, 0)
            self._reg_write(model.vars.MODEM_LONGRANGE4_LRCHPWRSH1, 0)
            self._reg_write(model.vars.MODEM_LONGRANGE4_LRCHPWRSH2, 0)
            self._reg_write(model.vars.MODEM_LONGRANGE4_LRCHPWRSH3, 0)
            self._reg_write(model.vars.MODEM_LONGRANGE4_LRCHPWRSH4, 0)
            self._reg_write(model.vars.MODEM_LONGRANGE5_LRCHPWRSH5, 0)
            self._reg_write(model.vars.MODEM_LONGRANGE5_LRCHPWRSH6, 0)
            self._reg_write(model.vars.MODEM_LONGRANGE5_LRCHPWRSH7, 0)
            self._reg_write(model.vars.MODEM_LONGRANGE5_LRCHPWRSH8, 0)
            self._reg_write(model.vars.MODEM_LONGRANGE5_LRCHPWRSH9, 0)
            self._reg_write(model.vars.MODEM_LONGRANGE5_LRCHPWRSH10, 0)
            self._reg_write(model.vars.MODEM_LONGRANGE5_LRCHPWRSH11, 0)
            self._reg_write(model.vars.MODEM_LONGRANGE6_LRCHPWRSH12, 0)

    def calc_antdiv_gainmode_reg(self, model):
        #Antenna diversity has its own AGC gain restore mode that determines how to set the gain index based on previous packet
        antdivmode = model.vars.antdivmode.value

        if antdivmode == model.vars.antdivmode.var_enum.DISABLE or \
                antdivmode == model.vars.antdivmode.var_enum.ANTENNA1:
            #The value of this field does matter even when antdiv is disabled
            gainmode = 0
        else:
            gainmode = 1  # : Always enable gain restore mode - clear stored gain from last packet when RX is turned off.

        self._reg_write(model.vars.AGC_ANTDIV_GAINMODE, gainmode)

    def calc_antdiv_debouncecntthd(self, model):
        """
        Number of clock cycles to wait after antenna switching before setting AGC gains.

        Args:
            model:

        Returns:

        """
        antdivmode = model.vars.antdivmode.value

        debouncecntthd = 40 # : Recommended value by He Gou

        if antdivmode == model.vars.antdivmode.var_enum.DISABLE or \
                antdivmode == model.vars.antdivmode.var_enum.ANTENNA1:
            self._reg_do_not_care(model.vars.AGC_ANTDIV_DEBOUNCECNTTHD)
        else:
            self._reg_write(model.vars.AGC_ANTDIV_DEBOUNCECNTTHD, debouncecntthd)

    def calc_gain_schedule_regs(self, model):
        lnaindexborder = 7
        pgaindexborder = 8

        self._reg_write(model.vars.AGC_GAINRANGE_LNAINDEXBORDER, lnaindexborder)
        self._reg_write(model.vars.AGC_GAINRANGE_PGAINDEXBORDER, pgaindexborder)