'''
Description: 查看生成的矩形
Author: Lirenjie
Date: 2021-12-21 21:22:08
LastEditors: Lirenjie
LastEditTime: 2021-12-22 15:55:23
'''
import cv2
import numpy as np
from numpy.linalg import norm
import sys
import os
import json

import turtle as t
import time

img = cv2.imread("C:\\Users\\lirenjie\\Desktop\\cAA662F.jpg")
oldimg = cv2.GaussianBlur(img, (3, 3), 0)  # 图片分辨率调整
pic_hight, pic_width = (576, 704)

box = [[114, 295], [115, 259], [208, 262], [207, 298]]

t.penup()
t.goto(box[0])

t.pendown()
t.goto(box[1])
t.goto(box[2])
t.goto(box[3])
t.goto(box[0])

t.down()
# time.sleep(3)

heigth_point = right_point = [0, 0]
left_point = low_point = [1000, 1000]
for point in box:
    if left_point[0] > point[0]:
        left_point = point
    if low_point[1] > point[1]:
        low_point = point
    if heigth_point[1] < point[1]:
        heigth_point = point
    if right_point[0] < point[0]:
        right_point = point

# 比较得到长边的斜率
length1 = (left_point[0] - low_point[0]) ** 2 + (left_point[1] - low_point[1]) ** 2
length2 = (right_point[0] - low_point[0]) ** 2 + (right_point[1] - low_point[1]) ** 2
if length1 > length2:
    k = (left_point[1] - low_point[1]) / (left_point[0] - low_point[0])
else:
    k = (right_point[1] - low_point[1]) / (right_point[0] - low_point[0])
print('k', k)

if k < 0:  # 正角度
    new_right_point = [right_point[0], heigth_point[1]]
    pts2 = np.float32([left_point, heigth_point, new_right_point])  # 字符只是高度需要改变
    pts1 = np.float32([left_point, heigth_point, right_point])
    M = cv2.getAffineTransform(pts1, pts2)  # 仿射变换
    dst = cv2.warpAffine(oldimg, M, (pic_width, pic_hight))
    cv2.imshow("card", dst)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
    # dst = cv2.warpAffine(oldimg, M, (pic_width, pic_hight))
    # cv2.imshow("card", dst)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()

    # point_limit(new_right_point)
    # point_limit(heigth_point)
    # point_limit(left_point)
    # card_img = dst[
    #     int(left_point[1]) : int(heigth_point[1]),
    #     int(left_point[0]) : int(new_right_point[0]),
    # ]

    # print(card_img)

    # card_imgs.append(card_img)
    # cv2.imshow("card", card_img)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()
elif k > 0:  # 负角度
    new_left_point = [left_point[0], heigth_point[1]]
    pts2 = np.float32([new_left_point, heigth_point, right_point])  # 字符只是高度需要改变
    pts1 = np.float32([left_point, heigth_point, right_point])
    M = cv2.getAffineTransform(pts1, pts2)
    dst = cv2.warpAffine(oldimg, M, (pic_width, pic_hight))
    cv2.imshow("card", dst)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
    # dst = cv2.warpAffine(oldimg, M, (pic_width, pic_hight))
    # cv2.imshow("card", dst)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()

    # point_limit(right_point)
    # point_limit(heigth_point)
    # point_limit(new_left_point)
    # card_img = dst[
    #     int(right_point[1]) : int(heigth_point[1]),
    #     int(new_left_point[0]) : int(right_point[0]),
    # ]

    # print(card_img)

    # card_imgs.append(card_img)
    # cv2.imshow("card", card_img)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()

print(M)


from typing import Callable, NamedTuple, Tuple

import jax
import jax.numpy as jnp

from blackjax.types import PyTree

class SMCInfo(NamedTuple):
    """Additional information on the tempered SMC step.

    weights: jnp.ndarray
        The weights after the MCMC pass.
    proposals: PyTree
        The particles that were proposed by the MCMC pass.
    ancestors: jnp.ndarray
        The index of the particles proposed by the MCMC pass that were selected
        by the resampling step.
    log_likelihood_increment: float
        The log-likelihood increment due to the current step of the SMC algorithm.
    """

    weights: jnp.ndarray
    proposals: PyTree
    ancestors: jnp.ndarray
    log_likelihood_increment: float

def kernel(
    mcmc_kernel_factory: Callable,
    mcmc_state_generator: Callable,
    resampling_fn: Callable,
    num_mcmc_iterations: int,
):
    """Build a generic SMC kernel.

    In Feynman-Kac equivalent terms, the algo goes roughly as follows:

    ```
        M_t = mcmc_kernel_factory(potential_fn)
        for i in range(num_mcmc_iterations):
            x_t^i = M_t(..., x_t^i)
        G_t = log_weights_fn
        log_weights = G_t(x_t)
        idx = resample(log_weights)
        x_t = x_t[idx]
    ```

    Parameters
    ----------
    mcmc_kernel_factory: Callable
        A function of the Markov potential that returns a mcmc_kernel.
    mcmc_state_generator: Callable
        A function that creates a new mcmc state from a position and a potential.
    resampling_fn: Callable
        A function that resamples the particles generated by the MCMC kernel,
        based of previously computed weights.
    num_mcmc_iterations: int
        Number of iterations of the MCMC kernel

    Returns
    -------
    A kernel that takes a PRNGKey, a set of particles, the log-likehood of the
    distribution and the Feynman-Kac potential at time `t`. The kernel returns
    a new set of particles.

    """

    def one_step(
        rng_key: jnp.ndarray,
        particles: PyTree,
        logprob_fn: Callable,
        log_weight_fn: Callable,
    ) -> Tuple[PyTree, SMCInfo]:
        """

        Parameters
        ----------
        rng_key: DeviceArray[int],
            JAX PRNGKey for randomness.
        particles: PyTree
            Current particles sample of the SMC algorithm.
        logprob_fn: Callable
            Log probability function we wish to sample from.
        log_weight_fn: Callable
            A function that represents the Feynman-Kac log potential at time t.

        Returns
        -------
        particles: PyTree,
            The updated set of particles.
        info: SMCInfo,
            Additional information on the SMC step
        """

        num_particles = jax.tree_flatten(particles)[0][0].shape[0]
        scan_key, resampling_key = jax.random.split(rng_key, 2)

        # First advance the particles using the MCMC kernel
        mcmc_kernel = mcmc_kernel_factory(logprob_fn)

        def mcmc_body_fn(curr_particles, curr_key):
            keys = jax.random.split(curr_key, num_particles)
            new_particles, _ = jax.vmap(mcmc_kernel, in_axes=(0, 0))(
                keys, curr_particles
            )
            return new_particles, None

        mcmc_state = jax.vmap(mcmc_state_generator, in_axes=(0, None))(
            particles, logprob_fn
        )
        keys = jax.random.split(scan_key, num_mcmc_iterations)
        proposed_states, _ = jax.lax.scan(mcmc_body_fn, mcmc_state, keys)
        proposed_particles = proposed_states.position

        # Resample the particles depending on their respective weights
        log_weights = jax.vmap(log_weight_fn, in_axes=(0,))(proposed_particles)
        weights, log_likelihood_increment = _normalize(log_weights)
        resampling_index = resampling_fn(weights, resampling_key)
        particles = jax.tree_map(lambda x: x[resampling_index], proposed_particles)

        info = SMCInfo(
            weights, proposed_particles, resampling_index, log_likelihood_increment
        )
        return particles, info

    return one_step

def _normalize(log_weights):
    """Normalize log-weights into weights and return resulting weights and log-likelihood increment."""
    n = log_weights.shape[0]
    max_logw = jnp.max(log_weights)
    w = jnp.exp(log_weights - max_logw)
    w_mean = w.mean()

    log_likelihood_increment = jnp.log(w_mean) + max_logw

    w = w / (n * w_mean)
    return w, log_likelihood_increment


"""
Notes
-----

Important attributes of continuous (order > 0) :class:`Field` and
:class:`SurfaceField` instances:

- `vertex_remap` : `econn[:, :n_vertex] = vertex_remap[conn]`
- `vertex_remap_i` : `conn = vertex_remap_i[econn[:, :n_vertex]]`

where `conn` is the mesh vertex connectivity, `econn` is the
region-local field connectivity.
"""
import numpy as nm

from sfepy.base.base import output, get_default, assert_
from sfepy.base.base import Struct
import fea
from sfepy.discrete.common.fields import parse_shape, Field
from sfepy.discrete.fem.mesh import Mesh
from sfepy.discrete.fem.meshio import convert_complex_output
from sfepy.discrete.fem.utils import (extend_cell_data, prepare_remap,
                                      invert_remap, get_min_value)
from sfepy.discrete.fem.fe_surface import FESurface
from sfepy.discrete.integrals import Integral
from sfepy.discrete.fem.linearizer import (get_eval_dofs, get_eval_coors,
                                           create_output)

def get_eval_expression(expression, ig,
                        fields, materials, variables,
                        functions=None, mode='eval', term_mode=None,
                        extra_args=None, verbose=True, kwargs=None):
    """
    Get the function for evaluating an expression given a list of elements,
    and reference element coordinates.
    """
    from sfepy.discrete.evaluate import eval_in_els_and_qp

    def _eval(iels, coors):
        val = eval_in_els_and_qp(expression, ig, iels, coors,
                                 fields, materials, variables,
                                 functions=functions, mode=mode,
                                 term_mode=term_mode,
                                 extra_args=extra_args, verbose=verbose,
                                 kwargs=kwargs)
        return val[..., 0]

    return _eval

def create_expression_output(expression, name, primary_field_name,
                             fields, materials, variables,
                             functions=None, mode='eval', term_mode=None,
                             extra_args=None, verbose=True, kwargs=None,
                             min_level=0, max_level=1, eps=1e-4):
    """
    Create output mesh and data for the expression using the adaptive
    linearizer.

    Parameters
    ----------
    expression : str
        The expression to evaluate.
    name : str
        The name of the data.
    primary_field_name : str
        The name of field that defines the element groups and polynomial
        spaces.
    fields : dict
        The dictionary of fields used in `variables`.
    materials : Materials instance
        The materials used in the expression.
    variables : Variables instance
        The variables used in the expression.
    functions : Functions instance, optional
        The user functions for materials etc.
    mode : one of 'eval', 'el_avg', 'qp'
        The evaluation mode - 'qp' requests the values in quadrature points,
        'el_avg' element averages and 'eval' means integration over
        each term region.
    term_mode : str
        The term call mode - some terms support different call modes
        and depending on the call mode different values are
        returned.
    extra_args : dict, optional
        Extra arguments to be passed to terms in the expression.
    verbose : bool
        If False, reduce verbosity.
    kwargs : dict, optional
        The variables (dictionary of (variable name) : (Variable
        instance)) to be used in the expression.
    min_level : int
        The minimum required level of mesh refinement.
    max_level : int
        The maximum level of mesh refinement.
    eps : float
        The relative tolerance parameter of mesh adaptivity.

    Returns
    -------
    out : dict
        The output dictionary.
    """
    field = fields[primary_field_name]
    vertex_coors = field.coors[:field.n_vertex_dof, :]

    coors = []
    vdofs = []
    conns = []
    mat_ids = []
    levels = []
    offset = 0
    for ig, ap in field.aps.iteritems():
        ps = ap.interp.poly_spaces['v']
        gps = ap.interp.gel.interp.poly_spaces['v']
        group = field.domain.groups[ig]
        vertex_conn = ap.econn[:, :group.shape.n_ep]

        eval_dofs = get_eval_expression(expression, ig,
                                        fields, materials, variables,
                                        functions=functions,
                                        mode=mode, extra_args=extra_args,
                                        verbose=verbose, kwargs=kwargs)
        eval_coors = get_eval_coors(vertex_coors, vertex_conn, gps)

        (level, _coors, conn,
         _vdofs, _mat_ids) = create_output(eval_dofs, eval_coors,
                                           group.shape.n_el, ps,
                                           min_level=min_level,
                                           max_level=max_level, eps=eps)

        _mat_ids[:] = field.domain.mesh.mat_ids[ig][0]

        coors.append(_coors)
        vdofs.append(_vdofs)
        conns.append(conn + offset)
        mat_ids.append(_mat_ids)
        levels.append(level)

        offset += _coors.shape[0]

    coors = nm.concatenate(coors, axis=0)
    vdofs = nm.concatenate(vdofs, axis=0)
    mesh = Mesh.from_data('linearized_mesh', coors, None, conns, mat_ids,
                          field.domain.mesh.descs)

    out = {}
    out[name] = Struct(name='output_data', mode='vertex',
                       data=vdofs, var_name=name, dofs=None,
                       mesh=mesh, levels=levels)

    out = convert_complex_output(out)

    return out

class FEField(Field):
    """
    Base class for finite element fields.

    Notes
    -----
    - Region can span over several groups -> different Aproximation
      instances
    - interps and hence node_descs are per region (must have single
      geometry!)
    - no two interps can be in a same group -> no two aps (with
      different regions) can be in a same group -> aps can be uniquely
      indexed with ig

    Field shape information:

    - ``shape`` - the shape of the base functions in a point
    - ``n_components`` - the number of DOFs per FE node
    - ``val_shape`` - the shape of field value (the product of DOFs and
      base functions) in a point
    """

    def __init__(self, name, dtype, shape, region, approx_order=1):
        """
        Create a finite element field.

        Parameters
        ----------
        name : str
            The field name.
        dtype : numpy.dtype
            The field data type: float64 or complex128.
        shape : int/tuple/str
            The field shape: 1 or (1,) or 'scalar', space dimension (2, or (2,)
            or 3 or (3,)) or 'vector', or a tuple. The field shape determines
            the shape of the FE base functions and is related to the number of
            components of variables and to the DOF per node count, depending
            on the field kind.
        region : Region
            The region where the field is defined.
        approx_order : int or tuple
            The FE approximation order. The tuple form is (order, has_bubble),
            e.g. (1, True) means order 1 with a bubble function.

        Notes
        -----
        Assumes one cell type for the whole region!
        """
        shape = parse_shape(shape, region.domain.shape.dim)
        if not self._check_region(region):
            raise ValueError('unsuitable region for field %s! (%s)' %
                             (name, region.name))

        Struct.__init__(self, name=name, dtype=dtype, shape=shape,
                        region=region)
        self.domain = self.region.domain
        self.igs = self.region.igs

        self._set_approx_order(approx_order)
        self._setup_geometry()
        self._setup_kind()
        self._setup_shape()

        self._create_interpolant()
        self._setup_approximations()
        self._setup_global_base()
        self.setup_coors()
        self.clear_mappings(clear_all=True)

    def _set_approx_order(self, approx_order):
        """
        Set a uniform approximation order.
        """
        if isinstance(approx_order, tuple):
            self.approx_order = approx_order[0]
            self.force_bubble = approx_order[1]

        else:
            self.approx_order = approx_order
            self.force_bubble = False

    def _create_interpolant(self):
        name = '%s_%s_%s_%d%s' % (self.gel.name, self.space,
                                  self.poly_space_base, self.approx_order,
                                  'B' * self.force_bubble)
        self.interp = fea.Interpolant(name, self.gel, self.space,
                                      self.poly_space_base, self.approx_order,
                                      self.force_bubble)

    def _setup_approximations(self):
        self.aps = {}
        self.aps_by_name = {}
        for ig in self.igs:
            name = self.interp.name + '_%s_ig%d' % (self.region.name, ig)
            ap = fea.Approximation(name, self.interp, self.region, ig)
            self.aps[ig] = ap
            self.aps_by_name[ap.name] = ap

    def get_true_order(self):
        """
        Get the true approximation order depending on the reference
        element geometry.

        For example, for P1 (linear) approximation the true order is 1,
        while for Q1 (bilinear) approximation in 2D the true order is 2.
        """
        gel = self.gel
        if (gel.dim + 1) == gel.n_vertex:
            order = self.approx_order

        else:
            order = gel.dim * self.approx_order

        if self.force_bubble:
            bubble_order = gel.dim + 1
            order = max(order, bubble_order)

        return order

    def is_higher_order(self):
        """
        Return True, if the field's approximation order is greater than one.
        """
        return self.force_bubble or (self.approx_order > 1)

    def _setup_global_base(self):
        """
        Setup global DOF/base functions, their indices and connectivity of the
        field. Called methods implemented in subclasses.
        """
        self._setup_facet_orientations()

        self._init_econn()

        self.n_vertex_dof, self.vertex_remap = self._setup_vertex_dofs()
        self.vertex_remap_i = invert_remap(self.vertex_remap)

        aux = self._setup_edge_dofs()
        self.n_edge_dof, self.edge_dofs, self.edge_remap = aux

        aux = self._setup_face_dofs()
        self.n_face_dof, self.face_dofs, self.face_remap = aux

        aux = self._setup_bubble_dofs()
        self.n_bubble_dof, self.bubble_dofs, self.bubble_remaps = aux

        self.n_nod = self.n_vertex_dof + self.n_edge_dof \
                     + self.n_face_dof + self.n_bubble_dof

        self._setup_esurface()

    def _setup_esurface(self):
        """
        Setup extended surface entities (edges in 2D, faces in 3D),
        i.e. indices of surface entities into the extended connectivity.
        """
        node_desc = self.node_desc

        for ig, ap in self.aps.iteritems():
            gel = ap.interp.gel
            ap.efaces = gel.get_surface_entities().copy()

            nd = node_desc.edge
            if nd is not None:
                efs = []
                for eof in gel.get_edges_per_face():
                    efs.append(nm.concatenate([nd[ie] for ie in eof]))
                efs = nm.array(efs).squeeze()

                if efs.ndim < 2:
                    efs = efs[:,nm.newaxis]
                ap.efaces = nm.hstack((ap.efaces, efs))

            efs = node_desc.face
            if efs is not None:
                efs = nm.array(efs).squeeze()

                if efs.ndim < 2:
                    efs = efs[:,nm.newaxis]
                ap.efaces = nm.hstack((ap.efaces, efs))

    def setup_coors(self, coors=None):
        """
        Setup coordinates of field nodes.
        """
        mesh = self.domain.mesh
        self.coors = nm.empty((self.n_nod, mesh.dim), nm.float64)

        if coors is None:
            coors = mesh.coors

        # Mesh vertex nodes.
        if self.n_vertex_dof:
            indx = self.vertex_remap_i
            self.coors[:self.n_vertex_dof] = nm.take(coors,
                                                     indx.astype(nm.int32),
                                                     axis=0)

        for ig, ap in self.aps.iteritems():
            ap.eval_extra_coor(self.coors, coors)

    def get_vertices(self):
        """
        Return indices of vertices belonging to the field region.
        """
        return self.vertex_remap_i

    def _get_facet_dofs(self, get_facets, remap, dofs, ig):
        gfacets = get_facets(ig)
        facets = remap[gfacets]

        return dofs[facets[facets >= 0]].ravel()

    def get_data_shape(self, ig, integral,
                       integration='volume', region_name=None):
        """
        Get element data dimensions.

        Parameters
        ----------
        ig : int
            The element group index.
        integral : Integral instance
            The integral describing used numerical quadrature.
        integration : 'volume', 'plate', 'surface', 'surface_extra' or 'point'
            The term integration type.
        region_name : str
            The name of surface region, required when `shape_kind` is
            'surface'.

        Returns
        -------
        data_shape : 4 ints
            The `(n_el, n_qp, dim, n_en)` for volume shape kind,
            `(n_fa, n_qp, dim, n_fn)` for surface shape kind and
            `(n_nod, 0, 0, 1)` for point shape kind.

        Notes
        -----
        - `n_el`, `n_fa` = number of elements/facets
        - `n_qp` = number of quadrature points per element/facet
        - `dim` = spatial dimension
        - `n_en`, `n_fn` = number of element/facet nodes
        - `n_nod` = number of element nodes
        """
        ap = self.aps[ig]

        region = self.domain.regions[region_name]
        shape = region.shape[ig]
        dim = region.dim

        if integration in ('surface', 'surface_extra'):
            sd = ap.surface_data[region_name]

            # This works also for surface fields.
            key = sd.face_type
            weights = ap.get_qp(key, integral).weights
            n_qp = weights.shape[0]

            if integration == 'surface':
                data_shape = (sd.n_fa, n_qp, dim, ap.n_ep[key])

            else:
                data_shape = (sd.n_fa, n_qp, dim, ap.n_ep['v'])

        elif integration in ('volume', 'plate'):
            _, weights = integral.get_qp(self.gel.name)
            n_qp = weights.shape[0]

            data_shape = (shape.n_cell, n_qp, dim, ap.n_ep['v'])

        elif integration == 'point':
            dofs = self.get_dofs_in_region(region, merge=True)
            data_shape = (dofs.shape[0], 0, 0, 1)

        else:
            raise NotImplementedError('unsupported integration! (%s)'
                                      % integration)

        return data_shape

    def get_dofs_in_region_group(self, region, ig, merge=True):
        """
        Return indices of DOFs that belong to the given region and group.
        """
        node_desc = self.node_desc

        dofs = []

        vdofs = nm.empty((0,), dtype=nm.int32)
        if node_desc.vertex is not None:
            ii = region.get_vertices(ig)
            vdofs = self.vertex_remap[ii]
            vdofs = vdofs[vdofs >= 0]
        dofs.append(vdofs)

        edofs = nm.empty((0,), dtype=nm.int32)
        if node_desc.edge is not None:
            edofs = self._get_facet_dofs(region.get_edges,
                                         self.edge_remap,
                                         self.edge_dofs, ig)
        dofs.append(edofs)

        fdofs = nm.empty((0,), dtype=nm.int32)
        if node_desc.face is not None:
            fdofs = self._get_facet_dofs(region.get_faces,
                                         self.face_remap,
                                         self.face_dofs, ig)
        dofs.append(fdofs)

        bdofs = nm.empty((0,), dtype=nm.int32)
        if (node_desc.bubble is not None) and region.has_cells():
            ii = region.get_cells(ig)
            group_els = self.bubble_remaps[ig][ii]
            bdofs = self.bubble_dofs[ig][group_els[group_els >= 0]].ravel()
        dofs.append(bdofs)

        if merge:
            dofs = nm.concatenate(dofs)

        return dofs

    def extend_dofs(self, dofs, fill_value=None):
        """
        Extend DOFs to the whole domain using the `fill_value`, or the
        smallest value in `dofs` if `fill_value` is None.
        """
        if fill_value is None:
            if nm.isrealobj(dofs):
                fill_value = get_min_value(dofs)

            else:
                # Complex values - treat real and imaginary parts separately.
                fill_value = get_min_value(dofs.real)
                fill_value += 1j * get_min_value(dofs.imag)

        if self.approx_order != 0:
            indx = self.get_vertices()

            n_nod = self.domain.shape.n_nod
            new_dofs = nm.empty((n_nod, dofs.shape[1]), dtype=self.dtype)
            new_dofs.fill(fill_value)
            new_dofs[indx] = dofs[:indx.size]

        else:
            new_dofs = extend_cell_data(dofs, self.domain, self.region,
                                        val=fill_value)

        return new_dofs

    def remove_extra_dofs(self, dofs):
        """
        Remove DOFs defined in higher order nodes (order > 1).
        """
        if self.approx_order != 0:
            new_dofs = dofs[:self.n_vertex_dof]

        else:
            new_dofs = dofs

        return new_dofs

    def linearize(self, dofs, min_level=0, max_level=1, eps=1e-4):
        """
        Linearize the solution for post-processing.

        Parameters
        ----------
        dofs : array, shape (n_nod, n_component)
            The array of DOFs reshaped so that each column corresponds
            to one component.
        min_level : int
            The minimum required level of mesh refinement.
        max_level : int
            The maximum level of mesh refinement.
        eps : float
            The relative tolerance parameter of mesh adaptivity.

        Returns
        -------
        mesh : Mesh instance
            The adapted, nonconforming, mesh.
        vdofs : array
            The DOFs defined in vertices of `mesh`.
        levels : array of ints
            The refinement level used for each element group.
        """
        assert_(dofs.ndim == 2)

        n_nod, dpn = dofs.shape

        assert_(n_nod == self.n_nod)
        assert_(dpn == self.shape[0])

        vertex_coors = self.coors[:self.n_vertex_dof, :]

        coors = []
        vdofs = []
        conns = []
        mat_ids = []
        levels = []
        offset = 0
        for ig, ap in self.aps.iteritems():
            ps = ap.interp.poly_spaces['v']
            gps = ap.interp.gel.interp.poly_spaces['v']
            group = self.domain.groups[ig]
            vertex_conn = ap.econn[:, :group.shape.n_ep]

            eval_dofs = get_eval_dofs(dofs, ap.econn, ps, ori=ap.ori)
            eval_coors = get_eval_coors(vertex_coors, vertex_conn, gps)

            (level, _coors, conn,
             _vdofs, _mat_ids) = create_output(eval_dofs, eval_coors,
                                               group.shape.n_el, ps,
                                               min_level=min_level,
                                               max_level=max_level, eps=eps)

            _mat_ids[:] = self.domain.mesh.mat_ids[ig][0]

            coors.append(_coors)
            vdofs.append(_vdofs)
            conns.append(conn + offset)
            mat_ids.append(_mat_ids)
            levels.append(level)

            offset += _coors.shape[0]

        coors = nm.concatenate(coors, axis=0)
        vdofs = nm.concatenate(vdofs, axis=0)
        mesh = Mesh.from_data('linearized_mesh', coors, None, conns, mat_ids,
                              self.domain.mesh.descs)

        return mesh, vdofs, levels

    def get_output_approx_order(self):
        """
        Get the approximation order used in the output file.
        """
        return min(self.approx_order, 1)

    def create_output(self, dofs, var_name, dof_names=None,
                      key=None, extend=True, fill_value=None,
                      linearization=None):
        """
        Convert the DOFs corresponding to the field to a dictionary of
        output data usable by Mesh.write().

        Parameters
        ----------
        dofs : array, shape (n_nod, n_component)
            The array of DOFs reshaped so that each column corresponds
            to one component.
        var_name : str
            The variable name corresponding to `dofs`.
        dof_names : tuple of str
            The names of DOF components.
        key : str, optional
            The key to be used in the output dictionary instead of the
            variable name.
        extend : bool
            Extend the DOF values to cover the whole domain.
        fill_value : float or complex
           The value used to fill the missing DOF values if `extend` is True.
        linearization : Struct or None
            The linearization configuration for higher order approximations.

        Returns
        -------
        out : dict
            The output dictionary.
        """
        linearization = get_default(linearization, Struct(kind='strip'))

        out = {}
        if linearization.kind is None:
            out[key] = Struct(name='output_data', mode='full',
                              data=dofs, var_name=var_name,
                              dofs=dof_names, field_name=self.name)

        elif ((not self.is_higher_order())
            or (linearization.kind == 'strip')):
            if extend:
                ext = self.extend_dofs(dofs, fill_value)

            else:
                ext = self.remove_extra_dofs(dofs)

            if ext is not None:
                approx_order = self.get_output_approx_order()

                if approx_order != 0:
                    # Has vertex data.
                    out[key] = Struct(name='output_data', mode='vertex',
                                      data=ext, var_name=var_name,
                                      dofs=dof_names)

                else:
                    ext.shape = (ext.shape[0], 1, ext.shape[1], 1)
                    out[key] = Struct(name='output_data', mode='cell',
                                      data=ext, var_name=var_name,
                                      dofs=dof_names)

        else:
            mesh, vdofs, levels = self.linearize(dofs,
                                                 linearization.min_level,
                                                 linearization.max_level,
                                                 linearization.eps)
            out[key] = Struct(name='output_data', mode='vertex',
                              data=vdofs, var_name=var_name, dofs=dof_names,
                              mesh=mesh, levels=levels)

        out = convert_complex_output(out)

        return out

    def create_mesh(self, extra_nodes=True):
        """
        Create a mesh from the field region, optionally including the field
        extra nodes.
        """
        mesh = self.domain.mesh

        if self.approx_order != 0:
            conns, mat_ids, descs = [], [], []
            for ig, ap in self.aps.iteritems():
                group = self.domain.groups[ig]
                if extra_nodes:
                    conn = ap.econn
                else:
                    offset = group.shape.n_ep
                    conn = ap.econn[:,:offset]
                conns.append(conn)
                mat_ids.append(mesh.mat_ids[ig])
                descs.append(mesh.descs[ig])

            if extra_nodes:
                coors = self.coors

            else:
                coors = self.coors[:self.n_vertex_dof]

            mesh = Mesh.from_data(self.name, coors, None, conns,
                                  mat_ids, descs)

        return mesh

    def interp_to_qp(self, dofs):
        """
        Interpolate DOFs into quadrature points.

        The quadrature order is given by the field approximation order.

        Parameters
        ----------
        dofs : array
            The array of DOF values of shape `(n_nod, n_component)`.

        Returns
        -------
        data_qp : array
            The values interpolated into the quadrature points.
        integral : Integral
            The corresponding integral defining the quadrature points.
        """
        integral = Integral('i', order=self.approx_order)

        data_qp = []
        for ig, ap in self.aps.iteritems():
            bf = ap.get_base('v', False, integral)
            bf = bf[:,0,:].copy()

            vals = nm.dot(bf, dofs[ap.econn])
            vals = nm.swapaxes(vals, 0, 1)
            vals.shape = vals.shape + (1,)

            data_qp.append(vals)

        data_qp = nm.concatenate(data_qp, axis=0)

        return data_qp, integral

    def get_coor(self, nods=None):
        """
        Get coordinates of the field nodes.

        Parameters
        ----------
        nods : array, optional
           The indices of the required nodes. If not given, the
           coordinates of all the nodes are returned.
        """
        if nods is None:
            return self.coors
        else:
            return self.coors[nods]

    def create_mapping(self, ig, region, integral, integration):
        """
        Create a new reference mapping.
        """
        ap = self.aps[ig]

        out = ap.describe_geometry(self, integration, region, integral,
                                   return_mapping=True)
        return out

class VolumeField(FEField):
    """
    Finite element field base class over volume elements (element dimension
    equals space dimension).
    """

    def _check_region(self, region):
        """
        Check whether the `region` can be used for the
        field. Non-surface fields require the region to span whole
        element groups.

        Returns
        -------
        ok : bool
            True if the region is usable for the field.
        """
        ok = True
        domain = region.domain
        for ig in region.igs:
            if domain.groups[ig].gel.dim != domain.shape.tdim:
                output('cells with a bad topological dimension! (%d == %d)'
                       % (domain.groups[ig].gel.dim, domain.shape.tdim))
                ok = False
                break
            shape = domain.groups[ig].shape
            if region.shape[ig].n_vertex < shape.n_vertex:
                output('region does not span a whole element group!')
                ok = False
                break

        return ok

    def _setup_geometry(self):
        """
        Setup the field region geometry.
        """
        ig = self.region.domain.cmesh.cell_groups[self.region.cells[0]]
        self.gel = self.domain.groups[ig].gel

        self.is_surface = False

    def _create_interpolant(self):
        name = '%s_%s_%s_%d%s' % (self.gel.name, self.space,
                                  self.poly_space_base, self.approx_order,
                                  'B' * self.force_bubble)
        self.interp = fea.Interpolant(name, self.gel, self.space,
                                      self.poly_space_base, self.approx_order,
                                      self.force_bubble)

    def _setup_approximations(self):
        self.aps = {}
        self.aps_by_name = {}
        for ig in self.igs:
            name = self.interp.name + '_%s_ig%d' % (self.region.name, ig)
            ap = fea.Approximation(name, self.interp, self.region, ig)
            self.aps[ig] = ap
            self.aps_by_name[ap.name] = ap

    def _init_econn(self):
        """
        Initialize the extended DOF connectivity.
        """
        for ig, ap in self.aps.iteritems():
            n_ep = ap.n_ep['v']
            n_cell = self.region.get_n_cells(ig)
            ap.econn = nm.zeros((n_cell, n_ep), nm.int32)

    def _setup_vertex_dofs(self):
        """
        Setup vertex DOF connectivity.
        """
        if self.node_desc.vertex is None:
            return 0, None

        region = self.region

        vertices = region.get_vertices_of_cells()
        remap = prepare_remap(vertices, region.n_v_max)
        n_dof = vertices.shape[0]

        ##
        # Remap vertex node connectivity to field-local numbering.
        for ig, ap in self.aps.iteritems():
            group = self.domain.groups[ig]
            offset = group.shape.n_ep
            cells = region.get_cells(ig)
            ap.econn[:,:offset] = nm.take(remap,
                                          nm.take(group.conn,
                                                  cells.astype(nm.int32),
                                                  axis=0))

        return n_dof, remap

    def setup_extra_data(self, geometry, info, is_trace):
        dct = info.dc_type.type

        if geometry != None:
            geometry_flag = 'surface' in geometry
        else:
            geometry_flag = False

        if (dct == 'surface') or (geometry_flag):
            reg = info.get_region()
            self.domain.create_surface_group(reg)
            self._setup_surface_data(reg, is_trace)

        elif dct == 'edge':
            raise NotImplementedError('dof connectivity type %s' % dct)

        elif dct == 'point':
            self._setup_point_data(self, info.region)

        elif dct not in ('volume', 'scalar', 'plate'):
            raise ValueError('unknown dof connectivity type! (%s)' % dct)

    def _setup_surface_data(self, region, is_trace=False):
        for ig, ap in self.aps.iteritems():
            if ig not in region.igs: continue
            if region.name not in ap.surface_data:
                ap.setup_surface_data(region)

        for ig, ap in self.aps.iteritems():
            if region.name in ap.surface_data and is_trace:
                sd = ap.surface_data[region.name]
                sd.setup_mirror_connectivity(region)

    def _setup_point_data(self, field, region):
        # Point data only in the first group to avoid multiple
        # assembling of nodes on group boundaries.
        ap = self.aps[self.igs[0]]
        if region.name not in ap.point_data:
            ap.setup_point_data(field, region)

    def get_econn(self, conn_type, region, ig, is_trace=False,
                  integration=None):
        """
        Get extended connectivity of the given type in the given region.
        """
        ct = conn_type.type if isinstance(conn_type, Struct) else conn_type

        if ((ig not in self.igs) or (ig not in region.igs)
            or (ct == 'point' and (ig > self.igs[0]))):
            # Point data only in the first group to avoid multiple
            # assembling of nodes on group boundaries.
            return None

        ap = self.aps[ig]

        if ct in ('volume', 'plate'):
            if region.name == self.region.name:
                conn = ap.econn

            else:
                aux = integration in ('volume', 'plate')
                cells = region.get_cells(ig, true_cells_only=aux)
                conn = nm.take(ap.econn, cells.astype(nm.int32), axis=0)

        elif ct == 'surface':
            sd = ap.surface_data[region.name]
            conn = sd.get_connectivity(is_trace=is_trace)

        elif ct == 'edge':
            raise NotImplementedError('connectivity type %s' % ct)

        elif ct == 'point':
            conn = ap.point_data[region.name]

        else:
            raise ValueError('unknown connectivity type! (%s)' % ct)

        return conn

    def get_full_econn(self, ig=None):
        if ig not in self.igs:
            msg = "The required index (%s) of element group is not available" \
                % (str(ig),)
            raise ValueError(msg)

        econn = self.aps[ig].econn
        n_econn = econn.shape[1]
        if self.n_basis == n_econn:
            full_econn = econn

        elif self.n_basis==self.n_components*n_econn:
            full_econn = nm.zeros([econn.shape[0], self.n_basis],
                                  dtype=nm.int32)
            for ii in nm.arange(self.n_components):
                slc = slice(n_econn*ii, n_econn*(ii+1))
                full_econn[:, slc] = econn*self.n_components + ii

        return full_econn

    def average_qp_to_vertices(self, data_qp, integral):
        """
        Average data given in quadrature points in region elements into
        region vertices.

        .. math::
           u_n = \sum_e (u_{e,avg} * volume_e) / \sum_e volume_e
               = \sum_e \int_{volume_e} u / \sum volume_e
        """
        region = self.region

        n_cells = region.get_n_cells()
        if n_cells != data_qp.shape[0]:
            msg = 'incomatible shape! (%d == %d)' % (n_cells,
                                                     data_qp.shape[0])
            raise ValueError(msg)

        n_vertex = self.n_vertex_dof
        nc = data_qp.shape[2]

        nod_vol = nm.zeros((n_vertex,), dtype=nm.float64)
        data_vertex = nm.zeros((n_vertex, nc), dtype=nm.float64)
        for ig, ap in self.aps.iteritems():
            vg = ap.describe_geometry(self, 'volume', ap.region, integral)

            volume = nm.squeeze(vg.volume)
            iels = ap.region.get_cells(ig)

            data_e = nm.zeros((volume.shape[0], 1, nc, 1), dtype=nm.float64)
            vg.integrate(data_e, data_qp[iels])

            ir = nm.arange(nc, dtype=nm.int32)

            conn = ap.econn[:, :self.gel.n_vertex]
            for ii, cc in enumerate(conn):
                # Assumes unique nodes in cc!
                ind2, ind1 = nm.meshgrid(ir, cc)
                data_vertex[ind1,ind2] += data_e[iels[ii],0,:,0]
                nod_vol[cc] += volume[ii]
        data_vertex /= nod_vol[:,nm.newaxis]

        return data_vertex

class SurfaceField(FEField):
    """
    Finite element field base class over surface (element dimension is one
    less than space dimension).
    """

    def _check_region(self, region):
        """
        Check whether the `region` can be used for the
        field.

        Returns
        -------
        ok : bool
            True if the region is usable for the field.
        """
        ok = True
        for ig in region.igs:
            n_cell = region.get_n_cells(ig, True)
            if n_cell == 0:
                ok = False
                break

        return ok

    def _setup_geometry(self):
        """
        Setup the field region geometry.
        """
        self.gel = self.domain.groups[self.region.igs[0]].gel.surface_facet
        if self.gel is None:
            raise ValueError('element group has no surface!')

        self.is_surface = True

    def _create_interpolant(self):
        name = '%s_%s_%s_%d%s' % (self.gel.name, self.space,
                                  self.poly_space_base, self.approx_order,
                                  'B' * self.force_bubble)
        self.interp = fea.SurfaceInterpolant(name, self.gel, self.space,
                                             self.poly_space_base,
                                             self.approx_order,
                                             self.force_bubble)

    def _setup_approximations(self):
        self.aps = {}
        self.aps_by_name = {}
        for ig in self.igs:
            name = self.interp.name + '_%s_ig%d' % (self.region.name, ig)
            ap = fea.SurfaceApproximation(name, self.interp, self.region, ig)
            self.aps[ig] = ap
            self.aps_by_name[ap.name] = ap

    def setup_extra_data(self, geometry, info, is_trace):
        dct = info.dc_type.type

        if dct != 'surface':
            msg = "dof connectivity type must be 'surface'! (%s)" % dct
            raise ValueError(msg)

        reg = info.get_region()

        for ig, ap in self.aps.iteritems():
            if ig not in reg.igs: continue

            if reg.name not in ap.surface_data:
                # Defined in setup_vertex_dofs()
                msg = 'no surface data of surface field! (%s)' % reg.name
                raise ValueError(msg)

        for ig, ap in self.aps.iteritems():
            if reg.name in ap.surface_data and is_trace:
                sd = ap.surface_data[reg.name]
                sd.setup_mirror_connectivity(reg)

    def _init_econn(self):
        """
        Initialize the extended DOF connectivity.
        """
        for ig, ap in self.aps.iteritems():
            n_ep = ap.n_ep['v']
            n_cell = self.region.get_n_cells(ig, True)
            ap.econn = nm.zeros((n_cell, n_ep), nm.int32)

    def _setup_vertex_dofs(self):
        """
        Setup vertex DOF connectivity.
        """
        if self.node_desc.vertex is None:
            return 0, None

        region = self.region

        remap = prepare_remap(region.vertices, region.n_v_max)
        n_dof = region.vertices.shape[0]

        ##
        # Remap vertex node connectivity to field-local numbering.
        for ig, ap in self.aps.iteritems():
            group = self.domain.groups[ig]
            faces = group.gel.get_surface_entities()
            aux = FESurface('aux', region, faces, group.conn, ig)
            ap.econn[:,:aux.n_fp] = aux.leconn
            ap.surface_data[region.name] = aux

        return n_dof, remap

    def _setup_bubble_dofs(self):
        """
        Setup bubble DOF connectivity.
        """
        return 0, None, None

    def get_econn(self, conn_type, region, ig, is_trace=False,
                  integration=None):
        """
        Get extended connectivity of the given type in the given region.
        """
        ct = conn_type.type if isinstance(conn_type, Struct) else conn_type

        if ct != 'surface':
            msg = 'connectivity type must be "surface"! (%s)' % ct
            raise ValueError(msg)

        ap = self.aps[ig]

        sd = ap.surface_data[region.name]
        conn = sd.get_connectivity(local=True, is_trace=is_trace)

        return conn

    def average_qp_to_vertices(self, data_qp, integral):
        """
        Average data given in quadrature points in region elements into
        region vertices.

        .. math::
           u_n = \sum_e (u_{e,avg} * area_e) / \sum_e area_e
               = \sum_e \int_{area_e} u / \sum area_e
        """
        region = self.region

        n_cells = region.get_n_cells(None, True)
        if n_cells != data_qp.shape[0]:
            msg = 'incomatible shape! (%d == %d)' % (n_cells,
                                                     data_qp.shape[0])
            raise ValueError(msg)

        n_vertex = len(region.vertices)
        nc = data_qp.shape[2]

        nod_vol = nm.zeros((n_vertex,), dtype=nm.float64)
        data_vertex = nm.zeros((n_vertex, nc), dtype=nm.float64)
        offset = 0
        for ig, ap in self.aps.iteritems():
            sg = ap.describe_geometry(self, 'surface', ap.region, integral)

            area = nm.squeeze(sg.volume)
            n_cells = region.get_n_cells(ig, True)
            iels = offset + nm.arange(n_cells, dtype=nm.int32)
            offset += n_cells

            data_e = nm.zeros((area.shape[0], 1, nc, 1), dtype=nm.float64)
            sg.integrate(data_e, data_qp[iels])

            ir = nm.arange(nc, dtype=nm.int32)

            sd = self.domain.surface_groups[ig][region.name]
            # Should be vertex connectivity!
            conn = sd.get_connectivity(local=True)
            for ii, cc in enumerate(conn):
                # Assumes unique nodes in cc!
                ind2, ind1 = nm.meshgrid(ir, cc)
                data_vertex[ind1,ind2] += data_e[iels[ii],0,:,0]
                nod_vol[cc] += area[ii]
        data_vertex /= nod_vol[:,nm.newaxis]

        return data_vertex

class H1Mixin(Struct):
    """
    Methods of fields specific to H1 space.
    """

    def _setup_shape(self):
        """
        Setup the field's shape-related attributes, see :class:`Field`.
        """
        self.n_components = nm.prod(self.shape)
        self.val_shape = self.shape


import os
import numpy

import _mypath
import chroma

data_dir = '../../../data/'

def load_LSST_filters():
    """ Read the LSST filter definitions from disk.  Also add a normalization filter that is
    only non-zero at 500nm, which is used to define the CatSim/PhoSim normalization.
    """
    filter_dir = data_dir+'filters/'
    filter_names = 'ugrizy'
    # establish wavelength array
    filters = {}
    for filter_name in filter_names:
        ffile = filter_dir+'LSST_{}.dat'.format(filter_name)
        wave, throughput = numpy.genfromtxt(ffile).T
        filters['LSST_'+filter_name] = {'wave':wave,
                                        'throughput':throughput}
    norm_throughput = numpy.zeros_like(wave)
    norm_throughput[wave==500] = 1.0
    filters['norm'] = {'wave':wave,
                       'throughput':norm_throughput}
    return filters

def load_Euclid_filters():
    """ Load Euclid filters as defined by Voigt+12, centered at 775nm with widths of 150, 250, 350,
    and 450 nm.  These are stored in the data directory."""

    filter_dir = data_dir+'filters/'
    filter_names = ['Euclid_{}'.format(width) for width in [150,250,350,450]]

    filters = {}
    for filter_name in filter_names:
        ffile = filter_dir+'{}.dat'.format(filter_name)
        wave, throughput = numpy.genfromtxt(ffile).T
        filters[filter_name] = {'wave':wave,
                                'throughput':throughput}
    norm_throughput = numpy.zeros_like(wave)
    norm_throughput[wave==500] = 1.0
    filters['norm'] = {'wave':wave,
                       'throughput':norm_throughput}
    return filters

def match_filter_wavelengths(filters, wave_match):
    for fname, f in filters.iteritems():
        new_throughput = numpy.interp(wave_match, f['wave'], f['throughput'])
        filters[fname] = {'wave':wave_match,
                          'throughput':new_throughput}
    return filters

def AB_zeropoints(filters):
    """Compute AB zeropoints for given filters.
    """
    # define AB source in flambda
    ABsource = 3631e-23 # 3631 Jy -> erg/s/Hz/cm^2
    c = 29979245800.0 # cm/s
    nm_to_cm = 1.0e-7

    zps = {}
    for filter_name, filter_ in filters.iteritems():
        fwave = filter_['wave']
        throughput = filter_['throughput']
        ABflambda = ABsource * c / fwave**2 / nm_to_cm # erg/s/Hz/cm^2*cm/s/nm^2 -> erg/s/cm^2/nm
        AB_photons = ABflambda * fwave * throughput
        dlambda = fwave[1] - fwave[0] # assuming linear wavelength bins!
        AB_sumphotons = (AB_photons * dlambda).sum()
        zps[filter_name] = -2.5 * numpy.log10(AB_sumphotons)
    return zps

def read_spec(specfile):
    wave, flambda = numpy.genfromtxt(specfile).T
    return {'wave':wave, 'flambda':flambda}

def match_wavelengths(spec, wave_match):
    """ Interpolate spectrum onto given wavelength array.
    """
    flux_i = numpy.interp(wave_match, spec['wave'], spec['flambda'])
    return {'wave':wave_match, 'flambda':flux_i}

def apply_redshift(spec, redshift):
    flux_i = numpy.interp(spec['wave'], spec['wave'] * (1.0 + redshift), spec['flambda'])
    return {'wave':spec['wave'], 'flambda':flux_i / (1.0 + redshift)}

def scale_spec(spec, target_mag, normfilter, normzp):
    """ Multiply spectrum flux such that the normalization magnitude matches the given target
    magnitude.
    """
    fwave = normfilter['wave']
    throughput = normfilter['throughput']
    flambda_i = numpy.interp(fwave, spec['wave'], spec['flambda'])
    photons = flambda_i * fwave * throughput
    dlambda = fwave[1] - fwave[0]
    sumphotons = (photons * dlambda).sum()
    current_mag = -2.5 * numpy.log10(sumphotons) - normzp
    multiplier = 10**(-0.4 * (target_mag - current_mag))
    spec['flambda'] *= multiplier
    return spec

def apply_extinction(spec, A_v, R_v=3.1):
    wave = spec['wave']
    valid = (wave > 92) & (wave < 1200)
    wave = wave[valid]
    ext = chroma.extinction.reddening(wave*10, a_v=A_v, r_v=R_v, model='f99')
    flambda = spec['flambda']
    flambda[valid] /= ext
    spec = {'wave':spec['wave'], 'flambda':flambda}
    return spec

def compute_mags(spec, filters, zps):
    """ Compute magnitudes from spectrum.  Assume that spectrum wavelengths are already matched to
    filter wavelengths.
    """
    mags = {}
    for filter_name, filter_ in filters.iteritems():
        fwave = filter_['wave']
        throughput = filter_['throughput']
        photons = spec['flambda'] * fwave * throughput
        dlambda = fwave[1] - fwave[0] # assuming linear wavelength bins!
        sumphotons = (photons * dlambda).sum()
        mags[filter_name] = -2.5 * numpy.log10(sumphotons) - zps[filter_name]
    return mags

def make_composite_spec(gal, filters, zps, wave_match):
    SED_dir = os.environ['CAT_SHARE_DATA']+'data/'

    if gal['sedPathBulge'] != 'None':
        bulge_spec = read_spec(SED_dir+gal['sedPathBulge'])
        bulge_spec = scale_spec(bulge_spec, gal['magNormBulge'],
                                filters['norm'], zps['norm'])
        bulge_spec = apply_extinction(bulge_spec,
                                      A_v=gal['internalAVBulge'],
                                      R_v=gal['internalRVBulge'])
        bulge_spec = apply_redshift(bulge_spec, gal['redshift'])
        bulge_spec = match_wavelengths(bulge_spec, wave_match)
    else:
        bulge_spec = {'wave':wave_match, 'flambda':numpy.zeros_like(wave_match)}
    if gal['sedPathDisk'] != 'None':
        disk_spec = read_spec(SED_dir+gal['sedPathDisk'])
        disk_spec = scale_spec(disk_spec, gal['magNormDisk'],
                                filters['norm'], zps['norm'])
        disk_spec = apply_extinction(disk_spec,
                                      A_v=gal['internalAVDisk'],
                                      R_v=gal['internalRVDisk'])
        disk_spec = apply_redshift(disk_spec, gal['redshift'])
        disk_spec = match_wavelengths(disk_spec, wave_match)
    else:
        disk_spec = {'wave':wave_match, 'flambda':numpy.zeros_like(wave_match)}
    if gal['sedPathAGN'] != 'None':
        AGN_spec = read_spec(SED_dir+gal['sedPathAGN'])
        AGN_spec = scale_spec(AGN_spec, gal['magNormAGN'],
                              filters['norm'], zps['norm'])
        AGN_spec = apply_redshift(AGN_spec, gal['redshift'])
        AGN_spec = match_wavelengths(AGN_spec, wave_match)
    else:
        AGN_spec = {'wave':wave_match, 'flambda':numpy.zeros_like(wave_match)}
    return {'wave':wave_match, 'flambda':(bulge_spec['flambda'] +
                                          disk_spec['flambda'] +
                                          AGN_spec['flambda'])}

def get_UV_flux(spec):
    # flambda in erg/s/cm^2/nm
    UV_flambda = numpy.interp(230.0, spec['wave'], spec['flambda'])
    # convert to erg/s/Hz
    c = 29979245800.0 # cm/s
    nm_to_cm = 1.0e-7
    UV_fnu = UV_flambda * (230.0)**2 / c * nm_to_cm
    return UV_fnu

def add_emission_lines(spec):
    wave = spec['wave'][:]
    flambda = spec['flambda'][:]
    UV_fnu = get_UV_flux(spec)
    lines = ['OII','OIII','Hbeta','Halpha','Lya']
    multipliers = numpy.array([1.0, 0.36, 0.61, 1.77, 2.0])*1.e13
    waves = [372.7, 500.7, 486.1, 656.3, 121.5]
    velocity = 200.0 # km/s
    for line, m, w in zip(lines, multipliers, waves):
        flux = UV_fnu * m
        sigma = velocity / 299792.458 * w # sigma in Angstroms
        amplitude = flux / sigma / numpy.sqrt(2.0 * numpy.pi)
        flambda += amplitude * numpy.exp(-(w-spec['wave'])**2/(2*sigma**2))
    return {'wave':wave, 'flambda':flambda}

def make_composite_spec_with_emission_lines(gal, filters, zps, wave_match):
    SED_dir = os.environ['CAT_SHARE_DATA']+'data/'

    if gal['sedPathBulge'] != 'None':
        bulge_spec = read_spec(SED_dir+gal['sedPathBulge'])
        bulge_spec = scale_spec(bulge_spec, gal['magNormBulge'],
                                filters['norm'], zps['norm'])
        bulge_spec = add_emission_lines(bulge_spec)
        bulge_spec = apply_extinction(bulge_spec,
                                      A_v=gal['internalAVBulge'],
                                      R_v=gal['internalRVBulge'])
        bulge_spec = apply_redshift(bulge_spec, gal['redshift'])
        bulge_spec = match_wavelengths(bulge_spec, wave_match)
    else:
        bulge_spec = {'wave':wave_match, 'flambda':numpy.zeros_like(wave_match)}
    if gal['sedPathDisk'] != 'None':
        disk_spec = read_spec(SED_dir+gal['sedPathDisk'])
        disk_spec = scale_spec(disk_spec, gal['magNormDisk'],
                                filters['norm'], zps['norm'])
        disk_spec = add_emission_lines(disk_spec)
        disk_spec = apply_extinction(disk_spec,
                                      A_v=gal['internalAVDisk'],
                                      R_v=gal['internalRVDisk'])
        disk_spec = apply_redshift(disk_spec, gal['redshift'])
        disk_spec = match_wavelengths(disk_spec, wave_match)
    else:
        disk_spec = {'wave':wave_match, 'flambda':numpy.zeros_like(wave_match)}
    if gal['sedPathAGN'] != 'None':
        AGN_spec = read_spec(SED_dir+gal['sedPathAGN'])
        AGN_spec = scale_spec(AGN_spec, gal['magNormAGN'],
                              filters['norm'], zps['norm'])
        AGN_spec = apply_redshift(AGN_spec, gal['redshift'])
        AGN_spec = match_wavelengths(AGN_spec, wave_match)
    else:
        AGN_spec = {'wave':wave_match, 'flambda':numpy.zeros_like(wave_match)}
    return {'wave':wave_match, 'flambda':(bulge_spec['flambda'] +
                                          disk_spec['flambda'] +
                                          AGN_spec['flambda'])}

def make_bulge_spec(gal, filters, zps, wave_match):
    SED_dir = os.environ['CAT_SHARE_DATA']+'data/'

    if gal['sedPathBulge'] != 'None':
        bulge_spec = read_spec(SED_dir+gal['sedPathBulge'])
        bulge_spec = scale_spec(bulge_spec, gal['magNormBulge'],
                                filters['norm'], zps['norm'])
        bulge_spec = apply_extinction(bulge_spec,
                                      A_v=gal['internalAVBulge'],
                                      R_v=gal['internalRVBulge'])
        bulge_spec = apply_redshift(bulge_spec, gal['redshift'])
        bulge_spec = match_wavelengths(bulge_spec, wave_match)
    else:
        return None
    return {'wave':wave_match, 'flambda':bulge_spec['flambda']}

def make_disk_spec(gal, filters, zps, wave_match):
    SED_dir = os.environ['CAT_SHARE_DATA']+'data/'

    if gal['sedPathDisk'] != 'None':
        disk_spec = read_spec(SED_dir+gal['sedPathDisk'])
        disk_spec = scale_spec(disk_spec, gal['magNormDisk'],
                                filters['norm'], zps['norm'])
        disk_spec = apply_extinction(disk_spec,
                                      A_v=gal['internalAVDisk'],
                                      R_v=gal['internalRVDisk'])
        disk_spec = apply_redshift(disk_spec, gal['redshift'])
        disk_spec = match_wavelengths(disk_spec, wave_match)
    else:
        return None
    return {'wave':wave_match, 'flambda':disk_spec['flambda']}


""" Minimum working example of an SME script
"""

import logging
import sys
from os.path import dirname, join, realpath

import numpy as np

from pysme import sme as SME
from pysme import util
from pysme.linelist.vald import ValdFile
from pysme.nso import load_solar_spectrum
from pysme.solve import solve

logger = logging.getLogger("pysme")
# logger.setLevel(logging.CRITICAL)

if __name__ == "__main__":
    # Define the location of all your files
    # this will put everything into the example dir
    target = "sun"
    examples_dir = dirname(realpath(__file__))
    # in_file = join(examples_dir, "../sun_6440_test.inp")
    # in_file = join(examples_dir, "gr8_HARPS_HD148816.inp")
    vald_file = join(examples_dir, "data/harps.lin")
    out_file = join(examples_dir, f"{target}.sme")
    plot_file = join(examples_dir, f"{target}.html")
    log_file = join(examples_dir, f"{target}.log")

    # Use HARPS spectral range
    wave, flux = load_solar_spectrum()
    wmin, wmax = 3780, 6910
    idx = (wave > wmin) & (wave < wmax)
    wave, flux = wave[idx], flux[idx]
    flux[flux < 0] = 0

    sme = SME.SME_Structure(wave=wave, sob=flux)
    sme.uncs = np.sqrt(flux)
    sme.mask = sme.mask_values["line"]

    sme.linelist = ValdFile(vald_file)

    sme.nmu = 7
    sme.vrad = 0
    sme.cscale = None
    sme.vrad_flag = "each"
    sme.cscale_flag = "linear"
    sme.cscale_type = "match"
    sme.atmo.source = "marcs2014.sav"
    sme.abund = "asplund2009"

    # elems = [
    #     "Al",
    #     "Ba",
    #     "Ca",
    #     "C",
    #     "H",
    #     "K",
    #     "Li",
    #     "Mg",
    #     "Mn",
    #     "Na",
    #     "N",
    #     "O",
    #     "Si",
    # ]
    # for elem in elems:
    #     sme.nlte.set_nlte(elem, f"nlte_{elem}_ama51_pysme.grd")
    # sme.nlte.set_nlte("Fe", "marcs2012_Fe2016.grd")

    # Load the parameters from the command line
    if len(sys.argv) > 1:
        param = sys.argv[1:]
    else:
        param = [5000, 4.0, 0]  # 5000.00, 4.00, 0.00
    sme.teff = float(param[0])
    sme.logg = float(param[1])
    sme.monh = float(param[2])

    # try:
    sme = solve(sme, ["teff", "logg", "monh"])
    print(f"{param[0]}, {param[1]}, {param[2]}, {sme.teff}, {sme.logg}, {sme.monh}")
    # except Exception as ex:
    #     print(f"{param[0]}, {param[1]}, {param[2]}, nan, nan, nan")


'''
Contains routines to read and manipulate particle size distribution data from the
2D-S and HVPS optical array probes during the IMPACTS experiment. Creation of 1 Hz
PSD data requires use of UIOOPS package (doi: 10.5281/zenodo.3976291).
Some of the arguments in the routines require radar data matched to the P-3
location to determine the optional degree of riming following Leinonen & Szyrmer (2015).

Copyright Joe Finlon, Univ. of Washington, 2022.
'''

import xarray as xr
import numpy as np
np.warnings.filterwarnings('ignore', category=np.VisibleDeprecationWarning)
np.warnings.filterwarnings('ignore', message='Mean of empty slice')
np.seterr(invalid='ignore')
from datetime import datetime, timedelta
from scipy.optimize import least_squares
try: # try importing the pytmatrix package
    from forward import *
except ImportError:
    print(
        'WARNING: The pytmatrix package cannot be installed for the psdread() function.'
    )

def psdread(
        twodsfile, hvpsfile, datestr, size_cutoff=1., minD=0.15, maxD=30.,
        qc=False, deadtime_thresh=0.6, verbose=True,
        start_time=None, end_time=None, tres=5.,
        compute_bulk=False, compute_fits=False, Z_interp=False,
        matchedZ_W=None, matchedZ_Ka=None, matchedZ_Ku=None, matchedZ_X=None):
    '''
    Load the 2DS and HVPS PSDs processed by UIOOPS and create n-second combined PSDs with optional bulk properties.
    Inputs:
        twodsfile: Path to the 2DS data
        hvpsfile: Path to the HVPS data
        datestr: YYYYMMDD [str]
        size_cutoff: Size [mm] for the 2DS-HVPS crossover
        minD: Minimum size [mm] to consider in the combined PSD
        maxD: Maximum size [mm] to consider in the combined PSD
        qc: Boolean to optionally ignore 1-Hz data from averaging when probe dead time > deadtime_thresh
        deadtime_thresh: Deadtime hreshold [0–1] to ignore 1-Hz data when qc is True
        verbose: Boolean to optionally print all data-related warnings (e.g., high probe dead time)
        start_time: Start time [YYYY-MM-DDTHH:MM:SS as str] to consider for the PSDs (optional)
        end_time: End time [YYYY-MM-DDTHH:MM:SS as str] to consider for the PSDs (optional)
        tres: Averaging interval [s]; tres=1. skips averaging routine
        compute_bulk: Boolean to optionally compute bulk statistics such as N, IWC, Dmm, rho_e
        compute_fits: Boolean to optionally compute gamma fit parameters N0, mu, lambda
        Z_interp: Boolean to optionally simulate Z for additional degrees of riming from Leinonen & Szyrmer (2015; LS15)
        matchedZ_Ka, ...: None (skips minimization) or masked array of matched Z values to perform LS15 m-D minimization
    '''
    p3psd = {}

    if (twodsfile is None) and (hvpsfile is not None):
        print('Only using the HVPS data for {}'.format(datestr))
        size_cutoff = 0.4 # start HVPS PSD at 0.4 mm
    elif (hvpsfile is None) and (twodsfile is not None):
        print('Only using the 2DS data for {}'.format(datestr))
        size_cutoff = 3.2 # end 2DS PSD at 3.2 mm
    elif (twodsfile is None) and (hvpsfile is None):
        print('No input files given...exiting')
        exit()

    # 2DS information
    if twodsfile is not None:
        ds1 = xr.open_dataset(twodsfile)
        time_raw = ds1['time'].values # HHMMSS from flight start date or numpy.datetime64
        if np.issubdtype(time_raw.dtype, np.datetime64): # numpy.datetime64
            time = np.array(time_raw, dtype='datetime64[s]')
        else: # native HHMMSS format (from UIOOPS SD file)
            time_dt = [
                datetime(int(datestr[0:4]), int(datestr[4:6]), int(datestr[6:]))
                + timedelta(
                    hours=int(str(int(time_raw[i])).zfill(6)[0:2]),
                    minutes=int(str(int(time_raw[i])).zfill(6)[2:4]),
                    seconds=int(str(int(time_raw[i])).zfill(6)[4:]))
                for i in range(len(time_hhmmss))
            ]
            time_str = [
                datetime.strftime(time_dt[i], '%Y-%m-%dT%H:%M:%S')
                for i in range(len(time_dt))
            ]
            time = np.array(time_str, dtype='datetime64[s]')
        bin_min_2ds = ds1['bin_min'].values # mm
        bin_max_2ds = ds1['bin_max'].values
        bin_inds = np.where((bin_min_2ds>=minD) & (bin_max_2ds<=size_cutoff))[0] # find bins within user-specified range
        bin_min_2ds = bin_min_2ds[bin_inds]; bin_max_2ds = bin_max_2ds[bin_inds]
        bin_width_2ds = ds1['bin_dD'].values[bin_inds] / 10. # cm
        bin_mid_2ds = bin_min_2ds + (bin_width_2ds * 10.) / 2.
        count_2ds = ds1['count'].values[:, bin_inds]
        sv_2ds = ds1['sample_vol'].values[:, bin_inds] # cm^3
        count_hab_2ds = ds1['habitsd'].values[:, bin_inds, :] * np.tile(np.reshape(sv_2ds, (sv_2ds.shape[0], sv_2ds.shape[1], 1)), (1, 1, 10)) * np.tile(
            np.reshape(bin_width_2ds, (1, len(bin_width_2ds), 1)), (sv_2ds.shape[0], 1, 10))
        ar_2ds = ds1['mean_area_ratio'].values[:, bin_inds] # mean area ratio (circular fit) per bin
        asr_2ds = ds1['mean_aspect_ratio_ellipse'].values[:, bin_inds] # mean aspect ratio (elliptical fit) per bin
        activetime_2ds = ds1['sum_IntArr'].values # s

        if hvpsfile is None:
            count = count_2ds; count_hab = count_hab_2ds; sv = sv_2ds; ar = ar_2ds; asr = asr_2ds; activetime_hvps = np.ones(count.shape[0])
            bin_min = bin_min_2ds; bin_mid = bin_mid_2ds; bin_max = bin_max_2ds; bin_width = bin_width_2ds

    # HVPS information
    if hvpsfile is not None:
        ds2 = xr.open_dataset(hvpsfile)
        bin_min_hvps = ds2['bin_min'].values # mm
        bin_max_hvps = ds2['bin_max'].values
        bin_inds = np.where((bin_min_hvps>=size_cutoff) & (bin_max_hvps<=maxD))[0] # find bins within user-specified range
        bin_min_hvps = bin_min_hvps[bin_inds]; bin_max_hvps = bin_max_hvps[bin_inds]
        bin_width_hvps = ds2['bin_dD'].values[bin_inds] / 10. # cm
        if size_cutoff==2.:
            bin_min_hvps = np.insert(bin_min_hvps, 0, 2.); bin_max_hvps = np.insert(bin_max_hvps, 0, 2.2); bin_width_hvps = np.insert(bin_width_hvps, 0, 0.02)
            bin_inds = np.insert(bin_inds, 0, bin_inds[0]-1)
        bin_mid_hvps = bin_min_hvps + (bin_width_hvps * 10.) / 2.
        count_hvps = ds2['count'].values[:, bin_inds]
        sv_hvps = ds2['sample_vol'].values[:, bin_inds] # cm^3
        count_hab_hvps = (ds2['habitsd'].values[:, bin_inds, :]) * np.tile(np.reshape(sv_hvps, (sv_hvps.shape[0], sv_hvps.shape[1], 1)), (1, 1, 10)) * np.tile(
            np.reshape(bin_width_hvps, (1, len(bin_width_hvps), 1)), (sv_hvps.shape[0], 1, 10))
        ar_hvps = ds2['mean_area_ratio'].values[:, bin_inds] # mean area ratio (circular fit) per bin
        asr_hvps = ds2['mean_aspect_ratio_ellipse'].values[:, bin_inds] # mean aspect ratio (elliptical fit) per bin
        activetime_hvps = ds2['sum_IntArr'].values # s
        if size_cutoff==2.: # normalize counts in first bin (1.8-2.2 mm, now only for 2-2.2 mm)
            count_hvps[:, 0] = count_hvps[:, 0] / 2.
            count_hab_hvps[:, 0, :] = count_hab_hvps[:, 0, :] / 2.

        if twodsfile is None:
            time_hhmmss = ds2['time'].values # HHMMSS from flight start date
            time_dt = [datetime(int(datestr[0:4]), int(datestr[4:6]), int(datestr[6:])) + timedelta(
                hours=int(str(int(time_hhmmss[i])).zfill(6)[0:2]), minutes=int(str(int(time_hhmmss[i])).zfill(6)[2:4]),
                seconds=int(str(int(time_hhmmss[i])).zfill(6)[4:])) for i in range(len(time_hhmmss))]
            time_str = [datetime.strftime(time_dt[i], '%Y-%m-%dT%H:%M:%S') for i in range(len(time_dt))]
            time = np.array(time_str, dtype='datetime64[s]')
            count = count_hvps; count_hab = count_hab_hvps; sv = sv_hvps; ar = ar_hvps; asr = asr_hvps; activetime_2ds = np.ones(count.shape[0])
            bin_min = bin_min_hvps; bin_mid = bin_mid_hvps; bin_max = bin_max_hvps; bin_width = bin_width_hvps

    # Combine the datasets
    if (twodsfile is not None) and (hvpsfile is not None):
        count = np.concatenate((count_2ds, count_hvps), axis=1)
        count_hab = np.concatenate((count_hab_2ds, count_hab_hvps), axis=1)
        sv = np.concatenate((sv_2ds, sv_hvps), axis=1)
        ar = np.concatenate((ar_2ds, ar_hvps), axis=1)
        asr = np.concatenate((asr_2ds, asr_hvps), axis=1)
        bin_min = np.concatenate((bin_min_2ds, bin_min_hvps))
        bin_mid = np.concatenate((bin_mid_2ds, bin_mid_hvps))
        bin_max = np.concatenate((bin_max_2ds, bin_max_hvps))
        bin_width = np.concatenate((bin_width_2ds, bin_width_hvps))

    # Average the data
    if start_time is None:
        start_dt64 = time[0]
    else:
        start_dt64 = np.datetime64(start_time)
    if end_time is None:
        end_dt64 = time[-1] if int(tres)>1 else time[-1]+np.timedelta64(1, 's')
    else:
        end_dt64 = np.datetime64(end_time) if int(tres)>1 else np.datetime64(end_time)+np.timedelta64(1, 's')
    dur = (end_dt64 - start_dt64) / np.timedelta64(1, 's') # dataset duration to consider [s]

    # Allocate arrays
    count_aver = np.zeros((int(dur/tres), len(bin_mid)))
    count_hab_aver = np.zeros((int(dur/tres), len(bin_mid), 8))
    sv_aver = np.zeros((int(dur/tres), len(bin_mid)))
    at_2ds_aver = np.ma.array(np.ones(int(dur/tres)), mask=False)
    at_hvps_aver = np.ma.array(np.ones(int(dur/tres)), mask=False)
    ND = np.zeros((int(dur/tres), len(bin_mid)))
    ar_aver = np.zeros((int(dur/tres), len(bin_mid)))
    asr_aver = np.zeros((int(dur/tres), len(bin_mid)))

    time_subset = start_dt64 # allocate time array of N-sec interval obs
    curr_time = start_dt64
    i = 0

    while curr_time+np.timedelta64(int(tres),'s')<=end_dt64:
        if curr_time>start_dt64:
            time_subset = np.append(time_subset, curr_time)
        time_inds = np.where((time>=curr_time) & (time<curr_time+np.timedelta64(int(tres), 's')))[0]
        if qc is True:
            activetime_thresh = 1. - deadtime_thresh
            time_inds = time_inds[(activetime_2ds[time_inds]>=activetime_thresh) & (activetime_hvps[time_inds]>=activetime_thresh)]
        if len(time_inds)>0:
            count_aver[i, :] = np.nansum(count[time_inds, :], axis=0)
            count_hab_aver[i, :, 0] = np.nansum(count_hab[time_inds, :, 3], axis=0) # tiny
            count_hab_aver[i, :, 1] = np.nansum(count_hab[time_inds, :, 0], axis=0) # spherical
            count_hab_aver[i, :, 2] = np.nansum(count_hab[time_inds, :, 1:3], axis=(0, 2)) # oriented + linear
            count_hab_aver[i, :, 3] = np.nansum(count_hab[time_inds, :, 4], axis=0) # hexagonal
            count_hab_aver[i, :, 4] = np.nansum(count_hab[time_inds, :, 5], axis=0) # irregular
            count_hab_aver[i, :, 5] = np.nansum(count_hab[time_inds, :, 6], axis=0) # graupel
            count_hab_aver[i, :, 6] = np.nansum(count_hab[time_inds, :, 7], axis=0) # dendrite
            count_hab_aver[i, :, 7] = np.nansum(count_hab[time_inds, :, 8], axis=0) # aggregate
            ar_aver[i, :] = np.nanmean(ar[time_inds, :], axis=0) # binned mean of area ratio
            asr_aver[i, :] = np.nanmean(asr[time_inds, :], axis=0) # binned mean of aspect ratio
            sv_aver[i, :] = np.nansum(sv[time_inds, :], axis=0)
            at_2ds_aver[i] = np.nansum(activetime_2ds[time_inds]) / len(time_inds)
            at_hvps_aver[i] = np.nansum(activetime_hvps[time_inds]) / len(time_inds)
            ND[i, :] = np.nanmean(count[time_inds, :]/sv[time_inds, :], axis=0) / bin_width # take N(D) for each sec, then average [cm**-4]
        else: # Mask data for current period if dead (active) time from either probe > 0.8*tres (< 0.2*tres) for all 1-Hz times
            if verbose is True:
                print('All 1-Hz data for the {}-s period beginning {} has high dead time. Masking data.'.format(str(tres), np.datetime_as_string(curr_time)))
            at_2ds_aver[i] = np.nansum(activetime_2ds[np.where((time>=curr_time) & (time<curr_time+np.timedelta64(int(tres), 's')))[0]]) / tres; at_2ds_aver.mask[i] = True
            at_hvps_aver[i] = np.nansum(activetime_hvps[np.where((time>=curr_time) & (time<curr_time+np.timedelta64(int(tres), 's')))[0]]) / tres; at_hvps_aver.mask[i] = True
            count_aver[i, :] = np.nan; count_hab_aver[i, :] = np.nan; sv_aver[i, :] = np.nan; ND[i, :] = np.nan; asr_aver[i, :] = np.nan
        i += 1
        curr_time += np.timedelta64(int(tres), 's')

    #ND = np.ma.masked_invalid(count_aver / sv_aver / np.tile(bin_width[np.newaxis, :], (int(dur/tres), 1))) # cm^-4

    # Mask arrays
    count_aver = np.ma.masked_where(np.isnan(count_aver), count_aver)
    count_hab_aver = np.ma.masked_where(np.isnan(count_hab_aver), count_hab_aver)
    sv_aver = np.ma.masked_where(np.isnan(sv_aver), sv_aver)
    ar_aver = np.ma.masked_invalid(ar_aver)
    asr_aver = np.ma.masked_invalid(asr_aver)
    ND[~np.isfinite(ND)] = 0.; ND = np.ma.masked_where(ND==0., ND)

    # Create dictionary
    p3psd['time'] = time_subset
    p3psd['count'] = count_aver
    p3psd['count_habit'] = count_hab_aver
    p3psd['sv'] = sv_aver
    p3psd['area_ratio'] = ar_aver
    p3psd['aspect_ratio'] = asr_aver
    p3psd['ND'] = ND
    p3psd['bin_min'] = bin_min
    p3psd['bin_mid'] = bin_mid
    p3psd['bin_max'] = bin_max
    p3psd['bin_width'] = bin_width
    p3psd['active_time_2ds'] = at_2ds_aver
    p3psd['active_time_hvps'] = at_hvps_aver

    if compute_bulk is True:
        # Compute Z for various degrees of riming and radar wavelengths
        # Based on work from Leionen and Szyrmer 2015 (LS15)
        # (https://agupubs.onlinelibrary.wiley.com/doi/pdf/10.1002/2015EA000102)
        # Follows https://github.com/dopplerchase/Leinonen_Python_Forward_Model
        # and uses forward.py and ess238-sup-0002-supinfo.tex in repo
        Z = forward_Z() #initialize class
        # get the PSD in the format to use in the routine (mks units)
        Z.set_PSD(PSD=ND*10.**8, D=bin_mid/1000., dD=bin_width/100., Z_interp=Z_interp)
        Z.load_split_L15() # Load the leinonen output
        Z.fit_sigmas(Z_interp) # Fit the backscatter cross-sections
        Z.fit_rimefrac(Z_interp) # Fit the riming fractions
        Z.calc_Z() # Calculate Z...outputs are Z.Z_x, Z.Z_ku, Z.Z_ka, Z.Z_w for the four radar wavelengths

        # Compute IWC and Dmm following Brown and Francis (1995), modified for a Dmax definition following Hogan et al.
        [
            N0_bf, N0_hy, mu_bf, mu_hy, lam_bf, lam_hy, iwc_bf, iwc_hy, iwc_hab,
            asr_nw, asr_bf, asr_hy, asr_hab, dmm_bf, dmm_hy, dmm_hab, dm_bf, dm_hy,
            dm_hab, rho_bf, rho_hy, rho_hab, rhoe_bf, rhoe_hy, rhoe_hab] = calc_bulk(
            count_aver, count_hab_aver, sv_aver, asr_aver, bin_mid, bin_width)

        # Add bulk variables to the dictionary
        if Z_interp is True: # consider additional degrees of riming from LS15
            p3psd['riming_mass_array'] = [
                0., 0.05, 0.1, 0.15, 0.2, 0.3, 0.4, 0.5, 0.75, 1., 2.]
        else:
            p3psd['riming_mass_array'] = [0., 0.1, 0.2, 0.5, 1., 2.]
        p3psd['a_coeff_array'] = Z.a_coeff
        p3psd['b_coeff_array'] = Z.b_coeff
        p3psd['dbz_W'] = Z.Z_w
        p3psd['dbz_Ka'] = Z.Z_ka
        p3psd['dbz_Ku'] = Z.Z_ku
        p3psd['dbz_X'] = Z.Z_x
        p3psd['N0_bf'] = N0_bf
        p3psd['N0_hy'] = N0_hy
        p3psd['mu_bf'] = mu_bf
        p3psd['mu_hy'] = mu_hy
        p3psd['lambda_bf'] = lam_bf
        p3psd['lambda_hy'] = lam_hy
        p3psd['iwc_bf'] = iwc_bf
        p3psd['iwc_hy'] = iwc_hy
        p3psd['iwc_hab'] = iwc_hab
        p3psd['mean_aspect_ratio'] = asr_nw
        p3psd['mean_aspect_ratio_bf'] = asr_bf
        p3psd['mean_aspect_ratio_hy'] = asr_hy
        p3psd['mean_aspect_ratio_habit'] = asr_hab
        p3psd['dmm_bf'] = dmm_bf
        p3psd['dmm_hy'] = dmm_hy
        p3psd['dmm_hab'] = dmm_hab
        p3psd['dm_bf'] = dm_bf
        p3psd['dm_hy'] = dm_hy
        p3psd['dm_hab'] = dm_hab
        p3psd['eff_density_bf'] = rhoe_bf
        p3psd['eff_density_hy'] = rhoe_hy
        p3psd['eff_density_hab'] = rhoe_hab
        p3psd['density_bf'] = rho_bf
        p3psd['density_hy'] = rho_hy
        p3psd['density_hab'] = rho_hab

        # Optionally constrain the matched Z at Ku- and Ka-band against PSDS to estimate bulk properties
        if (
                matchedZ_W is not None) or (matchedZ_Ka is not None) or (
                matchedZ_Ku is not None) or (matchedZ_X is not None):
            p3psd = calc_riming(
                p3psd, Z, matchedZ_W, matchedZ_Ka, matchedZ_Ku, matchedZ_X,
                compute_fits=compute_fits)

    return p3psd

def calc_bulk(particle_count, habit_count, sample_vol, aspect_ratio, bin_mid, bin_width):
    x0 = [1.e-1, -1., 5.] # initial guess for N0 [cm**-4], mu, lambda [cm**-1]

    # allocate arrays
    N0_bf = np.zeros(particle_count.shape[0])
    N0_hy = np.zeros(particle_count.shape[0])
    mu_bf = np.zeros(particle_count.shape[0])
    mu_hy = np.zeros(particle_count.shape[0])
    lam_bf = np.zeros(particle_count.shape[0])
    lam_hy = np.zeros(particle_count.shape[0])
    iwc_bf = np.zeros(particle_count.shape[0])
    iwc_hy = np.zeros(particle_count.shape[0])
    iwc_hab = np.zeros(particle_count.shape[0])
    asr_nw = np.zeros(particle_count.shape[0])
    asr_bf = np.zeros(particle_count.shape[0])
    asr_hy = np.zeros(particle_count.shape[0])
    asr_hab = np.zeros(particle_count.shape[0])
    dmm_bf = np.zeros(particle_count.shape[0])
    dmm_hy = np.zeros(particle_count.shape[0])
    dmm_hab = np.zeros(particle_count.shape[0])
    dm_bf = np.zeros(particle_count.shape[0])
    dm_hy = np.zeros(particle_count.shape[0])
    dm_hab = np.zeros(particle_count.shape[0])
    rhoe_bf = np.zeros(particle_count.shape[0])
    rhoe_hy = np.zeros(particle_count.shape[0])
    rhoe_hab = np.zeros(particle_count.shape[0])
    rho_bf = np.zeros((particle_count.shape[0], particle_count.shape[1]))
    rho_hy = np.zeros((particle_count.shape[0], particle_count.shape[1]))
    rho_hab = np.zeros((particle_count.shape[0], particle_count.shape[1]))

    # compute particle habit mass outside loop for speed
    a_coeff = np.array([1.96e-3, 1.96e-3, 1.666e-3, 7.39e-3, 1.96e-3, 4.9e-2, 5.16e-4, 1.96e-3])
    a_tile = np.tile(np.reshape(a_coeff, (1, len(a_coeff))), (habit_count.shape[1], 1))
    b_coeff = np.array([1.9, 1.9, 1.91, 2.45, 1.9, 2.8, 1.8, 1.9])
    b_tile = np.tile(np.reshape(b_coeff, (1, len(b_coeff))), (habit_count.shape[1], 1))
    D_tile = np.tile(np.reshape(bin_mid, (len(bin_mid), 1)), (1, habit_count.shape[2]))
    mass_tile = a_tile * (D_tile/10.) ** b_tile

    for time_ind in range(particle_count.shape[0]):
        if particle_count[time_ind, :].count()==particle_count.shape[1]: # time period is not masked...continue on
            Nt = 1000.*np.nansum(particle_count[time_ind, :]/sample_vol[time_ind, :]) # number concentratino [L**-1]

            # spherical volume from Chase et al. (2018) [cm**3 / cm**3]
            vol = (np.pi / 6.) * np.sum(0.6 * ((bin_mid/10.)**3.) * particle_count[time_ind, :] / sample_vol[time_ind, :])

            # number-weighted mean aspect rato
            asr_nw[time_ind] = np.nansum(aspect_ratio[time_ind, :] * particle_count[time_ind, :]) / np.nansum(particle_count[time_ind, :])

            # Brown & Francis products
            mass_particle = (0.00294/1.5) * (bin_mid/10.)**1.9 # particle mass [g]
            mass_bf = mass_particle * particle_count[time_ind, :] # g (binned)
            cumMass_bf = np.nancumsum(mass_bf)
            if cumMass_bf[-1]>0.:
                iwc_bf[time_ind] = 10.**6 * np.nansum(mass_bf / sample_vol[time_ind, :]) # g m^-3
                z_bf = 1.e12 * (0.174/0.93) * (6./np.pi/0.934)**2 * np.nansum(mass_particle**2*particle_count[time_ind, :]/sample_vol[time_ind, :]) # mm^6 m^-3
                sol = least_squares(calc_chisquare, x0, method='lm',ftol=1e-9,xtol=1e-9, max_nfev=int(1e6),\
                                    args=(Nt,iwc_bf[time_ind],z_bf,bin_mid,bin_width,0.00294/1.5,1.9)) # sove the gamma params using least squares minimziation
                N0_bf[time_ind] = sol.x[0]; mu_bf[time_ind] = sol.x[1]; lam_bf[time_ind] = sol.x[2]
                asr_bf[time_ind] = np.sum(aspect_ratio[time_ind, :] * mass_bf / sample_vol[time_ind, :]) / np.sum(mass_bf / sample_vol[time_ind, :]) # mass-weighted aspect ratio
                rhoe_bf[time_ind] = (iwc_bf[time_ind] / 10.**6) / vol # effective density from Chase et al. (2018) [g cm**-3]
                rho_bf[time_ind, :] = (mass_bf / particle_count[time_ind, :]) / (np.pi / 6.) / (bin_mid/10.)**3. # rho(D) following Heymsfield et al. (2003) [g cm**-3]
                dm_bf[time_ind] = 10. * np.sum((bin_mid/10.) * mass_bf / sample_vol[time_ind, :]) / np.sum(mass_bf / sample_vol[time_ind, :]) # mass-weighted mean D from Chase et al. (2020) [mm]
                if cumMass_bf[0]>=0.5*cumMass_bf[-1]:
                    dmm_bf[time_ind] = bin_mid[0]
                else:
                    dmm_bf[time_ind] = bin_mid[np.where(cumMass_bf>0.5*cumMass_bf[-1])[0][0]-1]

            # Heymsfield (2010) products [https://doi.org/10.1175/2010JAS3507.1]
            #mass_hy = (0.0061*(bin_mid/10.)**2.05) * particle_count[time_ind, :] # g (binned) H04 definition used in GPM NCAR files
            mass_particle = 0.00528 * (bin_mid/10.)**2.1 # particle mass [g]
            mass_hy = mass_particle * particle_count[time_ind, :] # g (binned)
            cumMass_hy = np.nancumsum(mass_hy)
            if cumMass_hy[-1]>0.:
                iwc_hy[time_ind] = 10.**6 * np.nansum(mass_hy / sample_vol[time_ind, :]) # g m^-3
                z_hy = 1.e12 * (0.174/0.93) * (6./np.pi/0.934)**2 * np.nansum(mass_particle**2*particle_count[time_ind, :]/sample_vol[time_ind, :]) # mm^6 m^-3
                sol = least_squares(calc_chisquare, x0, method='lm',ftol=1e-9,xtol=1e-9, max_nfev=int(1e6),\
                                    args=(Nt,iwc_hy[time_ind],z_hy,bin_mid,bin_width,0.00528,2.1)) # sove the gamma params using least squares minimziation
                N0_hy[time_ind] = sol.x[0]; mu_hy[time_ind] = sol.x[1]; lam_hy[time_ind] = sol.x[2]
                asr_hy[time_ind] = np.sum(aspect_ratio[time_ind, :] * mass_hy / sample_vol[time_ind, :]) / np.sum(mass_hy / sample_vol[time_ind, :]) # mass-weighted aspect ratio
                rhoe_hy[time_ind] = (iwc_hy[time_ind] / 10.**6) / vol # effective density from Chase et al. (2018) [g cm**-3]
                rho_hy[time_ind, :] = (mass_hy / particle_count[time_ind, :]) / (np.pi / 6.) / (bin_mid/10.)**3. # rho(D) following Heymsfield et al. (2003) [g cm**-3]
                dm_hy[time_ind] = 10. * np.sum((bin_mid/10.) * mass_hy / sample_vol[time_ind, :]) / np.sum(mass_hy / sample_vol[time_ind, :]) # mass-weighted mean D from Chase et al. (2020) [mm]
                if cumMass_hy[0]>=0.5*cumMass_hy[-1]:
                    dmm_hy[time_ind] = bin_mid[0]
                else:
                    dmm_hy[time_ind] = bin_mid[np.where(cumMass_hy>0.5*cumMass_hy[-1])[0][0]-1]

            # Habit-specific products
            mass_hab = np.sum(mass_tile * habit_count[time_ind, :, :], axis=1) # g (binned)
            cumMass_hab = np.nancumsum(mass_hab)
            if cumMass_hab[-1]>0.:
                if cumMass_hab[0]>=0.5*cumMass_hab[-1]:
                    dmm_hab[time_ind] = bin_mid[0]
                else:
                    dmm_hab[time_ind] = bin_mid[np.where(cumMass_hab>0.5*cumMass_hab[-1])[0][0]-1]
            iwc_hab[time_ind] = 10.**6 * np.nansum(mass_hab / sample_vol[time_ind, :]) # g m^-3
            asr_hab[time_ind] = np.sum(aspect_ratio[time_ind, :] * mass_hab / sample_vol[time_ind, :]) / np.sum(mass_hab / sample_vol[time_ind, :]) # mass-weighted aspect ratio
            rhoe_hab[time_ind] = (iwc_hab[time_ind] / 10.**6) / vol # effective density from Chase et al. (2018) [g cm**-3]
            rho_hab[time_ind, :] = (mass_hab / particle_count[time_ind, :]) / (np.pi / 6.) / (bin_mid/10.)**3. # rho(D) following Heymsfield et al. (2003) [g cm**-3]
            dm_hab[time_ind] = 10. * np.sum((bin_mid/10.) * mass_hab / sample_vol[time_ind, :]) / np.sum(mass_hab / sample_vol[time_ind, :]) # mass-weighted mean D from Chase et al. (2020) [mm]

    mu_bf = np.ma.masked_where(N0_bf==0., mu_bf)
    mu_hy = np.ma.masked_where(N0_hy==0., mu_hy)
    lam_bf = np.ma.masked_where(N0_bf==0., lam_bf)
    lam_hy = np.ma.masked_where(N0_hy==0., lam_hy)
    N0_bf = np.ma.masked_where(N0_bf==0., N0_bf)
    N0_hy = np.ma.masked_where(N0_hy==0., N0_hy)
    dmm_bf = np.ma.masked_where(dmm_bf==0., dmm_bf)
    dmm_hy = np.ma.masked_where(dmm_hy==0., dmm_hy)
    dmm_hab = np.ma.masked_where(dmm_hab==0., dmm_hab)
    dm_bf = np.ma.masked_where(dm_bf==0., dm_bf)
    dm_hy = np.ma.masked_where(dm_hy==0., dm_hy)
    dm_hab = np.ma.masked_where(dm_hab==0., dm_hab)
    asr_nw = np.ma.masked_where(np.ma.getmask(dmm_bf), asr_nw)
    asr_bf = np.ma.masked_where(np.ma.getmask(dmm_bf), asr_bf)
    asr_hy = np.ma.masked_where(np.ma.getmask(dmm_hy), asr_hy)
    asr_hab = np.ma.masked_where(np.ma.getmask(asr_hab), iwc_hab)
    rhoe_bf = np.ma.masked_where(np.ma.getmask(dmm_bf), rhoe_bf)
    rhoe_hy = np.ma.masked_where(np.ma.getmask(dmm_hy), rhoe_hy)
    rhoe_hab = np.ma.masked_where(np.ma.getmask(dmm_hab), rhoe_hab)
    iwc_bf = np.ma.masked_where(np.ma.getmask(dmm_bf), iwc_bf)
    iwc_hy = np.ma.masked_where(np.ma.getmask(dmm_hy), iwc_hy)
    iwc_hab = np.ma.masked_where(np.ma.getmask(dmm_hab), iwc_hab)
    rho_bf = np.ma.masked_where(rho_bf==0., rho_bf)
    rho_hy = np.ma.masked_where(rho_hy==0., rho_hy)
    rho_hab = np.ma.masked_where(rho_hab==0., rho_hab)

    return (N0_bf, N0_hy, mu_bf, mu_hy, lam_bf, lam_hy, iwc_bf, iwc_hy, iwc_hab, asr_nw, asr_bf, asr_hy, asr_hab, dmm_bf, dmm_hy, dmm_hab,\
            dm_bf, dm_hy, dm_hab, rho_bf, rho_hy, rho_hab, rhoe_bf, rhoe_hy, rhoe_hab)

def calc_riming(p3psd, Z, matchedZ_W, matchedZ_Ka, matchedZ_Ku, matchedZ_X, compute_fits=False):
    x0 = [1.e-1, -1., 5.] # initial guess for N0 [cm**-4], mu, lambda [cm**-1]

    rmass = np.zeros(len(p3psd['time']))
    rfrac = np.zeros(len(p3psd['time']))
    a_coeff = np.zeros(len(p3psd['time']))
    b_coeff = np.zeros(len(p3psd['time']))
    Nw = np.zeros(len(p3psd['time']))
    N0 = np.zeros(len(p3psd['time']))
    mu = np.zeros(len(p3psd['time']))
    lam = np.zeros(len(p3psd['time']))
    iwc = np.zeros(len(p3psd['time']))
    asr = np.zeros(len(p3psd['time']))
    dm = np.zeros(len(p3psd['time']))
    dmm = np.zeros(len(p3psd['time']))
    rho_eff = np.zeros(len(p3psd['time']))
    dfr_KuKa = np.zeros(len(p3psd['time']))
    error = np.zeros((len(p3psd['time']), len(p3psd['riming_mass_array'])))

    for i in range(len(p3psd['time'])):
        # loop through the different possible riming masses
        for j in range(len(p3psd['riming_mass_array'])):
            if (matchedZ_W is not None) and (np.ma.is_masked(matchedZ_W[i]) is False) and (np.ma.is_masked(p3psd['dbz_W'][i, :]) is False):
                error[i, j] = error[i, j] + np.abs(matchedZ_W[i] - p3psd['dbz_W'][i, j])
            if (matchedZ_Ka is not None) and (np.ma.is_masked(matchedZ_Ka[i]) is False) and (np.ma.is_masked(p3psd['dbz_Ka'][i, :]) is False):
                error[i, j] = error[i, j] + np.abs(matchedZ_Ka[i] - p3psd['dbz_Ka'][i, j])
            if (matchedZ_Ku is not None) and (np.ma.is_masked(matchedZ_Ku[i]) is False) and (np.ma.is_masked(p3psd['dbz_Ku'][i, :]) is False):
                error[i, j] = error[i, j] + np.abs(matchedZ_Ku[i] - p3psd['dbz_Ku'][i, j])
            if (matchedZ_X is not None) and (np.ma.is_masked(matchedZ_X[i]) is False) and (np.ma.is_masked(p3psd['dbz_X'][i, :]) is False):
                error[i, j] = error[i, j] + np.abs(matchedZ_X[i] - p3psd['dbz_X'][i, j])

        if np.sum(error[i, :])>0.:
            rmass[i] = p3psd['riming_mass_array'][np.argmin(error[i, :])]
            a_coeff[i] = p3psd['a_coeff_array'][np.argmin(error[i, :])]
            b_coeff[i] = p3psd['b_coeff_array'][np.argmin(error[i, :])]

            if p3psd['count'][i, :].count()==p3psd['count'].shape[1]: # time period is not masked...continue on
                Nt = 1000.*np.nansum(p3psd['count'][i, :]/p3psd['sv'][i, :]) # concentration [L**-1]
                mass_particle = a_coeff[i] * (p3psd['bin_mid']/10.)**b_coeff[i] # particle mass [g]
                mass = mass_particle * p3psd['count'][i, :] # g (binned)
                cumMass = np.nancumsum(mass)
                if cumMass[-1]>0.:
                    # Nw (follows Chase et al. 2021)
                    # [log10(m**-3 mm**-1)]
                    D_melt = ((6. * mass_particle) / (np.pi * 0.997))**(1./3.)
                    Nw[i] = np.log10((1e5) * (4.**4 / 6) * np.nansum(
                        D_melt**3 * p3psd['ND'][i, :] * p3psd['bin_width'])**5 / np.nansum(
                        D_melt**4 * p3psd['ND'][i, :] * p3psd['bin_width'])**4)

                    # IWC
                    iwc[i] = 10.**6 * np.nansum(mass / p3psd['sv'][i, :]) # g m^-3

                    # DFR
                    dfr_KuKa[i] = p3psd[
                        'dbz_Ku'][i, np.argmin(error[i, :])] - p3psd[
                        'dbz_Ka'][i, np.argmin(error[i, :])] # dB

                    # Optionally compute N0, mu, lambda
                    if compute_fits:
                        z = 10.**(p3psd['dbz_X'][i,np.argmin(error[i, :])]/10.) # mm^6 m^-3

                        # solve gamma params using least squares minimziation
                        sol = least_squares(
                            calc_chisquare, x0, method='lm', ftol=1e-9, xtol=1e-9,
                            max_nfev=int(1e6), args=(
                                Nt, iwc[i], z, p3psd['bin_mid'], p3psd['bin_width'],
                                a_coeff[i], b_coeff[i], np.argmin(error[i, :])))
                        N0[i] = sol.x[0]; mu[i] = sol.x[1]; lam[i] = sol.x[2]

                    # Mass-weighted mean aspect ratio
                    asr[i] = np.sum(
                        p3psd['aspect_ratio'][i, :] * mass / p3psd['sv'][i, :]) / np.sum(
                        mass / p3psd['sv'][i, :])

                    # Bulk riming fraction (see Eqn 1 of Morrison and Grabowski
                    # [2010, https://doi.org/10.1175/2010JAS3250.1] for binned version)
                    rfrac[i] = np.sum(
                        np.squeeze(Z.rimefrac[0, :, np.argmin(error[i, :])])
                        * mass / p3psd['sv'][i, :]) / np.nansum(
                        mass / p3psd['sv'][i, :]) # SUM(rimed mass conc)/iwc

                    # Effective density (follows Chase et al. 2018)
                    vol = (np.pi / 6.) * np.sum(
                        0.6 * ((p3psd['bin_mid']/10.)**3.) * p3psd['count'][i, :]
                        / p3psd['sv'][i, :]) # [cm**3 / cm**3]
                    rho_eff[i] = (iwc[i] / 10.**6) / vol # [g cm**-3]

                    # Mass-weighted mean diameter (follows Chase et al. 2020)
                    # M3/M2 if b==2, more generally M(b+1)/Mb
                    dm[i] = 10. * np.sum(
                        (p3psd['bin_mid']/10.) * mass / p3psd['sv'][i, :]) / np.sum(
                        mass / p3psd['sv'][i, :]) # [mm]

                    # Mass-weighted median diameter [mm]
                    if cumMass[0]>=0.5*cumMass[-1]:
                        dmm[i] = p3psd['bin_mid'][0]
                    else:
                        dmm[i] = p3psd[
                            'bin_mid'][np.where(cumMass>0.5*cumMass[-1])[0][0]-1]

    p3psd['sclwp'] = np.ma.masked_where(np.sum(error, axis=1)==0., rmass)
    p3psd['riming_frac'] = np.ma.masked_where(np.sum(error, axis=1)==0., rfrac)
    p3psd['a_coeff'] = np.ma.masked_where(np.sum(error, axis=1)==0., a_coeff)
    p3psd['b_coeff'] = np.ma.masked_where(np.sum(error, axis=1)==0., b_coeff)
    if compute_fits:
        p3psd['mu_ls'] = np.ma.masked_where(N0==0., mu)
        p3psd['lambda_ls'] = np.ma.masked_where(N0==0., lam)
        p3psd['N0_ls'] = np.ma.masked_where(N0==0., N0)
    p3psd['Nw_ls'] = np.ma.masked_where(np.sum(error, axis=1)==0., Nw)
    p3psd['iwc_ls'] = np.ma.masked_where(np.sum(error, axis=1)==0., iwc)
    p3psd['mean_aspect_ratio_ls'] = np.ma.masked_where(np.sum(error, axis=1)==0., asr)
    p3psd['dm_ls'] = np.ma.masked_where(np.sum(error, axis=1)==0., dm)
    p3psd['dmm_ls'] = np.ma.masked_where(np.sum(error, axis=1)==0., dmm)
    p3psd['eff_density_ls'] = np.ma.masked_where(np.sum(error, axis=1)==0., rho_eff)
    p3psd['dfr_KuKa_ls'] = np.ma.masked_where(np.sum(error, axis=1)==0., dfr_KuKa)

    return p3psd

def calc_chisquare(
    x, Nt_obs, iwc_obs, z_obs, bin_mid, bin_width, a_coefficient, b_coefficient,
    rime_ind=None, exponential=False):
    '''
    Compute gamma fit parameters for the PSD.
    Follows McFarquhar et al. (2015) by finding N0-mu-lambda minimizing first
    (Nt), third (mass), sixth (reflectivity) moments.
    Inputs:
        x: N0, mu, lambda to test on the minimization procedure
        Nt_obs: Observed number concentration [L^-1]
        iwc_obs: Observed IWC using an assumed m-D relation [g m**-3]
        z_obs: Observed Z (following Hogan et al. 2012 definition) using assumed m-D relation [mm**6 m**-3]
        bin_mid: Midpoints for the binned particle size [mm]
        bin_width: Bin width for the binned particle size [cm]
        a_coefficient: Prefactor component to the assumed m-D reltation [cm**-b]
        b_coefficient: Exponent component to the assumed m-D reltation
        rime_ind (optional, for LS products only): Riming category index to use for the reflectivity moment
        exponential: Boolean, True if setting mu=0 for the fit (exponential form)
    Outputs:
        chi_square: Chi-square value for the provided N0-mu-lambda configuration
    '''
    Dmax = bin_mid / 10. # midpoint in cm
    dD = bin_width # bin width in cm
    mass_particle = a_coefficient * Dmax**b_coefficient # binned particle mass [g]

    if exponential: # exponential form with mu=0
        ND_fit = x[0] * np.exp(-x[2]*Dmax)
    else: # traditional gamma function with variable mu
        ND_fit = x[0] * Dmax**x[1] * np.exp(-x[2]*Dmax)

    Nt_fit = 1000.*np.nansum(ND_fit*dD) # L**-1
    iwc_fit = 10.**6  * np.nansum(mass_particle*ND_fit*dD) # g m**-3
    if rime_ind is not None:
        Z_fit = forward_Z() #initialize class
        Z_fit.set_PSD(PSD=ND_fit[np.newaxis,:]*10.**8, D=Dmax/100., dD=dD/100., Z_interp=True) # get the PSD in the format to use in the routine (mks units)
        Z_fit.load_split_L15() # Load the leinonen output
        Z_fit.fit_sigmas(Z_interp=True) # Fit the backscatter cross-sections
        Z_fit.calc_Z() # Calculate Z...outputs are Z.Z_x, Z.Z_ku, Z.Z_ka, Z.Z_w for the four radar wavelengths
        z_fit = 10.**(Z_fit.Z_x[0, rime_ind] / 10.) # mm**6 m**-3
    else:
        z_fit = 1.e12 * (0.174/0.93) * (6./np.pi/0.934)**2 * np.nansum(mass_particle**2*ND_fit*dD) # mm**6 m**-3

    csq_Nt = ((Nt_obs-Nt_fit) / np.sqrt(Nt_obs*Nt_fit))**2
    csq_iwc = ((iwc_obs-iwc_fit) / np.sqrt(iwc_obs*iwc_fit))**2
    csq_z = ((z_obs-z_fit) / np.sqrt(z_obs*z_fit))**2
    chi_square = [csq_Nt, csq_iwc, csq_z]

    return chi_square