"""
Module: System
    This module shall be used to implement subclasses of system. It wraps all information needed and generated by a simulation.
"""

import numpy as np
import pandas as pd

pd.options.mode.use_inf_as_na = True

from ensembler.util import dataStructure as data
from ensembler.util.ensemblerTypes import samplerCls, conditionCls
from ensembler.util.ensemblerTypes import Union, Iterable, NoReturn, Number

from ensembler.potentials._basicPotentials import _potential1DClsPerturbed as _perturbedPotentialCls
from ensembler.potentials.OneD import linearCoupledPotentials
from ensembler.samplers.stochastic import metropolisMonteCarloIntegrator

from ensembler.system.basic_system import system

class perturbedSystem(system):
    """

    """
    name = "perturbed system"
    # Lambda Dependend Settings
    state = data.lambdaState
    currentState: data.lambdaState
    potential: _perturbedPotentialCls

    # current lambda
    _currentLambda: Number = np.nan
    _currentdHdLambda: Number = np.nan

    """
    Attributes
    """
    @property
    def lam(self) -> Number:
        return self._currentLambda

    @lam.setter
    def lam(self, lam: Number):
        if lam < 0.0 or lam > 1.0:
            raise ValueError(f"Variable lam = {lam}."\
                              "It cannot be lower than 0 or larger than 1.")
        self._currentLambda = lam
        self.potential.set_lambda(lam=self._currentLambda)
        self.update_current_state()

    def set_lambda(self, lam:Number):
        self.lam = lam

    """
    Magic
    """

    def __init__(self, potential: _perturbedPotentialCls=linearCoupledPotentials(), sampler: samplerCls=metropolisMonteCarloIntegrator(),
                 conditions: Iterable[conditionCls] = [],
                 temperature: float = 298.0, start_position: (Iterable[Number] or float) = None, lam: float = 0.0):
        """
            __init__
                construct a eds-System that can be used to manage a simulation.

        Parameters
        ----------
        potential:  pot.envelopedPotential, optional
            potential function class to be explored by sampling
        sampler: sampler, optional
            sampling method, that allows exploring the potential function
        conditions: Iterable[condition], optional
            conditions that shall be applied to the system.
        temperature: float, optional
            The temperature of the system (default: 298K)
        start_position:
            starting position for the simulation and setup of the system.
        lam: Number, optional
            the value of the copuling lambda
        """
        super().__init__(potential=potential, sampler=sampler, conditions=conditions, temperature=temperature,
                         start_position=start_position)

        self.lam = lam
        self.update_current_state()

    """
    Overwrite Functions to adapt to EDS
    """

    def set_current_state(self,
                          current_position: Union[Number, Iterable[Number]],
                          current_velocities: Union[Number, Iterable[Number]] = 0,
                          current_force: Union[Number, Iterable[Number]] = 0,
                          current_temperature: Union[Number, Iterable[Number]] = 298,
                          current_lambda: Union[Number, Iterable[Number]] = 0,
                          current_dHdLambda: Union[Number, Iterable[Number]] = 0):
        """
            set_current_state
                set s the current state to the given variables.

        Parameters
        ----------
        Parameters
        ----------
        current_position: Union[Number, Iterable[Number]]
            new current system position
        current_velocities: Union[Number, Iterable[Number]], optional
            new current system velocity. (default: 0)
        current_force: Union[Number, Iterable[Number]], optional
            new current system force. (default: 0)
        current_temperature: Union[Number, Iterable[Number]], optional
            new current system temperature. (default: 298)
        current_lam: Union[Number, Iterable[Number]],
            The new lambda value (default: 0)
        current_dHdLam: Union[Number, Iterable[Number]],
            The new dHdLam(default: 0)
        """
        self._currentPosition = current_position
        self._currentForce = current_force
        self._currentVelocities = current_force
        self._currentTemperature = current_temperature
        self._currentLambda = current_lambda
        self._currentdHdLambda = current_dHdLambda

        self._update_energies()
        self._update_dHdLambda()
        self.update_current_state()

    def update_system_properties(self) -> NoReturn:
        """
            updateSystemProperties
                update all system properties
        """
        self._update_energies()
        self._update_temperature()
        self._update_dHdLambda()

    def update_current_state(self):
        """
        updateCurrentState
                This function updates the current state from the _current Variables.

        """
        self._currentState = self.state(position=self._currentPosition, temperature=self._currentTemperature,
                                        total_system_energy=self._currentTotE,
                                        total_potential_energy=self._currentTotPot,
                                        total_kinetic_energy=self._currentTotKin,
                                        dhdpos=self._currentForce, velocity=self._currentVelocities,
                                        lam=self._currentLambda, dhdlam=self._currentdHdLambda)

    def append_state(self, new_position: Union[Number, Iterable[Number]], new_velocity: Union[Number, Iterable[Number]], new_forces: Union[Number, Iterable[Number]],
                     new_lambda: Number) -> NoReturn:
        """
            append_state
                Append a new state to the trajectory.

        Parameters
        ----------
        new_position: Union[Number, Iterable[Number]]
            new position for the system
        new_velocity: Union[Number, Iterable[Number]]
            new velocity for the system
        new_forces: Union[Number, Iterable[Number]]
            new forces for the system
        new_lambda: Union[Number, Iterable[Number]]
            new lambda for the system

        """
        self._currentPosition = new_position
        self._currentVelocities = new_velocity
        self._currentForce = new_forces
        self._currentLambda = new_lambda

        self._update_temperature()
        self._update_energies()
        self._update_dHdLambda()
        self.update_current_state()

        self._trajectory.append(self.current_state)

    """
    Functionality
    """

    def _update_dHdLambda(self) -> Number:
        """
            _update_dHdlambda
                update the current dHdLambda value

        Returns
        -------
        Number
            dHdlambda

        """
        self._currentdHdLambda = self.potential.dvdlam(self._currentPosition)
        self.update_current_state()
        return self._currentdHdLambda


# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""
Models for doing PSF/PRF fitting photometry on image data.
"""

from __future__ import division
import numpy as np
from astropy.table import Table
from astropy.modeling import models, Parameter, Fittable2DModel
from astropy.modeling.fitting import LevMarLSQFitter
from astropy.nddata.utils import subpixel_indices
from ..utils import mask_to_mirrored_num
from ..extern.nddata_compat import extract_array

__all__ = ['IntegratedGaussianPRF', 'PRFAdapter', 'prepare_psf_model',
           'get_grouped_psf_model']

class DiscretePRF(Fittable2DModel):
    """
    A discrete Pixel Response Function (PRF) model.

    The discrete PRF model stores images of the PRF at different
    subpixel positions or offsets as a lookup table. The resolution is
    given by the subsampling parameter, which states in how many
    subpixels a pixel is divided.

    In the typical case of wanting to create a PRF from an image with
    many point sources, use the `~DiscretePRF.create_from_image` method,
    rather than directly initializing this class.

    The discrete PRF model class in initialized with a 4 dimensional
    array, that contains the PRF images at different subpixel positions.
    The definition of the axes is as following:

        1. Axis: y subpixel position
        2. Axis: x subpixel position
        3. Axis: y direction of the PRF image
        4. Axis: x direction of the PRF image

    The total array therefore has the following shape
    (subsampling, subsampling, prf_size, prf_size)

    Parameters
    ----------
    prf_array : ndarray
        Array containing PRF images.
    normalize : bool
        Normalize PRF images to unity.  Equivalent to saying there is
        *no* flux outside the bounds of the PRF images.
    subsampling : int, optional
        Factor of subsampling. Default = 1.

    Notes
    -----
    See :ref:`psf-terminology` for more details on the distinction
    between PSF and PRF as used in this module.
    """

    flux = Parameter('flux')
    x_0 = Parameter('x_0')
    y_0 = Parameter('y_0')

    def __init__(self, prf_array, normalize=True, subsampling=1):
        # Array shape and dimension check
        if subsampling == 1:
            if prf_array.ndim == 2:
                prf_array = np.array([[prf_array]])
        if prf_array.ndim != 4:
            raise TypeError('Array must have 4 dimensions.')
        if prf_array.shape[:2] != (subsampling, subsampling):
            raise TypeError('Incompatible subsampling and array size')
        if np.isnan(prf_array).any():
            raise Exception("Array contains NaN values. Can't create PRF.")

        # Normalize if requested
        if normalize:
            for i in range(prf_array.shape[0]):
                for j in range(prf_array.shape[1]):
                    prf_array[i, j] /= prf_array[i, j].sum()

        # Set PRF asttributes
        self._prf_array = prf_array
        self.subsampling = subsampling

        constraints = {'fixed': {'x_0': True, 'y_0': True}}
        x_0 = 0
        y_0 = 0
        flux = 1
        super(DiscretePRF, self).__init__(n_models=1, x_0=x_0, y_0=y_0,
                                          flux=flux, **constraints)
        self.fitter = LevMarLSQFitter()

    @property
    def prf_shape(self):
        """Shape of the PRF image."""

        return self._prf_array.shape[-2:]

    def evaluate(self, x, y, flux, x_0, y_0):
        """
        Discrete PRF model evaluation.

        Given a certain position and flux the corresponding image of the
        PSF is chosen and scaled to the flux. If x and y are outside the
        boundaries of the image, zero will be returned.

        Parameters
        ----------
        x : float
            x coordinate array in pixel coordinates.
        y : float
            y coordinate array in pixel coordinates.
        flux : float
            Model flux.
        x_0 : float
            x position of the center of the PRF.
        y_0 : float
            y position of the center of the PRF.
        """

        # Convert x and y to index arrays
        x = (x - x_0 + 0.5 + self.prf_shape[1] // 2).astype('int')
        y = (y - y_0 + 0.5 + self.prf_shape[0] // 2).astype('int')

        # Get subpixel indices
        y_sub, x_sub = subpixel_indices((y_0, x_0), self.subsampling)

        # Out of boundary masks
        x_bound = np.logical_or(x < 0, x >= self.prf_shape[1])
        y_bound = np.logical_or(y < 0, y >= self.prf_shape[0])
        out_of_bounds = np.logical_or(x_bound, y_bound)

        # Set out of boundary indices to zero
        x[x_bound] = 0
        y[y_bound] = 0
        result = flux * self._prf_array[int(y_sub), int(x_sub)][y, x]

        # Set out of boundary values to zero
        result[out_of_bounds] = 0
        return result

    @classmethod
    def create_from_image(cls, imdata, positions, size, fluxes=None,
                          mask=None, mode='mean', subsampling=1,
                          fix_nan=False):
        """
        Create a discrete point response function (PRF) from image data.

        Given a list of positions and size this function estimates an
        image of the PRF by extracting and combining the individual PRFs
        from the given positions.

        NaN values are either ignored by passing a mask or can be
        replaced by the mirrored value with respect to the center of the
        PRF.

        Note that if fluxes are *not* specified explicitly, it will be
        flux estimated from an aperture of the same size as the PRF
        image. This does *not* account for aperture corrections so often
        will *not* be what you want for anything other than quick-look
        needs.

        Parameters
        ----------
        imdata : array
            Data array with the image to extract the PRF from
        positions : List or array or `~astropy.table.Table`
            List of pixel coordinate source positions to use in creating
            the PRF.  If this is a `~astropy.table.Table` it must have
            columns called ``x_0`` and ``y_0``.
        size : odd int
            Size of the quadratic PRF image in pixels.
        mask : bool array, optional
            Boolean array to mask out bad values.
        fluxes : array, optional
            Object fluxes to normalize extracted PRFs. If not given (or
            None), the flux is estimated from an aperture of the same
            size as the PRF image.
        mode : {'mean', 'median'}
            One of the following modes to combine the extracted PRFs:
                * 'mean':  Take the pixelwise mean of the extracted PRFs.
                * 'median':  Take the pixelwise median of the extracted PRFs.
        subsampling : int
            Factor of subsampling of the PRF (default = 1).
        fix_nan : bool
            Fix NaN values in the data by replacing it with the mirrored
            value. Assuming that the PRF is symmetrical.

        Returns
        -------
        prf : `photutils.psf.sandbox.DiscretePRF`
            Discrete PRF model estimated from data.
        """

        # Check input array type and dimension.
        if np.iscomplexobj(imdata):
            raise TypeError('Complex type not supported')
        if imdata.ndim != 2:
            raise ValueError('{0}-d array not supported. '
                             'Only 2-d arrays supported.'.format(imdata.ndim))
        if size % 2 == 0:
            raise TypeError("Size must be odd.")

        if fluxes is not None and len(fluxes) != len(positions):
            raise TypeError('Position and flux arrays must be of equal '
                            'length.')

        if mask is None:
            mask = np.isnan(imdata)

        if isinstance(positions, (list, tuple)):
            positions = np.array(positions)

        if isinstance(positions, Table) or \
            (isinstance(positions, np.ndarray) and
             positions.dtype.names is not None):
            # One can do clever things like
            # positions['x_0', 'y_0'].as_array().view((positions['x_0'].dtype,
            #                                          2))
            # but that requires positions['x_0'].dtype is
            # positions['y_0'].dtype.
            # Better do something simple to allow type promotion if required.
            pos = np.empty((len(positions), 2))
            pos[:, 0] = positions['x_0']
            pos[:, 1] = positions['y_0']
            positions = pos

        if isinstance(fluxes, (list, tuple)):
            fluxes = np.array(fluxes)

        if mode == 'mean':
            combine = np.ma.mean
        elif mode == 'median':
            combine = np.ma.median
        else:
            raise Exception('Invalid mode to combine prfs.')

        data_internal = np.ma.array(data=imdata, mask=mask)
        prf_model = np.ndarray(shape=(subsampling, subsampling, size, size))
        positions_subpixel_indices = \
            np.array([subpixel_indices(_, subsampling) for _ in positions],
                     dtype=np.int)

        for i in range(subsampling):
            for j in range(subsampling):
                extracted_sub_prfs = []
                sub_prf_indices = np.all(positions_subpixel_indices == [j, i],
                                         axis=1)
                positions_sub_prfs = positions[sub_prf_indices]
                for k, position in enumerate(positions_sub_prfs):
                    x, y = position
                    extracted_prf = extract_array(data_internal, (size, size),
                                                  (y, x))
                    # Check shape to exclude incomplete PRFs at the boundaries
                    # of the image
                    if (extracted_prf.shape == (size, size) and
                            np.ma.sum(extracted_prf) != 0):
                        # Replace NaN values by mirrored value, with respect
                        # to the prf's center
                        if fix_nan:
                            prf_nan = extracted_prf.mask
                            if prf_nan.any():
                                if (prf_nan.sum() > 3 or
                                        prf_nan[size // 2, size // 2]):
                                    continue
                                else:
                                    extracted_prf = mask_to_mirrored_num(
                                        extracted_prf, prf_nan,
                                        (size // 2, size // 2))
                        # Normalize and add extracted PRF to data cube
                        if fluxes is None:
                            extracted_prf_norm = (np.ma.copy(extracted_prf) /
                                                  np.ma.sum(extracted_prf))
                        else:
                            fluxes_sub_prfs = fluxes[sub_prf_indices]
                            extracted_prf_norm = (np.ma.copy(extracted_prf) /
                                                  fluxes_sub_prfs[k])
                        extracted_sub_prfs.append(extracted_prf_norm)
                    else:
                        continue
                prf_model[i, j] = np.ma.getdata(
                    combine(np.ma.dstack(extracted_sub_prfs), axis=2))
        return cls(prf_model, subsampling=subsampling)

class IntegratedGaussianPRF(Fittable2DModel):
    r"""
    Circular Gaussian model integrated over pixels. Because it is
    integrated, this model is considered a PRF, *not* a PSF (see
    :ref:`psf-terminology` for more about the terminology used here.)

    This model is a Gaussian *integrated* over an area of ``1`` (in
    units of the model input coordinates, e.g. 1 pixel).  This is in
    contrast to the apparently similar
    `astropy.modeling.functional_models.Gaussian2D`, which is the value
    of a 2D Gaussian *at* the input coordinates, with no integration.
    So this model is equivalent to assuming the PSF is Gaussian at a
    *sub-pixel* level.

    Parameters
    ----------
    sigma : float
        Width of the Gaussian PSF.
    flux : float (default 1)
        Total integrated flux over the entire PSF
    x_0 : float (default 0)
        Position of the peak in x direction.
    y_0 : float (default 0)
        Position of the peak in y direction.

    Notes
    -----
    This model is evaluated according to the following formula:

        .. math::

            f(x, y) =
                \frac{F}{4}
                \left[
                {\rm erf} \left(\frac{x - x_0 + 0.5}
                {\sqrt{2} \sigma} \right) -
                {\rm erf} \left(\frac{x - x_0 - 0.5}
                {\sqrt{2} \sigma} \right)
                \right]
                \left[
                {\rm erf} \left(\frac{y - y_0 + 0.5}
                {\sqrt{2} \sigma} \right) -
                {\rm erf} \left(\frac{y - y_0 - 0.5}
                {\sqrt{2} \sigma} \right)
                \right]

    where ``erf`` denotes the error function and ``F`` the total
    integrated flux.
    """

    flux = Parameter(default=1)
    x_0 = Parameter(default=0)
    y_0 = Parameter(default=0)
    sigma = Parameter(default=1, fixed=True)

    _erf = None
    fit_deriv = None

    @property
    def bounding_box(self):
        halfwidth = 4 * self.sigma
        return ((int(self.y_0 - halfwidth), int(self.y_0 + halfwidth)),
                (int(self.x_0 - halfwidth), int(self.x_0 + halfwidth)))

    def __init__(self, sigma=sigma.default,
                 x_0=x_0.default, y_0=y_0.default, flux=flux.default,
                 **kwargs):
        if self._erf is None:
            from scipy.special import erf
            self.__class__._erf = erf

        super(IntegratedGaussianPRF, self).__init__(n_models=1, sigma=sigma,
                                                    x_0=x_0, y_0=y_0,
                                                    flux=flux, **kwargs)

    def evaluate(self, x, y, flux, x_0, y_0, sigma):
        """Model function Gaussian PSF model."""

        return (flux / 4 *
                ((self._erf((x - x_0 + 0.5) / (np.sqrt(2) * sigma)) -
                  self._erf((x - x_0 - 0.5) / (np.sqrt(2) * sigma))) *
                 (self._erf((y - y_0 + 0.5) / (np.sqrt(2) * sigma)) -
                  self._erf((y - y_0 - 0.5) / (np.sqrt(2) * sigma)))))

class PRFAdapter(Fittable2DModel):
    """
    A model that adapts a supplied PSF model to act as a PRF. It
    integrates the PSF model over pixel "boxes".  A critical built-in
    assumption is that the PSF model scale and location parameters are
    in *pixel* units.

    Parameters
    ----------
    psfmodel : a 2D model
        The model to assume as representative of the PSF
    renormalize_psf : bool
        If True, the model will be integrated from -inf to inf and
        re-scaled so that the total integrates to 1.  Note that this
        renormalization only occurs *once*, so if the total flux of
        ``psfmodel`` depends on position, this will *not* be correct.
    xname : str or None
        The name of the ``psfmodel`` parameter that corresponds to the
        x-axis center of the PSF.  If None, the model will be assumed to
        be centered at x=0.
    yname : str or None
        The name of the ``psfmodel`` parameter that corresponds to the
        y-axis center of the PSF.  If None, the model will be assumed to
        be centered at y=0.
    fluxname : str or None
        The name of the ``psfmodel`` parameter that corresponds to the
        total flux of the star.  If None, a scaling factor will be
        applied by the ``PRFAdapter`` instead of modifying the
        ``psfmodel``.

    Notes
    -----
    This current implementation of this class (using numerical
    integration for each pixel) is extremely slow, and only suited for
    experimentation over relatively few small regions.
    """

    flux = Parameter(default=1)
    x_0 = Parameter(default=0)
    y_0 = Parameter(default=0)

    def __init__(self, psfmodel, renormalize_psf=True, flux=flux.default,
                 x_0=x_0.default, y_0=y_0.default, xname=None, yname=None,
                 fluxname=None, **kwargs):

        self.psfmodel = psfmodel.copy()

        if renormalize_psf:
            from scipy.integrate import dblquad
            self._psf_scale_factor = 1. / dblquad(self.psfmodel,
                                                  -np.inf, np.inf,
                                                  lambda x: -np.inf,
                                                  lambda x: np.inf)[0]
        else:
            self._psf_scale_factor = 1

        self.xname = xname
        self.yname = yname
        self.fluxname = fluxname

        # these can be used to adjust the integration behavior. Might be
        # used in the future to expose how the integration happens
        self._dblquadkwargs = {}

        super(PRFAdapter, self).__init__(n_models=1, x_0=x_0, y_0=y_0,
                                         flux=flux, **kwargs)

    def evaluate(self, x, y, flux, x_0, y_0):
        """The evaluation function for PRFAdapter."""

        if self.xname is None:
            dx = x - x_0
        else:
            dx = x
            setattr(self.psfmodel, self.xname, x_0)

        if self.xname is None:
            dy = y - y_0
        else:
            dy = y
            setattr(self.psfmodel, self.yname, y_0)

        if self.fluxname is None:
            return (flux * self._psf_scale_factor *
                    self._integrated_psfmodel(dx, dy))
        else:
            setattr(self.psfmodel, self.yname, flux * self._psf_scale_factor)
            return self._integrated_psfmodel(dx, dy)

    def _integrated_psfmodel(self, dx, dy):
        from scipy.integrate import dblquad

        # infer type/shape from the PSF model.  Seems wasteful, but the
        # integration step is a *lot* more expensive so its just peanuts
        out = np.empty_like(self.psfmodel(dx, dy))
        outravel = out.ravel()
        for i, (xi, yi) in enumerate(zip(dx.ravel(), dy.ravel())):
            outravel[i] = dblquad(self.psfmodel,
                                  xi-0.5, xi+0.5,
                                  lambda x: yi-0.5, lambda x: yi+0.5,
                                  **self._dblquadkwargs)[0]
        return out

def prepare_psf_model(psfmodel, xname=None, yname=None, fluxname=None,
                      renormalize_psf=True):
    """
    Convert a 2D PSF model to one suitable for use with
    `psf_photometry`.

    The resulting model may be a composite model, but should have only
    the x, y, and flux related parameters un-fixed.

    Parameters
    ----------
    psfmodel : a 2D model
        The model to assume as representative of the PSF.
    xname : str or None
        The name of the ``psfmodel`` parameter that corresponds to the
        x-axis center of the PSF.  If None, the model will be assumed to
        be centered at x=0, and a new parameter will be added for the
        offset.
    yname : str or None
        The name of the ``psfmodel`` parameter that corresponds to the
        y-axis center of the PSF.  If None, the model will be assumed to
        be centered at x=0, and a new parameter will be added for the
        offset.
    fluxname : str or None
        The name of the ``psfmodel`` parameter that corresponds to the
        total flux of the star.  If None, a scaling factor will be added
        to the model.
    renormalize_psf : bool
        If True, the model will be integrated from -inf to inf and
        re-scaled so that the total integrates to 1.  Note that this
        renormalization only occurs *once*, so if the total flux of
        ``psfmodel`` depends on position, this will *not* be correct.

    Returns
    -------
    outmod : a model
        A new model ready to be passed into `psf_photometry`.
    """

    if xname is None:
        xinmod = models.Shift(0, name='x_offset')
        xname = 'offset_0'
    else:
        xinmod = models.Identity(1)
        xname = xname + '_2'
    xinmod.fittable = True

    if yname is None:
        yinmod = models.Shift(0, name='y_offset')
        yname = 'offset_1'
    else:
        yinmod = models.Identity(1)
        yname = yname + '_2'
    yinmod.fittable = True

    outmod = (xinmod & yinmod) | psfmodel

    if fluxname is None:
        outmod = outmod * models.Const2D(1, name='flux_scaling')
        fluxname = 'amplitude_3'
    else:
        fluxname = fluxname + '_2'

    if renormalize_psf:
        # we do the import here because other machinery works w/o scipy
        from scipy import integrate

        integrand = integrate.dblquad(psfmodel, -np.inf, np.inf,
                                      lambda x: -np.inf, lambda x: np.inf)[0]
        normmod = models.Const2D(1./integrand, name='renormalize_scaling')
        outmod = outmod * normmod

    # final setup of the output model - fix all the non-offset/scale
    # parameters
    for pnm in outmod.param_names:
        outmod.fixed[pnm] = pnm not in (xname, yname, fluxname)

    # and set the names so that psf_photometry knows what to do
    outmod.xname = xname
    outmod.yname = yname
    outmod.fluxname = fluxname

    # now some convenience aliases if reasonable
    outmod.psfmodel = outmod[2]
    if 'x_0' not in outmod.param_names and 'y_0' not in outmod.param_names:
        outmod.x_0 = getattr(outmod, xname)
        outmod.y_0 = getattr(outmod, yname)
    if 'flux' not in outmod.param_names:
        outmod.flux = getattr(outmod, fluxname)

    return outmod

def get_grouped_psf_model(template_psf_model, star_group):
    """
    Construct a joint PSF model which consists of a sum of PSF's templated on
    a specific model, but whose parameters are given by a table of objects.

    Parameters
    ----------
    template_psf_model : `astropy.modeling.Fittable2DModel` instance
        The model to use for *individual* objects.  Must have parameters named
        ``x_0``, ``y_0``, and ``flux``.
    star_group : `~astropy.table.Table`
        Table of stars for which the compound PSF will be constructed.  It
        must have columns named ``x_0``, ``y_0``, and ``flux_0``.

    Returns
    -------
    group_psf
        An `astropy.modeling` ``CompoundModel`` instance which is a sum of the
        given PSF models.
    """

    group_psf = None
    for i in range(len(star_group)):
        psf_to_add = template_psf_model.copy()
        psf_to_add.flux = star_group['flux_0'][i]
        psf_to_add.x_0 = star_group['x_0'][i]
        psf_to_add.y_0 = star_group['y_0'][i]

        if group_psf is None:
            # this is the first one only
            group_psf = psf_to_add
        else:
            group_psf += psf_to_add

    return group_psf


"""
Python implementation of the LiNGAM algorithms.
The LiNGAM Project: https://sites.google.com/site/sshimizu06/lingam
"""

import itertools
import warnings
from abc import ABCMeta, abstractmethod

import numpy as np
from sklearn.linear_model import LinearRegression
from sklearn.utils import check_array

from .bootstrap import BootstrapMixin
from .hsic import hsic_test_gamma
from .utils import (get_exo_variables, get_sink_variables,
                    predict_adaptive_lasso)

class _BaseLiNGAM(BootstrapMixin, metaclass=ABCMeta):
    """Base class for all LiNGAM algorithms."""

    def __init__(self, random_state=None):
        """Construct a _BaseLiNGAM model.

        Parameters
        ----------
        random_state : int, optional (default=None)
            random_state is the seed used by the random number generator.
        """
        self._random_state = random_state
        self._causal_order = None
        self._adjacency_matrix = None

    @abstractmethod
    def fit(self, X):
        """Subclasses should implement this method!
        Fit the model to X.

        Parameters
        ----------
        X : array-like, shape (n_samples, n_features)
            Training data, where n_samples is the number of samples
            and n_features is the number of features.

        Returns
        -------
        self : object
            Returns the instance itself.
        """

    def estimate_total_effect(self, X, from_index, to_index):
        """Estimate total effect using causal model.

        Parameters
        ----------
        X : array-like, shape (n_samples, n_features)
            Original data, where n_samples is the number of samples
            and n_features is the number of features.
        from_index :
            Index of source variable to estimate total effect.
        to_index :
            Index of destination variable to estimate total effect.

        Returns
        -------
        total_effect : float
            Estimated total effect.
        """
        # Check parameters
        X = check_array(X)

        # Check from/to causal order
        from_order = self._causal_order.index(from_index)
        to_order = self._causal_order.index(to_index)
        if from_order > to_order:
            warnings.warn(f'The estimated causal effect may be incorrect because '
                          f'the causal order of the destination variable (to_index={to_index}) '
                          f'is earlier than the source variable (from_index={from_index}).')

        # from_index + parents indices
        parents = np.where(np.abs(self._adjacency_matrix[from_index]) > 0)[0]
        predictors = [from_index]
        predictors.extend(parents)

        # Estimate total effect
        coefs = predict_adaptive_lasso(X, predictors, to_index)

        return coefs[0]

    def get_error_independence_p_values(self, X):
        """Calculate the p-value matrix of independence between error variables.

        Parameters
        ----------
        X : array-like, shape (n_samples, n_features)
            Original data, where n_samples is the number of samples
            and n_features is the number of features.

        Returns
        -------
        independence_p_values : array-like, shape (n_features, n_features)
            p-value matrix of independence between error variables.
        """
        # Check parameters
        X = check_array(X)
        n_samples = X.shape[0]
        n_features = X.shape[1]

        E = X - np.dot(self._adjacency_matrix, X.T).T
        p_values = np.zeros([n_features, n_features])
        for i, j in itertools.combinations(range(n_features), 2):
            _, p_value = hsic_test_gamma(np.reshape(E[:, i], [n_samples, 1]),
                                         np.reshape(E[:, j], [n_samples, 1]))
            p_values[i, j] = p_value
            p_values[j, i] = p_value

        return p_values

    def _estimate_adjacency_matrix(self, X, prior_knowledge=None):
        """Estimate adjacency matrix by causal order.

        Parameters
        ----------
        X : array-like, shape (n_samples, n_features)
            Training data, where n_samples is the number of samples
            and n_features is the number of features.
        prior_knowledge : array-like, shape (n_variables, n_variables), optional (default=None)
            Prior knowledge matrix.

        Returns
        -------
        self : object
            Returns the instance itself.
        """
        sink_vars = get_sink_variables(prior_knowledge)
        exo_vars = get_exo_variables(prior_knowledge)

        B = np.zeros([X.shape[1], X.shape[1]], dtype='float64')
        for i in range(1, len(self._causal_order)):
            target = self._causal_order[i]
            predictors = self._causal_order[:i]

            # target is not used for prediction if it is included in exogenous variables
            if target in exo_vars:
                continue

            # sink variables are not used as predictors
            predictors = [v for v in predictors if v not in sink_vars]

            B[target, predictors] = predict_adaptive_lasso(
                X, predictors, target)

        self._adjacency_matrix = B
        return self

    @property
    def causal_order_(self):
        """Estimated causal ordering.

        Returns
        -------
        causal_order_ : array-like, shape (n_features)
            The causal order of fitted model, where
            n_features is the number of features.
        """
        return self._causal_order

    @property
    def adjacency_matrix_(self):
        """Estimated adjacency matrix.

        Returns
        -------
        adjacency_matrix_ : array-like, shape (n_features, n_features)
            The adjacency matrix B of fitted model, where
            n_features is the number of features.
        """
        return self._adjacency_matrix


import functools
import multiprocessing
import os
import pathlib
import shelve

import numpy
import pandas
import seaborn
import ujson
import urllib

from collections import defaultdict

from IPython.display import Image, display

from astropy.coordinates import SkyCoord
from astropy import units as u
import astropy.io.fits as fits
import astropy.utils.data
from astropy.table import vstack
from astropy.timeseries import TimeSeries
from astropy.stats import sigma_clip

from astroquery.vizier import Vizier

from matplotlib import pyplot

astropy.utils.data.Conf.remote_timeout.set(60)

DATA_LOCATION = os.path.join(os.path.expanduser('~'), 'Documents', 'superwasp-data')
CACHE_LOCATION = os.path.join(DATA_LOCATION, 'cache')

SECONDS_PER_DAY = 60 * 60 * 24

MAIN_WORKFLOW = 7534
JUNK_WORKFLOW = 17313

if not os.path.exists(DATA_LOCATION):
    os.mkdir(DATA_LOCATION)
if not os.path.exists(CACHE_LOCATION):
    os.mkdir(CACHE_LOCATION)

def batches(i, batch_size=100):
    for x in range(int(len(i)/batch_size)+1):
        subset = i[x*batch_size:(x+1)*batch_size]
        if len(subset) == 0:
            return
        yield subset

class PandasDFWrapper(object):
    def limit(self, limit):
        return self.__class__(df=self.df[:limit])

    def cached_pandas_load(self, filename):
        cache_file_path = pathlib.Path(os.path.join(CACHE_LOCATION, '{}.pickle'.format(filename)))
        orig_file_path = pathlib.Path(os.path.join(DATA_LOCATION, filename))
        if cache_file_path.exists() and (
            not orig_file_path.exists() or
            cache_file_path.stat().st_mtime > orig_file_path.stat().st_mtime
        ):
            return (pandas.read_pickle(cache_file_path), cache_file_path)
        return (None, cache_file_path)

    def _mpapply(self, func, df=None, axis='rows'):
        """
        Splits the DataFrame and applies the function across a pool of worker processes.

        On experiment this can actually add a lot of overhead to execution time.
        But it might be worth it on really long running operations.
        """
        if df is None:
            df = self.df

        split_df = numpy.array_split(df, multiprocessing.cpu_count())

        results = []
        with multiprocessing.Pool() as pool:
            for part in split_df:
                results.append(pool.apply_async(part.apply, args=(func, axis)))
            return pandas.concat([r.get() for r in results])

class CoordinatesMixin(object):
    VSX_MAG_AMPLITUDE_FLAG = '('
    VSX_PERIOD_THRESHOLD = 0.1
    VSX_SEARCH_RADIUS = 2 * u.arcsec
    VSX_MAGNITUDE_LIMIT = 15

    @property
    def coords(self):
        return SkyCoord(self.df['SWASP ID'].replace(r'^1SWASP', '', regex=True).values, unit=(u.hour, u.deg))

    @property
    def fits_urls(self):
        return self.df['SWASP ID'].apply(lambda s: 'http://wasp.warwick.ac.uk/lcextract?{}'.format(
            urllib.parse.urlencode(
                {'objid': s.replace('1SWASP', '1SWASP ')},
                quote_via=urllib.parse.quote,
            )
        ))

    @property
    def fits(self):
        for swasp_id, url in zip(self.df['SWASP ID'], self.fits_urls):
            yield fits.open(url)

    @property
    def timeseries(self):
        for fits_file in self.fits:
            hjd_col = fits.Column(name='HJD', format='D', array=fits_file[1].data['TMID'] / 86400 + 2453005.5)
            lc_data = fits.BinTableHDU.from_columns(fits_file[1].data.columns + fits.ColDefs([hjd_col]))
            yield TimeSeries.read(lc_data, time_column='HJD', time_format='jd')

    @property
    def timeseries_folded(self):
        for period, timeseries in zip(self.df['Period'], self.timeseries):
            yield timeseries.fold(
                period=period * u.second,
            )

    def add_coords(self):
        if 'Coords' not in self.df:
            coords = self.coords
            self.df['_RAJ2000'] = coords.ra
            self.df['_DEJ2000'] = coords.dec

    def add_fits_urls(self):
        if 'FITS URL' not in self.df:
            self.df['FITS URL'] = self.fits_urls

    def _extend_epochs(self, ts, epochs=1):
        epoch_length = ts['time'].max() - ts['time'].min()
        ts_out = [ts]
        for i in range(epochs):
            ts_new = ts.copy()
            ts_new['time'] = ts_new['time'] + epoch_length * (i + 1)
            ts_out.append(ts_new)
        return vstack(ts_out)

    def plot(self, folded=False, clip=False, sigma=4, hue=None):
        if folded:
            self.add_classification_labels()
            if 'Period' not in self.df:
                self.df = self.df.merge(FoldedLightcurves().df, how='left')
            ts_iter = self.timeseries_folded
        else:
            plotted_ids = set()
            ts_iter = self.timeseries

        for (subject_id, row), ts in zip(self.df.iterrows(), ts_iter):
            if folded:
                ts = self._extend_epochs(ts)
            else:
                if row['SWASP ID'] in plotted_ids:
                    continue
                plotted_ids.add(row['SWASP ID'])

            if clip:
                ts_flux = sigma_clip(ts['TAMFLUX2'], sigma=sigma)
            else:
                ts_flux = ts['TAMFLUX2']

            ts_data = {
                'time': ts.time.jd,
                'flux': ts_flux,
                'camera': ts['CAMERA_ID'],
            }
            pyplot.figure()
            plot = seaborn.scatterplot(
                data=ts_data,
                x='time',
                y='flux',
                hue=hue,
                alpha=0.5,
                s=1,
                palette='Set2',
            )
            if folded:
                plot.set_title('{} Period {}s ({})'.format(
                    row['SWASP ID'],
                    row['Period'],
                    row['Classification Label'],
                ))
            else:
                plot.set_title(row['SWASP ID'])

    def _query_vsx_for_coord(self, coord, cache):
        coord_str = coord.to_string()
        if coord_str not in cache:
            cache[coord_str] = Vizier.query_region(
                coord,
                radius=self.VSX_SEARCH_RADIUS,
                catalog='B/vsx/vsx',
            )

        return cache[coord_str]

    def _coords_for_row(self, row, cache):
        if row['SWASP ID'] not in cache:
            cache[row['SWASP ID']] = SkyCoord(
                row['SWASP ID'].replace('1SWASP', ''),
                unit=(u.hour, u.deg)
            )

        return cache[row['SWASP ID']]

    def add_vsx_types(self):
        if self.df.index.name:
            orig_index_name = self.df.index.name
            self.df.reset_index(inplace=True)
        else:
            orig_index_name = None

        vsx_types, vsx_types_cache_file = self.cached_pandas_load('vsx_types')
        if vsx_types is None:
            vsx_results_dict = defaultdict(list)
            batch_size = 100
            result_map = {
                'VSX Period': 'Period',
                'VSX Type': 'Type',
                'VSX Name': 'Name',
                'VSX Mag Max': 'max',
                'VSX Mag Min': 'min',
                'VSX Mag Format': 'f_min',
            }

            with shelve.open(os.path.join(CACHE_LOCATION, 'vsx_cache')) as vsx_cache:
                with shelve.open(os.path.join(CACHE_LOCATION, 'coord_cache')) as coord_cache:
                    for i, (_, row) in enumerate(self.df.iterrows(), start=1):
                        if i % 100 == 0:
                            print('Processing row: {}'.format(i), end='\r')
                        vsx_query = self._query_vsx_for_coord(
                            self._coords_for_row(row, coord_cache),
                            vsx_cache
                        )
                        if vsx_query is None:
                            continue

                        period_min = (row['Period'] / SECONDS_PER_DAY) * (1 - self.VSX_PERIOD_THRESHOLD)
                        period_max = (row['Period'] / SECONDS_PER_DAY) * (1 + self.VSX_PERIOD_THRESHOLD)

                        for vsx_table in vsx_query:
                            for vsx_row in vsx_table:
                                if vsx_row['Period'] < period_min:
                                    continue
                                if vsx_row['Period'] > period_max:
                                    continue

                                # When max is actually a mean and min is actually an amplitude
                                if vsx_row['f_min'] == self.VSX_MAG_AMPLITUDE_FLAG:
                                    if (vsx_row['max'] + vsx_row['min']) > self.VSX_MAGNITUDE_LIMIT:
                                        continue
                                else:
                                    if vsx_row['min'] > self.VSX_MAGNITUDE_LIMIT:
                                        continue

                                vsx_results_dict['subject_id'].append(row['subject_id'])
                                for result_key, vsx_key in result_map.items():
                                    vsx_results_dict[result_key].append(vsx_row[vsx_key])

            vsx_types = pandas.DataFrame(vsx_results_dict)
            vsx_types.to_pickle(vsx_types_cache_file)

        if len(vsx_types.index) > 0:
            vsx_types['VSX Period'] = vsx_types['VSX Period'] * SECONDS_PER_DAY
            self.df = self.df.merge(
                vsx_types,
                left_on='subject_id',
                right_on='subject_id',
                how='left',
            )

        if orig_index_name:
            self.df.set_index('subject_id', inplace=True)

class ZooLookupMixin(object):
    @property
    def zoo_lookup(self):
        zoo_lookup, cache_file = self.cached_pandas_load('lookup.dat')
        if zoo_lookup is not None:
            return zoo_lookup

        zoo_lookup = pandas.read_csv(
            os.path.join(DATA_LOCATION, 'lookup.dat'),
            delim_whitespace=True,
            header=None,
        )
        zoo_lookup.columns = [
            'subject_id',
            'SWASP ID',
            'Period',
            'Period Number',
        ]
        # Period in this file is rounded differently to the others
        # So drop it here so it doesn't stop us from merging later
        zoo_lookup.drop('Period', 'columns', inplace=True)
        zoo_lookup.to_pickle(cache_file)
        return zoo_lookup

    def merge_zoo_lookup(self):
        if self.df.index.name:
            orig_index_name = self.df.index.name
            self.df.reset_index(inplace=True)
        else:
            orig_index_name = None

        self.df = self.df.merge(
            self.zoo_lookup,
            how='left',
        )
        if orig_index_name:
            self.df.set_index(orig_index_name, inplace=True)

class ZooniverseSubjects(PandasDFWrapper, ZooLookupMixin):
    def __init__(self, df=None):
        if df is not None:
            self.df = df
            return

        self.df, self.cache_file = self.cached_pandas_load('superwasp-variable-stars-subjects.csv')
        if self.df is not None:
            return

        self.df = pandas.read_csv(
            os.path.join(DATA_LOCATION, 'superwasp-variable-stars-subjects.csv'),
            index_col='subject_id',
        )
        self.df.to_pickle(self.cache_file)

    @property
    def subject_sets(self):
        return { set_id: self.get_subject_set(set_id) for set_id in set(self.df['subject_set_id']) }

    @property
    def workflows(self):
        return {
            workflow_id: self.get_workflow(workflow_id)
            for workflow_id in set(self.df[self.df['workflow_id'].notna()]['workflow_id'])
        }

    @property
    def retired(self):
        return self.__class__(df=self.df[self.df['retired_at'].notna()])

    @property
    def active(self):
        return self.__class__(df=self.df[self.df['retired_at'].isna()])

    @property
    def distinct(self):
        new_df = self.df.reset_index('subject_id')
        new_df.drop_duplicates('subject_id', inplace=True)
        new_df.set_index('subject_id', inplace=True)
        return self.__class__(df=new_df)

    def get_subject_set(self, set_id):
        return self.__class__(df=self.df[self.df['subject_set_id'] == set_id])

    def get_workflow(self, workflow_id):
        return self.__class__(df=self.df[self.df['workflow_id'] == workflow_id])

    def decode_locations(self, index=0, target='lightcurve'):
        self.df = self.df.copy()
        self.df[target] = self.df['locations'].apply(
            lambda s: ujson.loads(s)[str(index)]
        )

    def display_lightcurves(self, col='lightcurve', start=0, end=None):
        if col not in self.df:
            self.decode_locations(target=col)

        self.df[col][start:end].apply(
            lambda s: display(Image(url=s, width=500, height=500))
        )

class ZooniverseClassifications(PandasDFWrapper):
    ANNOTATION_PREFIX = 'annotation_'

    def __init__(self, df=None, drop_duplicates=False, duplicate_columns=('subject_ids', 'user_id')):
        if df is not None:
            self.df = df
            return

        try:
            self.df, self.cache_file = self.cached_pandas_load('superwasp-variable-stars-classifications.csv')
            if self.df is not None:
                return

            self.df = pandas.read_csv(
                os.path.join(DATA_LOCATION, 'superwasp-variable-stars-classifications.csv'),
                index_col='classification_id',
            )
            self.df.to_pickle(self.cache_file)
        finally:
            if drop_duplicates:
                self.df.drop_duplicates(duplicate_columns, inplace=True)

    @property
    def workflows(self):
        return {
            workflow_id: self.get_workflow(workflow_id)
            for workflow_id in set(self.df[self.df['workflow_id'].notna()]['workflow_id'])
        }

    @property
    def annotations(self):
        self.decode_annotations()
        return self.df[['subject_ids', 'user_id'] + self.annotation_keys]

    @property
    def annotation_keys(self):
        self.decode_annotations()
        return [col for col in self.df.keys() if col.startswith(self.ANNOTATION_PREFIX)]

    def get_workflow(self, workflow_id):
        return ZooniverseClassifications(df=self.df[self.df['workflow_id'] == workflow_id])

    def get_subjects(self, subject_ids):
        return ZooniverseClassifications(df=self.df[self.df['subject_ids'].isin(subject_ids)])

    def get_users(self, user_names):
        return ZooniverseClassifications(df=self.df[self.df['user_name'].isin(user_names)])

    def decode_annotations(self):
        if not 'annotations' in self.df.keys():
            return
        self.df = self.df.copy()

        for classification_id, annotations in self.df['annotations'].items():
            for annotation in ujson.loads(annotations):
                annotation_col = self.ANNOTATION_PREFIX + annotation['task']
                if annotation_col not in self.df:
                    self.df[annotation_col] = pandas.Series([], dtype=str)
                self.df.at[classification_id, annotation_col] = annotation['value']
        self.df.drop('annotations', 'columns', inplace=True)

    def count_annotations(self, col=None, drop_duplicates=True):
        self.decode_annotations()
        if not col:
            col = self.annotation_keys[0]

        df = self.annotations.reset_index()
        if drop_duplicates:
            df.drop_duplicates(['user_id', 'subject_ids'], inplace=True)

        return pandas.pivot_table(
            df,
            index='subject_ids',
            values='classification_id',
            columns=col,
            aggfunc=lambda x: len(x.unique()),
            fill_value=0,
        )

class FoldedLightcurves(PandasDFWrapper, CoordinatesMixin, ZooLookupMixin):
    def __init__(self, min_period=0, df=None):
        self.min_period = min_period

        if df is not None:
            self.df = df
            return

        self.df, self.cache_file = self.cached_pandas_load('results_total.dat')
        if self.df is not None:
            return

        self.df = pandas.read_csv(
            os.path.join(DATA_LOCATION, 'results_total.dat'),
            delim_whitespace=True,
            header=None,
        )
        self.df.columns = [
            'Camera Number',
            'SWASP',
            'ID',
            'Period Number',
            'Period',
            'Sigma',
            'Chi Squared',
            'Period Flag'
        ]
        self.df = self.df[(self.df['Period Flag'] == 0) & (self.df['Period'] >= min_period)]
        self.df['SWASP ID'] = self.df['SWASP'] + self.df['ID']
        self.df.drop(['Period Flag', 'Camera Number', 'SWASP', 'ID'], 'columns', inplace=True)
        self.df.to_pickle(self.cache_file)

    def get_siblings(self, swasp_id):
        return self.__class__(df=self.df[self.df['SWASP ID'] == swasp_id], min_period=self.min_period)

class AggregatedClassifications(PandasDFWrapper, CoordinatesMixin):
    PULSATOR = 1
    EA_EB = 2
    EW = 3
    ROTATOR = 4
    UNKNOWN = 5
    JUNK = 6
    CLASSIFICATION_LABELS = {
        PULSATOR: 'Pulsator',
        EA_EB: 'EA/EB',
        EW: 'EW',
        ROTATOR: 'Rotator',
        UNKNOWN: 'Unknown',
        JUNK: 'Junk',
    }

    def __init__(self, df=None):
        if df is not None:
            self.df = df
            return

        self.df, self.cache_file = self.cached_pandas_load('class_top.csv')
        if self.df is not None:
            return

        self.df = pandas.read_csv(
            os.path.join(DATA_LOCATION, 'class_top.csv'),
            delim_whitespace=True,
            header=None,
        )
        self.df.columns = [
            'subject_id',
            'SWASP ID',
            'Period Number',
            'Period',
            'Classification',
            'Period Uncertainty',
            'Classification Count',
        ]
        # Period in this file is rounded differently to the others
        # So drop it here so it doesn't stop us from merging later
        self.df.drop('Period', 'columns', inplace=True)
        self.df.set_index('subject_id', inplace=True)
        self.df.to_pickle(self.cache_file)

    def add_classification_labels(self):
        self.df = self.df.copy()
        self.df['Classification Label'] = self.get_classification_labels(self.df['Classification'])

    def get_classification_labels(self, series):
        return series.apply(lambda c: self.CLASSIFICATION_LABELS.get(c, None))

    def get_class(self, classification):
        return self.__class__(df=self.df[self.df['Classification'] == classification])

    def remove_class(self, classification):
        return self.__class__(df=self.df[self.df['Classification'] != classification])

    @property
    def pulsators(self):
        return self.get_class(self.PULSATOR)

    @property
    def eaebs(self):
        return self.get_class(self.EA_EB)

    @property
    def ews(self):
        return self.get_class(self.EW)

    @property
    def rotators(self):
        return self.get_class(self.ROTATOR)

    @property
    def unknowns(self):
        return self.get_class(self.UNKNOWN)

    @property
    def junk(self):
        return self.get_class(self.JUNK)

    @property
    def real(self):
        return self.remove_class(self.JUNK)

class UnifiedSubjects(ZooniverseSubjects, FoldedLightcurves, AggregatedClassifications):
    def __init__(
        self,
        zooniverse_subjects=None,
        folded_lightcurves=None,
        aggregated_classifications=None,
        df=None,
        min_period=0
    ):
        self.min_period = min_period

        if df is not None:
            self.df = df
            return

        if not zooniverse_subjects:
            zooniverse_subjects = ZooniverseSubjects()
        if not folded_lightcurves:
            folded_lightcurves = FoldedLightcurves(min_period=min_period)
        if not aggregated_classifications:
            aggregated_classifications = AggregatedClassifications()

        self.df = zooniverse_subjects.df
        self.merge_zoo_lookup()
        self.df = self.df.reset_index().merge(
            aggregated_classifications.df.reset_index(),
            how='left',
        )
        self.df = self.df.merge(
            folded_lightcurves.df,
            how='left',
        )
        self.df.set_index('subject_id', inplace=True)

    def get_siblings(self, obj_id):
        if type(obj_id) == int:
            swasp_id = self.df[self.df.index == obj_id].iloc[0]['SWASP ID']
        else:
            swasp_id = obj_id

        return super().get_siblings(swasp_id)

import numpy,import astropy,from astropy


"""
Module used to mimic observations.

----

.. include license and copyright
.. include:: ../include/copy.rst

----

.. include common links, assuming primary doc root is up one directory
.. include:: ../include/links.rst
"""

#Extraction:
#    - fiber-extraction aperture (pixels, optimal?)
#    - fiber PSF FWHM on detector (pixels)
#
#    - spectral resolution (R)
#    - dispersion on detector (A per pixel)
#
#Detector:
#    - arcsec / pixel (spatial)
#    - angstrom / pixel (spectral)
#    - detector readnoise (e-)
#    - detector darkcurrent (e-/hour)
#
#Throughput
#    - spectrograph throughput
#        - detector QE
#        - camera efficiency
#        - grating efficiency
#        - focal-plane to grating tansmission efficiency
#            - foreoptics, lenslet, coupling, fiber transmission, FRD,
#              collimator, dichroic(s)
#
#    - top-of-telescope throughput
#        - telescope efficiency
#        - spectrograph efficiency
#
#Point-source Aperture losses
#    - fiber diameter (arcsec, mm)
#    - focal-plane plate scale (arcsec/mm)
#    - seeing (arcsec)
#    - pointing error (arcsec)
#
#*Source:
#    - source spectrum (1e-17 erg/s/cm^2/angstrom)
#
#    - point source
#        - source mag (at wavelength/in band above)
#
#    - extended source
#        - source surface brightness (at wavelength/in band above)
#        - source size (kpc/arcsec)
#        - cosmology
#
#    - source velocity/redshift
#
#*Sky:
#    - sky spectrum (1e-17 erg/s/cm^2/angstrom)
#
#Observation:
#    - obs date
#    - lunar phase/illumination (days, fraction)
#    - sky coordinates
#    - wavelength for calculation (angstroms)
#    - band for calculation
#    - sky surface brightness (mag/arcsec^2, at wavelength/in band above)
#    - airmass
#    - exposure time (s)
#
#Telescope:
#    - location
#    - telescope diameter (or effective) (m^2)
#    - central obstruction

import os
import warnings

from IPython import embed

import numpy

from scipy import signal

from matplotlib import pyplot

from . import source, spectrum, util

class Observation:
    """
    Observation of the sky with or without a source.

    Args:
        telescope (:class:`~enyo.etc.telescopes.Telescope`):
            Telescope used for the observation.
        sky_spectrum (:class:`~enyo.etc.spectrum.Spectrum`):
            Sky spectrum
        spec_aperture (:class:`~enyo.etc.aperture.Aperture`):
            On-sky entrance aperture.
        exposure_time (:obj:`float`):
            Exposure time in seconds.
        detector (:class:`~enyo.etc.detector.Detector`):
            Instrument detector object.
        system_throughput (:class:`~enyo.etc.efficiency.Efficiency`, optional):
            System throughput; i.e., ratio of the number of photons
            detected to the total number incident on the telescope
            primary.  If None, system throughput is unity.
        atmospheric_throughput (:class:`~enyo.etc.efficiency.Efficiency`, optional):
            Atmospheric throughput; i.e., the ratio of the number of
            the photons incident on the telescope primary to the
            number of photons hitting the top of the atmosphere. If
            None, assume there is no atmosphere.
        airmass (:obj:`float`, optional):
            Airmass of the observation. If None, assume
            ``atmospheric_throughput`` is at the correct airmass or
            there is no atmosphere.
        onsky_source_distribution (:class:`~enyo.etc.source.Source`, optional):
            On-sky source surface-brightness distribution, including
            any seeing effects. If None, assume observation is only
            of the sky. Normalization is irrelevant; used to
            calculate aperture losses.
        source_spectrum (:class:`~enyo.etc.spectrum.Spectrum`, optional):
            Spectrum of the source with the flux normalized to the
            *total* source flux.
        extraction (:class:`~enyo.etc.extract.Extraction`, optional):
            Object used to calculate the extraction losses and tally
            the number of pixels included in the extraction to get
            the read-noise hit. TODO: Currently *cannot* be None;
            make a required argument.
        snr_units (:obj:`str`, optional):
            Units of the S/N calculation. Must be ``'pixel'``,
            ``'angstrom'``, or ``'resolution'`` for S/N per pixel,
            per angstrom, or per resolution element.

    Raises:
        ValueError:
            Raised if the requested S/N units cannot be calculated
            (i.e., the provided spectrum does not define the
            resolution meaning that the S/N per resolution element is
            undefined) or if the unit string for the S/N is not
            recognized.

    """
    def __init__(self, telescope, sky_spectrum, spec_aperture, exposure_time, detector,
                 system_throughput=None, atmospheric_throughput=None, airmass=None,
                 onsky_source_distribution=None, source_spectrum=None, extraction=None,
                 snr_units='pixel'):

        # Save the input or use defaults
        self.source = onsky_source_distribution

        # Sky spectrum is expected to be independent of position within
        # the aperture and be the sky flux density per unit area, where
        # the unit area is defined by the aperture object (arcsec^2)
        # i.e., the units are, e.g., erg/s/cm^2/angstrom/arcsec^2
        self.source_spectrum = source_spectrum

        # Match the sampling of the sky and source spectrum, if possible
        self.sky_spectrum = sky_spectrum if self.source_spectrum is None else \
                                spectrum.Spectrum(self.source_spectrum.wave,
                                                  sky_spectrum.interp(self.source_spectrum.wave),
                                                  log=self.source_spectrum.log)

        # In the current implementation, the source spectrum is expected
        # to be:
        #   - independent of position within the source
        #   - the source flux density integrated over the full source
        #     distribution; i.e., units are, e.g., erg/s/cm^2/angstrom
        self.wave = self.sky_spectrum.wave.copy()
        self.sres = self.sky_spectrum.sres.copy() if sky_spectrum.sres is not None \
                        else (None if self.source_spectrum is None
                                else self.source_spectrum.sres.copy())

        self.atmospheric_throughput = atmospheric_throughput
        self.telescope = telescope
        self.aperture = spec_aperture
        self.system_throughput = system_throughput
        self.detector = detector
        self.exptime = exposure_time
        self.extraction = extraction

        # Get the "aperture factor". If the source distribution is not
        # provided, the source surface brightness is assumed to be
        # uniform within the aperture (like the sky). In this case, the
        # aperture factor is the area of the aperture itself so that
        # the object flux is the integral of the surface brightness
        # over the aperture size (like the sky).
        self.aperture_factor = self.aperture.area if self.source is None \
                                     else self.aperture.integrate_over_source(self.source) \
                                                / self.source.integral

        # Get the total object flux incident on the focal plane in
        # electrons per second per angstrom
        _object_flux = numpy.zeros(self.wave.size, dtype=float) if self.source_spectrum is None \
                            else self.source_spectrum.photon_flux(inplace=False) \
                                    * self.telescope.area * self.aperture_factor \
                                    * self.detector(self.wave)
        if self.atmospheric_throughput is not None:
            _object_flux *= self.atmospheric_throughput(self.wave)
        if self.system_throughput is not None:
            _object_flux *= self.system_throughput(self.wave)

        # Total sky flux in electrons per second per angstrom; the
        # provided sky spectrum is always assumed to be uniform over the
        # aperture
        _sky_flux = self.sky_spectrum.photon_flux(inplace=False) \
                        * self.telescope.area * self.aperture.area \
                        * self.detector(self.wave)
        if self.system_throughput is not None:
            _sky_flux *= self.system_throughput(self.wave)

        # Set the units for the output:
        dw = self.sky_spectrum.wavelength_step()
        if snr_units == 'pixel':
            _object_flux *= dw
            _sky_flux *= dw
            spectral_width = 1.
        elif snr_units == 'angstrom':
            spectral_width = 1./dw
        elif snr_units == 'resolution':
            if self.sres is None:
                raise ValueError('Cannot compute S/N per resolution element without resolution '
                                 'vector.')
            _object_flux *= self.wave/self.sres
            _sky_flux *= self.wave/self.sres
            spectral_width = self.wave/self.sres/dw
        else:
            raise ValueError('Unknown S/N units requested.')

        # Observe and extract the source
        # TODO: Set spectral_pixels...
        self.object_flux, self.obj_shot_var, self.sky_flux, self.sky_shot_var, self.read_var \
                = self.extraction.sum_signal_and_noise(_object_flux, _sky_flux, self.exptime,
                                                       spectral_width=spectral_width)

    def simulate(self, sky_only=False, sky_sub=False, sky_err=0.1):
        """
        Return a simulated spectrum.

        Args:
            sky_only (:obj:`bool`, optional):
                Only include the sky flux in the simulated spectrum
                (no object flux)
            sky_sub (:obj:`bool`, optional):
                Provide the object spectrum only, but include the sky
                flux shot noise and additional error from the sky
                subtraction.
            sky_err (:obj:`float`, optional):
                The fraction of the total sky error incurred due to
                the sky subtraction. Should be between 0 and 1; 0
                means no additional error is incurred, 1 means that
                the sky noise from the sky subtration is the same as
                the sky noise from the observation itself.

        Returns:
            :class:`~enyo.etc.spectrum.Spectrum`: The simulated
            spectrum.
        """
        if sky_only:
            shot_var = self.sky_shot_var
            flux = self.sky_flux
        elif sky_sub:
            shot_var = self.obj_shot_var + (1 + numpy.square(sky_err))*self.sky_shot_var
            flux = self.object_flux
        else:
            shot_var = self.obj_shot_var + self.sky_shot_var
            flux = self.object_flux + self.sky_flux

#        error=numpy.sqrt(shot_var + self.read_var)
#        draw = numpy.random.normal(scale=error)
#        return spectrum.Spectrum(self.wave, flux + draw, error=error,
#                                 log=self.sky_spectrum.log if self.source_spectrum is None
#                                         else self.source_spectrum.log)

        # Draw from a Poisson distribution for the shot noise,
        # subtracted the expectation value of the distribution so that
        # only the noise is added
        shot_draw = numpy.random.poisson(lam=shot_var)-shot_var
        # Draw from a Gaussian distribution for the read noise
        read_draw = numpy.random.normal(scale=numpy.sqrt(self.read_var))
        return spectrum.Spectrum(self.wave, flux + shot_draw + read_draw,
                                 error=numpy.sqrt(shot_var + self.read_var),
                                 log=self.sky_spectrum.log if self.source_spectrum is None
                                         else self.source_spectrum.log)

    def snr(self, sky_sub=False, sky_err=0.1):
        """
        Calculate the S/N.

        Args:
            sky_sub (:obj:`bool`, optional):
                Provide the object spectrum only, but include the sky
                flux shot noise and additional error from the sky
                subtraction.
            sky_err (:obj:`float`, optional):
                The fraction of the total sky error incurred due to
                the sky subtraction. Should be between 0 and 1; 0
                means no additional error is incurred, 1 means that
                the sky noise from the sky subtration is the same as
                the sky noise from the observation itself.

        Returns:
            :class:`~enyo.etc.spectrum.Spectrum`: The S/N spectrum.
            WARNING: Here, :class:`~enyo.etc.spectrum.Spectrum` is
            used as a container class for the S/N vector; some
            functionality of the class will not be valid!
        """
        flux = self.object_flux + self.sky_flux
        var = self.obj_shot_var + self.sky_shot_var + self.read_var
        if sky_sub:
            flux -= self.sky_flux
            var += numpy.square(sky_err)*self.sky_shot_var
        # TODO: add additional noise from sky subtraction
        return spectrum.Spectrum(self.wave, flux / numpy.sqrt(var),
                                 log=self.sky_spectrum.log if self.source_spectrum is None
                                         else self.source_spectrum.log)

def monochromatic_image(sky, spec_aperture, spec_kernel, platescale, pixelsize, onsky_source=None,
                        scramble=False):
    """
    Generate a monochromatic image of the sky, with or with out a
    source, taken by a spectrograph through an aperture.

    .. warning::

        - May resample the `source` and `spec_kernel` maps.

    .. todo::

        - Add effect of differential atmospheric refraction
        - Allow a force map size and pixel sampling

    Args:
        sky (:class:`enyo.etc.source.OnSkySource`):
            Sky flux distribution.
        spec_aperture (:class:`enyo.etc.aperture.Aperture`):
            Spectrograph aperture. Aperture is expected to be
            oriented with the dispersion along the first axis (e.g.,
            the slit width is along the abcissa).
        spec_kernel (:class:`enyo.etc.kernel.SpectrographGaussianKernel`):
            Convolution kernel describing the point-spread function
            of the spectrograph.
        platescale (:obj:`float`):
            Platescale in mm/arcsec at the detector
        pixelsize (:obj:`float`):
            Size of the detector pixels in mm.
        onsky_source (:class:`enyo.etc.source.OnSkySource`, optional):
            On-sky distribution of the source flux. If None, only sky
            is observed through the aperture.
        scramble (:obj:`bool`, optional):
            Fully scramble the source light passing through the
            aperture. This should be False for slit observations. For
            fiber observations, this should be True and makes the
            nominal assumption that the focal plane incident on the
            fiber face is perfectly scrambled.
    """
    # Check input
    if onsky_source is not None:
        if onsky_source.sampling is None or onsky_source.size is None:
            warnings.warn('Source was not provided with an initial map sampling; doing so now '
                          'with default sampling and size.')
            onsky_source.make_map()

    # Detector pixel scale in arcsec/pixel
    pixelscale = pixelsize/platescale

    # Assume the sampling of the source is provided with the maximum
    # allowed pixel size. Determine a pixel size that is no more than
    # this, up to an integer number of detector pixels. Use an integer
    # number specifically so that the result of the kernel convolution
    # can be simply rebinned to match the detector pixel size.
    # `sampling` is in arcsec per pixel
    sampling = pixelscale if onsky_source is None else min(onsky_source.sampling, pixelscale)
    oversample = int(pixelscale/sampling)+1 if sampling < pixelscale else 1
    sampling = pixelscale/oversample

    # Assume the size of the image properly samples the source.
    # Determine a map size that at least encompasses the input source
    # and the input aperture. The factor of 1.5 is ad hoc; could likely
    # be lower. `size` is in arcsec
    dx, dy = 1.5*numpy.diff(numpy.asarray(spec_aperture.bounds).reshape(2,-1), axis=0).ravel()
    size = max(dx,dy) #if onsky_source is None else max(onsky_source.size, dx, dy)

    # TODO: Below alters `source` and `spec_kernel`. Should maybe
    # instead save the old sampling and size and then resample back to
    # the input before returning.

    # Resample the distribution maps. Classes OnSkySource and Aperture
    # use arcsecond units.
    if onsky_source is not None:
        onsky_source.make_map(sampling=sampling, size=size)
    sky.make_map(sampling=sampling, size=size)
    ap_img = spec_aperture.response(sky.x, sky.y)

    # However, SpectrographGaussianKernel uses mm, so we need to
    # convert sampling from arcsec/pixel to mm/pixel using the
    # platescale.
    spec_kernel.resample(pixelscale=platescale*sampling)

    # Construct the image. Note that scrambling simply scales the
    # aperture image by the flux that enters it; otherwise, the source
    # image is just attenuated by the aperture response map.
    source_img = sky.data if onsky_source is None else onsky_source.data + sky.data
    input_img = ap_img*numpy.sum(source_img*ap_img)/numpy.sum(ap_img) if scramble \
                        else source_img*ap_img

    # Convolve it with the spectrograph imaging kernel
    mono_img = signal.fftconvolve(input_img, spec_kernel.array, mode='same')
    # Return the image, downsampling if necessary
    return util.boxcar_average(mono_img, oversample) if oversample > 1 else mono_img

def twod_spectrum(sky_spectrum, spec_aperture, spec_kernel, platescale, linear_dispersion,
                  pixelsize, source_distribution=None, source_spectrum=None, thresh=None,
                  scramble=False, wave_lim=None, field_coo=None, opticalmodel=None):
    """
    Documentation TBW.

    if optical model is not provided:
        - ignore field_coo
        - return rectilinear 2D spectrum

    platescale is in mm/arcsec
    linear_dispersion is in A/mm
    pixelsize is in mm

    """
    # Ensure that the sky spectrum and source spectrum will have the same wavelength limits
    if source_spectrum is not None and wave_lim is None:
        wave_lim = source_spectrum.wave[[0,-1]]

    # Get the sky-only monochromatic image
    sky = source.OnSkyConstant(1.0)
    sky_img = monochromatic_image(sky, spec_aperture, spec_kernel, platescale, pixelsize,
                                  scramble=scramble)

    # Renormalize the sky slit image such that the integral is the area of the aperture.
    sky_img *= spec_aperture.area / numpy.sum(sky_img)/numpy.square(sky.sampling)

    source_img = None
    if source_distribution is not None and source_spectrum is not None:
        # Reset the source distribution map
        source_distribution.reset_map()

        # Get the source-only monochromatic image
        sky = source.OnSkyConstant(0.0)
        source_img = monochromatic_image(sky, spec_aperture, spec_kernel, platescale, pixelsize,
                                         onsky_source=source_distribution, scramble=scramble)
        # Renormalize the source image by the integral of the onsky-source;
        # this maintains the effects of aperture losses
        source_img /= source_distribution.integral

    s = numpy.array([0,0])
    e = numpy.array([*sky_img.shape])
    if thresh is not None:
        indx = sky_img > thresh
        s, e = numpy.append(numpy.where(numpy.any(indx, axis=1))[0][[0,-1]],
                            numpy.where(numpy.any(indx, axis=0))[0][[0,-1]]).reshape(2,2).T
        if source_img is not None:
            indx = source_img > thresh
            _s, _e = numpy.append(numpy.where(numpy.any(indx, axis=1))[0][[0,-1]],
                                  numpy.where(numpy.any(indx, axis=0))[0][[0,-1]]).reshape(2,2).T
            s = numpy.minimum(s, _s)
            e = numpy.maximum(e, _e)

    sky_img = sky_img[s[0]:e[0],s[1]:e[1]]
    if source_img is not None:
        source_img = source_img[s[0]:e[0],s[1]:e[1]]

    # Get the 2D spectrum
    dispscale = linear_dispersion*pixelsize
    wave0, sky_2d_spec = rectilinear_twod_spectrum(sky_spectrum, sky_img, dispscale,
                                                   wave_lim=wave_lim)

    if source_distribution is None or source_spectrum is None:
        return sky_2d_spec if opticalmodel is None \
                    else opticalmodel.project_2d_spectrum(sky_2d_spec, platescale,
                                                          linear_dispersion, pixelsize, wave0,
                                                          field_coo=field_coo)

    # Get the 2D spectrum
    wave0, source_2d_spec = rectilinear_twod_spectrum(source_spectrum, source_img, dispscale,
                                                      wave_lim=wave_lim)

    return sky_2d_spec + source_2d_spec if opticalmodel is None \
                    else opticalmodel.project_2d_spectrum(sky_2d_spec + source_2d_spec, platescale,
                                                          linear_dispersion, pixelsize, wave0,
                                                          field_coo=field_coo)

#def rectilinear_twod_spectrum(spectrum, aperture_image, dispscale, wave_lim=None, oversample=1):
#    """
#    Construct a rectilinear 2D spectrum.
#
#    spectral dimension of aperture_image is along the first axis
#
#    aperture_image has to be odd?
#    """
#
#    # TODO: Let dispersion scale be non-linear?
#
#    # Resample the spectrum to the appropriate dispersion scale up to some constant
#    if wave_lim is None:
#        # TODO: This should be the pixel boundaries, not the pixel centers!
#        wave_lim = spectrum.wave[[0,-1]]
#    resamp_wave = numpy.arange(wave_lim[0], wave_lim[1] + dispscale/oversample,
#                               dispscale/oversample)
#    resampled_spectrum = spectrum.resample(resamp_wave)
##    pyplot.plot(spectrum.wave, spectrum.flux)
##    pyplot.plot(resampled_spectrum.wave, resampled_spectrum.flux)
##    pyplot.show()
#
#    # Oversample the image spectrally
#    _aperture_image = util.block_replicate(aperture_image, (oversample,1)) \
#                            if oversample > 1 else aperture_image
#
#    # Number of spectral and spatial channels in the aperture image
#    nspec, nspat = _aperture_image.shape
#    width = nspec//2
#    # Length of the spectrum
#    npix = len(resampled_spectrum)
#
#    # Pad the spectrum with zeros and create one shifted copy per spectral channel
#    flux = numpy.zeros((nspec,npix), dtype=float)
#    for i in range(nspec):
#        flux[i,width:-width] = resampled_spectrum.flux[i:npix-nspec+i+1]
#
#    import time
#    t = time.perf_counter()
#    flux = numpy.tile(flux, (nspat,1))
#    print('Tile: {0}s'.format(time.perf_counter()-t))
#
#    # Do the convolution
#    t = time.perf_counter()
#    twodspec = flux[:,:] * _aperture_image.ravel()[:,None]
#    indx = numpy.arange(0,nspec*nspat,nspec)
#    twodspec = numpy.add.reduceat(twodspec, indx, axis=0)
#    print('Mult: {0}s'.format(time.perf_counter()-t))
#
#    return util.block_average(twodspec, (oversample,1)) if oversample > 1 else twodspec

def rectilinear_twod_spectrum(spectrum, aperture_image, dispscale, wave_lim=None, oversample=1):
    """
    Construct a rectilinear 2D spectrum.

    Documentation TBW.

    spectral dimension of aperture_image is along the first axis

    aperture_image has to be odd?
    dispscale is A/pixel

    remove rows columns with no pixels above thresh

    """

    # TODO: Let dispersion scale be non-linear?

    # Resample the spectrum to the appropriate dispersion scale up to some constant
    if wave_lim is None:
        # TODO: This should be the pixel boundaries, not the pixel centers!
        wave_lim = spectrum.wave[[0,-1]]
    resamp_wave = numpy.arange(wave_lim[0], wave_lim[1] + dispscale/oversample,
                               dispscale/oversample)
    resampled_spectrum = spectrum.resample(resamp_wave)
    wave0 = numpy.mean(resampled_spectrum.wave[:oversample])
#    pyplot.plot(spectrum.wave, spectrum.flux)
#    pyplot.plot(resampled_spectrum.wave, resampled_spectrum.flux)
#    pyplot.show()

    # Oversample the image spectrally
    _aperture_image = util.block_replicate(aperture_image, (oversample,1)) \
                            if oversample > 1 else aperture_image.copy()

    # Number of spectral and spatial channels in the aperture image;
    # number of convolution kernels is nspat
    nspat, nspec = _aperture_image.shape

    twodspec = numpy.zeros((len(resampled_spectrum),nspat), dtype=float)
    for i in range(nspat):
        twodspec[:,i] = signal.fftconvolve(resampled_spectrum.flux, _aperture_image[i],
                                           mode='same')

    return wave0, (util.block_average(twodspec, (oversample,1)) if oversample > 1 else twodspec)


from __future__ import annotations
import numpy as np
import math
from pyscses.constants import fundamental_charge, boltzmann_eV
from pyscses.grid_point import GridPoint
from pyscses.defect_species import DefectSpecies
from typing import List, Optional, Dict
from pyscses.defect_at_site import DefectAtSite
import warnings

class LabelError(Exception):
    pass

class Site:
    """The Site class contains all the information about a given site and the defects occupying that site.
    This class contains functions for the calculations which correspond to each individual site, rather than the system as a whole.

    Attributes:
        label (str): Reference label for this site. i.e. 'O' for an oxygen site.
        x (float): x coordinate of the site.
        defect_energies (list): List of segregation energies for all defects present at the site.
        defect_species (list): List of defect species for all defects present at the site.
        defects (list): List of DefectAtSite objects, containing the properties of all individual defects at the site.
        scaling (float): A scaling factor that can be applied in the charge calculation.
        valence (float): The charge of the defect present at the site (in atomic units).
        saturation_parameter (float): Optional saturation parameter as described in
            `Hendricks et al. Sol. Stat. Ionics (2002)`_
            and `Swift et al. Nature Comp. Sci. (2021)`_.
            Setting `saturation_parameter` < `1.0` sets some proportion of excluded sites that are
            unavailable for occupation by any explicit defects.
            Default value is `1.0`, i.e., 100% of sites may be occupied.
        defects (list): List of Defect_Species objects for all defects present at the site.
        sites (list): List containing all x coordinates and corresponding  defect segregation energies.

    .. _Hendricks et al. Sol. Stat. Ionics (2002):
       https://doi.org/10.1016/S0167-2738(02)00484-8

    .. _Swift et al. Nature Comp. Sci. (2021):
       https://doi.org/10.1038/s43588-021-00041-y

    """

    def __init__(self,
                 label: str,
                 x: float,
                 defect_species: List[DefectSpecies],
                 defect_energies: List[float],
                 scaling: Optional[np.ndarray] = None,
                 valence: float = 0.0,
                 saturation_parameter: float = 1.0) -> None:
        """Initialise a Site object.

        Args:
            label (str): Reference label for this site.
            x (float): x coordinate of this site.
            defect_species (list(DefectSpecies)): List of `DefectSpecies` objects (one for each defect species that can occupy this site).
            defect_energies (list(float)): List of defect segregation energies for each defect species at this site.
            scaling (optional, list(float): Optional list of scaling factors for the net charge at this site. Default scaling for each defect species is 1.0.
            valence (optional, float): Optional formal valence for this site in the absence of any defects. Default is 0.0.
            saturation_parameter (optional, float): Optional saturation parameter as described in
                Hendricks et al. Sol. Stat. Ionics (2002) [#HendricksEtAl_SolStatIonics2002]_
                and Swift et al. Nature Comp. Sci. (2021). [#SwiftEtAl_NatureCompSci2021]_.
                A saturation parameter < 1.0 introduces some proportion of excluded sites that are
                unavailable for occupation by any explicit defects. Default is 1.0.

        Raises:
            ValueError if the number of DefectSpecies != the number of defect segregation energies != the number of scaling factors (if passed).

        .. [HendricksEtAl_SolStatIonics2002]:
           https://doi.org/10.1016/S0167-2738(02)00484-8

        .. [SwiftEtAl_NatureCompSci2021]:
           https://doi.org/10.1038/s43588-021-00041-y

        """
        if len(defect_species) != len(defect_energies):
            raise ValueError("len(defect_species) must be equal to len(defect_energies)")
        if scaling:
            if len(defect_species) != len(scaling):
                raise ValueError("len(defect_species) must be equal to len(scaling)")
        self.label = label
        self.x = x
        self.defect_energies = defect_energies
        self.defect_species = defect_species
        self.defects = [DefectAtSite(label=d.label,
                                     valence=d.valence,
                                     mole_fraction=d.mole_fraction,
                                     mobility=d.mobility,
                                     energy=e,
                                     site=self,
                                     fixed=d.fixed)
            for d, e in zip(defect_species, defect_energies)]
        if scaling:
            self.scaling = scaling
        else:
            self.scaling = np.ones_like(defect_energies, dtype=float)
        self.grid_point: Optional[GridPoint] = None
        self.valence = valence
        self.saturation_parameter = saturation_parameter
        self.fixed_defects = tuple(d for d in self.defects if d.fixed)
        self.mobile_defects = tuple(d for d in self.defects if not d.fixed)
        self.alpha = self.saturation_parameter - sum((d.mole_fraction for d in self.fixed_defects))

    def competing_defect_species(self) -> Dict[str, int]:
        """Returns a dictionary reporting the number of fixed and / or mobile defect species that can occupy this site.

        Args:
            None

        Returns
            Dict(str, int): Dictionary {'fixed': n_fixed, 'mobile': n_mobile}

        """
        pass

    def defect_with_label(self,
                          label: str) -> DefectAtSite:
        """Select a defect at this site by the species label.

        Args:
            label (str): Label to identify defect species.

        Returns:
                DefectAtSite: The DefectAtSite that matches the label.

        """
        if not label in (d.label for d in self.defects):
            raise LabelError(f"\"{label}\" does not match any of the defect species labels for this site.")
        else:
            return next(d for d in self.defects if d.label == label)

    def energies(self) -> List[float]:
        """Returns a list of the segregation energies for each defect from self.defects """
        return [d.energy for d in self.defects]

    def average_local_energy(self,
                             method: str = 'mean') -> Optional[np.ndarray]:
        """
        Returns the average local segregation energy for each site based on a specified method.

        Args:
            method (str): The method in which the average segregation energies will be calculated.
                          'mean' - Returns the sum of all values at that site divided by the number of values at that site.
                          'min' - Returns the minimum segregation energy value for that site (appropriate for low temperature calculations).

        Returns:
            numpy.array: Average segregation energies on the site coordinates grid.

        """
        if self.grid_point is not None:
            return self.grid_point.average_site_energy(method)
        else:
            raise ValueError("TODO")

    def probabilities(self,
                      phi: float,
                      temp: float) -> Dict[str, float]:
        """Calculates the probabilities of this site being occupied by each defect species.

        Args:
            phi (float): Electrostatic potential at this site in Volts.
            temp (float): Temperature in Kelvin.

        Returns:
            dict(str, float): Probabilities of site occupation for each defect species.

        """
        probabilities_dict = {}
        boltzmann_factors = {d.label: d.boltzmann_factor(phi, temp) for d in self.mobile_defects}
        denominator = (self.alpha +
                       sum([d.mole_fraction * (boltzmann_factors[d.label] - 1.0)
                            for d in self.mobile_defects]))
        for defect in self.defects:
            if defect.fixed:
                probabilities_dict[defect.label] = defect.mole_fraction
            else:
                numerator = self.alpha * defect.mole_fraction * boltzmann_factors[defect.label]
                probabilities_dict[defect.label] = numerator / denominator
        return probabilities_dict

    def probabilities_as_list(self,
                              phi: float,
                              temp: float) -> List[float]:
        """Calculates the probabilities of this site being occupied by each defect species.

        Legacy interface that returns a list of site-occupation probabilities
        in the same order as `Site.defects`.

            Args:
            phi (float): Electrostatic potential at this site in Volts.
            temp (float): Temperature in Kelvin.

        Returns:
            list(float): Probabilities of site occupation for each defect species.

        """
        warnings.warn("Site.probabilities_as_list() is deprecated and targeted for removal. Please use Site.probabilities() instead.", DeprecationWarning)
        probabilities_dict = self.probabilities(phi=phi, temp=temp)
        return [probabilities_dict[d.label] for d in self.defects]

    def defect_valences(self) -> np.ndarray:
        """Returns an array of valences for each defect in `self.defects`"""
        return np.array([d.valence for d in self.defects])

    def charge(self,
               phi: float,
               temp: float) -> float:
        """
        Charge at this site (in Coulombs).

        Args:
            phi (float):  Electrostatic potential at this site in Volts.
            temp (float): Temperature in Kelvin.

        Returns:
            float: The charge at this site.

        """
        defect_probabilities = self.probabilities(phi=phi, temp=temp)
        charge = sum([defect_probabilities[d.label] * d.valence
                      for d in self.defects]) * self.scaling
        charge += self.valence
        charge *= fundamental_charge
        return float(charge)