import numpy as np
from numbers import Number
import typing
import warnings

from .steps import _ModelStep
from ._model_utils import get_gap_from_model
from ._step_utils import HeightOptimisationFunction, make_interpolation_func

__all__ = ['QuasiStaticStep']

class QuasiStaticStep(_ModelStep):
    """
    A model step for quasi static relative movement

    To be used for loading, unloading and sliding as well as quasi static impacts, can be used for rolling-sliding
    with appropriate sub models for tangential behaviour.

    Parameters
    ----------

    step_name: str
        An identifying name for the step used for errors and outputs
    number_of_steps: int
        The number of sub steps the problem will be split into, together with the time_period this controls the duration
        of the time steps used
    no_time: bool, optional (False)
        Set to true if there is no time dependence and the time steps can be solved in any order (no permanent changes
        between steps such as plastic deformation or heat generation), if True the model will be solved more efficiently
    time_period: float, optional (1.0)
        The total time period of this model step, used for solving sub-models and writing outputs
    off_set_x, off_set_y: float or Sequence of float, optional (0.0)
        The off set between the surfaces in the x and y directions, this can be a relative off set or an absolute off
        set, controlled by the relative_loading parameter:

        * A constant (float) value, indicating a constant offset between the surfaces (no relative movement of profiles)
        * A two element sequence of floats, indicating the the start and finish offsets, if this is used, the
          movement_interpolation_mode will be used to generate intermediate values
        * A 2 by n array of n absolute position values and n time values normalised to a 0-1 scale. array[0] should be
          position values and array[1] should be time values, time values must be between 0 and 1. The
          movement_interpolation_mode will be used to generate intermediate values
    interference, normal_load: float or Sequence of float, optional (None)
        The interference and normal load between the surfaces, only one of these can be set (the other will be solved
        for) setting neither keeps the interference as it is at the start of this model step. As above for the off sets,
        either of these parameters can be:

        * A constant (float) value, indicating a constant load/ interference between the surfaces.
        * A two element sequence of floats, indicating the the start and finish load. interference, if this is used, the
          movement_interpolation_mode will be used to generate intermediate values
        * A 2 by n array of n absolute position values and n time values normalised to a 0-1 scale. array[0] should be
          position values and array[1] should be time values, time values must be between 0 and 1. The
          movement_interpolation_mode will be used to generate intermediate values
    relative_loading: bool, optional (False)
        If True the load or displacement and off set will be applied relative to the value at the start of the step,
        otherwise the absolute value will be used. eg, if the previous step ended with a load of 10N and this step ramps
        from 0 to 10N setting relative_loading to True will ramp the total load form 10 to 20N over this step.
    adhesion: bool, optional (True)
        If True the adhesion model set for the contact model will be used, If set to false this step will ignore the
        adhesion model (typically used for loading steps)
    unloading: bool, optional (False)
        If True the contact nodes will be constrained to be a sub set of those found in the previous time step.
    fast_ld: bool, optional (False)
        If True the load and displacements can be swapped to give a faster simulation, eg: if 100 equally spaced load
        steps are specified by a start and finish load, the displacement at the maximum load will be found and instead
        100 equally spaced displacement controlled steps will solved between just touching and the maximum deflection.
        This swapping makes the computation much faster as it removes an optimisation loop.
    impact_properties: list, optional (None)
        This is currently not supported, providing a value will cause an error
    movement_interpolation_mode: str or int, optional ('linear')
        Any valid input to scipy.interpolate.interp1d as the 'kind' parameter, using 'nearest', 'previous' or 'next'
        will cause a warning. This parameter controls how the offset and loading is interpolated over the step
    profile_interpolation_mode: {'nearest', 'linear'}, optional ('nearest')
        Used to generate the grid points for the second surface at the location of the grid points for the first
        surface, nearest ensures compatibility with sub models which change the profile, if the grid spacings of the
        surfaces match
    periodic_geometry: bool, optional (False)
        If True the surface profile will warp when applying the off set between the surfaces
    periodic_axes: tuple, optional ((False, False))
        For each True value the corresponding axis will be solved by circular convolution, meaning the result is
        periodic in that direction
    periodic_im_repeats: tuple, optional (1,1)
        The number of times the influence matrix should be wrapped along periodic dimensions, only used if at least one
        of periodic axes is True. This is necessary to ensure truly periodic behaviour, no physical limit exists
    method: {'auto', 'pk', 'double'}, optional ('auto')
        The method by which the normal contact is solved, only used for load controlled contact.
        'pk' uses the Polonsky and Keer algorithm for elastic contact.
        'double' uses a double iteration procedure, suitable for elastic contact with a maximum pressure.
        'auto' automatically selects 'pk' if there is no maximum pressure and 'double' if there is.
    max_it_interference: int, optional (100)
        The maximum number of iterations used to find the interference between the surfaces, only used if
        a total normal load is specified (Not used if contact is displacement controlled)
    rtol_interference: float, optional (1e-3)
        The relative tolerance on the load used as the convergence criteria for the interference optimisation loop, only
        used if a total normal load is specified (Not used if contact is displacement controlled)
    max_it_displacement: int, optional (100)
        The maximum number of iterations used to find the surface pressures from the interference, used for all IM
        based materials
    rtol_displacement: float, optional (1e-4)
        The norm of the residual used to declare convergence of the bccg iterations
    no_update_warning: bool, optional (True)
        Change to False to suppress warning given when no movement or loading changes are specified
    upper: float, optional (4.0)
        For load controlled contact the upper bound for the interference between the bodies will be this factor
        multiplied by the largest just touching gap. 4 is suitable for flat on flat contacts but for ball on flat
        contacts a lower value will give faster converging solutions

    Examples
    --------
    In the following example we model the contact between a rough cylinder and a flat plane.
    Both surface are elastic. This code could be used to generate load displacement curves.

    >>> import slippy.surface as s
    >>> import slippy.contact as c
    >>> # define contact geometry
    >>> cylinder = s.RoundSurface((1 ,np.inf, 1), shape=(256, 256), grid_spacing=0.001)
    >>> roughness = s.HurstFractalSurface(1, 0.2, 1000, shape=(256, 256), grid_spacing=0.001,
    >>>                                   generate = True)
    >>> combined = cylinder + roughness * 0.00001
    >>> flat = s.FlatSurface(shape=(256, 256), grid_spacing=0.001, generate = True)
    >>> # define material behaviour and assign to surfaces
    >>> material = c.Elastic('steel', properties = {'E':200e9, 'v':0.3})
    >>> combined.material = material
    >>> flat.material = material
    >>>
    >>> # make a contact model
    >>> my_model = c.ContactModel('qss_test', combined, flat)
    >>>
    >>> # make a modelling step to describe the problem
    >>> max_int = 0.002
    >>> n_time_steps = 20
    >>> my_step = c.QuasiStaticStep('loading', n_time_steps, no_time=True,
    >>>                             interference = [max_int*0.001, max_int],
    >>>                             periodic_geometry=True, periodic_axes = (False, True))
    >>> # add the steps to the model
    >>> my_model.add_step(my_step)
    >>> # add output requests
    >>> output_request = c.OutputRequest('Output-1',
    >>>                                  ['interference', 'total_normal_load',
    >>>                                   'loads_z', 'total_displacement',
    >>>                                   'converged'])
    >>> my_step.add_output(output_request)
    >>> # solve the model
    >>> final_result = my_model.solve()
    """
    _just_touching_gap = None
    _adhesion_model = None
    _initial_contact_nodes = None
    _upper = None

    def __init__(self, step_name: str, number_of_steps: int, no_time: bool = False,
                 time_period: float = 1.0,
                 off_set_x: typing.Union[float, typing.Sequence[float]] = 0.0,
                 off_set_y: typing.Union[float, typing.Sequence[float]] = 0.0,
                 interference: typing.Union[float, typing.Sequence[float]] = None,
                 normal_load: typing.Union[float, typing.Sequence[float]] = None,
                 relative_loading: bool = False,
                 adhesion: bool = True,
                 unloading: bool = False,
                 fast_ld: bool = False,
                 impact_properties: dict = None,
                 movement_interpolation_mode: str = 'linear',
                 profile_interpolation_mode: str = 'nearest',
                 periodic_geometry: bool = False, periodic_axes: tuple = (False, False),
                 periodic_im_repeats: tuple = (1, 1), method: str = 'auto',
                 max_it_interference: int = 100, rtol_interference=1e-3,
                 max_it_displacement: int = None, rtol_displacement=1e-5, no_update_warning: bool = True,
                 upper: float = 4.0):

        # movement interpolation mode sort out movement interpolation mode make array of values
        if impact_properties is not None:
            raise NotImplementedError("Impacts are not yet implemented")
        # work out the time step if needed:
        self._no_time = no_time
        self.total_time = time_period
        if not no_time and fast_ld:
            raise ValueError("Cannot have time dependence and fast_ld, either set no_time True or set fast_ld False")
        self._periodic_im_repeats = periodic_im_repeats
        self._fast_ld = fast_ld
        self._relative_loading = relative_loading
        self.profile_interpolation_mode = profile_interpolation_mode
        self._periodic_profile = periodic_geometry
        self._periodic_axes = periodic_axes
        self._max_it_interference = max_it_interference
        self._rtol_interference = rtol_interference
        self._max_it_displacement = max_it_displacement
        self._rtol_displacement = rtol_displacement
        self.number_of_steps = number_of_steps

        if method not in {'auto', 'pk', 'double'}:
            raise ValueError(f"Unrecognised method for step {step_name}: {method}")

        self._method = method
        self._height_optimisation_func = None
        self._adhesion = adhesion
        self._unloading = unloading
        self._upper_factor = upper

        self.time_step = time_period / number_of_steps

        # check that something is actually changing
        self.update = set()

        if not isinstance(off_set_x, Number) or not isinstance(off_set_y, Number):
            if no_time:
                raise ValueError("Can not have no time dependence and sliding contact")
            off_set_x = [off_set_x] * 2 if isinstance(off_set_x, Number) else off_set_x
            off_set_y = [off_set_y] * 2 if isinstance(off_set_y, Number) else off_set_y
            off_set_x_func = make_interpolation_func(off_set_x, movement_interpolation_mode, 'relative_off_set_x')
            off_set_y_func = make_interpolation_func(off_set_y, movement_interpolation_mode, 'relative_off_set_y')
            self._off_set_upd = lambda time: np.array([off_set_x_func(time), off_set_y_func(time)])
            self.update.add('off_set')
            self.off_set = None
        else:
            self.off_set = np.array([off_set_x, off_set_y])

        if normal_load is not None and interference is not None:
            raise ValueError("Both normal_load and interference are set, only one of these can be set")
        if normal_load is None and interference is None:
            if relative_loading:
                interference = 0
            else:
                raise ValueError("Cannot have no set load or interference and not relative loading, set either the"
                                 "normal load, normal interference or change relative_loading to True")

        if normal_load is not None:
            if isinstance(normal_load, Number):
                self.normal_load = normal_load
            else:
                self.normal_load = None
                self._normal_load_upd = make_interpolation_func(normal_load, movement_interpolation_mode,
                                                                'normal_load')
                self.update.add('normal_load')
            self.load_controlled = True
        else:
            self.normal_load = None

        if interference is not None:
            if isinstance(interference, Number):
                self.interference = interference
            else:
                self.interference = None
                self._interference_upd = make_interpolation_func(interference, movement_interpolation_mode,
                                                                 'interference')
                self.update.add('interference')
            self.load_controlled = False

        if not self.update and no_update_warning:
            warnings.warn("Nothing set to update")

        provides = {'off_set', 'loads_z', 'surface_1_displacement_z', 'surface_2_displacement_z',
                    'total_displacement_z', 'interference', 'just_touching_gap', 'surface_1_points', 'contact_nodes',
                    'surface_2_points', 'time', 'time_step', 'new_step', 'converged', 'gap', 'total_normal_load'}
        super().__init__(step_name, time_period, provides)

    def solve(self, previous_state, output_file):
        current_time = previous_state['time']
        if self._fast_ld:
            # solve the normal contact problem
            raise NotImplementedError("TODO")
            # TODO
            # change to disp controlled, set the displacement variable

        for s in self.sub_models:
            s.no_time = self._no_time

        if self.load_controlled:
            update_func = self._solve_load_controlled
        else:  # displacement controlled
            update_func = self._solve_displacement_controlled

        relative_time = np.linspace(0, 1, self.number_of_steps + 1)[1:]
        just_touching_gap = None

        original = dict()

        if self._relative_loading:
            original['normal_load'] = previous_state['total_normal_load'] if 'total_normal_load' in previous_state \
                else 0
            original['interference'] = previous_state['interference'] if 'interference' in previous_state else 0
            original['off_set'] = np.array(previous_state['off_set']) if 'off_set' in previous_state else \
                np.array([0, 0])

        self._adhesion_model = self.model.adhesion if self._adhesion else None

        current_state = dict()

        for i in range(self.number_of_steps):
            self.update_movement(relative_time[i], original)
            # find overlapping nodes
            if 'off_set' in self.update or just_touching_gap is None or not self._no_time:
                just_touching_gap, surface_1_points, surface_2_points \
                    = get_gap_from_model(self.model, interference=0, off_set=self.off_set,
                                         mode=self.profile_interpolation_mode, periodic=self._periodic_profile)
                self._just_touching_gap = just_touching_gap
                self._upper = None
            current_state = dict(just_touching_gap=just_touching_gap, surface_1_points=surface_1_points,
                                 surface_2_points=surface_2_points, off_set=self.off_set,
                                 time_step=self.time_step)
            if i == 0:
                current_state['new_step'] = True
            else:
                current_state['new_step'] = False
            # solve contact
            if self._unloading and 'contact_nodes' in previous_state:
                initial_contact_nodes = previous_state['contact_nodes']
            else:
                initial_contact_nodes = None

            self._initial_contact_nodes = initial_contact_nodes
            print('#####################################################\nTime step:', i,
                  '\n#####################################################')
            print('Set load:', self.normal_load)

            results = update_func(current_state)
            current_state.update(results)
            current_state['gap'] = (just_touching_gap - current_state['interference'] +
                                    current_state['total_displacement_z'])
            current_time += self.time_step
            current_state['time'] = current_time
            # solve sub models
            self.solve_sub_models(current_state)
            self.save_outputs(current_state, output_file)

        return current_state

    @property
    def upper(self):
        if self._upper is None:
            self._upper = np.max(self._just_touching_gap) * self._upper_factor
        return self._upper

    def update_movement(self, relative_time, original):
        for name in self.update:
            if self._relative_loading:
                self.__setattr__(name, original[name] + self.__getattribute__(f'_{name}_upd')(relative_time))
            else:
                self.__setattr__(name, self.__getattribute__(f'_{name}_upd')(relative_time))

    def _solve_load_controlled(self, current_state) -> dict:
        # if there is time dependence or we don't already have one, make a new height optimiser
        if not self._no_time or self._height_optimisation_func is None:
            opt_func = HeightOptimisationFunction(just_touching_gap=self._just_touching_gap,
                                                  model=self.model,
                                                  adhesion_model=self._adhesion_model,
                                                  initial_contact_nodes=self._initial_contact_nodes,
                                                  max_it_inner=self._max_it_displacement,
                                                  tol_inner=self._rtol_displacement,
                                                  material_options=dict(),
                                                  max_set_load=self.normal_load,
                                                  tolerance=self._rtol_interference,
                                                  periodic_axes=self._periodic_axes,
                                                  periodic_im_repeats=self._periodic_im_repeats)
            self._height_optimisation_func = opt_func
        else:
            opt_func = self._height_optimisation_func

        if self._unloading and 'contact_nodes' in current_state:
            contact_nodes = current_state['contact_nodes']
        else:
            contact_nodes = None
            # contact_nodes = np.ones(self._just_touching_gap.shape, dtype=np.bool)
        if self._method == 'auto':
            if np.isinf(opt_func.max_pressure):
                self._method = 'pk'
            else:
                self._method = 'double'

        if self._method == 'pk':
            opt_func.contact_nodes = None
            opt_func.p_and_k(self.normal_load)
        else:
            opt_func.change_load(self.normal_load, contact_nodes)
            # need to set bounds and pick a sensible starting point
            upper = self.upper
            print(f'upper bound set at: {upper}')
            if self._no_time:
                brackets = opt_func.get_bounds_from_cache(0, upper)
            else:
                brackets = (0, upper)
            print(f'Bounds adjusted using cache to: {brackets}')
            print(f'Interference tolerance set to {self._rtol_interference} Relative')
            opt_func.brent(0, upper, r_tol=self._rtol_interference, max_iter=self._max_it_interference)

        # noinspection PyProtectedMember

        results = opt_func.results
        load_conv = (np.abs(results['total_normal_load'] - self.normal_load) / self.normal_load) < 0.05
        results['converged'] = bool(load_conv) and not opt_func.last_call_failed
        return results

    def _solve_displacement_controlled(self, current_state):
        if not self._no_time or self._height_optimisation_func is None:
            h_opt_func = HeightOptimisationFunction(just_touching_gap=self._just_touching_gap,
                                                    model=self.model,
                                                    adhesion_model=self._adhesion_model,
                                                    initial_contact_nodes=self._initial_contact_nodes,
                                                    max_it_inner=self._max_it_displacement,
                                                    tol_inner=self._rtol_displacement,
                                                    material_options=dict(),
                                                    max_set_load=1,
                                                    tolerance=self._rtol_interference,
                                                    periodic_axes=self._periodic_axes,
                                                    periodic_im_repeats=self._periodic_im_repeats)
            self._height_optimisation_func = h_opt_func
        else:
            h_opt_func = self._height_optimisation_func

        if self._unloading and 'contact_nodes' in current_state:
            contact_nodes = current_state['contact_nodes']
        else:
            contact_nodes = None
            # contact_nodes = np.ones(self._just_touching_gap.shape, dtype=np.bool)

        h_opt_func.change_load(1, contact_nodes)
        _ = h_opt_func(self.interference, current_state)
        results = h_opt_func.results
        results['interference'] = self.interference
        results['converged'] = not h_opt_func.last_call_failed
        return results

    def __repr__(self):
        return f'{self.name}: QuasiStaticStep'


#!/usr/bin/env python
# -*- coding: utf-8 -*-

from __future__ import division
from __future__ import print_function
from __future__ import absolute_import
from __future__ import unicode_literals
from builtins import zip
from builtins import next
from builtins import str
from builtins import range
from builtins import object
__all__ = ["getColors", "nColors", "allmods"]
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as mpl_colors
from matplotlib import cm as cmx
import fsps
from .generalTools import calcQ, air_to_vac, getEmis
from .astrodata import dopita, sdss, vanzee, kewley
import pkg_resources

c = 2.9979e18 #ang/s
lsun = 3.839e33 #erg/s
planck = 6.626e-27
pc_to_cm = 3.08568e18

def getColors(vals, cname='CMRmap', minv=0.05, maxv=0.8, cmap=None,
              set_bad_vals=False, return_cNorm=False,
              logNorm=False, Ncol=100, return_cmap=False):
    '''
    sM = getColors(arr, cname='jet', minv=0.0, maxv=1.0)
    sM,cNorm = getColors(arr, cmap=cubehelix.cmap(), return_cNorm=True)
    '''
    if cmap is None:
        cmap = plt.get_cmap(cname)
    new_cmap = mpl_colors.LinearSegmentedColormap.from_list('trunc({0}, {1:.2f}, {2:.2f})'.format(cmap.name, minv, maxv), cmap(np.linspace(minv, maxv, Ncol)))
    if set_bad_vals:
        new_cmap.set_bad('white', alpha=1.0)
    if logNorm:
        cNorm = mpl_colors.LogNorm(vmin=vals.min(), vmax=vals.max())
    else:
        cNorm = mpl_colors.Normalize(vmin=vals.min(), vmax=vals.max())
    scalarMap = cmx.ScalarMappable(norm=cNorm, cmap=new_cmap)
    if return_cNorm:
        return scalarMap, cNorm
    if return_cmap:
        color_list = new_cmap(np.linspace(minv, maxv, Ncol))
        cmap_name = new_cmap.name+str(Ncol)
        return new_cmap.from_list(cmap_name, color_list, Ncol)
    else:
        scalarMap.set_array(vals)
        return scalarMap

def nColors(n, **kwargs):
    '''
    n_colors(50, minv=0.5, maxv=1.0, cname='Blues')
    '''
    sM = getColors(np.linspace(0.0, 1.0), **kwargs)
    colors = [sM.to_rgba(v) for v in np.linspace(0.0, 1.0, n)]
    return colors

def sextract(text, par1=None, par2=None):
    '''
    to extract stuff from cloudy output
    '''
    if np.size(text) == 1:
        if type(par1) is int:
            str1 = text[par1::]
        elif type(par1) is str or type(par1) is str:
            str1 = text.split(par1)
            if len(str1) == 1:
                return ''
            else:
                str1 = str1[-1]
        else:
            str1 = text
        if type(par2) is int:
            str2 = str1[0:par2]
        elif type(par2) is str or type(par2) is str:
            str2 = str1.split(par2)
            if len(str2) == 1:
                return ''
            else:
                str2 = str2[0]
        else:
            str2 = str1
        return str2
    else:
        res = []
        for subtext in text:
            res1 = sextract(subtext, par1=par1, par2=par2)
            if res1 != '':
                res.append(res1)
        return res

class modObj(object):
    '''
    '''
    def __init__(self, dir_, prefix, parline, read_out=False, read_rad=False,
                 read_cont=False, use_doublet=False, read_emis=False,
                 read_heat=False, read_cool=False, **kwargs):
        '''
        this needs to be called from other class or given
        a line from a ".pars" file
        [0]modnum; [1]logZ; [2]age; [3]logU; [4]logR; [5]logQ
        '''
        self.modnum = int(parline[0])
        self.logZ = parline[1]
        self.age = parline[2]
        self.logU = parline[3]
        self.logR = parline[4]
        self.logQ = parline[5]
        self.nH = parline[6]
        try:
            self.efrac = parline[7]
        except IndexError:
            self.efrac = -1.0
        try:
            self.fbhb = parline[8]
        except IndexError:
            self.fbhb = 0.0

        self.logq = np.log10((10.0**self.logQ)/(np.pi*4.0*self.nH*(10.0**self.logR)**2.0))
        self.fl = '{}{}{}'.format(dir_, prefix, self.modnum)
        self.load_lines(use_doublet=use_doublet)
        if read_out:
            self._read_out()
        if read_cont:
            self._load_cont()
        if read_rad:
            self._dat = dict()
            self._init_rad()
            eles = ['H', 'He', 'C', 'N', 'O', 'S', 'Si', 'Fe']
            self.ion_names, self.n_ions, self.ion_arr = dict(), dict(), dict()
            for ele in eles:
                self._init_ele(ele)
            self._init_phys()
        if read_emis:
            self._init_emis()
        if read_heat:
            self._init_heat()
        if read_cool:
            self._init_cool()
        return
    def add_lines(self, lines):
        line_info = np.genfromtxt(self.fl+'.lineflux')
        lam, flu = line_info[:,0], line_info[:,1]
        for name, wav in list(lines.items()):
            matchind = np.argmin(np.abs(lam-wav))
            self.__setattr__(name, flu[matchind])
        return
    def load_lines(self, use_doublet=False, **kwargs):
        names, vacwavs = getEmis()
        self.lines = dict(names=names, wavs=vacwavs)
        lines = {'Lya':1215.67,
                 'Ha':6564.60,
                 'HeI':5877.243,
                 'HeII':4687.015,
                 'HeIIu':1640.42,
                 'Hb':4862.71,
                 'Hg':4341.692,
                 'Hd':4102.892,
                 'OIIIa':4960.295,
                 'OIIIb':5008.240,
                 'NIIa':6549.86,
                 'NIIb':6585.27,
                 'OIIa':3727.10,
                 'OIIb':3729.86,
                 'SIIa':6718.294,
                 'SIIb':6732.673,
                 'OI':6302.046,
                 'NeIIIb':3869.86,
                 'NeIIIa':3968.59,
                 'SIII':6313.81,
                 'ArIII':7137.77}
        line_info = np.genfromtxt(self.fl+'.lineflux')
        lam, flu = line_info[:,0], line_info[:,1]
        for name, wav in list(lines.items()):
            matchind = np.argmin(np.abs(lam-wav))
            self.__setattr__(name, flu[matchind])
        self.HaHb = self.Ha/self.Hb
        def logify(a,b):
            return np.log10(a/b)
        def logHa(x):
            return np.log10(x/self.Ha)
        def logHb(x):
            return np.log10(x/self.Hb)
        self.log_NII_Ha = logHa(self.NIIa+self.NIIb)
        self.log_SII_Ha = logHa(self.SIIa+self.SIIb)
        self.log_OIII_Hb = logHb(self.OIIIa+self.OIIIb)
        #
        self.log_NIIa_Ha = logHa(self.NIIa)
        self.log_NIIb_Ha = logHa(self.NIIb)
        self.log_SIIa_Ha = logHa(self.SIIa)
        self.log_SIIb_Ha = logHa(self.SIIb)
        #
        self.log_OIIIa_Hb = logHb(self.OIIIa)
        self.log_OIIIb_Hb = logHb(self.OIIIb)
        #
        self.log_OIII_OII = logify(self.OIIIa+self.OIIIb, self.OIIa+self.OIIb)
        self.log_OIIIa_OII = logify(self.OIIIa, self.OIIa+self.OIIb)
        self.log_OIIIb_OII = logify(self.OIIIb, self.OIIa+self.OIIb)
        #
        self.log_OI_Ha = logHa(self.OI)
        self.log_NII_OII = logify(self.NIIa+self.NIIb, self.OIIa+self.OIIb)
        self.R23 = logHb(self.OIIa+self.OIIb+self.OIIIa+self.OIIIb)
        return
    def _load_cont(self, dist_corr=False, output_units=False, **kwargs):
        cont_info = np.genfromtxt(self.fl+'.contflux', skip_header=1)
        # erg / s / cm2
        self.lam, self.nebflu = cont_info[:,0], cont_info[:,3]
        self.incflu, self.attflu = cont_info[:,1], cont_info[:,2]
        self.spec_Q = calcQ(self.lam, self.incflu, f_nu=True)
        if dist_corr: # erg/s
            self.nebflu *= self.dist_fact
            self.attflu *= self.dist_fact
            self.incflu *= self.dist_fact
            self.spec_Q = calcQ(self.lam, self.incflu*c/self.lam, f_nu=True)
        elif output_units: # Lsun/Hz
            self.nebflu *= self.dist_fact/lsun * self.lam / c
            self.attflu *= self.dist_fact/lsun * self.lam / c
            self.incflu *= self.dist_fact/lsun * self.lam / c
            self.spec_Q = calcQ(self.lam, self.incflu*lsun, f_nu=True)
        return
    def get_fsps_spec(self, **kwargs):
        sp = fsps.StellarPopulation(zcontinuous=1)
        sp.params['logzsol'] = self.logZ
        lam, spec = sp.get_spectrum(tage=self.age*1.0e-9)
        self.__setattr__('fsps_lam', lam)
        self.__setattr__('fsps_spec', spec)
        self.__setattr__('fsps_Q', calcQ(lam, spec*lsun, f_nu=True))
        return
    def _read_f(self, key, delimiter='\t', comments=';', names=True, **kwargs):
        '''
        self._read_f('.rad')
        '''
        file_ = self.fl+key
        try:
            return np.genfromtxt(file_,delimiter=delimiter,
                                 comments=comments,
                                 names=names, **kwargs)
        except IOError:
            return None

    def _init_rad(self):
        '''
        self._init_rad()
        attributes:
            n_zones
            zones
            depth
            thickness (cm)
            radius_all (cm)
            rad_pc (pc)
            dr_all (cm)
            dv_all (cm)
            r_in (cm)
            r_out (cm)
        '''
        self._dat['rad'] = self._read_f('.rad')
        if self._dat['rad'] is not None:
            self.n_zones = self._dat['rad'].size
            self.zones = np.arange(self.n_zones)
            self.depth = self._dat['rad']['depth']
            self.thickness = self.depth[-1]
            self.radius_all = self._dat['rad']['radius']
            self.rad_pc = self.radius_all/pc_to_cm
            self.dr_all = self._dat['rad']['dr']
            self.dv_all = 4.*np.pi*self.radius_all**2*self.dr_all
            self.r_in = self.radius_all[0] - self.dr_all[0]/2.
            self.r_out = self.radius_all[-1] + self.dr_all[0]/2.
            if self.Phi0 == 0.0:
                self.Phiarr = self.Qarr/(4.*np.pi*self.r_in**2.)
                self.Phi0 = self.Phiarr.sum()
        return
    def _init_emis(self):
        '''
        '''
        self._dat['emis'] = self._read_f('.emis')
        if self._dat['emis'] is not None:
            emis = self._dat['emis']
            self.depth = emis['depth']/pc_to_cm
            self.frac_depth = self.depth/(self.rad_pc[-1]-self.rad_pc[0])
            self.emis_labels = np.asarray(emis.dtype.names[1::])
            self.n_emis = np.size(self.emis_labels)
            self.emis_full = np.zeros((self.n_emis, np.size(emis)))
            trans_emis = lambda x: pow(10., x)
            for i, label in enumerate(self.emis_labels):
                self.emis_full[i] = trans_emis(emis[label])
            self.indHe = np.argmin(np.abs(self.ion_arr['He'][1]-0.5))
            self.indH = np.argmin(np.abs(self.ion_arr['H'][1]-0.5))
            return
    def _init_phys(self):
        '''
        self._init_phys()
        adds attributes:
            ne_all: electron density (cm^-3)
            nH_all: hydrogen density (cm^-3)
            Te: electron temperature (K)
        '''
        key = 'phys'
        self._dat[key] =  self._read_f('.phys')
        if self._dat[key] is not None:
            self.ne_all = self._dat[key]['ne']
            self.nH_all = self._dat[key]['nH']
            self.nenH = self.ne_all*self.nH_all
            self.Te = self._dat[key]['Te']
            self.ff_all = self._dat[key]['fillfac']
        return
    def _init_cool(self):
        key = 'cool'
        try:
            self._dat[key] = self._read_f('.'+key, comments='#')
        except:
            self._dat[key] = None
            print("ERROR")
            print(self.fl)
        if self._dat[key] is not None:
            attkeys = dict(cool_RecMet='hvFB', # heavy elem recomb cooling
                           cool_ColMet='Hvin', # heavy elem collis ionization
                           cool_FFC='FFcm', #F-F cooling (non e-)
                           cool_H='H',
                           cool_He='He',
                           cool_N='N',
                           cool_O='O',
                           cool_S='S',
                           cool_Ne='Ne',
                           cool_C='C',
                           cool_Ar='Ar',
                           cool_Fe='Fe',
                           cool_Al='Al',
                           cool_Si='Si',
                           cool_HminFB='Hfb')
            self.cool = self._dat[key]['Ctotergcm3s']
            self.Ctot = self._vol_integ(self.cool)
            for att,keyname in list(attkeys.items()):
                vals = self._dat[key][keyname]
                self.__setattr__(att, vals)
                self.__setattr__('frac_'+att, self._vol_integ(vals)/self.Ctot)
            CE_names = self._dat[key].dtype.names[5:33]
            self.cool_CE = np.sum([self._vol_integ(self._dat[key][elname]) for elname in CE_names])
            self.frac_cool_CE = np.sum([self._vol_integ(self._dat[key][elname])/self.Ctot for elname in CE_names])
            self.frac_cool_FF_FB = self.Cool_HFFc+self.Cool_HFBc
        return
    def _init_heat(self):
        htots, outfs = [], []
        heat_labs = []
        fname = self.fl+'.heat'
        file_ = open(fname, 'r')
        for line in file_:
            if line[0] == '#':
                continue
            line = line.split('\n')[0]
            lps = line.split('\t')
            htots.append(lps[2])
            x = lps[4::]
            labs = x[0::2]
            fracs = x[1::2]
            od = {}
            [od.__setitem__(key, float(val)) for key, val in zip(labs,fracs)]
            outfs.append(od)
            for lab in labs:
                if lab not in heat_labs:
                    heat_labs.append(lab)
        hf = {}
        [hf.__setitem__(lab,
                        np.array([outf[lab]
                                  if lab in list(outf.keys())
                                  else 0.0 for outf in outfs]))
         for lab in heat_labs]
        Htots = np.array([float(htot) for htot in htots])
        Htot = self._vol_integ(Htots)
        hr = {}
        for key, arr in list(hf.items()):
            hr.__setitem__(key, np.sum(self._vol_integ(arr*Htots))/Htot)
        self.__setattr__('heatfracs', hr)
        return
    def _init_ele(self, key):
        '''
        keys = [H, He, C, N, O, S, Si, Fe]
        attributes:
            ion_names['C'] = C__1, C__2,...C__n
            n_ions['C'] = n
            ion_arr['C'][0] = f1, f2,...,f_r
        '''
        self._dat[key] = self._read_f('.ele_'+key)
        if self._dat[key] is not None:
            ion_names = self._dat[key].dtype.names[1:]
            n_ions = np.size(ion_names)
            ion_arr = np.zeros((n_ions, self.n_zones))
            for i, ion in enumerate(ion_names):
                ion_arr[i,:] = self._dat[key][ion]
            self.ion_names[key] = ion_names
            self.n_ions[key] = n_ions
            self.ion_arr[key] = ion_arr
        return

    @property
    def dvff(self):
        try:
            return self.dv_all*self.ff_all
        except:
            return None

    def _quiet_div(self, a, b):
        if a is None or b is None:
            to_return = None
        else:
            np.seterr(all="ignore")
            to_return = a/b
            np.seterr(all=None)
        return to_return
    def _vol_cum(self, a):
        if a is None or self.dvff is None:
            return None
        else:
            return (a*self.dvff).cumsum()
    def _vol_integ(self, a):
        if a is None or self.dvff is None:
            return None
        else:
            return (a*self.dvff).sum()
    def _dV(self, a):
        if a is None or self.dvff is None:
            return None
        else:
            return (a*self.dvff)
    def _vol_mean(self, a, b=1.):
        return self._quiet_div(self._vol_integ(a*b), self._vol_integ(b))

    @property
    def T0(self):
        try:
            return self._vol_mean(self.Te, self.nenH)
        except:
            return None

    @property
    def Tpiem(self):
       try:
           return self._vol_mean((self.Te - self.T0)**2./self.nenH, self.T0**2)
       except:
           return None
    def _i_emis(self, ref):
        '''
        get index for emissivity by label or index
        '''
        if type(ref) is str or type(ref) is np.string_:
            if ref in self.emis_labels:
                to_return = np.squeeze(np.where(self.emis_labels == ref)).item()
            else:
                return None
        elif type(ref) is int or type(ref) is np.int32:
            to_return = ref
        else:
            to_return = None
        return to_return
    def get_emis(self, ref):
        '''
        return emissivity.
        '''
        if self._i_emis(ref) is not None:
            return self.emis_full[self._i_emis(ref)]
        else:
            return None
    def get_cumVol_emis(self, ref):
        if self._i_emis(ref) is not None:
            return
        else:
            return None
    def get_emis_vol(self, ref):
        '''
        return integration of the emissivity on the volume
        '''
        return self._vol_integ(self.get_emis(ref))
    def get_frac_emis(self, ref):
        '''
        emiss dV / tot emiss
        '''
        return self._dV(self.get_emis(ref))/self.get_emis_vol(ref)
    def get_emis_rad(self, ref):
        '''
        return integration of the emissivity on the radius
        '''
        return self._rad_integ(self.get_emis(ref))
    def _read_out(self):
        '''
        self._read_out()
        attributes:
            dist_fact: 4 pi Rinner^2 (cm^2)
            Phi0: ionizing photon flux (s^-1 cm^-2)
            clogQ: Phi0 * dist_fact = Q (s-1)
            gasC, gasN, gasO: n relative to H
            DGR: dust to gas ratio
            Av_ex: extinction from extended source
            Av_pt: extinction from pt source
        '''
        filename = self.fl+'.out'
        self.out = {}
        file_ = open(filename, 'r')
        for line in file_:
            line = line.split('\n')[0]
            if line[0:8] == ' ####  1':
                self.out['###First'] = line
            elif line[0:5] == ' ###':
                self.out['###Last'] = line
            elif 'Hi-Con' in line:
                for i in range(7):
                    self.out['SED' + str(i+1)] = next(file_)
            elif line[0:15] == ' IONIZE PARMET:':
                self.out['INZ'] = line
            elif 'H :' in line:
                self.out['gascomp'] = line
            elif 'Dust to gas ratio' in line:
                self.out['dust'] = line
            elif 'ENERGY BUDGET' in line:
                self.out['energy'] = line
            elif 'Cooling:' in line:
                self.out['cool'] = line
            elif 'Heating:' in line:
                self.out['heat'] = line
            elif line[0:5] == ' HFBc':
                self.out['HFBc'] = line
            elif 'The geometry is' in line:
                self.out['geometry'] = line
        file_.close()
        self.dist_fact = 4.0*np.pi*(10.0**self.logR)**2.0
        try:
            self.H_Rec_Lum = float(sextract(sextract(self.out['HFBc'], 'HFBc', 18), 9, 8))
        except:
            self.H_Rec_Lum = 0.0
            #print(sextract(sextract(self.out['HFBc'], 'HFBc', 18), 9, 8))
        #self.Phi0 = float(sextract(self.out['SED2'], 'Ion pht flx:'))
        # Ion pht flx: phi(H) = Q/4piR2
        #self.clogQ = np.log10(self.Phi0*self.dist_fact)
        self.strom_logU = float(sextract(self.out['INZ'], 'U(sp):', 'Q(ion):'))
        # Q(ion) is exiting
        self.Qarr = np.zeros(4)
        self.Phiarr = np.zeros(4)
        try:
            self.Qarr[0] = float(sextract(self.out['SED2'], 'Q(1.0-1.8):', 'Q(1.8-4.0):'))
            self.Qarr[1] = float(sextract(self.out['SED2'], 'Q(1.8-4.0):', 'Q(4.0-20):'))
            self.Qarr[2] = float(sextract(self.out['SED2'], 'Q(4.0-20):', 'Q(20--):'))
            self.Qarr[3] = float(sextract(self.out['SED2'], 'Q(20--):', 'Ion pht'))
            self.Qarr = pow(10., self.Qarr)
            self.input_lum = True
        except:
            pass
        self.Qh = self.Qarr.sum()
        try:
            self.Phiarr[0] = float(sextract(self.out['SED2'], 'phi(1.0-1.8):', 'phi(1.8-4.0):'))
            self.Phiarr[1] = float(sextract(self.out['SED2'], 'phi(1.8-4.0):', 'phi(4.0-20):'))
            self.Phiarr[2] = float(sextract(self.out['SED2'], 'phi(4.0-20):', 'phi(20--):'))
            self.Phiarr[3] = float(sextract(self.out['SED2'], 'phi(20--):', 'Ion pht'))
            self.Phiarr = pow(10., self.Phiarr)
            self.input_lum = False
        except:
            pass
        self.Phi0 = self.Phiarr.sum()
        if self.Qh == 0.0:
            self.Qarr = self.Phiarr*self.dist_fact
            self.Qh = self.Qarr.sum()
        self.Qhe = self.Qarr[1::].sum()
        self.QhQhe = np.log10(self.Qh) - np.log10(self.Qhe)
        self._set_lineCool()
        try:
            self.Heat_BF = float(sextract(sextract(self.out['heat'], 'BFH1', 14), ':', 5))
        except:
            self.Heat_BF = 0.0
        self.Heat = float(sextract(self.out['energy'], 'Heat:', 'Coolg:'))
        self.Cool = float(sextract(self.out['energy'], 'Coolg:', 'Error:'))
        self.RecLin = float(sextract(self.out['energy'], 'Rec Lin:', 8))
        self.gasC = float(sextract(self.out['gascomp'], 'C :', 8))
        self.gasN = float(sextract(self.out['gascomp'], 'N :', 8))
        self.gasO = float(sextract(self.out['gascomp'], 'O :', 8))
        if 'dust' in self.out:
            self.DGR = float(sextract(self.out['dust'], '(by mass):',',' ))
            self.Av_ex = float(sextract(self.out['dust'], 'AV(ext):', '(pnt)'))
            self.Av_pt = float(sextract(self.out['dust'], ' (pnt):'))
        return
    def _set_lineCool(self):
        keys = ['HFBc', 'HFFc', 'Clin 912.000A:', 'N  2 6584.00A:',
                ' S II 6731.00A:','S II 6716.00A:', 'TOTL 3727.00A:',
                'S  3 9532.00A:', 'O  3 5007.00A:', 'O  3 4959.00A:']
        keyattrs = ['HFBc', 'HFFc', 'Ly', 'NII', 'SIIb', 'SIIa',
                    'O_3727', 'SIII', 'O_5007', 'O_4959']
        self.cool_frac = {}
        for key, keyattr in zip(keys, keyattrs):
            if key[0] == 'H':
                try:
                    val = float(sextract(sextract(self.out['cool'],
                                                  key, 14),
                                         ':', 5))
                except:
                    val = 0.0
            else:
                try:
                    val = float(sextract(self.out['cool'], key, 5))
                except:
                    val = 0.0
            self.__setattr__('Cool_'+keyattr, val)
        self.Cool_Otot = np.sum([self.Cool_O_3727,
                                 self.Cool_O_5007,
                                 self.Cool_O_4959])
        self.Cool_Stot = np.sum([self.Cool_SIIa,
                                 self.Cool_SIIb,
                                 self.Cool_SIII])
        return

class allmods(object):
    '''
    mods = outobj.allmods(dir, prefix, read_out=True, read_rad=False)
    '''
    def __init__(self, dir_, prefix, **kwargs):
        self.modpars = np.genfromtxt('{}{}.pars'.format(dir_, prefix))
        self.load_mods(dir_, prefix, **kwargs)
        self.set_pars()
        self.set_arrs()
        read_out = kwargs.get('read_out', False)
        read_rad = kwargs.get('read_rad', False)
        if read_out:
            self.add_arrs('gasC', 'gasN', 'gasO')
            if hasattr(self.mods[0], 'DGR'):
                self.add_arrs('DGR', 'Av_ex', 'Av_pt')
            if read_rad:
                self.add_arrs('Te')

    def load_mods(self, dir_, prefix, **kwargs):
        mods = []
        for par in self.modpars:
            mod = modObj(dir_, prefix, par, **kwargs)
            mods.append(mod)
        self.__setattr__('mods', mods)
        self.__setattr__('nmods', len(mods))
        return
    def set_pars(self):
        self.logZ_vals = np.unique(self.modpars[:,1])
        self.age_vals = np.unique(self.modpars[:,2])
        self.logU_vals = np.unique(self.modpars[:,3])
        self.logR_vals = np.unique(self.modpars[:,4])
        self.logQ_vals = np.unique(self.modpars[:,5])
        self.nH_vals = np.unique(self.modpars[:,6])
        try:
            self.efrac_vals = np.unique(self.modpars[:,7])
        except IndexError:
            self.efrac_vals = np.array([-1.0])
        try:
            self.fbhb_vals = np.unique(self.modpars[:,8])
        except IndexError:
            self.fbhb_vals = np.array([0.0])
    def set_arrs(self):
        iterstrings = ['logZ', 'age', 'logU', 'logR', 'logQ', 'nH',
                       'efrac','fbhb',
                       'log_NII_Ha','log_NIIa_Ha','log_NIIb_Ha',
                       'log_OIII_Hb','log_OIIIa_Hb','log_OIIIb_Hb',
                       'log_SII_Ha','log_SIIa_Ha','log_SIIb_Ha',
                       'HaHb', 'R23','log_NII_OII', 'log_OIII_OII']
        for i in iterstrings:
            vals = np.array([mod.__getattribute__(i) for mod in self.mods])
            self.__setattr__(i, vals)
    def add_arrs(self, *args):
        for item in args:
            try:
                vals = np.array([mod.__getattribute__(item) for mod in self.mods])
                self.__setattr__(item, vals)
            except AttributeError:
                continue
        return
    def add_lines(self, linedict={}):
        '''
        self.add_lines(linedict={'O3':1666.0})
        '''
        [mod.add_lines(linedict) for mod in self.mods]
        return
    def makeBPT(self, ax=None, plot_data=True, line_ratio='NIIb',
                gridnames=None, bpt_inds=None, axlabs=None, varsize=22,
                plt_pars={}, data_only=False, **kwargs):
        '''
        mo.makeBPT(ax=ax, const1='age', val1=0.5e6, const2=logR, val2=19.0,
                   const3='nH', val3=10.0)
        default line_ratio is NII_b / Ha
        line_ratio = ['NII', 'SII', 'OII', 'OI', 'R23']
        or bpt_inds=['log_OIb_Ha', 'log_OIIIb_Hb']
        '''
        if axlabs is None:
            xlabel = r'\textbf{log [N II] $\lambda 6584$ / H$\alpha$}'
            ylabel = r'\textbf{log [O III] $\lambda 5007$ / H$\beta$}'
        else:
            xlabel=axlabs[0]
            ylabel=ylabs[0]
        if bpt_inds is None:
            xlabel = r'\textbf{log [N II] $\lambda 6584$ / H$\alpha$}'
            ylabel = r'\textbf{log [O III] $\lambda 5007$ / H$\beta$}'
            bpt_inds = ['log_NIIb_Ha', 'log_OIIIb_Hb']
            if line_ratio == 'NII':
                xlabel = r'\textbf{log [N II] $\lambda 6548,6584$ / H$\alpha$}'
                ylabel = r'\textbf{log [O III] $\lambda 4959,5007$ / H$\beta$}'
                bpt_inds = ['log_NII_Ha', 'log_OIII_Hb']
            if line_ratio == 'SII':
                bpt_inds[0] = 'log_SII_Ha'
                xlabel = r'\textbf{log [S II] $\lambda 6716,6731$ / H$\alpha$}'
            if line_ratio == 'OI':
                bpt_inds[0] = 'log_OI_Ha'
                xlabel = r'\textbf{log [O I] $\lambda 6300$ / H$\alpha$}'
            if line_ratio == 'OII':
                bpt_inds = ['log_NII_OII', 'log_OIII_OII']
                ylabel = r'\textbf{log [O III] $\lambda 4959,5007$ / [O II] $\lambda 3726,3727$}'
                xlabel = r'\textbf{log [N II] $\lambda 6548,6584$ / [O II] $\lambda 3726,3727$}'
            if line_ratio == 'R23':
                bpt_inds = ['R23','log_OIII_OII']
                ylabel = r'\textbf{log [O III] $\lambda 4959,5007$ / [O II] $\lambda 3726,3727$}'
                xlabel = r'\textbf{(log [O II] $\lambda 3726,3727$ + [O III] $\lambda 4959,5007$) / H$\beta$}'
        if ax is None:
            fig = plt.figure()
            ax = fig.add_subplot(111)
        if data_only:
            vanzee.plot_bpt(plot_data, line_ratio=line_ratio, ax=ax, **plt_pars)
            sdss.plot_bpt(plot_data, line_ratio=line_ratio, ax=ax, **plt_pars)
            return
        pd = {'const1':'age',
              'val1':0.5e6,
              'const2':'logR',
              'val2':19.0,
              'const3':'nH',
              'val3':100.0,
              'const4':None,
              'val4':0.0,
              'const5':None,
              'val5':0.0}
        for key, val in list(kwargs.items()):
            pd[key] = val
        allvars = ['nH', 'logZ', 'logR', 'logU', 'age', 'efrac', 'fbhb']
        [allvars.remove(x) for x in [pd['const1'], pd['const2'], pd['const3'], pd['const4'], pd['const5']] if x is not None]
        if gridnames is None:
            x_name, y_name = allvars[0], allvars[1]
        else:
            x_name, y_name = gridnames[0], gridnames[1]
        grid_x = self.__getattribute__(x_name+'_vals')
        grid_y = self.__getattribute__(y_name+'_vals')

        cut_z = kwargs.get('cut_z', None)
        if cut_z is not None:
            logZmin= cut_z[0]
            logZmax=cut_z[1]
            if x_name == 'logZ':
                grid_x = grid_x[(grid_x >= logZmin) & (grid_x <= logZmax)]
            if y_name == 'logZ':
                grid_y = grid_y[(grid_y >= logZmin) & (grid_y <= logZmax)]

        use_mods = [mod for mod in self.mods
                    if (mod.__getattribute__(pd['const1']) == pd['val1'])
                    & (mod.__getattribute__(pd['const2']) == pd['val2'])
                    & (mod.__getattribute__(pd['const3']) == pd['val3'])]
        if pd['const4'] is not None:
            use_mods = [mod for mod in self.mods
                        if (mod.__getattribute__(pd['const1']) == pd['val1'])
                        & (mod.__getattribute__(pd['const2']) == pd['val2'])
                        & (mod.__getattribute__(pd['const3']) == pd['val3'])
                        & (mod.__getattribute__(pd['const4']) == pd['val4'])]

        gshape = (len(grid_y), len(grid_x))
        X, Y = np.meshgrid(grid_x, grid_y, indexing='xy')
        Zx = np.zeros(gshape)
        Zy = np.zeros(gshape)
        nrows = gshape[0]
        ncols = gshape[1]
        for ind, val in np.ndenumerate(X):
            arr = [mod for mod in use_mods if (mod.__getattribute__(x_name) == val) & (mod.__getattribute__(y_name) == Y[ind])]
            Zx[ind] = arr[0].__getattribute__(bpt_inds[0])
            Zy[ind] = arr[0].__getattribute__(bpt_inds[1])
        if plot_data:
            vanzee.plot_bpt(plot_data, line_ratio=line_ratio, ax=ax)
            sdss.plot_bpt(plot_data, line_ratio=line_ratio, ax=ax, **plt_pars)
        ax.set_xlabel(xlabel, fontsize=16)
        ax.set_ylabel(ylabel, fontsize=16)
        for i in range(nrows):
            color = kwargs.get('color', 'k')
            lw = kwargs.get('lw', 2)
            alpha = kwargs.get('alpha', 0.95)
            if i == 0:
                par_label = kwargs.get('par_label', '__nolegend__')
                ax.plot(Zx[i,:], Zy[i,:], color=color, lw=lw,
                        alpha=alpha, label=par_label)
            else:
                ax.plot(Zx[i,:], Zy[i,:], color=color, lw=lw, alpha=alpha,
                        label='__nolegend__')
        row_labs = [(Zx[i,0], Zy[i,0], r'${0:.1f}$'.format(float(np.unique(Y[i,:])))) for i in range(gshape[0])]
        for i in range(ncols):
            ax.plot(Zx[:,i], Zy[:,i], color=color, lw=lw, alpha=alpha,
                    label='__nolegend__')
        col_labs = [(Zx[0, i], Zy[0, i], r'${0:.1f}$'.format(float(np.unique(X[:,i])))) for i in range(gshape[1])]

        var_label = kwargs.get('var_label', '__nolegend__')
        if var_label:
            for lab in col_labs:
                ax.annotate(lab[-1],
                            xy=(lab[0], lab[1]), xycoords='data',
                            xytext=(0, -10), textcoords='offset points',
                            size=varsize,
                            horizontalalignment='left',
                            verticalalignment='top')
            ax.annotate(r'$\log \mathrm{Z}/\mathrm{Z}_{\odot}$',
                        xy=(col_labs[2][0], col_labs[2][1]),
                        xycoords='data', xytext=(0, -50),
                        textcoords='offset points', size=varsize,
                        horizontalalignment='left',
                        verticalalignment='top')
            for lab in row_labs:
                ax.annotate(lab[-1],
                            xy=(lab[0], lab[1]), xycoords='data',
                            xytext=(-10, 0), textcoords='offset points',
                            size=varsize,
                            horizontalalignment='right',
                            verticalalignment='bottom')
            ax.annotate(r'$\log \mathcal{U}_0$',
                        xy=(row_labs[1][0], row_labs[1][1]),
                        xycoords='data', xytext=(-50, 15),
                        textcoords='offset points', size=varsize,
                        horizontalalignment='right',
                        verticalalignment='bottom')
        plt.legend(numpoints=1)
        return
    def group_mods(self, xval='logZ', yval='age', zval='NIIb',
                   const='logU', cval=-2.0, make_cut=False, **kwargs):
        grid_x = self.__getattribute__(xval+'_vals')
        grid_y = self.__getattribute__(yval+'_vals')
        if make_cut:
            xlims = kwargs.get('xlims', (-1.0, 0.1))
            ylims = kwargs.get('ylims', (0.0, 10.e6))
            grid_x = grid_x[(grid_x >= xlims[0]) & (grid_x <= xlims[1])]
            grid_y = grid_y[(grid_y >= ylims[0]) & (grid_y <= ylims[1])]
        X, Y = np.meshgrid(grid_x, grid_y)
        Z = np.zeros_like(X)
        for index, x in np.ndenumerate(Z):
            mind = [i for i in range(self.nmods)
                    if (self.mods[i].__getattribute__(xval) == X[index]
                        and self.mods[i].__getattribute__(yval) == Y[index]
                        and self.mods[i].__getattribute__(const) == cval)]
            try:
                Z[index] = self.mods[mind[0]].__getattribute__(zval)
            except AttributeError:
                print('not a valid attribute.')
        if xval == 'age':
            X*=1.0e-6
        if yval == 'age':
            Y*=1.0e-6
        return X,Y,Z
    def pxl_plot(self, xval='logZ', yval='age', zval='log_OIII_Hb',
                 const='logU', cval=-2.0, ax=None, cname='CMRmap',
                 cmap=None, sM=None, cNorm=None, cb_arr=None,
                 no_cbar=False, show_grid=False, aspect='auto',
                 xlabels=None, ylabels=None,**kwargs):
        '''
        mods.pxl_plot(xval='logZ', yval='age', zval='log_OIII_Hb',
                      const='logR', cval=18, clab='log R (cm)')
        '''
        X, Y, Z = self.group_mods(xval=xval, yval=yval, zval=zval,
                                  const=const, cval=cval, **kwargs)
        masked_array = np.ma.array(Z, mask=np.isnan(Z))
        if (sM is None) or (cNorm is None):
            if cb_arr is not None:
                arr_in = cb_arr
            else:
                arr_in = masked_array
            if cmap is not None:
                sM, cNorm = getColors(arr_in, return_cNorm=True,
                                      set_bad_vals=True, cmap=cmap)
            else:
                sM, cNorm = getColors(arr_in, return_cNorm=True,
                                      set_bad_vals=True, cname=cname)
            cmap = plt.get_cmap(cname)
        if ax is None:
            fig = plt.figure()
            ax = fig.add_subplot(111)
        pf = ax.imshow(Z, norm=cNorm, aspect=aspect,
                       interpolation='nearest', origin='lower',
                       cmap=cmap)
        xCenters = X[0]
        yCenters = Y[:,0]
        if xlabels is None:
            xlabels = ['{0:.1f}'.format(x) for x in xCenters]
        if ylabels is None:
            ylabels = ['{0:.1f}'.format(y) for y in yCenters]
        for axis, labels in zip([ax.xaxis, ax.yaxis], [xlabels, ylabels]):
            locs = np.arange(len(labels))
            axis.set_ticks(locs + 0.5, minor=True)
            axis.set(ticks=locs, ticklabels=labels)
        if show_grid:
            ax.grid(True, which='minor')
        xlab = kwargs.get('xlab', None)
        ylab = kwargs.get('ylab', None)
        clab = kwargs.get('clab', None)
        if xlab is None:
            xlab = xval
            ylab = yval
        ax.set_xlabel(xlab)
        ax.set_ylabel(ylab)
        if no_cbar:
            return pf
        else:
            cb = plt.colorbar(pf, ax=ax)
            if clab is not None:
                clab_size=kwargs.get('clab_size', 14)
                cb.set_label(clab, size=clab_size)
            return ax

def nice_lines(key):
    lines = {'ha':[6562.50, r'H\alpha', r'\lambda6563'],
             'hb':[4861.36, r'H\beta', r'\lambda4861'],
             'o3':[5007.00, r'O III', r'\lambda5007'],
             'n2':[6584.00, r'N II', r'\lambda6584'],
             'an2':[6548.00, r'N II', r'\lambda6548'],
             'ao3':[4959.00, r'O III', r'\lambda4959'],
             'o2':[3727.00, r'O II', r'\lambdalambda3726,9'],
             's2a':[6716.00, r'S II', r'\lambda6716'],
             's2b':[6731.00, r'S II', r'\lambda6731'],
             'o1':[6300, r'O I', r'\lambda6300']}
    return lines[key]

def calc_dim(X,Y,Z):
    extent = [np.min(X), np.max(X), np.min(Y), np.max(Y)]
    dx = (extent[1] - extent[0]) / float(Z.shape[1])
    dy = (extent[3] - extent[2]) / float(Z.shape[0])
    return extent, dx/dy

def add_dopita(**kwargs):
    dopita.plot_bpt(**kwargs)
    return


"""
Pathfinder/CHIME Telescope Model

A model for both CHIME and Pathfinder telescopes. This attempts to query the
configuration db (:mod:`~ch_analysis.pathfinder.configdb`) for the details of
the feeds and their positions.
"""

import logging

import numpy as np
import h5py
import healpy
from scipy.interpolate import RectBivariateSpline

from caput import config, mpiutil

from cora.util import coord, hputil

from drift.core import telescope
from drift.telescope import cylbeam

from draco.core.containers import ContainerBase, GridBeam, HEALPixBeam

from ch_util import ephemeris, tools
from caput.cache import cached_property

# Get the logger for the module
logger = logging.getLogger(__name__)

class CHIME(telescope.PolarisedTelescope):
    """Model telescope for the CHIME/Pathfinder.

    This class currently uses a simple Gaussian model for the primary beams.

    Attributes
    ----------
    layout : datetime or int
        Specify which layout to use.
    correlator : string
        Restrict to a specific correlator.
    skip_non_chime : boolean
        Ignore non CHIME feeds in the BeamTransfers.
    stack_type : string, optional
        Stacking type.
        `redundant`: feeds of same polarization have same beam class (default).
        `redundant_cyl`: feeds of same polarization and cylinder have same beam
        class.
        `unique`: Each feed has a unique beam class.
    use_pathfinder_freq: boolean
        Use the pathfinder channelization of 1024 frequencies between 400 and
        800 MHz.  Setting this to True also enables the specification of a
        subset of these frequencies through the four attributes below.  Default
        is True.
    channel_bin : int, optional
        Number of channels to bin together. Must exactly divide the total
        number. Binning is performed prior to selection of any subset. Default
        is 1.
    freq_physical : list, optional
        Select subset of frequencies using a list of physical frequencies in
        MHz. Finds the closests pathfinder channel.
    channel_range : list, optional
        Select subset of frequencies using a range of frequency channel indices,
        either [start, stop, step], [start, stop], or [stop] is acceptable.
    channel_index : list, optional
        Select subset of frequencies using a list of frequency channel indices.
    input_sel : list, optional
        Select a reduced set of feeds to use. Useful for generating small
        subsets of the data.
    baseline_masking_type : string, optional
        Select a subset of baselines.  `total_length` selects baselines according to
        their total length. Need to specify `minlength` and `maxlength` properties
        (defined in baseclass).  `individual_length` selects baselines according to
        their seperation in the North-South (specify `minlength_ns` and `maxlength_ns`)
        or the East-West (specify `minlength_ew` and `maxlength_ew`).
    minlength_ns, maxlength_ns : float
        Minimum and maximum North-South baseline lengths to include (in metres)
    minlength_ew, maxlength_ew: float
        Minimum and maximum East-West baseline lengths to include (in metres)
    dec_normalized: float, optional
        Normalize the beam by its magnitude at transit at this declination
        in degrees.
    skip_pol_pair : list
        List of antenna polarisation pairs to skip. Valid entries are "XX", "XY", "YX"
        or "YY". Like the skipped frequencies these pol pairs will have entries
        generated but their beam transfer matrices are implicitly zero and thus not
        calculated.
    """

    # Configure which feeds and layout to use
    layout = config.Property(default=None)
    correlator = config.Property(proptype=str, default=None)
    skip_non_chime = config.Property(proptype=bool, default=False)

    # Redundancy settings
    stack_type = config.enum(
        ["redundant", "redundant_cyl", "unique"], default="redundant"
    )

    # Configure frequency properties
    use_pathfinder_freq = config.Property(proptype=bool, default=True)
    channel_bin = config.Property(proptype=int, default=1)
    freq_physical = config.Property(proptype=list, default=[])
    channel_range = config.Property(proptype=list, default=[])
    channel_index = config.Property(proptype=list, default=[])

    # Input selection
    input_sel = config.Property(proptype=list, default=None)

    # Baseline masking options
    baseline_masking_type = config.enum(
        ["total_length", "individual_length"], default="individual_length"
    )
    minlength_ew = config.Property(proptype=float, default=0.0)
    maxlength_ew = config.Property(proptype=float, default=1.0e7)
    minlength_ns = config.Property(proptype=float, default=0.0)
    maxlength_ns = config.Property(proptype=float, default=1.0e7)

    # Auto-correlations setting (overriding default in baseclass)
    auto_correlations = config.Property(proptype=bool, default=True)

    # Beam normalization
    dec_normalized = config.Property(proptype=float, default=None)
    # Skipping frequency/baseline parameters
    skip_pol_pair = config.list_type(type_=str, maxlength=4, default=[])

    # Fix base properties
    cylinder_width = 20.0
    cylinder_spacing = tools._PF_SPACE

    _exwidth = [0.7]
    _eywidth = _exwidth

    _hxwidth = [1.2]
    _hywidth = _hxwidth

    _pickle_keys = ["_feeds"]

    #
    # === Initialisation routines ===
    #

    def __init__(self, feeds=None):
        import datetime

        self._feeds = feeds

        # Set location properties
        self.latitude = ephemeris.CHIMELATITUDE
        self.longitude = ephemeris.CHIMELONGITUDE
        self.altitude = ephemeris.CHIMEALTITUDE

        # Set the LSD start epoch (i.e. CHIME/Pathfinder first light)
        self.lsd_start_day = datetime.datetime(2013, 11, 15)

        # Set the overall normalization of the beam
        self._set_beam_normalization()

    @classmethod
    def from_layout(cls, layout, correlator=None, skip=False):
        """Create a CHIME/Pathfinder telescope description for the specified layout.

        Parameters
        ----------
        layout : integer or datetime
            Layout id number (corresponding to one in database), or datetime
        correlator : string, optional
            Name of the specific correlator. Needed to return a unique config
            in some cases.
        skip : boolean, optional
            Whether to skip non-CHIME antennas. If False, leave them in but
            set them to infinite noise (unsupported at the moment).

        Returns
        -------
        tel : CHIME
        """

        tel = cls()

        tel.layout = layout
        tel.correlator = correlator
        tel.skip_non_chime = skip
        tel._load_layout()

        return tel

    def _load_layout(self):
        """Load the CHIME/Pathfinder layout from the database.

        Generally this routine shouldn't be called directly. Use
        :method:`CHIME.from_layout` or configure from a YAML file.
        """
        if self.layout is None:
            raise Exception("Layout attributes not set.")

        # Fetch feed layout from database
        feeds = tools.get_correlator_inputs(self.layout, self.correlator)

        if mpiutil.size > 1:
            feeds = mpiutil.world.bcast(feeds, root=0)

        if self.skip_non_chime:
            raise Exception("Not supported.")

        self._feeds = feeds

    def _finalise_config(self):
        # Override base method to implement automatic loading of layout when
        # configuring from YAML.

        if self.layout is not None:
            logger.debug("Loading layout: %s", str(self.layout))
            self._load_layout()

        # Set the overall normalization of the beam
        self._set_beam_normalization()

    #
    # === Redefine properties of the base class ===
    #

    # Tweak the following two properties to change the beam width
    @cached_property
    def fwhm_ex(self):
        """Full width half max of the E-plane antenna beam for X polarization."""

        return np.polyval(np.array(self._exwidth) * 2.0 * np.pi / 3.0, self.frequencies)

    @cached_property
    def fwhm_hx(self):
        """Full width half max of the H-plane antenna beam for X polarization."""

        return np.polyval(np.array(self._hxwidth) * 2.0 * np.pi / 3.0, self.frequencies)

    @cached_property
    def fwhm_ey(self):
        """Full width half max of the E-plane antenna beam for Y polarization."""

        return np.polyval(np.array(self._eywidth) * 2.0 * np.pi / 3.0, self.frequencies)

    @cached_property
    def fwhm_hy(self):
        """Full width half max of the H-plane antenna beam for Y polarization."""

        return np.polyval(np.array(self._hywidth) * 2.0 * np.pi / 3.0, self.frequencies)

    # Set the approximate uv feed sizes
    @property
    def u_width(self):
        return self.cylinder_width

    # v-width property override
    @property
    def v_width(self):
        return 1.0

    # Set non-zero rotation angle for pathfinder and chime
    @property
    def rotation_angle(self):
        if self.correlator == "pathfinder":
            return tools._PF_ROT
        elif self.correlator == "chime":
            return tools._CHIME_ROT
        else:
            return 0.0

    def calculate_frequencies(self):
        """Override default version to give support for specifying by frequency
        channel number.
        """
        if self.use_pathfinder_freq:
            # Use pathfinder channelization of 1024 bins between 400 and 800 MHz.
            basefreq = np.linspace(800.0, 400.0, 1024, endpoint=False)

            # Bin the channels together
            if len(basefreq) % self.channel_bin != 0:
                raise Exception(
                    "Channel binning must exactly divide the total number of channels"
                )

            basefreq = basefreq.reshape(-1, self.channel_bin).mean(axis=-1)

            # If requested, select subset of frequencies.
            if self.freq_physical:
                basefreq = basefreq[
                    [np.argmin(np.abs(basefreq - freq)) for freq in self.freq_physical]
                ]

            elif self.channel_range and (len(self.channel_range) <= 3):
                basefreq = basefreq[slice(*self.channel_range)]

            elif self.channel_index:
                basefreq = basefreq[self.channel_index]

            # Save to object
            self._frequencies = np.unique(basefreq)[::-1]

        else:
            # Otherwise use the standard method
            telescope.TransitTelescope.calculate_frequencies(self)

    @property
    def feeds(self):
        """Return a description of the feeds as a list of :class:`tools.CorrInput` instances."""

        if self.input_sel is None:
            feeds = self._feeds
        else:
            feeds = [self._feeds[fi] for fi in self.input_sel]

        return feeds

    @property
    def input_index(self):
        """An index_map describing the inputs known to the telescope. Useful
        for generating synthetic datasets.
        """
        # Extract lists of channel ID and serial numbers
        channels, feed_sn = list(
            zip(*[(feed.id, feed.input_sn) for feed in self.feeds])
        )

        # Create an input index map and return it.
        from ch_util import andata

        return andata._generate_input_map(feed_sn, channels)

    _pos = None

    @property
    def feedpositions(self):
        """The set of feed positions on *all* cylinders.

        This is constructed for the given layout and includes all rotations of
        the cylinder axis.

        Returns
        -------
        feedpositions : np.ndarray
            The positions in the telescope plane of the receivers. Packed as
            [[u1, v1], [u2, v2], ...].
        """

        if self._pos is None:
            # Fetch cylinder relative positions
            pos = tools.get_feed_positions(self.feeds)

            # The above routine returns NaNs for non CHIME feeds. This is a bit
            # messy, so turn them into zeros.
            self._pos = np.nan_to_num(pos)

        return self._pos

    @property
    def beamclass(self):
        """Beam class definition for the CHIME/Pathfinder.

        When `self.stack_type` is `redundant`, the X-polarisation feeds get
        `beamclass = 0`, and the Y-polarisation gets `beamclass = 1`.
        When `self.stack_type` is `redundant_cyl`, feeds of same polarisation
        and cylinder have same beam class. The beam class is given by
        `beamclass = 2*cyl + pol` where `cyl` is the cylinder number according to
        `ch_util.tools` convention and `pol` is the polarisation (0 for X and 1
        for Y polarisation)
        When `self.stack_type` is `unique`, then the feeds are just given an
        increasing unique class.
        In all cases, any other type of feed gets set to `-1` and should be
        ignored.
        """
        # Make beam class just channel number.

        def _feedclass(f, redundant_cyl=False):
            if tools.is_array(f):
                if tools.is_array_x(f):  # feed is X polarisation
                    pol = 0
                else:  # feed is Y polarisation
                    pol = 1

                if redundant_cyl:
                    return 2 * f.cyl + pol
                else:
                    return pol
            return -1

        if self.stack_type == "redundant":
            return np.array([_feedclass(f) for f in self.feeds])
        elif self.stack_type == "redundant_cyl":
            return np.array([_feedclass(f, redundant_cyl=True) for f in self.feeds])
        else:
            beamclass = [
                fi if tools.is_array(feed) else -1 for fi, feed in enumerate(self.feeds)
            ]
            return np.array(beamclass)

    @property
    def polarisation(self):
        """
        Polarisation map.

        Returns
        -------
        pol : np.ndarray
            One-dimensional array with the polarization for each feed ('X' or 'Y').
        """

        def _pol(f):
            if tools.is_array(f):
                if tools.is_array_x(f):  # feed is X polarisation
                    return "X"
                else:  # feed is Y polarisation
                    return "Y"
            return "N"

        return np.asarray([_pol(f) for f in self.feeds], dtype=np.str)

    #
    # === Setup the primary beams ===
    #

    def beam(self, feed, freq, angpos=None):
        """Primary beam implementation for the CHIME/Pathfinder.

        This only supports normal CHIME cylinder antennas. Asking for the beams
        for other types of inputs will cause an exception to be thrown. The
        beams from this routine are rotated by `self.rotation_angle` to account
        for the CHIME/Pathfinder rotation.

        Parameters
        ----------
        feed : int
            Index for the feed.
        freq : int
            Index for the frequency.
        angpos : np.ndarray[nposition, 2], optional
            Angular position on the sky (in radians).
            If not provided, default to the _angpos
            class attribute.

        Returns
        -------
        beam : np.ndarray[nposition, 2]
            Amplitude vector of beam at each position on the sky.
        """
        # # Fetch beam parameters out of config database.

        feed_obj = self.feeds[feed]

        # Check that feed exists and is a CHIME cylinder antenna
        if feed_obj is None:
            raise ValueError("Craziness. The requested feed doesn't seem to exist.")

        if not tools.is_array(feed_obj):
            raise ValueError("Requested feed is not a CHIME antenna.")

        # If the angular position was not provided, then use the values in the
        # class attribute.
        if angpos is None:
            angpos = self._angpos

        # Get the beam rotation parameters.
        yaw = -self.rotation_angle
        pitch = 0.0
        roll = 0.0

        rot = np.radians([yaw, pitch, roll])

        # We can only support feeds angled parallel or perp to the cylinder
        # axis. Check for these and throw exception for anything else.
        if tools.is_array_y(feed_obj):
            beam = cylbeam.beam_y(
                angpos,
                self.zenith,
                self.cylinder_width / self.wavelengths[freq],
                self.fwhm_ey[freq],
                self.fwhm_hy[freq],
                rot=rot,
            )
        elif tools.is_array_x(feed_obj):
            beam = cylbeam.beam_x(
                angpos,
                self.zenith,
                self.cylinder_width / self.wavelengths[freq],
                self.fwhm_ex[freq],
                self.fwhm_hx[freq],
                rot=rot,
            )
        else:
            raise RuntimeError(
                "Given polarisation (feed.pol=%s) not supported." % feed_obj.pol
            )

        # Normalize the beam
        if self._beam_normalization is not None:
            beam *= self._beam_normalization[freq, feed, np.newaxis, :]

        return beam

    #
    # === Override methods determining the feed pairs we should calculate ===
    #
    # These should probably get ported back into `driftscan` as options.

    def _sort_pairs(self):
        # Reimplemented sort pairs to ensure that returned array is in
        # channel order.

        # Create mask of included pairs, that are not conjugated
        tmask = np.logical_and(self._feedmask, np.logical_not(self._feedconj))
        uniq = telescope._get_indices(self._feedmap, mask=tmask)

        # Get channel id for each feed in the pair, this will be used for the sort
        ci, cj = np.array([(self.feeds[fi].id, self.feeds[fj].id) for fi, fj in uniq]).T

        # # Sort by constructing a numpy array with the keys as fields, and use
        # # np.argsort to get the indices

        # Create array of keys to sort
        dt = np.dtype("i4,i4")
        sort_arr = np.zeros(ci.size, dtype=dt)
        sort_arr["f0"] = ci
        sort_arr["f1"] = cj

        # Get map which sorts
        sort_ind = np.argsort(sort_arr)

        # Invert mapping
        tmp_sort_ind = sort_ind.copy()
        sort_ind[tmp_sort_ind] = np.arange(sort_ind.size)

        # Remap feedmap entries
        fm_copy = self._feedmap.copy()
        wmask = np.where(self._feedmask)
        fm_copy[wmask] = sort_ind[self._feedmap[wmask]]

        self._feedmap = fm_copy

    def _make_ew(self):
        # # Reimplemented to make sure entries we always pick the upper
        # # triangle (and do not reorder to make EW baselines)
        if self.stack_type != "unique":
            super(CHIME, self)._make_ew()

    def _unique_baselines(self):
        # Reimplement unique baselines in order to mask out either according to total
        # baseline length or maximum North-South and East-West baseline seperation.

        from drift.core import telescope

        # Construct array of indices
        fshape = [self.nfeed, self.nfeed]
        f_ind = np.indices(fshape)

        # Construct array of baseline separations
        bl1 = self.feedpositions[f_ind[0]] - self.feedpositions[f_ind[1]]
        bl2 = np.around(bl1[..., 0] + 1.0j * bl1[..., 1], self._bl_tol)

        # Construct array of baseline lengths
        blen = np.sum(bl1 ** 2, axis=-1) ** 0.5

        if self.baseline_masking_type == "total_length":
            # Create mask of included baselines
            mask = np.logical_and(blen >= self.minlength, blen <= self.maxlength)
        else:
            mask_ew = np.logical_and(
                abs(bl1[..., 0]) >= self.minlength_ew,
                abs(bl1[..., 0]) <= self.maxlength_ew,
            )
            mask_ns = np.logical_and(
                abs(bl1[..., 1]) >= self.minlength_ns,
                abs(bl1[..., 1]) <= self.maxlength_ns,
            )
            mask = np.logical_and(mask_ew, mask_ns)

        # Remove the auto correlated baselines between all polarisations
        if not self.auto_correlations:
            mask = np.logical_and(blen > 0.0, mask)

        return telescope._remap_keyarray(bl2, mask), mask

    def _unique_beams(self):
        # Override to mask out any feed where the beamclass is less than zero.
        # This is used to get exclude feeds which are not normal CHIME cylinder
        # feeds

        beam_map, beam_mask = telescope.TransitTelescope._unique_beams(self)

        # Construct a mask including only the feeds where the beam class is
        # greater than zero
        bc_mask = self.beamclass >= 0
        bc_mask = np.logical_and(bc_mask[:, np.newaxis], bc_mask[np.newaxis, :])

        beam_mask = np.logical_and(beam_mask, bc_mask)

        return beam_map, beam_mask

    def _set_beam_normalization(self):
        """Determine the beam normalization for each feed and frequency.

        The beam will be normalized by its value at transit at the declination
        provided in the dec_normalized config parameter.  If this config parameter
        is set to None, then there is no additional normalization applied.
        """

        self._beam_normalization = None

        if self.dec_normalized is not None:

            angpos = np.array(
                [(0.5 * np.pi - np.radians(self.dec_normalized)), 0.0]
            ).reshape(1, -1)

            beam = np.ones((self.nfreq, self.nfeed, 2), dtype=np.float64)

            beam_lookup = {}

            for fe, feed in enumerate(self.feeds):

                if not tools.is_array(feed):
                    continue

                beamclass = self.beamclass[fe]

                if beamclass not in beam_lookup:

                    beam_lookup[beamclass] = np.ones((self.nfreq, 2), dtype=np.float64)
                    for fr in range(self.nfreq):
                        beam_lookup[beamclass][fr] = self.beam(fe, fr, angpos)[0]

                beam[:, fe, :] = beam_lookup[beamclass]

            self._beam_normalization = tools.invert_no_zero(
                np.sqrt(np.sum(beam ** 2, axis=-1))
            )

    def _skip_baseline(self, bl_ind):
        """Override to skip baselines based on which polarisation pair they are."""

        # Pull in baseline skip choice from parent class
        skip_bl = super()._skip_baseline(bl_ind)

        pol_i, pol_j = self.polarisation[self.uniquepairs[bl_ind]]
        pol_pair = pol_i + pol_j

        skip_pol = pol_pair in self.skip_pol_pair

        return skip_bl or skip_pol

def _flat_top_gauss6(x, A, sig, x0):
    """Flat-top gaussian. Power of 6."""
    return A * np.exp(-abs((x - x0) / sig) ** 6)

def _flat_top_gauss3(x, A, sig, x0):
    """Flat-top gaussian. Power of 3."""
    return A * np.exp(-abs((x - x0) / sig) ** 3)

class CHIMEParameterizedBeam(CHIME):
    """CHIME telescope that uses a parameterized fit to the driftscan beam.

    This speeds up evaluation of the beam model.
    """

    SIGMA_EW = [14.87857614, 9.95746878]

    FUNC_NS = [_flat_top_gauss6, _flat_top_gauss3]
    PARAM_NS = np.array(
        [[9.97981768e-01, 1.29544939e00, 0.0], [9.86421047e-01, 8.10213326e-01, 0.0]]
    )

    def _sigma(self, pol_index, freq_index, dec):
        """Width of the power beam in the EW direction."""
        return self.SIGMA_EW[pol_index] / self.frequencies[freq_index] / np.cos(dec)

    def _beam_amplitude(self, pol_index, dec):
        """Amplitude of the power beam at meridian."""
        return self.FUNC_NS[pol_index](
            dec - np.radians(self.latitude), *self.PARAM_NS[pol_index]
        )

    def beam(self, feed, freq, angpos=None):
        """Parameterized fit to driftscan cylinder beam model for CHIME telescope.

        Parameters
        ----------
        feed : int
            Index for the feed.
        freq : int
            Index for the frequency.
        angpos : np.ndarray[nposition, 2], optional
            Angular position on the sky (in radians).
            If not provided, default to the _angpos
            class attribute.

        Returns
        -------
        beam : np.ndarray[nposition, 2]
            Amplitude vector of beam at each position on the sky.
        """

        feed_obj = self.feeds[feed]

        # Check that feed exists and is a CHIME cylinder antenna
        if feed_obj is None:
            raise ValueError("Craziness. The requested feed doesn't seem to exist.")

        if not tools.is_array(feed_obj):
            raise ValueError("Requested feed is not a CHIME antenna.")

        # If the angular position was not provided, then use the values in the
        # class attribute.
        if angpos is None:
            angpos = self._angpos

        dec = 0.5 * np.pi - angpos[:, 0]
        ha = angpos[:, 1]

        # We can only support feeds angled parallel or perp to the cylinder
        # axis. Check for these and throw exception for anything else.
        if tools.is_array_x(feed_obj):
            pol = 0
        elif tools.is_array_y(feed_obj):
            pol = 1
        else:
            raise RuntimeError(
                "Given polarisation (feed.pol=%s) not supported." % feed_obj.pol
            )

        beam = np.zeros((angpos.shape[0], 2), dtype=np.float64)
        beam[:, 0] = np.sqrt(
            self._beam_amplitude(pol, dec)
            * np.exp(-((ha / self._sigma(pol, freq, dec)) ** 2))
        )

        # Normalize the beam
        if self._beam_normalization is not None:
            beam *= self._beam_normalization[freq, feed, np.newaxis, :]

        return beam

class CHIMEFitBeam(CHIME):
    """Driftscan model with revised FWHM for north-south beam.

    Point source beam model was fit to a flat Gaussian at each frequency.
    Best-fit FWHM as a function of frequency was fit with a cubic
    polynomial. This class revises coefficients of FWHM from the
    fit. Detailed comparisons are documented in:
    https://bao.chimenet.ca/doc/documents/1448

    Note that the optimization was carried out in 585-800 MHz.
    This model will produce large extrapolation errors if used below that.
    """

    _eywidth = (
        3.0
        / (2 * np.pi)
        * np.array([1.15310483e-07, -2.30462590e-04, 1.50451290e-01, -3.07440520e01])
    )

    _hxwidth = (
        3.0
        / (2 * np.pi)
        * np.array([2.97495306e-07, -6.00582101e-04, 3.99949759e-01, -8.66733249e01])
    )

class CHIMEExternalBeam(CHIME):
    """Model telescope for the CHIME.

    This class uses an external beam model that is read in from a file.

    Attributes
    ----------
    primary_beam_filename : str
        Path to the file containing the primary beam. Can either be a Healpix beam or a
        GridBeam.
    freq_interp_beam : bool, optional
        Interpolate between neighbouring frequencies if we don't have a beam for every
        frequency channel.
    force_real_beam : bool, optional
        Ensure the output beam is real, regardless of what the datatype of the beam file
        is. This can help save memory if the saved beam is complex but you know the
        imaginary part is zero.
    """

    primary_beam_filename = config.Property(proptype=str)
    freq_interp_beam = config.Property(proptype=bool, default=False)
    force_real_beam = config.Property(proptype=bool, default=False)

    def _finalise_config(self):
        """Get the beam file object."""

        logger.debug("Reading beam model from {}...".format(self.primary_beam_filename))
        self._primary_beam = ContainerBase.from_file(
            self.primary_beam_filename, mode="r", distributed=False, ondisk=True
        )

        self._is_grid_beam = isinstance(self._primary_beam, GridBeam)

        # cache axes
        self._beam_freq = self._primary_beam.freq[:]
        self._beam_nside = None if self._is_grid_beam else self._primary_beam.nside

        # TODO must use bytestring here because conversion doesn't work with ondisk=True
        if self._is_grid_beam:
            self._beam_pol_map = {
                "X": list(self._primary_beam.pol[:]).index(b"XX"),
                "Y": list(self._primary_beam.pol[:]).index(b"YY"),
            }
        else:
            self._beam_pol_map = {
                "X": list(self._primary_beam.pol[:]).index(b"X"),
                "Y": list(self._primary_beam.pol[:]).index(b"Y"),
            }

        if len(self._primary_beam.input) > 1:
            raise ValueError("Per-feed beam model not supported for now.")

        complex_beam = np.issubclass_(
            self._primary_beam.beam.dtype.type, np.complexfloating
        )
        self._output_dtype = (
            np.complex128 if complex_beam and not self.force_real_beam else np.float64
        )

        super()._finalise_config()

    def beam(self, feed, freq_id):
        """Get the beam pattern.

        Parameters
        ----------
        feed : int
            Feed index.
        freq_id : int
            Frequency ID.

        Returns
        -------
        beam : np.ndarray[pixel, pol]
            Return the vector beam response at each point in the Healpix grid. This
            array is of type `np.float64` if the input beam pattern is real, of
            `force_real_beam` is set, otherwise is is on type `np.complex128`.
        """
        feed_obj = self.feeds[feed]
        tel_freq = self.frequencies
        nside = self._nside
        npix = healpy.nside2npix(nside)

        if feed_obj is None:
            raise ValueError("The requested feed doesn't seem to exist.")

        if tools.is_array_x(feed_obj):
            pol_ind = self._beam_pol_map["X"]
        elif tools.is_array_y(feed_obj):
            pol_ind = self._beam_pol_map["Y"]
        else:
            raise ValueError("Polarisation not supported by this feed", feed_obj)

        # find nearest frequency
        freq_sel = _nearest_freq(
            tel_freq, self._beam_freq, freq_id, single=(not self.freq_interp_beam)
        )
        # Raise an error if we can't find any suitable frequency
        if len(freq_sel) == 0:
            raise ValueError(f"No beam model spans frequency {tel_freq[freq_id]}.")

        if self._is_grid_beam:
            # Either we haven't set up interpolation coords yet, or the nside has
            # changed
            if (self._beam_nside is None) or (self._beam_nside != nside):
                self._beam_nside = nside
                self._setup_gridbeam_interpolation()

            # interpolate gridbeam onto HEALPix
            beam_map = self._interpolate_gridbeam(freq_sel, pol_ind)

        else:  # Healpix input beam just need to change to the required resolution
            beam_map = self._primary_beam.beam[freq_sel, pol_ind, 0, :]

            # Check resolution and resample to a better resolution if needed
            if nside != self._beam_nside:

                if nside > self._beam_nside:
                    logger.warning(
                        f"Requested nside={nside} higher than that of "
                        f"beam {self._beam_nside}"
                    )

                logger.debug(
                    "Resampling external beam from nside {:d} to {:d}".format(
                        self._beam_nside,
                        nside,
                    )
                )
                beam_map_new = np.zeros((len(freq_sel), npix), dtype=beam_map.dtype)
                beam_map_new["Et"] = healpy.ud_grade(beam_map["Et"], nside)
                beam_map_new["Ep"] = healpy.ud_grade(beam_map["Ep"], nside)
                beam_map = beam_map_new

        map_out = np.empty((npix, 2), dtype=self._output_dtype)

        # Pull out the real part of the beam if we are forcing a conversion. This should
        # do nothing if the array is already real
        def _conv_real(x):
            if self.force_real_beam:
                x = x.real
            return x

        if len(freq_sel) == 1:
            # exact match
            map_out[:, 0] = _conv_real(beam_map["Et"][0])
            map_out[:, 1] = _conv_real(beam_map["Ep"][0])
        else:
            # interpolate between pair of frequencies
            freq_high = self._beam_freq[freq_sel[1]]
            freq_low = self._beam_freq[freq_sel[0]]
            freq_int = tel_freq[freq_id]

            alpha = (freq_high - freq_int) / (freq_high - freq_low)
            beta = (freq_int - freq_low) / (freq_high - freq_low)

            map_out[:, 0] = _conv_real(
                beam_map["Et"][0] * alpha + beam_map["Et"][1] * beta
            )
            map_out[:, 0] = _conv_real(
                beam_map["Ep"][0] * alpha + beam_map["Ep"][1] * beta
            )

        return map_out

    def _setup_gridbeam_interpolation(self):
        # grid beam coordinates
        self._x_grid = self._primary_beam.phi[:]
        self._y_grid = self._primary_beam.theta[:]

        # celestial coordinates
        angpos = hputil.ang_positions(self._nside)
        x_cel = coord.sph_to_cart(angpos).T

        # rotate to telescope coords
        # first align y with N, then polar axis with NCP
        self._x_tel = cylbeam.rotate_ypr(
            (1.5 * np.pi, np.radians(90.0 - self.latitude), 0), *x_cel
        )

        # mask any pixels outside grid
        x_t, y_t, z_t = self._x_tel
        self._pix_mask = (
            (z_t > 0)
            & (np.abs(x_t) < np.abs(self._x_grid.max()))
            & (np.abs(y_t) < np.abs(self._y_grid.max()))
        )

        # pre-compute polarisation pattern
        # taken from driftscan
        zenith = np.array([np.pi / 2.0 - np.radians(self.latitude), 0.0])
        that, phat = coord.thetaphi_plane_cart(zenith)
        xhat, yhat, zhat = cylbeam.rotate_ypr(
            [-self.rotation_angle, 0.0, 0.0], phat, -that, coord.sph_to_cart(zenith)
        )

        self._pvec_x = cylbeam.polpattern(angpos, xhat)
        self._pvec_y = cylbeam.polpattern(angpos, yhat)

    def _interpolate_gridbeam(self, f_sel, p_ind):
        x, y = self._x_grid, self._y_grid
        x_t, y_t, z_t = self._x_tel
        mask = self._pix_mask

        # interpolation routine requires increasing axes
        reverse_x = (np.diff(self._x_grid) < 0).any()
        if reverse_x:
            x = x[::-1]

        npix = healpy.nside2npix(self._nside)
        beam_out = np.zeros(
            (len(f_sel), npix), dtype=HEALPixBeam._dataset_spec["beam"]["dtype"]
        )
        for i, fi in enumerate(f_sel):
            # For now we just use the magnitude. Assumes input is power beam
            beam = self._primary_beam.beam[fi, p_ind, 0]
            if reverse_x:
                beam = beam[:, ::-1]
            beam_spline = RectBivariateSpline(y, x, np.sqrt(np.abs(beam)))

            # beam amplitude
            amp = np.zeros(npix, dtype=beam.real.dtype)
            amp[mask] = beam_spline(y_t[mask], x_t[mask], grid=False)

            # polarisation projection
            pvec = self._pvec_x if self._beam_pol_map["X"] == p_ind else self._pvec_y

            beam_out[i]["Et"] = amp * pvec[:, 0]
            beam_out[i]["Ep"] = amp * pvec[:, 1]

        return beam_out

def _nearest_freq(tel_freq, map_freq, freq_id, single=False):
    """Find nearest neighbor frequencies. Assumes map frequencies
    are uniformly spaced.

    Parameters
    ----------
    tel_freq : float
        frequencies from telescope object.
    map_freq : float
        frequencies from beam map file.
    freq_id : int
        frequency selection.
    single : bool
        Only return the single nearest neighbour.

    Returns
    -------
    freq_ind : list of neighboring map frequencies matched to tel_freq.

    """

    diff_freq = abs(map_freq - tel_freq[freq_id])
    if single:
        return np.array([np.argmin(diff_freq)])

    map_freq_width = abs(map_freq[1] - map_freq[0])
    match_mask = diff_freq < map_freq_width

    freq_ind = np.nonzero(match_mask)[0]

    return freq_ind