#!/usr/bin/env python3

#
# https://github.com/mseitzer/pytorch-fid
#

"""Calculates the Frechet Inception Distance (FID) to evalulate GANs

The FID metric calculates the distance between two distributions of images.
Typically, we have summary statistics (mean & covariance matrix) of one
of these distributions, while the 2nd distribution is given by a GAN.

When run as a stand-alone program, it compares the distribution of
images that are stored as PNG/JPEG at a specified location with a
distribution given by summary statistics (in pickle format).

The FID is calculated by assuming that X_1 and X_2 are the activations of
the pool_3 layer of the inception net for generated samples and real world
samples respectively.

See --help to see further details.

Code apapted from https://github.com/bioinf-jku/TTUR to use PyTorch instead
of Tensorflow

Copyright 2018 Institute of Bioinformatics, JKU Linz

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

   http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
import os
import pathlib
from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter

import numpy as np
import torch
from scipy import linalg
from torch.nn.functional import adaptive_avg_pool2d

from PIL import Image

try:
	from tqdm import tqdm
except ImportError:
	# If not tqdm is not available, provide a mock version of it
	def tqdm(x): return x

from .inceptionV3 import InceptionV3

cuda = True if torch.cuda.is_available() else False

parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)
parser.add_argument('path', type=str, nargs=2,
					help=('Path to the generated images or '
						  'to .npz statistic files'))
parser.add_argument('--batch-size', type=int, default=50,
					help='Batch size to use')
parser.add_argument('--dims', type=int, default=2048,
					choices=list(InceptionV3.BLOCK_INDEX_BY_DIM),
					help=('Dimensionality of Inception features to use. '
						  'By default, uses pool3 features'))
parser.add_argument('-c', '--gpu', default='', type=str,
					help='GPU to use (leave blank for CPU only)')

def calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6):
	"""Numpy implementation of the Frechet Distance.
	The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1)
	and X_2 ~ N(mu_2, C_2) is
			d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)).

	Stable version by Dougal J. Sutherland.

	Params:
	-- mu1   : Numpy array containing the activations of a layer of the
			   inception net (like returned by the function 'get_predictions')
			   for generated samples.
	-- mu2   : The sample mean over activations, precalculated on an
			   representative data set.
	-- sigma1: The covariance matrix over activations for generated samples.
	-- sigma2: The covariance matrix over activations, precalculated on an
			   representative data set.

	Returns:
	--   : The Frechet Distance.
	"""

	mu1 = np.atleast_1d(mu1)
	mu2 = np.atleast_1d(mu2)

	sigma1 = np.atleast_2d(sigma1)
	sigma2 = np.atleast_2d(sigma2)

	assert mu1.shape == mu2.shape, \
		'Training and test mean vectors have different lengths'
	assert sigma1.shape == sigma2.shape, \
		'Training and test covariances have different dimensions'

	diff = mu1 - mu2

	# Product might be almost singular
	covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
	if not np.isfinite(covmean).all():
		msg = ('fid calculation produces singular product; '
			   'adding %s to diagonal of cov estimates') % eps
		print(msg)
		offset = np.eye(sigma1.shape[0]) * eps
		covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))

	# Numerical error might give slight imaginary component
	if np.iscomplexobj(covmean):
		if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
			m = np.max(np.abs(covmean.imag))
			raise ValueError('Imaginary component {}'.format(m))
		covmean = covmean.real

	tr_covmean = np.trace(covmean)

	return (diff.dot(diff) + np.trace(sigma1) +
			np.trace(sigma2) - 2 * tr_covmean)

def get_activations(batch, model, batch_size=50, dims=2048):

	model.eval()

	if batch_size > batch.shape[0]:
		print(('Warning: batch size is bigger than the data size. '
			   'Setting batch size to data size'))
		batch_size = batch.shape[0]

	pred_arr = np.empty((batch.shape[0], dims))

	for i in tqdm(range(0, batch.shape[0], batch_size)): #tqdm: progressive bar

		start = i
		end = i + batch_size

		data = batch[start:end]

		if cuda:
			data=data.type(torch.cuda.FloatTensor)

		#pred = model(batch)[0]
		pred = model(data)[0]

		# If model output is not scalar, apply global spatial average pooling.
		# This happens if you choose a dimensionality not equal 2048.
		if pred.size(2) != 1 or pred.size(3) != 1:
			pred = adaptive_avg_pool2d(pred, output_size=(1, 1))

		pred_arr[start:end] = pred.cpu().data.numpy().reshape(pred.size(0), -1)

	return pred_arr

def calculate_activation_statistics(batch, model, batch_size=50, dims=2048):

	act = get_activations(batch, model, batch_size, dims)
	mu = np.mean(act, axis=0)
	sigma = np.cov(act, rowvar=False)
	return mu, sigma

def _compute_statistics_of_batch(batch, model, batch_size, dims):

	m, s = calculate_activation_statistics(batch, model, batch_size, dims)

	return m, s

#def calculate_fid_given_paths(paths, batch_size, cuda, dims):
def calculate_fid_given_batches(batch1, batch2, batch_size, dims=2048):

	# gray image
	if batch1.shape[1] < 3:  batch1 = batch1.expand(-1,3,-1,-1)
	if batch2.shape[1] < 3:  batch2 = batch2.expand(-1,3,-1,-1)

	block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[dims]

	model = InceptionV3([block_idx])

	if cuda:
		model.cuda()

	m1, s1 = _compute_statistics_of_batch(batch1, model, batch_size, dims)
	m2, s2 = _compute_statistics_of_batch(batch2, model, batch_size, dims)
	fid_value = calculate_frechet_distance(m1, s1, m2, s2)

	return fid_value

if __name__ == '__main__':
	args = parser.parse_args()
	os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu

	fid_value = calculate_fid_given_paths(args.path, args.batch_size, args.gpu != '', args.dims)
	print('FID: ', fid_value)


"""Eigensolvers with a few extra computation methods

The :class:`.Solver` class is the main interface for dealing with eigenvalue problems. It
is made to work specifically with pybinding's :class:`.Model` objects, but it may use any
eigensolver algorithm under the hood.

A few different algorithms are provided out of the box: the :func:`.lapack`, :func:`.arpack`
and :func:`.feast` functions return concrete :class:`.Solver` implementation using the LAPACK,
ARPACK and FEAST algorithms, respectively.

The :class:`.Solver` may easily be extended with new eigensolver algorithms. All that is
required is a function which takes a Hamiltonian matrix and returns the computed
eigenvalues and eigenvectors. See :class:`._SolverPythonImpl` for example.
"""
import time
import math

import numpy as np

from . import _cpp
from . import results
from .model import Model
from .system import System

__all__ = ['Solver', 'arpack', 'feast', 'lapack']

class Solver:
    """Computes the eigenvalues and eigenvectors of a Hamiltonian matrix

    This the common interface for various eigensolver implementations. It should not
    be created directly, but via the specific functions: :func:`.lapack`, :func:`.arpack`
    and :func:`.feast`. Those functions will set up their specific solver strategy and
    return a properly configured :class:`.Solver` object.
    """
    def __init__(self, impl: _cpp.Solver):
        self.impl = impl

    @property
    def model(self) -> Model:
        """The tight-binding model attached to this solver"""
        return self.impl.model

    @model.setter
    def model(self, model):
        self.impl.model = model

    @property
    def system(self) -> System:
        """The tight-binding system attached to this solver (shortcut for Solver.model.system)"""
        return System(self.impl.system, self.model.lattice)

    @property
    def eigenvalues(self) -> np.ndarray:
        """1D array of computed energy states"""
        return self.impl.eigenvalues

    @property
    def eigenvectors(self) -> np.ndarray:
        """2D array where each column represents a wave function

        eigenvectors.shape == (system.num_sites, eigenvalues.size)
        """
        return self.impl.eigenvectors

    def solve(self):
        """Explicitly solve the eigenvalue problem right now

        This method is usually not needed because the main result properties,
        :attr:`.eigenvalues` and :attr:`.eigenvectors`, will call this implicitly
        the first time they are accessed. However, since the :meth:`solve()` routine
        may be computationally expensive, it is useful to have the ability to call it
        ahead of time as needed.
        """
        self.impl.solve()

    def clear(self):
        """Clear the computed results and start over"""
        self.impl.clear()

    def report(self, shortform=False) -> str:
        """Return a report of the last :meth:`solve()` computation

        Parameters
        ----------
        shortform : bool, optional
            Return a short one line version of the report
        """
        return self.impl.report(shortform)

    def set_wave_vector(self, k):
        """Set the wave vector for periodic models

        Parameters
        ----------
        k : array_like
            Wave vector in reciprocal space.
        """
        self.clear()
        self.model.set_wave_vector(k)

    def calc_eigenvalues(self, map_probability_at=None):
        """Return an :class:`.Eigenvalues` result object with an optional probability colormap

        While the :attr:`.eigenvalues` property returns the raw values array, this
        method returns a result object with more data. In addition to the energy
        states, this result may show a colormap of the probability density for each
        state at a single position.

        Parameters
        ----------
        map_probability_at : array_like, optional
            Cartesian position where the probability density of each energy state
            should be calculated.

        Returns
        -------
        :class:`~pybinding.Eigenvalues`
        """
        if not map_probability_at:
            return results.Eigenvalues(self.eigenvalues)
        else:
            site_idx = self.system.find_nearest(position=map_probability_at)
            probability = abs(self.eigenvectors[site_idx, :])**2

            # sum probabilities of degenerate states
            for idx in self.find_degenerate_states(self.eigenvalues):
                probability[idx] = np.sum(probability[idx]) / len(idx)

            return results.Eigenvalues(self.eigenvalues, probability)

    def calc_probability(self, n, reduce=1e-5):
        r"""Calculate the spatial probability density

        .. math::
            \text{P}(r) = |\Psi_n(r)|^2

        for each position :math:`r` in `system.positions` where :math:`\Psi_n(r)`
        is `eigenvectors[:, n]`.

        Parameters
        ----------
        n : int or array_like
            Index of the desired eigenstate. If an array of indices is given, the
            probability will be calculated at each one and a sum will be returned.
        reduce : float, optional
            Reduce degenerate states by summing their probabilities. Neighboring
            states are considered degenerate if their energy is difference is lower
            than the value of `reduce`. This is disabled by passing `reduce=0`.

        Returns
        -------
        :class:`~pybinding.StructureMap`
        """
        if reduce and np.isscalar(n):
            n = np.flatnonzero(abs(self.eigenvalues[n] - self.eigenvalues) < reduce)

        probability = abs(self.eigenvectors[:, n]) ** 2
        if probability.ndim > 1:
            probability = np.sum(probability, axis=1)
        return self.system.with_data(probability)

    def calc_dos(self, energies, broadening):
        r"""Calculate the density of states as a function of energy

        .. math::
            \text{DOS}(E) = \frac{1}{c \sqrt{2\pi}}
                            \sum_n{e^{-\frac{(E_n - E)^2}{2 c^2}}}

        for each :math:`E` in `energies`, where :math:`c` is `broadening` and
        :math:`E_n` is `eigenvalues[n]`.

        Parameters
        ----------
        energies : array_like
            Values for which the DOS is calculated.
        broadening : float
            Controls the width of the Gaussian broadening applied to the DOS.

        Returns
        -------
        :class:`~pybinding.Series`
        """
        if hasattr(self.impl, 'calc_dos'):
            dos = self.impl.calc_dos(energies, broadening)
        else:
            scale = 1 / (broadening * math.sqrt(2 * math.pi))
            delta = self.eigenvalues[:, np.newaxis] - energies
            dos = scale * np.sum(np.exp(-0.5 * delta**2 / broadening**2), axis=0)
        return results.Series(energies, dos, labels=dict(variable="E (eV)", data="DOS"))

    def calc_ldos(self, energies, broadening, position, sublattice="", reduce=True):
        r"""Calculate the local density of states as a function of energy at the given position

        .. math::
            \text{LDOS}(E) = \frac{1}{c \sqrt{2\pi}}
                             \sum_n{|\Psi_n(r)|^2 e^{-\frac{(E_n - E)^2}{2 c^2}}}

        for each :math:`E` in `energies`, where :math:`c` is `broadening`,
        :math:`E_n` is `eigenvalues[n]` and :math:`r` is a single site position
        determined by the arguments `position` and `sublattice`.

        Parameters
        ----------
        energies : array_like
            Values for which the DOS is calculated.
        broadening : float
            Controls the width of the Gaussian broadening applied to the DOS.
        position : array_like
            Cartesian position of the lattice site for which the LDOS is calculated.
            Doesn't need to be exact: the method will find the actual site which is
            closest to the given position.
        sublattice : str
            Only look for sites of a specific sublattice, closest to `position`.
            The default value considers any sublattice.
        reduce : bool
            This option is only relevant for multi-orbital models. If true, the
            resulting LDOS will summed over all the orbitals at the target site
            and the result will be a 1D array. If false, the individual orbital
            results will be preserved and the result will be a 2D array with
            `shape == (energy.size, num_orbitals)`.

        Returns
        -------
        :class:`~pybinding.Series`
        """
        if hasattr(self.impl, 'calc_ldos'):
            ldos = self.impl.calc_ldos(energies, broadening, position, sublattice)
        else:
            delta = self.eigenvalues[:, np.newaxis] - energies
            gaussian = np.exp(-0.5 * delta**2 / broadening**2)
            scale = 1 / (broadening * math.sqrt(2 * math.pi))

            sys_idx = self.system.find_nearest(position, sublattice)
            ham_idx = self.system.to_hamiltonian_indices(sys_idx)

            def calc_single(index):
                psi2 = np.abs(self.eigenvectors[index])**2
                return scale * np.sum(psi2[:, np.newaxis] * gaussian, axis=0)

            ldos = np.array([calc_single(i) for i in ham_idx]).T
            if reduce:
                ldos = np.sum(ldos, axis=1)

        return results.Series(energies, ldos.squeeze(), labels=dict(variable="E (eV)", data="LDOS",
                                                                    columns="orbitals"))

    def calc_spatial_ldos(self, energy, broadening):
        r"""Calculate the spatial local density of states at the given energy

        .. math::
            \text{LDOS}(r) = \frac{1}{c \sqrt{2\pi}}
                             \sum_n{|\Psi_n(r)|^2 e^{-\frac{(E_n - E)^2}{2 c^2}}}

        for each position :math:`r` in `system.positions`, where :math:`E` is `energy`,
        :math:`c` is `broadening`, :math:`E_n` is `eigenvalues[n]` and :math:`\Psi_n(r)`
        is `eigenvectors[:, n]`.

        Parameters
        ----------
        energy : float
            The energy value for which the spatial LDOS is calculated.
        broadening : float
            Controls the width of the Gaussian broadening applied to the DOS.

        Returns
        -------
        :class:`~pybinding.StructureMap`
        """
        if hasattr(self.impl, 'calc_spatial_ldos'):
            ldos = self.impl.calc_spatial_ldos(energy, broadening)
        else:
            scale = 1 / (broadening * math.sqrt(2 * math.pi))
            gaussian = np.exp(-0.5 * (self.eigenvalues - energy)**2 / broadening**2)
            psi2 = np.abs(self.eigenvectors)**2
            ldos = scale * np.sum(psi2 * gaussian, axis=1)

        return self.system.with_data(ldos)

    def calc_bands(self, k0, k1, *ks, step=0.1):
        """Calculate the band structure on a path in reciprocal space

        Parameters
        ----------
        k0, k1, *ks : array_like
            Points in reciprocal space which form the path for the band calculation.
            At least two points are required.
        step : float, optional
            Calculation step length in reciprocal space units. Lower `step` values
            will return more detailed results.

        Returns
        -------
        :class:`~pybinding.Bands`
        """
        k_points = [np.atleast_1d(k) for k in (k0, k1) + ks]
        k_path = results.make_path(*k_points, step=step)

        bands = []
        for k in k_path:
            self.set_wave_vector(k)
            bands.append(self.eigenvalues)

        return results.Bands(k_path, np.vstack(bands))

    @staticmethod
    def find_degenerate_states(energies, abs_tolerance=1e-5):
        """Return groups of indices which belong to degenerate states

        Parameters
        ----------
        energies : array_like
        abs_tolerance : float, optional

        Examples
        --------
        >>> energies = np.array([0.1, 0.1, 0.2, 0.5, 0.5, 0.5, 0.7, 0.8, 0.8])
        >>> Solver.find_degenerate_states(energies)
        [[0, 1], [3, 4, 5], [7, 8]]

        >>> energies = np.array([0.1, 0.2, 0.5, 0.7])
        >>> Solver.find_degenerate_states(energies)
        []
        """
        # when:   energy == [0.1, 0.1, 0.2, 0.5, 0.5, 0.5, 0.7, 0.8, 0.8]
        # ...     idx == [0, 3, 4, 7]
        idx = np.flatnonzero(abs(np.diff(energies)) < abs_tolerance)
        if idx.size == 0:
            return []
        groups = np.split(idx, np.flatnonzero(np.diff(idx) != 1) + 1)
        # ...     groups == [[0], [3, 4], [7]]
        # return: [[0, 1], [3, 4, 5], [7, 8]]
        return [list(g) + [g[-1] + 1] for g in groups]

class _SolverPythonImpl:
    """Python eigensolver implementation

    This is intended to make use of scipy's LAPACK and ARPACK solvers.
    """
    def __init__(self, solve_func, model, **kwargs):
        self.solve_func = solve_func
        self._model = model

        self.kwargs = kwargs
        self.vals = np.empty(0)
        self.vecs = np.empty(0)
        self.compute_time = .0

    def clear(self):
        self.vals = np.empty(0)
        self.vecs = np.empty(0)
        self.compute_time = .0

    @property
    def model(self):
        return self._model

    @model.setter
    def model(self, model):
        self.clear()
        self._model = model

    @property
    def system(self):
        return self.model.system.impl

    @property
    def eigenvalues(self) -> np.ndarray:
        self.solve()
        return self.vals

    @property
    def eigenvectors(self) -> np.ndarray:
        self.solve()
        return self.vecs

    def solve(self):
        if len(self.vals):
            return

        start_time = time.time()

        self.vals, self.vecs = self.solve_func(self.model.hamiltonian, **self.kwargs)
        idx = self.vals.argsort()
        self.vals = self.vals[idx]
        self.vecs = self.vecs[:, idx]

        self.compute_time = time.time() - start_time

    def report(self, _=False):
        from .utils.time import pretty_duration
        return "Converged in " + pretty_duration(self.compute_time)

def lapack(model, **kwargs):
    """LAPACK :class:`.Solver` implementation for dense matrices

    This solver is intended for small models which are best represented by
    dense matrices. Always solves for all the eigenvalues and eigenvectors.
    Internally this solver uses the :func:`scipy.linalg.eigh` function for
    dense Hermitian matrices.

    Parameters
    ----------
    model : Model
        Model which will provide the Hamiltonian matrix.
    **kwargs
        Advanced arguments: forwarded to :func:`scipy.linalg.eigh`.

    Returns
    -------
    :class:`~pybinding.solver.Solver`
    """
    def solver_func(hamiltonian, **kw):
        from scipy.linalg import eigh
        return eigh(hamiltonian.toarray(), **kw)

    return Solver(_SolverPythonImpl(solver_func, model, **kwargs))

def arpack(model, k, sigma=0, **kwargs):
    """ARPACK :class:`.Solver` implementation for sparse matrices

    This solver is intended for large models with sparse Hamiltonian matrices.
    It only computes a small targeted subset of eigenvalues and eigenvectors.
    Internally this solver uses the :func:`scipy.sparse.linalg.eigsh` function
    for sparse Hermitian matrices.

    Parameters
    ----------
    model : Model
        Model which will provide the Hamiltonian matrix.
    k : int
        The desired number of eigenvalues and eigenvectors. This number must be smaller
        than the size of the matrix, preferably much smaller for optimal performance.
        The computed eigenvalues are the ones closest to `sigma`.
    sigma : float, optional
        Look for eigenvalues near `sigma`.
    **kwargs
        Advanced arguments: forwarded to :func:`scipy.sparse.linalg.eigsh`.

    Returns
    -------
    :class:`~pybinding.solver.Solver`
    """
    from scipy.sparse.linalg import eigsh
    if sigma == 0:
        # eigsh can cause problems when sigma is exactly zero
        sigma = np.finfo(model.hamiltonian.dtype).eps
    return Solver(_SolverPythonImpl(eigsh, model, k=k, sigma=sigma, **kwargs))

def feast(model, energy_range, initial_size_guess, recycle_subspace=False, is_verbose=False):
    """FEAST :class:`.Solver` implementation for sparse matrices

    This solver is only available if the C++ extension module was compiled with FEAST.

    Parameters
    ----------
    model : Model
        Model which will provide the Hamiltonian matrix.
    energy_range : tuple of float
        The lowest and highest eigenvalue between which to compute the solutions.
    initial_size_guess : int
        Initial user guess for number of eigenvalues which will be found in the given
        `energy_range`. This value may be completely wrong - the solver will auto-correct
        as needed. However, for optimal performance the estimate should be as close to
        1.5 * actual_size as possible.
    recycle_subspace : bool, optional
        Reuse previously computed values as a starting point for the next computation.
        This improves performance when subsequent computations differ only slightly, as
        is the case for the band structure of periodic systems where the results change
        gradually as a function of the wave vector. It may hurt performance otherwise.
    is_verbose : bool, optional
        Show the raw output from the FEAST routine.

    Returns
    -------
    :class:`~pybinding.solver.Solver`
    """
    try:
        # noinspection PyUnresolvedReferences
        return Solver(_cpp.FEAST(model, energy_range, initial_size_guess,
                                 recycle_subspace, is_verbose))
    except AttributeError:
        raise Exception("The module was compiled without the FEAST solver.\n"
                        "Use a different solver or recompile the module with FEAST.")


import torch
import os
import shutil
import numpy as np
import logging
import salem
import rasterio
from scipy.optimize import minimize
import matplotlib.pyplot as plt
from combine2d.core import data_logging
from combine2d.core.data_logging import DataLogger
from combine2d.core.arithmetics import RMSE, mean_BIAS
from combine2d.core.utils import NonRGIGlacierDirectory
from combine2d.core.cost_function import create_cost_func
from combine2d.core.data_logging import write_pickle, load_pickle
# -------------------------------
# Further initialization / extended import tasks
# Module logger
log = logging.getLogger(__name__)

class InversionDirectory(object):

    def __init__(self, gdir: NonRGIGlacierDirectory):
        self.gdir = gdir
        self.inv_settings = gdir.inversion_settings
        self.true_bed = None
        self.first_guessed_bed = None
        self.surf_noise = None
        self.ref_surf = None
        self.ice_mask = None
        self.minimize_log = ''
        self.cost_func = None
        self.data_logger = None
        self.bed_measurements = None
        if not 'minimize_bounds' in self.inv_settings:
            self.inv_settings['minimize_bounds'] = None

    def iteration_info_callback(self, x0):
        i = len(self.data_logger.costs) - 1
        if i >= 0:
            dl = self.data_logger
            dl.step_indices.append(i)
            b = self.true_bed
            log_entry = '''
            ----------------------------------------------
            Function Call: {func_call:d}
            Iteration: {iteration:d}
            Cost: {cost:g}
            Bed RMSE: {bed_rmse:g}
            Bed Bias: {bed_bias:g}
            Bed Max_diff: {bed_maxdiff:g}
            Surface RMSE: {surf_rmse:g}
            Surface Max_diff: {surf_maxdiff:g}
            '''
            myargs = {
                'func_call': i,
                'iteration': len(dl.step_indices),
                'cost': dl.costs[i],
                'bed_rmse': RMSE(dl.beds[i], b),
                'bed_bias': mean_BIAS(dl.beds[i], b, np.sum(self.ice_mask)),
                'bed_maxdiff': np.max(np.abs(dl.beds[i] - b)),
                'surf_rmse': RMSE(dl.surfs[i], self.ref_surf),
                'surf_maxdiff': np.max(np.abs(dl.surfs[i] - self.ref_surf))
            }

            if self.surf_noise is not None:
                log_entry += 'RMSE to perturbed surf: {:g}\n'.format(
                    RMSE(dl.surfs[i], self.ref_surf + self.surf_noise))
            log_entry = log_entry.format(**myargs)
            print(log_entry)
            self.minimize_log += log_entry

    def write_string_to_file(self, filename, text):
        dir = self.get_current_basedir()
        if not os.path.exists(dir):
            os.makedirs(dir, exist_ok=True)
        with open(os.path.join(dir, filename), 'w') as f:
            f.write(text)

    def get_current_basedir(self):
        return os.path.join(self.gdir.dir,
                            self.inv_settings['inversion_subdir'])

    def clear_dir(self, dir):
        if os.path.exists(dir):
            for f in os.listdir(dir):
                if (not str.endswith(f, '.py')) and (
                        not os.path.isdir(os.path.join(dir, f))):
                    os.remove(os.path.join(dir, f))
                elif os.path.isdir(os.path.join(dir, f)):
                    shutil.rmtree(os.path.join(dir, f))
        else:
            if not os.path.exists(dir):
                os.makedirs(dir, exist_ok=True)

    def _read_all_data(self):
        """
        Reads all necessary information from files in gdir for
        minimization/optimization and logging.
        """
        self.true_bed = salem.GeoTiff(
            self.gdir.get_filepath('dem')).get_vardata()
        self.ref_surf = salem.GeoTiff(
            self.gdir.get_filepath('ref_dem')).get_vardata()
        self.first_guessed_bed = salem.GeoTiff(
            self.get_subdir_filepath('first_guessed_bed')).get_vardata()
        self.ice_mask = np.load(self.gdir.get_filepath('ref_ice_mask'))
        if os.path.exists(self.gdir.get_filepath('dem_noise')): #TODO: once
            # surface noise is present, it cant get rid off ...
            shutil.copy(self.gdir.get_filepath('dem_noise'),
                        self.get_subdir_filepath('dem_noise'))
            self.surf_noise = np.load(self.get_subdir_filepath('dem_noise'))
        else:
            self.surf_noise = None

        if os.path.exists(self.gdir.get_filepath('bed_measurements')):
            shutil.copy(self.gdir.get_filepath('bed_measurements'),
                        self.get_subdir_filepath('bed_measurements'))
            self.bed_measurements = np.load(self.get_subdir_filepath(
                'bed_measurements'))
        else:
            self.bed_measurements = None

    def get_subdir_filepath(self, filename, filesuffix=None):
        """
        Gets the filepath for a file with a given name (without extension).
        Works as and is based on get_filepath in GlacierDirectory,
        but returns filepath in this inversion directory.

        Parameters
        ----------
        filename: str
            name of the file
        filesuffix: str
            optional filesuffix to the filename

        Returns
        -------
        Entire path to this file in this inversion directory

        """
        original_path = self.gdir.get_filepath(filename, filesuffix=filesuffix)
        original_path = os.path.split(original_path)
        return os.path.join(self.get_current_basedir(), original_path[1])

    def get_bounds(self):
        """
        Creates bounds for the minimization on the current domain. If
        'bounds_min_max' in inversion settings is None, no bounds are set.
        Else, in Areas without ice, upper and lower bound are exactly as the
        observed surface, otherwise min and max values for ice thickness are
        taken from 'bounds_min_max' and give bounds in glacierized areas. (
        min ice thickness determines upper bound and max ice thickness
        determines lower bound)

        Returns
        -------
        bounds for this domain and this inversion settings
        """
        bounds = None
        if self.inv_settings['bounds_min_max'] is not None:
            surf = self.ref_surf
            if self.surf_noise is not None:
                surf += self.surf_noise

            upper_bounds = surf.copy()
            lower_bounds = surf.copy()

            min_ice_thickness = self.inv_settings['bounds_min_max'][0]
            max_ice_thickness = self.inv_settings['bounds_min_max'][1]
            if min_ice_thickness is not None:
                upper_bounds = upper_bounds - min_ice_thickness * self.ice_mask
            else:
                upper_bounds = np.where(self.ice_mask, None, upper_bounds)

            if max_ice_thickness is not None:
                lower_bounds = lower_bounds - max_ice_thickness * self.ice_mask
            else:
                lower_bounds = np.where(self.ice_mask, None, lower_bounds)

            bounds = np.c_[lower_bounds.flatten(), upper_bounds.flatten()]
        return bounds

    def run_minimize(self):
        """
        Here the actual minimization of the cost_function is done via
        scipy.optimize.minimize.
        First, data from the glacier directory is read and optionally a
        DataLogger is created. The inversion settings used for this
        particular inversion are saved in this subdirectory. Bounds for the
        minimization are derived. Then the cost function is created and the
        minimization of this cost function started. In the end, the result is
        written to disk and optionally, further information is written to disk.

        The whole process is dominated by the set inversion settings

        Returns
        -------
        Result of minimization as scipy.optimize.minimize returns (res.x
        gives flattened ndarray with bed, needs to be reshaped)

        """

        # Copy first_guessed_bed to inversion directory
        if self.inv_settings['log_minimize_steps']:
            # TODO: really useful? -> respect reset argument in gdir?
            self.clear_dir(self.get_current_basedir())

        with rasterio.open(self.gdir.get_filepath('first_guessed_bed')) as src:
            profile = src.profile
            data = src.read(1)
        with rasterio.open(self.get_subdir_filepath('first_guessed_bed'),
                           'w', **profile) as dst:
            dst.write(data, 1)
        if os.path.exists(self.gdir.get_filepath('first_guessed_bed_noise')):
            shutil.copy(self.gdir.get_filepath('first_guessed_bed_noise'),
                        self.get_subdir_filepath('first_guessed_bed_noise'))

        write_pickle(self.inv_settings,
                     self.get_subdir_filepath('inversion_settings'))
        # Write out reg_parameters to check easier later on
        self.write_string_to_file(self.get_subdir_filepath('reg_parameters'),
                                  str(self.inv_settings['reg_parameters']))
        self.inv_settings = load_pickle(
            self.get_subdir_filepath('inversion_settings'))
        self._read_all_data()
        self.minimize_log = ''
        self.data_logger = None
        callback = None

        if self.inv_settings['log_minimize_steps']:
            dl = DataLogger(self)
            self.data_logger = dl
            callback = self.iteration_info_callback

        # ----------------------------------------------------------------------
        # Core: things are happening here:
        bounds = self.get_bounds()

        self.cost_func = create_cost_func(self.gdir, self.data_logger,
                                          self.surf_noise,
                                          self.bed_measurements)
        res = None
        try:
            res = minimize(fun=self.cost_func,
                           x0=self.first_guessed_bed.astype(np.float64).flatten(),
                           method=self.inv_settings['solver'], jac=True,
                           bounds=bounds,
                           options=self.inv_settings['minimize_options'],
                           callback=callback)

            inverted_bed = res.x.reshape(self.first_guessed_bed.shape)
            # ----------------------------------------------------------------------

            profile['dtype'] = 'float64'
            with rasterio.open(self.get_subdir_filepath('inverted_bed'),
                               'w', **profile) as dst:
                dst.write(inverted_bed, 1)

        except MemoryError as me:
            self.write_string_to_file(os.path.join(self.get_current_basedir(),
                                                   'warning.txt'),
                                      'Error during iteration: ' + str(me))

        if self.inv_settings['log_minimize_steps']:
            self.write_string_to_file('log.txt', self.minimize_log)
            dir = self.get_current_basedir()
            dl.filter_data_from_optimization()  # Optional, if we want to
            data_logging.write_pickle(dl,
                                      self.get_subdir_filepath('data_logger'))
            #dl.plot_all(dir)
            #plt.close('all')

        return res


import numpy as np

# This section of code is copied from ProFET (Ofer & Linial, DOI: 10.1093/bioinformatics/btv345)
# Original comment by the ProFET authors: 'Acquired from georgiev's paper of
# AAscales using helper script "GetTextData.py". + RegEx cleaning DOI: 10.1089/cmb.2008.0173'
gg_1 = {'Q': -2.54, 'L': 2.72, 'T': -0.65, 'C': 2.66, 'I': 3.1, 'G': 0.15, 'V': 2.64, 'K': -3.89, 'M': 1.89, 'F': 3.12, 'N': -2.02, 'R': -2.8, 'H': -0.39, 'E': -3.08, 'W': 1.89, 'A': 0.57, 'D': -2.46, 'Y': 0.79, 'S': -1.1, 'P': -0.58}
gg_2 = {'Q': 1.82, 'L': 1.88, 'T': -1.6, 'C': -1.52, 'I': 0.37, 'G': -3.49, 'V': 0.03, 'K': 1.47, 'M': 3.88, 'F': 0.68, 'N': -1.92, 'R': 0.31, 'H': 1, 'E': 3.45, 'W': -0.09, 'A': 3.37, 'D': -0.66, 'Y': -2.62, 'S': -2.05, 'P': -4.33}
gg_3 = {'Q': -0.82, 'L': 1.92, 'T': -1.39, 'C': -3.29, 'I': 0.26, 'G': -2.97, 'V': -0.67, 'K': 1.95, 'M': -1.57, 'F': 2.4, 'N': 0.04, 'R': 2.84, 'H': -0.63, 'E': 0.05, 'W': 4.21, 'A': -3.66, 'D': -0.57, 'Y': 4.11, 'S': -2.19, 'P': -0.02}
gg_4 = {'Q': -1.85, 'L': 5.33, 'T': 0.63, 'C': -3.77, 'I': 1.04, 'G': 2.06, 'V': 2.34, 'K': 1.17, 'M': -3.58, 'F': -0.35, 'N': -0.65, 'R': 0.25, 'H': -3.49, 'E': 0.62, 'W': -2.77, 'A': 2.34, 'D': 0.14, 'Y': -0.63, 'S': 1.36, 'P': -0.21}
gg_5 = {'Q': 0.09, 'L': 0.08, 'T': 1.35, 'C': 2.96, 'I': -0.05, 'G': 0.7, 'V': 0.64, 'K': 0.53, 'M': -2.55, 'F': -0.88, 'N': 1.61, 'R': 0.2, 'H': 0.05, 'E': -0.49, 'W': 0.72, 'A': -1.07, 'D': 0.75, 'Y': 1.89, 'S': 1.78, 'P': -8.31}
gg_6 = {'Q': 0.6, 'L': 0.09, 'T': -2.45, 'C': -2.23, 'I': -1.18, 'G': 7.47, 'V': -2.01, 'K': 0.1, 'M': 2.07, 'F': 1.62, 'N': 2.08, 'R': -0.37, 'H': 0.41, 'E': 0, 'W': 0.86, 'A': -0.4, 'D': 0.24, 'Y': -0.53, 'S': -3.36, 'P': -1.82}
gg_7 = {'Q': 0.25, 'L': 0.27, 'T': -0.65, 'C': 0.44, 'I': -0.21, 'G': 0.41, 'V': -0.33, 'K': 4.01, 'M': 0.84, 'F': -0.15, 'N': 0.4, 'R': 3.81, 'H': 1.61, 'E': -5.66, 'W': -1.07, 'A': 1.23, 'D': -5.15, 'Y': -1.3, 'S': 1.39, 'P': -0.12}
gg_8 = {'Q': 2.11, 'L': -4.06, 'T': 3.43, 'C': -3.49, 'I': 3.45, 'G': 1.62, 'V': 3.93, 'K': -0.01, 'M': 1.85, 'F': -0.41, 'N': -2.47, 'R': 0.98, 'H': -0.6, 'E': -0.11, 'W': -1.66, 'A': -2.32, 'D': -1.17, 'Y': 1.31, 'S': -1.21, 'P': -1.18}
gg_9 = {'Q': -1.92, 'L': 0.43, 'T': 0.34, 'C': 2.22, 'I': 0.86, 'G': -0.47, 'V': -0.21, 'K': -0.26, 'M': -2.05, 'F': 4.2, 'N': -0.07, 'R': 2.43, 'H': 3.55, 'E': 1.49, 'W': -5.87, 'A': -2.01, 'D': 0.73, 'Y': -0.56, 'S': -2.83, 'P': 0}
gg_10 = {'Q': -1.67, 'L': -1.2, 'T': 0.24, 'C': -3.78, 'I': 1.98, 'G': -2.9, 'V': 1.27, 'K': -1.66, 'M': 0.78, 'F': 0.73, 'N': 7.02, 'R': -0.99, 'H': 1.52, 'E': -2.26, 'W': -0.66, 'A': 1.31, 'D': 1.5, 'Y': -0.95, 'S': 0.39, 'P': -0.66}
gg_11 = {'Q': 0.7, 'L': 0.67, 'T': -0.53, 'C': 1.98, 'I': 0.89, 'G': -0.98, 'V': 0.43, 'K': 5.86, 'M': 1.53, 'F': -0.56, 'N': 1.32, 'R': -4.9, 'H': -2.28, 'E': -1.62, 'W': -2.49, 'A': -1.14, 'D': 1.51, 'Y': 1.91, 'S': -2.92, 'P': 0.64}
gg_12 = {'Q': -0.27, 'L': -0.29, 'T': 1.91, 'C': -0.43, 'I': -1.67, 'G': -0.62, 'V': -1.71, 'K': -0.06, 'M': 2.44, 'F': 3.54, 'N': -2.44, 'R': 2.09, 'H': -3.12, 'E': -3.97, 'W': -0.3, 'A': 0.19, 'D': 5.61, 'Y': -1.26, 'S': 1.27, 'P': -0.92}
gg_13 = {'Q': -0.99, 'L': -2.47, 'T': 2.66, 'C': -1.03, 'I': -1.02, 'G': -0.11, 'V': -2.93, 'K': 1.38, 'M': -0.26, 'F': 5.25, 'N': 0.37, 'R': -3.08, 'H': -1.45, 'E': 2.3, 'W': -0.5, 'A': 1.66, 'D': -3.85, 'Y': 1.57, 'S': 2.86, 'P': -0.37}
gg_14 = {'Q': -1.56, 'L': -4.79, 'T': -3.07, 'C': 0.93, 'I': -1.21, 'G': 0.15, 'V': 4.22, 'K': 1.78, 'M': -3.09, 'F': 1.73, 'N': -0.89, 'R': 0.82, 'H': -0.77, 'E': -0.06, 'W': 1.64, 'A': 4.39, 'D': 1.28, 'Y': 0.2, 'S': -1.88, 'P': 0.17}
gg_15 = {'Q': 6.22, 'L': 0.8, 'T': 0.2, 'C': 1.43, 'I': -1.78, 'G': -0.53, 'V': 1.06, 'K': -2.71, 'M': -1.39, 'F': 2.14, 'N': 3.13, 'R': 1.32, 'H': -4.18, 'E': -0.35, 'W': -0.72, 'A': 0.18, 'D': -1.98, 'Y': -0.76, 'S': -2.42, 'P': 0.36}
gg_16 = {'Q': -0.18, 'L': -1.43, 'T': -2.2, 'C': 1.45, 'I': 5.71, 'G': 0.35, 'V': -1.31, 'K': 1.62, 'M': -1.02, 'F': 1.1, 'N': 0.79, 'R': 0.69, 'H': -2.91, 'E': 1.51, 'W': 1.75, 'A': -2.6, 'D': 0.05, 'Y': -5.19, 'S': 1.75, 'P': 0.08}
gg_17 = {'Q': 2.72, 'L': 0.63, 'T': 3.73, 'C': -1.15, 'I': 1.54, 'G': 0.3, 'V': -1.97, 'K': 0.96, 'M': -4.32, 'F': 0.68, 'N': -1.54, 'R': -2.62, 'H': 3.37, 'E': -2.29, 'W': 2.73, 'A': 1.49, 'D': 0.9, 'Y': -2.56, 'S': -2.77, 'P': 0.16}
gg_18 = {'Q': 4.35, 'L': -0.24, 'T': -5.46, 'C': -1.64, 'I': 2.11, 'G': 0.32, 'V': -1.21, 'K': -1.09, 'M': -1.34, 'F': 1.46, 'N': -1.71, 'R': -1.49, 'H': 1.87, 'E': -1.47, 'W': -2.2, 'A': 0.46, 'D': 1.38, 'Y': 2.87, 'S': 3.36, 'P': -0.34}
gg_19 = {'Q': 0.92, 'L': 1.01, 'T': -0.73, 'C': -1.05, 'I': -4.18, 'G': 0.05, 'V': 4.77, 'K': 1.36, 'M': 0.09, 'F': 2.33, 'N': -0.25, 'R': -2.57, 'H': 2.17, 'E': 0.15, 'W': 0.9, 'A': -4.22, 'D': -0.03, 'Y': -3.43, 'S': 2.67, 'P': 0.04}

# Package all georgiev parameters
georgiev_parameters = [gg_1, gg_2, gg_3, gg_4, gg_5, gg_6, gg_7, gg_8, gg_9,
                       gg_10, gg_11, gg_12, gg_13, gg_14, gg_15, gg_16, gg_17,
                       gg_18, gg_19]

def get_georgiev_params_for_aa(aa):
    return [gg[aa] for gg in georgiev_parameters]

def get_georgiev_params_for_seq(s):
    return np.concatenate([get_georgiev_params_for_aa(aa) for aa in s])

def seqs_to_georgiev(seqs):
    return np.stack([get_georgiev_params_for_seq(s) for s in seqs])


# Copyright 2014-2019 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Samragni Banerjee <samragnibanerjee4@gmail.com>
#         Alexander Sokolov <alexander.y.sokolov@gmail.com>
#

import numpy as np
import pyscf.ao2mo

### Integral transformation ###
def transform_integrals(myadc):
    occ_a = myadc.mo_coeff[0][:,:myadc._nocc[0]]
    occ_b = myadc.mo_coeff[1][:,:myadc._nocc[1]]
    vir_a = myadc.mo_coeff[0][:,myadc._nocc[0]:]
    vir_b = myadc.mo_coeff[1][:,myadc._nocc[1]:]

    occ = occ_a, occ_b
    vir = vir_a, vir_b

    eris = lambda:None

    eris.oovv = transform_antisymmetrize_integrals(myadc._scf._eri, (occ,occ,vir,vir))
    eris.vvvv = transform_antisymmetrize_integrals(myadc._scf._eri, (vir,vir,vir,vir))
    eris.oooo = transform_antisymmetrize_integrals(myadc._scf._eri, (occ,occ,occ,occ))
    eris.voov = transform_antisymmetrize_integrals(myadc._scf._eri, (vir,occ,occ,vir))
    eris.ooov = transform_antisymmetrize_integrals(myadc._scf._eri, (occ,occ,occ,vir))
    eris.vovv = transform_antisymmetrize_integrals(myadc._scf._eri, (vir,occ,vir,vir))
    eris.vvoo = transform_antisymmetrize_integrals(myadc._scf._eri, (vir,vir,occ,occ))
    eris.vvvo = transform_antisymmetrize_integrals(myadc._scf._eri, (vir,vir,vir,occ))
    eris.ovoo = transform_antisymmetrize_integrals(myadc._scf._eri, (occ,vir,occ,occ))
    eris.ovov = transform_antisymmetrize_integrals(myadc._scf._eri, (occ,vir,occ,vir))
    eris.vooo = transform_antisymmetrize_integrals(myadc._scf._eri, (vir,occ,occ,occ))
    eris.oovo = transform_antisymmetrize_integrals(myadc._scf._eri, (occ,occ,vir,occ))
    eris.vovo = transform_antisymmetrize_integrals(myadc._scf._eri, (vir,occ,vir,occ))
    eris.vvov = transform_antisymmetrize_integrals(myadc._scf._eri, (vir,vir,occ,vir))
    eris.ovvo = transform_antisymmetrize_integrals(myadc._scf._eri, (occ,vir,vir,occ))
    eris.ovvv = transform_antisymmetrize_integrals(myadc._scf._eri, (occ,vir,vir,vir))

    return eris

# TODO: disk flag
def transform_antisymmetrize_integrals(v2e_ao, mo, disk = False):

    mo_1, mo_2, mo_3, mo_4 = mo

    mo_1_a, mo_1_b = mo_1
    mo_2_a, mo_2_b = mo_2
    mo_3_a, mo_3_b = mo_3
    mo_4_a, mo_4_b = mo_4

    v2e_a = None
    v2e_a = pyscf.ao2mo.general(v2e_ao, (mo_1_a, mo_3_a, mo_2_a, mo_4_a), compact=False)
    v2e_a = v2e_a.reshape(mo_1_a.shape[1], mo_3_a.shape[1], mo_2_a.shape[1], mo_4_a.shape[1])
    v2e_a = v2e_a.transpose(0,2,1,3).copy()

    if (mo_1_a is mo_2_a):
        v2e_a -= v2e_a.transpose(1,0,2,3).copy()
    elif (mo_3_a is mo_4_a):
        v2e_a -= v2e_a.transpose(0,1,3,2).copy()
    else:
        v2e_temp = None
        v2e_temp = pyscf.ao2mo.general(v2e_ao, (mo_1_a, mo_4_a, mo_2_a, mo_3_a), compact=False)
        v2e_temp = v2e_temp.reshape(mo_1_a.shape[1], mo_4_a.shape[1], mo_2_a.shape[1], mo_3_a.shape[1])
        v2e_a -= v2e_temp.transpose(0,2,3,1).copy()
        del v2e_temp

    v2e_a = disk_helper.dataset(v2e_a) if disk else v2e_a

    v2e_b = None
    v2e_b = pyscf.ao2mo.general(v2e_ao, (mo_1_b, mo_3_b, mo_2_b, mo_4_b), compact=False)
    v2e_b = v2e_b.reshape(mo_1_b.shape[1], mo_3_b.shape[1], mo_2_b.shape[1], mo_4_b.shape[1])
    v2e_b = v2e_b.transpose(0,2,1,3).copy()

    if (mo_1_b is mo_2_b):
        v2e_b -= v2e_b.transpose(1,0,2,3).copy()
    elif (mo_3_b is mo_4_b):
        v2e_b -= v2e_b.transpose(0,1,3,2).copy()
    else:
        v2e_temp = None
        v2e_temp = pyscf.ao2mo.general(v2e_ao, (mo_1_b, mo_4_b, mo_2_b, mo_3_b), compact=False)
        v2e_temp = v2e_temp.reshape(mo_1_b.shape[1], mo_4_b.shape[1], mo_2_b.shape[1], mo_3_b.shape[1])
        v2e_b -= v2e_temp.transpose(0,2,3,1).copy()
        del v2e_temp

    v2e_b = disk_helper.dataset(v2e_b) if disk else v2e_b

    v2e_ab = None
    v2e_ab = pyscf.ao2mo.general(v2e_ao, (mo_1_a, mo_3_a, mo_2_b, mo_4_b), compact=False)
    v2e_ab = v2e_ab.reshape(mo_1_a.shape[1], mo_3_a.shape[1], mo_2_b.shape[1], mo_4_b.shape[1])
    v2e_ab = v2e_ab.transpose(0,2,1,3).copy()

    v2e_ab = disk_helper.dataset(v2e_ab) if disk else v2e_ab

    return (v2e_a, v2e_ab, v2e_b)


#!/usr/bin/env python
"""
Classy data
"""
__author__ = "Sid Mau"

# Python libraries
import os
import glob
import numpy as np
import healpy as hp
import fitsio as fits
import scipy.ndimage

# Ugali libraries
import ugali.utils.healpix
import ugali.utils.projector

# Simple libraries
import simple.filters
import simple.objects.result

# TODO:
# - use point
# - make derived classes for each survey

########################################################################

class Data:
    """
    Class object for analyzing photometric data.
    """
    def __init__(self, survey, nside, datadir, fracdet, band_1, band_2, mag, mag_err, mag_dered, basis_1, basis_2, mag_max):
        self.survey      = survey
        self.nside       = nside
        self.datadir     = datadir
        self.fracdet     = fracdet
        self.band_1      = band_1
        self.band_2      = band_2
        self.mag_1       = mag.format(band_1.upper())
        self.mag_2       = mag.format(band_2.upper())
        self.mag_err_1   = mag_err.format(band_1.upper())
        self.mag_err_2   = mag_err.format(band_2.upper())
        self.mag_dered_1 = mag_dered.format(band_1.upper())
        self.mag_dered_2 = mag_dered.format(band_2.upper())
        self.basis_1     = basis_1
        self.basis_2     = basis_2
        self.mag_max     = mag_max

    def quality_filter(self, data):
        """
        Return cut on high quality objects.
        """
        return simple.filters.quality_filter(self.survey, data)

    def star_filter(self, data):
        """
        Return cut on star-like objects.
        """
        return simple.filters.star_filter(self.survey, data)

    def galaxy_filter(self, data):
        """
        Return cut on galaxy-like objects.
        """
        return simple.filters.galaxy_filter(self.survey, data)

    def color_filter(self, data):
        """
        Return cut on blue objects.
        """
        return simple.filters.color_filter(self.survey, data)

    def dered_mag(self, data):
        """
        Deredden magnitudes.
        """
        return simple.filters.dered_mag(self.survey, data)

    @property
    def load_fracdet(self):
        """
        Load fracdet if given.
        """
        if self.fracdet is not None:
            return fits.read(self.fracdet)
        else:
            msg = "No fracdet found at {}".format(self.fracdet)
            raise Exception(msg)
            return None

    def load_local_data(self, ra, dec):
        """
        Load data corresponding to the 8 nearest neighbors of the
        centroid at (ra, dec) into memory.
        """
        hpixel  = ugali.utils.healpix.angToPix(self.nside, ra, dec)
        hpixels = np.concatenate([[hpixel], hp.get_all_neighbours(self.nside, hpixel)])

        data_array = []
        for hpix in hpixels:
            inlist = glob.glob('{}/*_{:05d}.fits'.format(self.dirname, hpix))
            for infile in inlist:
                if not os.path.exists(infile):
                    continue
                data_array.append(fits.read(infile))
        data = np.concatenate(data_array)

        return data

    def compute_char_density(self, ra, dec):
        """
        Compute the characteristic density of a region.
        """

        data = self.load_local_data(ra, dec)

        mag_cut = (data[self.mag_1] < self.mag_max)

        proj = ugali.utils.projector.Projector(ra, dec)
        x, y = proj.sphereToImage(data[self.basis_1][mag_cut], data[self.basis_2][mag_cut]) # Trimmed magnitude range for hotspot finding
        delta_x = 0.01
        area = delta_x**2
        smoothing = 2. / 60. # Was 3 arcmin
        bins = np.arange(-8., 8. + 1.e-10, delta_x)
        centers = 0.5 * (bins[0: -1] + bins[1:])
        yy, xx = np.meshgrid(centers, centers)

        h = np.histogram2d(x, y, bins=[bins, bins])[0]

        h_g = scipy.ndimage.filters.gaussian_filter(h, smoothing / delta_x)

        delta_x_coverage = 0.1
        area_coverage = (delta_x_coverage)**2
        bins_coverage = np.arange(-5., 5. + 1.e-10, delta_x_coverage)
        h_coverage = np.histogram2d(x, y, bins=[bins_coverage, bins_coverage])[0]
        h_goodcoverage = np.histogram2d(x, y, bins=[bins_coverage, bins_coverage])[0]

        n_goodcoverage = h_coverage[h_goodcoverage > 0].flatten()

        characteristic_density = np.median(n_goodcoverage) / area_coverage # per square degree
        print('Characteristic density = {:0.1f} deg^-2').format(characteristic_density)

        # Use pixels with fracdet ~1.0 to estimate the characteristic density
        if self.fracdet is not None:
            fracdet = fits.read(self.fracdet)
            fracdet_zero = np.tile(0., len(fracdet))
            cut = (fracdet != hp.UNSEEN)
            fracdet_zero[cut] = fracdet[cut]

            nside_fracdet = hp.npix2nside(len(fracdet))

            subpix_region_array = []
            for pix in np.unique(ugali.utils.healpix.angToPix(self.nside, data[self.basis_1], data[self.basis_2])):
                subpix_region_array.append(ugali.utils.healpix.subpixel(pix, self.nside, nside_fracdet))
            subpix_region_array = np.concatenate(subpix_region_array)

            # Compute mean fracdet in the region so that this is available as a correction factor
            cut = (fracdet[subpix_region_array] != hp.UNSEEN)
            mean_fracdet = np.mean(fracdet[subpix_region_array[cut]])

            # Correct the characteristic density by the mean fracdet value
            characteristic_density_raw = 1. * characteristic_density
            characteristic_density /= mean_fracdet
            print('Characteristic density (fracdet corrected) = {:0.1f} deg^-2').format(characteristic_density)

        return characteristic_density

    def compute_local_char_density(self, ra, dec, x_peak, y_peak, angsep_peak):
        """
        Compute the local characteristic density of a region.
        """

        data = self.load_local_data(ra, dec)

        characteristic_density = self.compute_char_density(ra, dec)

        mag_cut = (data[mag_dered_1] < self.mag_max)

        proj = ugali.utils.projector.Projector(ra, dec)
        x, y = proj.sphereToImage(data[basis_1][mag_cut], data[basis_2][mag_cut]) # Trimmed magnitude range for hotspot finding

        # If fracdet map is available, use that information to either compute local density,
        # or in regions of spotty coverage, use the typical density of the region
        if self.fracdet is not None:
            fracdet = fits.read(self.fracdet)
            fracdet_zero = np.tile(0., len(fracdet))
            cut = (fracdet != hp.UNSEEN)
            fracdet_zero[cut] = fracdet[cut]

            nside_fracdet = hp.npix2nside(len(fracdet))

            subpix_region_array = []
            for pix in np.unique(ugali.utils.healpix.angToPix(nside, data[basis_1], data[basis_2])):
                subpix_region_array.append(ugali.utils.healpix.subpixel(pix, self.nside, nside_fracdet))
            subpix_region_array = np.concatenate(subpix_region_array)

            # Compute mean fracdet in the region so that this is available as a correction factor
            cut = (fracdet[subpix_region_array] != hp.UNSEEN)
            mean_fracdet = np.mean(fracdet[subpix_region_array[cut]])

            subpix_region_array = subpix_region_array[fracdet[subpix_region_array] > 0.99]
            subpix = ugali.utils.healpix.angToPix(nside_fracdet,
                                                  data[self.basis_1][mag_cut],
                                                  data[self.basis_2][mag_cut]) # Remember to apply mag threshold to objects

            # This is where the local computation begins
            ra_peak, dec_peak = proj.imageToSphere(x_peak, y_peak)
            subpix_all = ugali.utils.healpix.angToDisc(nside_fracdet, ra_peak, dec_peak, 0.5)
            subpix_inner = ugali.utils.healpix.angToDisc(nside_fracdet, ra_peak, dec_peak, 0.3)
            subpix_annulus = subpix_all[~np.in1d(subpix_all, subpix_inner)]
            mean_fracdet = np.mean(fracdet_zero[subpix_annulus])
            print('mean_fracdet {}'.format(mean_fracdet))
            if mean_fracdet < 0.5:
                characteristic_density_local = characteristic_density
                print('characteristic_density_local baseline {}').format(characteristic_density_local)
            else:
                # Check pixels in annulus with complete coverage
                subpix_annulus_region = np.intersect1d(subpix_region_array, subpix_annulus)
                print('{} percent pixels with complete coverage'.format(float(len(subpix_annulus_region)) / len(subpix_annulus)))
                if (float(len(subpix_annulus_region)) / len(subpix_annulus)) < 0.25:
                    characteristic_density_local = characteristic_density
                    print('characteristic_density_local spotty {}'.format(characteristic_density_local))
                else:
                    characteristic_density_local = float(np.sum(np.in1d(subpix, subpix_annulus_region))) \
                                                   / (hp.nside2pixarea(nside_fracdet, degrees=True) * len(subpix_annulus_region)) # deg^-2
                    print('characteristic_density_local cleaned up {}'.format(characteristic_density_local))
        else:
            # Compute the local characteristic density
            area_field = np.pi * (0.5**2 - 0.3**2)
            n_field = np.sum((angsep_peak > 0.3) & (angsep_peak < 0.5))
            characteristic_density_local = n_field / area_field

            # If not good azimuthal coverage, revert
            cut_annulus = (angsep_peak > 0.3) & (angsep_peak < 0.5)
            #phi = np.degrees(np.arctan2(y_full[cut_annulus] - y_peak, x_full[cut_annulus] - x_peak)) # Use full magnitude range, NOT TESTED!!!
            phi = np.degrees(np.arctan2(y[cut_annulus] - y_peak, x[cut_annulus] - x_peak)) # Impose magnitude threshold
            h = np.histogram(phi, bins=np.linspace(-180., 180., 13))[0]
            if np.sum(h > 0) < 10 or np.sum(h > 0.5 * np.median(h)) < 10:
                #angsep_peak = np.sqrt((x - x_peak)**2 + (y - y_peak)**2)
                characteristic_density_local = characteristic_density

        print('Characteristic density local = {:0.1f} deg^-2 = {:0.3f} arcmin^-2'.format(characteristic_density_local, characteristic_density_local / 60.**2))

        return characteristic_density_local

    def find_peaks(self, ra, dec, distance_modulus):
        """
        Convolve field to find characteristic density and peaks within the selected pixel.
        """

        pix_nside_select = ugali.utils.healpix.angToPix(self.nside, ra, dec)
        data = self.load_local_data(ra, dec)
        characteristic_density = self.compute_char_density(ra, dec)
        mag_cut = (data[mag_dered_1] < self.mag_max)

        # convolve field and find peaks
        proj = ugali.utils.projector.Projector(ra, dec)
        x, y = proj.sphereToImage(data[self.basis_1][mag_cut], data[self.basis_2][mag_cut]) # Trimmed magnitude range for hotspot finding
        delta_x = 0.01
        area = delta_x**2
        smoothing = 2. / 60. # Was 3 arcmin
        bins = np.arange(-8., 8. + 1.e-10, delta_x)
        centers = 0.5 * (bins[0: -1] + bins[1:])
        yy, xx = np.meshgrid(centers, centers)

        h = np.histogram2d(x, y, bins=[bins, bins])[0]

        h_g = scipy.ndimage.filters.gaussian_filter(h, smoothing / delta_x)

        factor_array = np.arange(1., 5., 0.05)
        rara, decdec = proj.imageToSphere(xx.flatten(), yy.flatten())
        cutcut = (ugali.utils.healpix.angToPix(self.nside, rara, decdec) == pix_nside_select).reshape(xx.shape)
        threshold_density = 5 * characteristic_density * area
        for factor in factor_array:
            h_region, n_region = scipy.ndimage.measurements.label((h_g * cutcut) > (area * characteristic_density * factor))
            #print 'factor', factor, n_region, n_region < 10
            if n_region < 10:
                threshold_density = area * characteristic_density * factor
                break

        h_region, n_region = scipy.ndimage.measurements.label((h_g * cutcut) > threshold_density)
        h_region = np.ma.array(h_region, mask=(h_region < 1))

        x_peak_array = []
        y_peak_array = []
        angsep_peak_array = []

        for index in range(1, n_region + 1): # loop over peaks
            index_peak = np.argmax(h_g * (h_region == index))
            x_peak, y_peak = xx.flatten()[index_peak], yy.flatten()[index_peak]
            #print index, np.max(h_g * (h_region == index))

            #angsep_peak = np.sqrt((x_full - x_peak)**2 + (y_full - y_peak)**2) # Use full magnitude range, NOT TESTED!!!
            angsep_peak = np.sqrt((x - x_peak)**2 + (y - y_peak)**2) # Impose magnitude threshold

            x_peak_array.append(x_peak)
            y_peak_array.append(y_peak)
            angsep_peak_array.append(angsep_peak)

        return x_peak_array, y_peak_array, angsep_peak_array

    def fit_aperture(self, ra, dec, proj, distance_modulus, x_peak, y_peak, angsep_peak):
        """
        Fit aperture by varing radius and computing the significance.
        """

        # use result.Result()
        ra_peak_array          = []
        dec_peak_array         = []
        r_peak_array           = []
        sig_peak_array         = []
        distance_modulus_array = []
        n_obs_peak_array       = []
        n_obs_half_peak_array  = []
        n_model_peak_array     = []

        size_array = np.arange(0.01, 0.3, 0.01)
        sig_array = np.tile(0., len(size_array))

        size_array_zero = np.concatenate([[0.], size_array])
        area_array = np.pi * (size_array_zero[1:]**2 - size_array_zero[0:-1]**2)

        characteristic_density_local = self.compute_local_char_density(ra, dec, x_peak, y_peak, angsep_peak)

        n_obs_array = np.tile(0, len(size_array))
        n_model_array = np.tile(0., len(size_array))
        for ii in range(0, len(size_array)):
            n_obs = np.sum(angsep_peak < size_array[ii])
            n_model = characteristic_density_local * (np.pi * size_array[ii]**2)
            sig_array[ii] = np.clip(scipy.stats.norm.isf(scipy.stats.poisson.sf(n_obs, n_model)), 0., 37.5) # Clip at 37.5
            n_obs_array[ii] = n_obs
            n_model_array[ii] = n_model

        ra_peak, dec_peak = proj.imageToSphere(x_peak, y_peak)

        index_peak = np.argmax(sig_array)
        r_peak = size_array[index_peak]
        #if np.max(sig_array) >= 37.5:
        #    r_peak = 0.5
        n_obs_peak = n_obs_array[index_peak]
        n_model_peak = n_model_array[index_peak]
        n_obs_half_peak = np.sum(angsep_peak < (0.5 * r_peak))

        # Compile resilts
        print('Candidate: x_peak: {:12.3f}, y_peak: {:12.3f}, r_peak: {:12.3f}, sig: {:12.3f}, ra_peak: {:12.3f}, dec_peak: {:12.3f}'.format(x_peak, y_peak, r_peak, np.max(sig_array), ra_peak, dec_peak))
        ra_peak_array.append(ra_peak)
        dec_peak_array.append(dec_peak)
        r_peak_array.append(r_peak)
        #sig_peak_array.append(np.max(sig_array))
        sig_peak_array.append(sig_array[index_peak])
        distance_modulus_array.append(distance_modulus)
        n_obs_peak_array.append(n_obs_peak)
        n_obs_half_peak_array.append(n_obs_half_peak)
        n_model_peak_array.append(n_model_peak)

        return (ra_peak_array, dec_peak_array, r_peak_array, sig_peak_array, distance_modulus_array, n_obs_peak_array, n_obs_half_peak_array, n_model_peak_array)


import numpy as np

def provide_PSF_2D(x=None,y=None,PSF_version=None):
    """ Provide 2D PSF at any position in the detector plane
        This a version which takes a finite nubmer of pregenerated PSF and \
        creates the interpolated version at required position
        (Future: version which takes interpolated values for Zernike \
        coefficients and generates image on the fly?)

    This version with

    @param[in] x            x-coordinate
    @param[in] y            y-coordinate
    @param[in] PSF_version  version of the PSF input files

    @returns                numpy array, (if PSF input file is  Apr15_v2 is 189x189, oversampled 9 times, \
                            corresponding to 21x21 physical pixels (315x315 microns))
    """

    # on tiger this directory is at:
    #DATA_DIRECTORY='/tigress/ncaplar/PIPE2D-521/'
    DATA_DIRECTORY='/Users/nevencaplar/Documents/PFS/Tickets/PIPE2D-521/'

    if PSF_version is None:
        PSF_version='Apr15_v2'

    positions_of_simulation=np.load(DATA_DIRECTORY+'positions_of_simulation_00_from_'+PSF_version+'.npy',allow_pickle=True)
    array_of_simulation=np.load(DATA_DIRECTORY+'array_of_simulation_00_from_'+PSF_version+'.npy',allow_pickle=True)

    # x and y position with simulated PSFs
    x_positions_of_simulation=positions_of_simulation[:,1]
    y_positions_of_simulation=positions_of_simulation[:,2]

    # This is a simple code that finds the closest avaliable PSFs, given the x and y position
    # This will have to be improved in order when we get to work with the full populated dectector plane

    # how far in x-dimension are you willing to search for suitable simulated PSFs
    x_search_distance=20
    # positions of all simulated PSFs in that range
    positions_of_simulation_in_acceptable_x_range=\
    positions_of_simulation[(x_positions_of_simulation<(x+x_search_distance))\
                            &(x_positions_of_simulation>(x-x_search_distance))]

    # if there are no simulated PSF avaliable in the specified x-range we are not able to provide the solution
    if len(positions_of_simulation_in_acceptable_x_range)<2:
        print('No simulated PSFs are avaliable in this x-area of the detector,')
        print('probably because this fiber has not been illuminated;')
        print('returning the closest avaliable PSFs, BUT that is probably not what you want')
        distances=np.sqrt(((x-x_positions_of_simulation)**2+\
                           (y-y_positions_of_simulation)**2).astype(float))
        index_of_closest_distance=np.where(distances[distances==\
                                                     np.min(distances)])[0][0]

        # ! change here on April 22, 2020!
        #       changed so the output is more similar (output as a list) to more general case below
        return [array_of_simulation[index_of_closest_distance]]
        # ! end of change here on April 22, 2020

    # y-distance from the requested positions for all of the suitable simulated PSFs
    distances_of_y_requested_position_from_avaliable=\
    y-positions_of_simulation_in_acceptable_x_range[:,2]

    # if you request for a spot exactly at the position of a input from the array
    if np.min(np.abs(distances_of_y_requested_position_from_avaliable))==0:
        index_of_simulated_psf=\
        np.where(distances_of_y_requested_position_from_avaliable==0)[0][0]

        y1_distance=0
        y2_distance=0

        # where are that exact PSF in the initial table
        index_of_1st_closest_simulated_psf_in_positions_of_simulation=\
        np.where(np.sum(positions_of_simulation,axis=1)==\
                 np.sum(positions_of_simulation_in_acceptable_x_range[index_of_simulated_psf]))[0][0]
        print(index_of_1st_closest_simulated_psf_in_positions_of_simulation)
            # extract the 2 simulated PSFs
        first_array_simulation=\
        array_of_simulation[index_of_1st_closest_simulated_psf_in_positions_of_simulation]
        second_array_simulation=\
        array_of_simulation[index_of_1st_closest_simulated_psf_in_positions_of_simulation]

    else:
        # ! change here on April 22, 2020!
        # separate the distances into distances which are for the spots above the requested position and below the requested position
        distances_of_y_requested_position_from_avaliable_which_are_above_the_spot=distances_of_y_requested_position_from_avaliable[distances_of_y_requested_position_from_avaliable<0]
        distances_of_y_requested_position_from_avaliable_which_are_below_the_spot=distances_of_y_requested_position_from_avaliable[distances_of_y_requested_position_from_avaliable>0]

        # if there are no sources above, do not do interpolation and select the nearest source
        if len(distances_of_y_requested_position_from_avaliable_which_are_above_the_spot)==0:
            index_of_1st_closest_above_simulated_psf=-99
        else:
            index_of_1st_closest_above_simulated_psf=\
            np.where(distances_of_y_requested_position_from_avaliable==\
                     np.max(distances_of_y_requested_position_from_avaliable_which_are_above_the_spot))[0][0]

        # if there are no sources below, do not do interpolation and select the nearest source
        if len(distances_of_y_requested_position_from_avaliable_which_are_below_the_spot)==0:
            index_of_1st_closest_below_simulated_psf=-99
        else:
            index_of_1st_closest_below_simulated_psf=\
            np.where(distances_of_y_requested_position_from_avaliable==\
                     np.min(distances_of_y_requested_position_from_avaliable_which_are_below_the_spot))[0][0]

        if index_of_1st_closest_below_simulated_psf==-99:
            index_of_1st_closest_below_simulated_psf=index_of_1st_closest_above_simulated_psf
        if index_of_1st_closest_above_simulated_psf==-99:
            index_of_1st_closest_above_simulated_psf=index_of_1st_closest_below_simulated_psf
        # ! end of change here on April 22, 2020

        # where are these 2 closest PSF in the initial table
        index_of_1st_closest_simulated_psf_in_positions_of_simulation=\
        np.where(np.sum(positions_of_simulation,axis=1)==\
                 np.sum(positions_of_simulation_in_acceptable_x_range[index_of_1st_closest_above_simulated_psf]))[0][0]
        index_of_2nd_closest_simulated_psf_in_positions_of_simulation=\
        np.where(np.sum(positions_of_simulation,axis=1)==\
                 np.sum(positions_of_simulation_in_acceptable_x_range[index_of_1st_closest_below_simulated_psf]))[0][0]

        # extract the 2 simulated PSFs
        first_array_simulation=\
        array_of_simulation[index_of_1st_closest_simulated_psf_in_positions_of_simulation]
        second_array_simulation=\
        array_of_simulation[index_of_2nd_closest_simulated_psf_in_positions_of_simulation]

        #print('1st:'+str(positions_of_simulation[index_of_1st_closest_simulated_psf_in_positions_of_simulation]),\
        #      '2nd:'+str(positions_of_simulation[index_of_2nd_closest_simulated_psf_in_positions_of_simulation]))

        # distance of each PSF from the proposed position
        #print(positions_of_simulation[index_of_1st_closest_simulated_psf_in_positions_of_simulation],positions_of_simulation[index_of_2nd_closest_simulated_psf_in_positions_of_simulation])
        y1_distance=\
        y-positions_of_simulation[index_of_1st_closest_simulated_psf_in_positions_of_simulation][2]
        y2_distance=\
        y-positions_of_simulation[index_of_2nd_closest_simulated_psf_in_positions_of_simulation][2]

    # ! change here on April 22, 2020!
    #       if you requested psf at the exact position of existing PSF use that one OR
    #       if you are outside of the range covered by spots, use the last avaliable image
    if y1_distance==0 or y1_distance==y2_distance:

        #    changed so the output is equivalent as in more general case
        return first_array_simulation,first_array_simulation,second_array_simulation,y1_distance,y2_distance
        # ! end of change here on April 22, 2020
    else:
        # create the predicted PSF as a linear interpolation of these two PSFs
        predicted_psf=(second_array_simulation-first_array_simulation*(y2_distance/y1_distance))/(1-y2_distance/y1_distance)
        return predicted_psf,first_array_simulation,second_array_simulation,y1_distance,y2_distance


#!/usr/bin/env python
# Started: Jan 2015 (KDG)
#    Revised to read in all the data from files: Mar 2016 (KDG)
#       observed data defines the wavelength grids to fit on
"""
ObsData class
  observed data that will be used to constrain the dust model
"""
from __future__ import print_function

import numpy as np
from astropy.table import Table
from astropy.io import fits

__all__ = ["ObsData"]

# Object for the observed dust data
class ObsData():
    """
    ObsData Class

    Parameters
    ----------
    ext_filenames: list of 'string'
        filenames with the observed extincction curve

    avnhi_filenames: list of 'string'
        filename with the observed A(V)/N(HI) value + unc

    abund_filename: 'string'
        filename with the observed atomic abundances

    ir_emis_filename: 'string'
        filename with the observed infrared dust emission

    dust_scat_filename: 'string'
        filename with the observed dust scattering (a, g) parameters
        [currently not used - hard coded for MW diffuse - need to change]

    ext_tags : list of 'string'
        list of tags identifying the origin of the
        dust extinction curve segments

    Attributes
    ----------
    alnhi : float
        A(lamda)/N(HI) value for extinction curve

    alnhi_unc : float
        uncertainty in A(lamda)/N(HI) value for extinction curve

    ext_waves : 'numpy.ndarray'
        wavelengths for the extinction curve

    ext_alav : 'numpy.ndarray'
        extinction curve in A(lambda)/A(V) units

    ext_alav_unc : 'numpy.ndarray'
        extinction curve uncertainties in A(lambda)/A(V) units

    ext_alnhi : 'numpy.ndarray'
        extinction curve in A(lambda)/N(HI) units

    ext_alnhi_unc : 'numpy.ndarray'
        extinction curve uncertainties in A(lambda)/N(HI) units

    ext_tags : 'numpy.ndarray'
        string tags identifying the origin of the extinction curve measurement

    """

    # read in the data from files
    def __init__(self, ext_filenames, avnhi_filename,
                 abund_filename, ir_emis_filename,
                 dust_scat_filename, ext_tags=None,
                 scat_path="./"):

        # extinction curve
        self.fit_extinction = True
        self.ext_waves = np.empty((0))
        self.ext_alav = np.empty((0))
        self.ext_alav_unc = np.empty((0))
        self.ext_tags = []

        if isinstance(ext_filenames, (list, tuple)):
            for i, filename in enumerate(ext_filenames):
                t = Table.read(filename, format='ascii.commented_header')
                self.ext_waves = np.concatenate([self.ext_waves,
                                                 1.0/t['wave']])
                self.ext_alav = np.concatenate([self.ext_alav,
                                                t['A(l)/A(V)']])
                self.ext_alav_unc = np.concatenate([self.ext_alav_unc,
                                                    t['unc']])
                if ext_tags is not None:
                    cur_tag = ext_tags[i]
                else:
                    cur_tag = 'Tag' + str(i+1)
                self.ext_tags = self.ext_tags + len(t['wave'])*[cur_tag]
        else:
            # assume it is a FITS file (need to add checks)
            hdulist = fits.open(ext_filenames)
            for i in range(1, len(hdulist)):
                t = hdulist[i].data
                # hack to get AzV 215 to work
                #  need to get a better file format for FITS extinction curves
                #  units, etc.
                trv = 3.65
                ext = (t['EXT']/trv) + 1
                ext_unc = t['UNC']/trv

                # only keep positive measurements
                gindxs, = np.where(ext > 0.0)
                self.ext_waves = np.concatenate([self.ext_waves,
                                                 t['WAVELENGTH'][gindxs]])
                self.ext_alav = np.concatenate([self.ext_alav, ext[gindxs]])
                self.ext_alav_unc = np.concatenate([self.ext_alav_unc,
                                                    ext_unc[gindxs]])
                self.ext_tags = self.ext_tags + \
                    len(t['WAVELENGTH'])*[hdulist[i].header['EXTNAME']]

            hdulist.close()

        # sort
        sindxs = np.argsort(self.ext_waves)
        self.ext_waves = self.ext_waves[sindxs]
        self.ext_alav = self.ext_alav[sindxs]
        self.ext_alav_unc = self.ext_alav_unc[sindxs]
        self.ext_tags = np.array(self.ext_tags)[sindxs]

        # normalization from A(V) to N(HI)
        t = Table.read(avnhi_filename,
                       format='ascii.commented_header',
                       header_start=-1)
        self.avnhi = t['Av_to_NHI'][0]
        self.avnhi_unc = t['unc'][0]

        # change the extinction normalization from A(V) to N(HI)
        self.ext_alnhi = self.ext_alav*self.avnhi
        self.ext_alnhi_unc = (np.square(self.ext_alav_unc/self.ext_alav)
                              + np.square(self.avnhi_unc/self.avnhi))
        self.ext_alnhi_unc = self.ext_alnhi*np.sqrt(self.ext_alnhi_unc)

        # dust abundances
        self.fit_abundance = False
        if abund_filename is not None:
            self.fit_abundance = True
            t = Table.read(abund_filename, format='ascii.commented_header')
            self.abundance = {}
            self.total_abundance = {}
            for i in range(len(t)):
                self.abundance[t['atom'][i]] = (t['abund'][i],
                                                t['abund_unc'][i])
                self.total_abundance[t['atom'][i]] = (t['total_abund'][i],
                                                      t['total_abund_unc'][i])

        # diffuse IR emission spectrum
        self.fit_ir_emission = False
        if ir_emis_filename is not None:
            self.fit_ir_emission = True
            t = Table.read(ir_emis_filename, format='ascii.commented_header')
            self.ir_emission_waves = np.array(t['WAVE'])
            self.ir_emission = np.array(t['SPEC'])/1e20
            self.ir_emission_unc = np.array(t['ERROR'])/1e20
            # check if any uncs are zero
            gindxs, = np.where(self.ir_emission_unc == 0.0)
            if len(gindxs) > 0:
                self.ir_emission_unc[gindxs] = 0.1*self.ir_emission[gindxs]

            # sort
            sindxs = np.argsort(self.ir_emission_waves)
            self.ir_emission_waves = self.ir_emission_waves[sindxs]
            self.ir_emission = self.ir_emission[sindxs]
            self.ir_emission_unc = self.ir_emission_unc[sindxs]

        # dust albedo (Gordon et al. AoD proceedings)
        self.fit_scat_a = False
        self.fit_scat_g = False
        if dust_scat_filename is not None:
            self.fit_scat_a = True
            self.fit_scat_g = True
            files_dgl = ["mathis73", "morgan76", "lillie76", "toller81",
                         "murthy93", "murthy95", "petersohn97", "witt97",
                         "schiminovich01", "shalima04", "sujatha05",
                         "sujatha07", "sujatha10"]

            scat_waves = []
            scat_albedo = []
            scat_albedo_unc = []
            scat_g = []
            scat_g_unc = []
            scat_ref = []
            for sfile in files_dgl:
                f = open(scat_path + sfile + '.dat', 'r')
                ref = f.readline().rstrip()
                f.close()

                t = Table.read(scat_path+sfile+'.dat',
                               format='ascii',
                               header_start=1)
                for k in range(len(t)):
                    scat_waves.append(t['wave,'][k])
                    scat_albedo.append(t['albedo,'][k])
                    scat_albedo_unc.append(t['delta,'][k])
                    scat_g.append(t['g,'][k])
                    scat_g_unc.append(t['delta'][k])
                    scat_ref.append(ref)

            # remove all the measurements with zero uncertainty
            gindxs, = np.where(np.array(scat_albedo_unc) > 0.0)
            self.scat_a_waves = np.array(scat_waves)[gindxs]*1e-4
            self.scat_albedo = np.array(scat_albedo)[gindxs]
            self.scat_albedo_unc = np.array(scat_albedo_unc)[gindxs]
            self.scat_a_ref = np.array(scat_ref)[gindxs]

            # sort
            sindxs = np.argsort(self.scat_a_waves)
            self.scat_a_waves = self.scat_a_waves[sindxs]
            self.scat_albedo = self.scat_albedo[sindxs]
            self.scat_albedo_unc = self.scat_albedo_unc[sindxs]
            self.scat_a_ref = self.scat_a_ref[sindxs]

            # remove all the measurements with zero uncertainty
            gindxs, = np.where(np.array(scat_g_unc) > 0.0)
            self.scat_g_waves = np.array(scat_waves)[gindxs]*1e-4
            self.scat_g = np.array(scat_g)[gindxs]
            self.scat_g_unc = np.array(scat_g_unc)[gindxs]
            self.scat_g_ref = np.array(scat_ref)[gindxs]

            # sort
            sindxs = np.argsort(self.scat_g_waves)
            self.scat_g_waves = self.scat_g_waves[sindxs]
            self.scat_g = self.scat_g[sindxs]
            self.scat_g_unc = self.scat_g_unc[sindxs]
            self.scat_g_ref = self.scat_g_ref[sindxs]


import numpy as np
import pygame
import time

pygame.init() # iniciamos el pygame

ancho, alto = 600, 600
pantalla = pygame.display.set_mode((ancho,alto)) # creamos la pantalla
fondoPantalla = 35,35,35 # de color casi negro

pantalla.fill(fondoPantalla) # se pinta el fondo de la pantalla

# numero de celdas en cada eje
celdasX, celdasY = 50,50

anchoCeldas = ancho / celdasX # ancho de la celda
altoCeldas = alto / celdasY # alto de la celda

pauseExecution = True # control de la ejecución del juego

# celda a 1 --> viva || celda a 0 --> muerta
estadoTablero = np.zeros((celdasX,celdasY)) # matriz de tamaño del tablero

# inicialización del tablero
# palo
estadoTablero[5,3] = 1
estadoTablero[5,4] = 1
estadoTablero[5,5] = 1

# automata que se mueve por pantalla
estadoTablero[21,21] = 1
estadoTablero[22,22] = 1
estadoTablero[22,23] = 1
estadoTablero[21,23] = 1
estadoTablero[20,23] = 1

while True:
    # para no sobreescribir el tablero puesto que las comparaciones se hacen con el estado inicial y no con el actualizado en el momento
    copiaTablero = np.copy(estadoTablero)

    # limpiamos la pantalla
    pantalla.fill(fondoPantalla) # se pinta el fondo de la pantalla

    # delay
    time.sleep(0.1)

    ev = pygame.event.get()

    for event in ev:
        if event.type == pygame.KEYDOWN: # al pulsar el teclado
            pauseExecution = not pauseExecution

        mouseClick = pygame.mouse.get_pressed() # nos da el boton izq, rueda o boton derecho pulsado

        if sum(mouseClick) > 0:
            posX, posY = pygame.mouse.get_pos() # lo devuelve en pixeles
            celX, celY = int(np.floor(posX / anchoCeldas)), int(np.floor(posY / altoCeldas)) # pos celda
            copiaTablero[celX,celY] = not mouseClick[2] # si pulso el izq pon 0, si pulso otro pon 1

    # recorremos el tablero
    for y in range (0,celdasY):
        for x in range (0,celdasX):

            if  not pauseExecution:
                # vecinos cercanos a cada x,y
                # con el modulo conseguimos que el vecino del tablero de un borde actue como un toroide [de la izquierda se pasa a la derecha y de arriba a abajo]
                vecinos = estadoTablero[(x-1) % celdasX,(y-1) % celdasY] + \
                        estadoTablero[(x-1) % celdasX,(y) % celdasY] + \
                        estadoTablero[(x-1) % celdasX,(y+1) % celdasY] + \
                        estadoTablero[(x) % celdasX,(y-1) % celdasY] + \
                        estadoTablero[(x) % celdasX,(y+1) % celdasY] + \
                        estadoTablero[(x+1) % celdasX,(y-1) % celdasY] + \
                        estadoTablero[(x+1) % celdasX,(y) % celdasY] + \
                        estadoTablero[(x+1) % celdasX,(y+1) % celdasY]

                # regla 1 del juego de la vida: si muerto y tres vecinos vivos, entonces revive
                if estadoTablero[x,y] == 0 and vecinos == 3:
                    copiaTablero[x,y] = 1

                # regla 2 del juego de la vida: si vivo y vecinos menor que 2 o mayor que 3, entonces muere
                elif estadoTablero[x,y] == 1 and (vecinos < 2 or vecinos > 3):
                    copiaTablero[x,y] = 0

            # polígono
            poligono = [((x)*anchoCeldas, (y)*altoCeldas),
                        ((x+1)*anchoCeldas, (y)*altoCeldas),
                        ((x+1)*anchoCeldas, (y+1)*altoCeldas),
                        ((x)*anchoCeldas, (y+1)*altoCeldas)]

            # el ultimo término es el grosor del poligono
            if copiaTablero[x,y] == 0: # muerto en negro
                pygame.draw.polygon(pantalla, (128,128,128), poligono, 1)
            else: # vivo en blanco
                pygame.draw.polygon(pantalla, (255,255,255), poligono, 0)

    estadoTablero = np.copy(copiaTablero)

    pygame.display.flip() # actualizo los fotogramas