"""spectral Electrical Impedance Tomography (sEIT) container
"""
# import functools
import os
from numbers import Number

import numpy as np

from reda.main.logger import LoggingClass
from reda.containers.ERT import ImportersBase
import reda.importers.eit_fzj as eit_fzj
import reda.importers.radic_sip256c as reda_sip256c
import reda.importers.crtomo as reda_crtomo_exporter
import reda.utils.eit_fzj_utils as eit_fzj_utils
import reda.utils.geometric_factors as geometric_factors
from reda.utils.fix_sign_with_K import fix_sign_with_K
import reda.eis.plots as eis_plot

from reda.utils.decorators_and_managers import append_doc_of
from reda.utils.decorators_and_managers import LogDataChanges

import reda.exporters.crtomo as exporter_crtomo

import reda.plotters.pseudoplots as PS
import reda.plotters.histograms as HS

import reda.utils.mpl
plt, mpl = reda.utils.mpl.setup()

class importers(ImportersBase):
    """This class provides wrappers for most of the importer functions, and is
    meant to be inherited by the data containers
    """
    @append_doc_of(reda_crtomo_exporter.load_seit_data)
    def import_crtomo(self, directory, frequency_file='frequencies.dat',
                      data_prefix='volt_', **kwargs):
        """CRTomo importer"""
        timestep = kwargs.get('timestep', None)
        if 'timestep' in kwargs:
            del (kwargs['timestep'])

        # we get no electrode positions (dummy1) and no topography data
        # (dummy2)
        data, dummy1, dumm2 = reda_crtomo_exporter.load_seit_data(
            directory, frequency_file, data_prefix, **kwargs)
        if timestep is not None:
            data['timestep'] = timestep
        self._add_to_container(data)

        print('Summary:')
        self._describe_data(data)

    def import_sip256c(self, filename, settings=None, reciprocal=None,
                       **kwargs):
        """Radic SIP256c data import"""
        timestep = kwargs.get('timestep', None)
        if 'timestep' in kwargs:
            del (kwargs['timestep'])
        if settings is None:
            settings = {}
        # we get not electrode positions (dummy1) and no topography data
        # (dummy2)
        data, dummy1, dummy2 = reda_sip256c.parse_radic_file(
            filename, settings, reciprocal=reciprocal, **kwargs)
        if timestep is not None:
            data['timestep'] = timestep
        self._add_to_container(data)

        print('Summary:')
        self._describe_data(data)

    def import_eit_fzj(self, filename, configfile, correction_file=None,
                       timestep=None, **kwargs):
        """EIT data import for FZJ Medusa systems"""
        # we get not electrode positions (dummy1) and no topography data
        # (dummy2)
        df_emd, dummy1, dummy2 = eit_fzj.read_3p_data(
            filename,
            configfile,
            **kwargs
        )
        if correction_file is not None:
            eit_fzj_utils.apply_correction_factors(df_emd, correction_file)

        if timestep is not None:
            df_emd['timestep'] = timestep

        self._add_to_container(df_emd)

        print('Summary:')
        self._describe_data(df_emd)

class sEIT(LoggingClass, importers):

    def __init__(self, dataframe=None):
        self.setup_logger()
        if dataframe is not None:
            self.check_dataframe(dataframe)
        # normal data (or full data, if reciprocals are not sorted
        self.data = dataframe

        self.required_data_columns = ['r', 'rpha']

    def check_dataframe(self, dataframe):
        """Check the given dataframe for the required columns
        """
        required_columns = (
            'a',
            'b',
            'm',
            'n',
            'r',
        )
        for column in required_columns:
            if column not in dataframe:
                raise Exception('Required column not in dataframe: {0}'.format(
                    column
                ))

    @property
    def abmn(self):
        return self.data.groupby(['a', 'b', 'm', 'n'])

    def subquery(self, subset, filter, inplace=True):
        """

        Examples
        --------

        ::

            subquery(
                'timestep == 2',
                'R > 4',
            )

        """
        # build the full query
        full_query = ''.join((
            'not (',
            subset,
            ') or not (',
            filter,
            ')',
        ))
        result = self.data.query(full_query, inplace=inplace)
        return result

    def query(self, query, inplace=True):
        """State what you want to keep

        """
        # TODO: add to queue
        result = self.data.query(query, inplace=inplace)
        return result

    def filter(self, query, inplace=True):
        """Use a query statement to filter data. Note that you specify the data
        to be removed!

        Parameters
        ----------
        query : string
            The query string to be evaluated. Is directly provided to
            pandas.DataFrame.query
        inplace : bool
            if True, change the container dataframe in place (defaults to True)

        Returns
        -------
        result : :py:class:`pandas.DataFrame`
            DataFrame that contains the result of the filter application

        """
        with LogDataChanges(self, filter_action='filter', filter_query=query):
            result = self.data.query(
                'not ({0})'.format(query),
                inplace=inplace,
            )
        return result

    def remove_frequencies(self, fmin, fmax):
        """Remove frequencies from the dataset
        """
        self.data.query(
            'frequency > {0} and frequency < {1}'.format(fmin, fmax),
            inplace=True
        )
        g = self.data.groupby('frequency')
        print('Remaining frequencies:')
        print(sorted(g.groups.keys()))

    def compute_K_analytical(self, spacing):
        """Assuming an equal electrode spacing, compute the K-factor over a
        homogeneous half-space.

        For more complex grids, please refer to the module:
        reda.utils.geometric_factors

        Parameters
        ----------
        spacing: float
            Electrode spacing

        """
        assert isinstance(spacing, Number)
        K = geometric_factors.compute_K_analytical(self.data, spacing)
        self.data = geometric_factors.apply_K(self.data, K)
        fix_sign_with_K(self.data)

    @append_doc_of(fix_sign_with_K)
    def fix_sign_with_K(self):
        """ """
        fix_sign_with_K(self.data)

    def scatter_norrec(self, filename=None, individual=False):
        """Create a scatter plot for all diff pairs

        Parameters
        ----------

        filename : string, optional
            if given, save plot to file
        individual : bool, optional
            if set to True, return one figure for each row

        Returns
        -------
        fig : matplotlib.Figure or list of :py:class:`matplotlib.Figure.Figure`
            objects the figure object
        axes : list of matplotlib.axes
            the individual axes

        """
        # if not otherwise specified, use these column pairs:
        std_diff_labels = {
            'r': 'rdiff',
            'rpha': 'rphadiff',
        }

        diff_labels = std_diff_labels

        # check which columns are present in the data
        labels_to_use = {}
        for key, item in diff_labels.items():
            # only use if BOTH columns are present
            if key in self.data.columns and item in self.data.columns:
                labels_to_use[key] = item

        g_freq = self.data.groupby('frequency')
        frequencies = list(sorted(g_freq.groups.keys()))

        if individual:
            figures = {}
            axes_all = {}
        else:
            Nx = len(labels_to_use.keys())
            Ny = len(frequencies)
            fig, axes = plt.subplots(
                Ny, Nx,
                figsize=(Nx * 2.5, Ny * 2.5)
            )

        for row, (name, item) in enumerate(g_freq):
            if individual:
                fig, axes_row = plt.subplots(
                    1, 2, figsize=(16 / 2.54, 6 / 2.54))
            else:
                axes_row = axes[row, :]
            # loop over the various columns
            for col_nr, (key, diff_column) in enumerate(
                    sorted(labels_to_use.items())):
                indices = np.where(~np.isnan(item[diff_column]))[0]
                ax = axes_row[col_nr]
                ax.scatter(
                    item[key],
                    item[diff_column],
                )
                ax.set_xlabel(key)
                ax.set_ylabel(diff_column)
                ax.set_title('N: {}'.format(len(indices)))
            if individual:
                fig.tight_layout()
                figures[name] = fig
                axes_all[name] = axes_row

        if individual:
            return figures, axes_all
        else:
            fig.tight_layout()
            return fig, axes

    def filter_incomplete_spectra(self, flimit=1000, percAccept=85):
        """Remove all data points that belong to spectra that did not retain at
        least **percAccept** percent of the number of data points.

        ..warning::

            This function does not honor additional dimensions (e.g.,
            timesteps) yet!

        """
        assert percAccept > 0 and percAccept < 100

        def _retain_only_complete_spectra(item, fmax, acceptN):
            """Function called using pd.filter, applied to all spectra in the
            data set. Return true if the number of data points <= **fmax** in
            item is equal, or larger, than **acceptN**.

            Parameters
            ----------
            item : :py:class:`pandas.DataFrame`
                dataframe containing one spectrum
            fmax : float
                maximum frequency up to which data points are counted
            acceptN : int
                the number of data points required to pass this test

            Returns
            -------
            true : bool
                if enough data points are present
            false : bool
                if not enough data points are present
            """
            frequencies = item['frequency'].loc[item['frequency'] < fmax]
            fN = frequencies.size
            if fN >= acceptN:
                return True
            return False

        group_abmn = self.data.groupby(['a', 'b', 'm', 'n'])
        frequencies = np.array(
            list(sorted(self.data.groupby('frequency').groups.keys()))
        )
        assert frequencies.size > 0
        assert flimit >= frequencies.min() and flimit <= frequencies.max()
        Nlimit = len(np.where(frequencies <= flimit)[0])
        Naccept = np.ceil(Nlimit * percAccept / 100.0)
        self.data = group_abmn.filter(
            _retain_only_complete_spectra, fmax=flimit, acceptN=Naccept
        ).copy()

    def get_spectrum(self, nr_id=None, abmn=None, plot_filename=None):
        """Return a spectrum and its reciprocal counter part, if present in the
        dataset. Optimally, refer to the spectrum by its normal-reciprocal id.

        Returns
        -------
        spectrum_nor : :py:class:`reda.eis.plots.sip_response`
            Normal spectrum. None if no normal spectrum is available
        spectrum_rec : :py:class:`reda.eis.plots.sip_response` or None
            Reciprocal spectrum. None if no reciprocal spectrum is available
        fig : :py:class:`matplotlib.Figure.Figure` , optional
            Figure object (only if plot_filename is set)

        """
        assert nr_id is None or abmn is None
        # determine nr_id for given abmn tuple
        if abmn is not None:
            subdata = self.data.query(
                'a == {} and b == {} and m == {} and n == {}'.format(*abmn)
            ).sort_values('frequency')
            if subdata.shape[0] == 0:
                return None, None

            # determine the norrec-id of this spectrum
            nr_id = subdata['id'].iloc[0]

        # get spectra
        subdata_nor = self.data.query(
            'id == {} and norrec=="nor"'.format(nr_id)
        ).sort_values('frequency')

        subdata_rec = self.data.query(
            'id == {} and norrec=="rec"'.format(nr_id)
        ).sort_values('frequency')

        # create spectrum objects
        spectrum_nor = None
        spectrum_rec = None

        if subdata_nor.shape[0] > 0:
            spectrum_nor = eis_plot.sip_response(
                frequencies=subdata_nor['frequency'].values,
                rmag=subdata_nor['r'],
                rpha=subdata_nor['rpha'],
            )
        if subdata_rec.shape[0] > 0:
            spectrum_rec = eis_plot.sip_response(
                frequencies=subdata_rec['frequency'].values,
                rmag=subdata_rec['r'],
                rpha=subdata_rec['rpha'],
            )
        if plot_filename is not None:
            if spectrum_nor is not None:
                fig = spectrum_nor.plot(
                    plot_filename,
                    reciprocal=spectrum_rec,
                    return_fig=True,
                    title='a: {} b: {} m: {}: n: {}'.format(
                        *subdata_nor[['a', 'b', 'm', 'n']].values[0, :]
                    )
                )
                return spectrum_nor, spectrum_rec, fig
        return spectrum_nor, spectrum_rec

    def plot_all_spectra(self, outdir):
        r"""This is a convenience function to plot ALL spectra currently
        stored in the container. It is useful to asses whether data filters
        do perform correctly.

        Note that the function just iterates over all ids and plots the
        corresponding spectra, thus it is slow.

        Spectra a named using the format: \%.2i_spectrum_id_\{\}.png.

        Parameters
        ----------
        outdir : string
            Output directory to store spectra in. Created if it does not
            exist.
        """
        os.makedirs(outdir, exist_ok=True)

        g = self.data.groupby('id')
        for nr, (name, item) in enumerate(g):
            print(
                'Plotting spectrum with id {} ({} / {})'.format(
                    name, nr, len(g.groups.keys()))
            )
            plot_filename = ''.join((
                outdir + os.sep,
                '{:04}_spectrum_id_{}.png'.format(nr, name)
            ))
            spec_nor, spec_rec, spec_fig = self.get_spectrum(
                nr_id=name,
                plot_filename=plot_filename
            )
            plt.close(spec_fig)

    def plot_pseudosections(self, column, filename=None, return_fig=False):
        """Create a multi-plot with one pseudosection for each frequency.

        Parameters
        ----------
        column : string
            which column to plot
        filename : None|string
            output filename. If set to None, do not write to file. Default:
            None
        return_fig : bool
            if True, return the generated figure object, also if filename is
            set. Default: False

        Returns
        -------
        fig : None|matplotlib.Figure
            if return_fig is set to True or filename is None, return the
            generated Figure object
        """
        assert column in self.data.columns

        g = self.data.groupby('frequency')
        fig, axes = plt.subplots(
            4, 2,
            figsize=(15 / 2.54, 20 / 2.54),
            sharex=True, sharey=True
        )
        for ax, (key, item) in zip(axes.flat, g):
            fig, ax, cb = PS.plot_pseudosection_type2(
                item, ax=ax, column=column
            )
            ax.set_title('f: {:.3f} Hz'.format(key))
        fig.tight_layout()
        if filename is not None:
            fig.savefig(filename, dpi=300)

        if return_fig or filename is None:
            return fig
        else:
            plt.close(fig)

    def export_to_crtomo_multi_frequency(self, directory, norrec='norrec'):
        """Export the sEIT data into data files that can be read by CRTomo.

        Parameters
        ----------
        directory : string
            output directory. will be created if required
        norrec : string (nor|rec|norrec)
            Which data to export. Default: norrec

        """
        exporter_crtomo.write_files_to_directory(
            self.data, directory, norrec=norrec
        )

    def export_to_crtomo_one_frequency(
            self, volt_file, frequency, norrec='norrec'):
        """Export one frequency into a CRTomo volt.dat file

        Parameters
        ----------
        volt_file : string
            output file. Will be overwritten if it exists
        frequency : float
            frequency to export
        norrec : str (nor|rec|norrec)
            Which data to export. Default: norrec
        """
        assert isinstance(frequency, float)
        frequency_data = self.data.query('frequency == {}'.format(frequency))
        exporter_crtomo.save_block_to_crt(
            volt_file, frequency_data, norrec=norrec
        )

    def export_to_crtomo_seit_manager(self, grid, norrec='norrec'):
        """Return a ready-initialized seit-manager object from the CRTomo
        tools. This function only works if the crtomo_tools are installed.

        WARNING: Not timestep aware!

        Parameters
        ----------
        grid : crtomo.crt_grid
            A CRTomo grid instance
        norrec : str (nor|rec|norrec)
            Which data to export. Default: norrec (all)

        """
        import crtomo
        subdata = self.data.query('norrec == "{}"'.format(norrec))
        g = subdata.groupby('frequency')
        seit_data = {}
        for name, item in g:
            print(name, item.shape, item.size)
            if item.shape[0] > 0:
                seit_data[name] = item[
                    ['a', 'b', 'm', 'n', 'r', 'rpha']
                ].values
        seit = crtomo.eitMan(grid=grid, seit_data=seit_data)
        return seit

    def export_to_crtomo_td_manager(self, grid, frequency, norrec='norrec'):
        """Return a ready-initialized tdman object from the CRTomo tools. Use
        the given frequency data to initialize it.

        WARNING: Not timestep aware!

        Parameters
        ----------
        grid : crtomo.crt_grid
            A CRTomo grid instance
        frequency : float
            The frequency to export data for
        norrec : str (nor|rec|norrec)
            Which data to export. Default: norrec (all)
        """
        subdata = self.data.query('norrec == "{}"'.format(norrec))
        import crtomo
        data = subdata.query('frequency == {}'.format(frequency))[
            ['a', 'b', 'm', 'n', 'r', 'rpha']
        ]
        tdman = crtomo.tdMan(grid=grid, volt_data=data)
        return tdman

    def plot_histograms(
            self, column='r', primary_dim=None, filename=None, **kwargs):
        """Plot a histograms for all frequencies of one data column

        Parameters
        ----------
        column : str, optional
            data column to plot. defaults to "r" for resistance
        primary_dim : None|str
            ???
        filename : None|str
            Prefix for filename
        **kwargs : dict
            ???

        TODO: Check saving to file for more than one secondary dimension
        Parameters
        ----------
        """
        dict_dimension, figs = HS.plot_histograms_extra_dims(
            self.data, column, primary_dim, **kwargs)
        if filename is not None:
            for key, item in figs.items():
                item.savefig(
                    filename + '_{}.jpg'.format(key.replace('_', '-')), dpi=300
                )
        return dict_dimension, figs

    @property
    def nr_frequencies(self):
        """Return the number of frequencies in the data set"""
        if self.data is None:
            return 0
        group_f = self.data.groupby('frequency')
        return group_f.ngroups

    @property
    def Nf(self):
        """Shortcut for self.nr_frequencies"""
        return self.nr_frequencies()

    @property
    def frequencies(self):
        """Return the frequencies contained in the data set"""
        if self.data is None:
            return 0
        frequencies = sorted(self.data.groupby('frequency').groups.keys())
        return frequencies

    @property
    def nr_timesteps(self):
        """Return the number of timesteps registered with this container"""
        if self.data is None or 'timestep' not in self.data:
            return 0
        group_ts = self.data.groupby('timestep')
        return group_ts.ngroups


#! /usr/bin/python3
# -*- coding: utf-8 -*-

__author__ = 'Alan Loh, Julien Girard'
__copyright__ = 'Copyright 2019, nenupytv'
__credits__ = ['Alan Loh', 'Julien Girard']
__maintainer__ = 'Alan'
__email__ = 'alan.loh@obspm.fr'
__status__ = 'Production'
__all__ = [
    'nenufar_loc',
    'lst',
    'lha',
    'ho_zenith',
    'eq_zenith',
    'radec',
    'radec_hd',
    'eq_coord',
    'to_radec',
    'to_altaz',
    'to_gal',
    'to_lmn',
    'rephase',
    'rotz',
    'wavelength',
    'ref_location',
    'radio_sources',
    'ateam',
    'AstroPlot',
    'astro_image'
    ]

from os import path
import numpy as np
from astropy.time import Time
from astropy import units as u
from astropy.coordinates import (
    EarthLocation,
    Angle,
    SkyCoord,
    AltAz,
    Galactic,
    ICRS,
    get_body,
    solar_system_ephemeris
)
from astropy.constants import c as lspeed
from astropy.wcs import WCS
from astropy.io import fits

import nenupytv

# ============================================================= #
# ------------------------ nenufar_loc ------------------------ #
# ============================================================= #
def nenufar_loc():
    """
    """
    return EarthLocation(
        lat=47.375944 * u.deg,
        lon=2.193361 * u.deg,
        height=136.195 * u.m
    )
# ============================================================= #

# ============================================================= #
# ---------------------------- lst ---------------------------- #
# ============================================================= #
def lst(time, location=None):
    """ Local sidereal time

        Parameters
        ----------
        time : `astropy.time.Time`
            UTC time
        location : `astropy.coord.EarthLocation`
            Location of the instrument

        Returns
        -------
        lst : float
            Local sidereal time in degrees
    """
    if not isinstance(time, Time):
        raise TypeError(
            'time is not an astropy Time.'
            )
    if location is None:
        location = nenufar_loc()
    if not isinstance(location, EarthLocation):
        raise TypeError(
            'time is not an astropy EarthLocation.'
            )
    lon = location.to_geodetic().lon
    lst = time.sidereal_time('apparent', lon)
    return lst.deg#.hourangle
# ============================================================= #

# ============================================================= #
# ---------------------------- lha ---------------------------- #
# ============================================================= #
def lha(time, ra, location=None):
    """ Local hour angle of an object in the observer's sky

        Parameters
        ----------
        time : `astropy.time.Time`
            UTC time
        ra : float
            Right Ascension in degrees
        location : `astropy.coord.EarthLocation`, optional
            Location of the instrument

        Returns
        -------
        lha : float
            Local hour angle in degrees
    """
    ra = Angle(ra * u.deg).deg#.hourangle
    ha = lst(time, location) - ra
    if ha < 0:
        ha += 360.
    elif ha > 360:
        ha -= 360.
    return ha
# ============================================================= #

# ============================================================= #
# ------------------------- ho_zenith ------------------------- #
# ============================================================= #
def ho_zenith(time, location=None):
    """ Horizontal coordinates of zenith
    """
    if location is None:
        location = nenufar_loc()
    if not isinstance(location, EarthLocation):
        raise TypeError(
            'time is not an astropy EarthLocation.'
            )
    altaz = AltAz(
        az=0.*u.deg,
        alt=90.*u.deg,
        location=location,
        obstime=time
    )
    return altaz
# ============================================================= #

# ============================================================= #
# ------------------------- eq_zenith ------------------------- #
# ============================================================= #
def eq_zenith(time, location=None):
    """ Get the ra dec coordinates of the zenith

        Parameters
        ----------
        time : `astropy.time.Time`
            UTC time
        location : `astropy.coord.EarthLocation`
            Location of the instrument

        Returns
        -------
        ra : float
            Right Ascension in degrees
        dec : float
            Declination in degrees
    """
    if location is None:
        location = nenufar_loc()
    if not isinstance(location, EarthLocation):
        raise TypeError(
            'time is not an astropy EarthLocation.'
            )

    zen_alt = 90*u.deg
    zen_az = 0*u.deg
    azel = SkyCoord(
        alt=zen_alt,
        az=zen_az,
        obstime=time,
        location=location,
        frame='altaz'
        )
    eq = azel.icrs
    return eq.ra.deg, eq.dec.deg
# ============================================================= #

# ============================================================= #
# --------------------------- radec --------------------------- #
# ============================================================= #
def radec(ra, dec):
    """ Equatorial coordinates

        :param ra:
            Right ascension in degrees
        :type ra: float
        :param dec:
            Declination in degrees
        :type dec: float

        :returns: :class:`astropy.coordinates.ICRS` object
        :rtype: :class:`astropy.coordinates.ICRS`

        :Example:

        >>> from nenupysim.astro import eq_coord
        >>> radec = eq_coord(
                ra=51,
                dec=39,
            )
    """
    return ICRS(
        ra=ra*u.deg,
        dec=dec*u.deg
    )
# ============================================================= #

# ============================================================= #
# ------------------------- radec_hd -------------------------- #
# ============================================================= #
def radec_hd(hms_dms):
    """
    """
    return SkyCoord(
        hms_dms,
        unit=[u.hourangle, u.deg]
    )
# ============================================================= #

# ============================================================= #
# ------------------------- eq_coord -------------------------- #
# ============================================================= #
def eq_coord(ra, dec):
    """ Equatorial coordinates

        :param ra:
            Right ascension in degrees
        :type ra: float
        :param dec:
            Declination in degrees
        :type dec: float

        :returns: :class:`astropy.coordinates.ICRS` object
        :rtype: :class:`astropy.coordinates.ICRS`

        :Example:

        >>> from nenupysim.astro import eq_coord
        >>> radec = eq_coord(
                ra=51,
                dec=39,
            )
    """
    eq = radec(
        ra=ra*u.deg,
        dec=dec*u.deg
    )
    return eq.ra.deg, eq.dec.deg
# ============================================================= #

# ============================================================= #
# ------------------------- eq_zenith ------------------------- #
# ============================================================= #
def to_radec(alt, az, time, location=None):
    """ Get the ra dec coordinates of the a altaz pointing

        Parameters
        ----------
        alt : float
            Elevation in degrees
        az : float
            Azimuth in degrees
        time : `astropy.time.Time`
            UTC time
        location : `astropy.coord.EarthLocation`
            Location of the instrument

        Returns
        -------
        ra : float
            Right Ascension in degrees
        dec : float
            Declination in degrees
    """
    if location is None:
        location = nenufar_loc()
    if not isinstance(location, EarthLocation):
        raise TypeError(
            'time is not an astropy EarthLocation.'
            )

    zen_alt = alt*u.deg
    zen_az = az*u.deg
    azel = SkyCoord(
        alt=zen_alt,
        az=zen_az,
        obstime=time,
        location=location,
        frame='altaz'
        )
    eq = azel.icrs
    return eq.ra.deg, eq.dec.deg
# ============================================================= #

# ============================================================= #
# ------------------------- to_altaz -------------------------- #
# ============================================================= #
def to_altaz(ra, dec, time, location=None):
    """ Transform altaz coordinates to ICRS equatorial system

        :param radec:
            Equatorial coordinates
        :type altaz: :class:`astropy.coordinates.ICRS`
        :param time:
            Time at which the local coordinates should be
            computed. It can either be provided as an
            :class:`astropy.time.Time` object or a string in ISO
            or ISOT format.
        :type time: str, :class:`astropy.time.Time`

        :returns: :class:`astropy.coordinates.AltAz` object
        :rtype: :class:`astropy.coordinates.AltAz`

        :Example:

        >>> from nenupysim.astro import eq_coord
        >>> radec = eq_coord(
                ra=51,
                dec=39,
            )
    """
    if location is None:
        location = nenufar_loc()
    if not isinstance(location, EarthLocation):
        raise TypeError(
            'time is not an astropy EarthLocation.'
            )

    altaz_frame = AltAz(
        obstime=time,
        location=location
    )
    eq = ICRS(
        ra=ra*u.deg,
        dec=dec*u.deg
    )
    altaz = eq.transform_to(altaz_frame)
    return altaz
# ============================================================= #

# ============================================================= #
# -------------------------- to_gal --------------------------- #
# ============================================================= #
def to_gal(src):
    """ Convert an astorpy source to galactic cooridnates
    """
    return src.transform_to(Galactic)
# ============================================================= #

# ============================================================= #
# -------------------------- to_gal --------------------------- #
# ============================================================= #
def to_lmn(ra, dec, ra_0, dec_0):
    """
    """
    ra = np.radians(ra)
    dec = np.radians(dec)
    ra_0 = np.radians(ra_0)
    dec_0 = np.radians(dec_0)
    ra_delta = ra - ra_0
    l = np.cos(dec)*np.sin(ra_delta)
    m = np.sin(dec)*np.cos(dec_0) - np.cos(dec)*np.sin(dec_0)*np.cos(ra_delta)
    n = np.sqrt(1 - l**2 - m**2)
    return l, m, n
# ============================================================= #

# ============================================================= #
# -------------------------- rephase -------------------------- #
# ============================================================= #
def rephase(ra, dec, time, loc, dw=False):
    """
    """
    raz, decz = eq_zenith(
        time=time,
        location=loc
        )
    def rotMatrix(r, d):
        """ r: ra in radians
            d: dec in radians
        """
        w = np.array([
            [  np.sin(r)*np.cos(d) ,  np.cos(r)*np.cos(d) , np.sin(d) ]
        ]).T
        v = np.array([
            [ -np.sin(r)*np.sin(d) , -np.cos(r)*np.sin(d) , np.cos(d) ]
        ]).T
        u = np.array([
            [  np.cos(r)           , -np.sin(r)           , 0.        ]
        ]).T
        rot_matrix = np.concatenate([u, v, w], axis=-1)
        return rot_matrix, w
    final_trans, wnew = rotMatrix(
        r=np.radians(ra),
        d=np.radians(dec)
    )
    original_trans, wold = rotMatrix(
        r=np.radians(raz),
        d=np.radians(decz)
    )
    total_trans = np.dot(final_trans.T, original_trans)

    if dw:
        return total_trans, original_trans, final_trans, wold-wnew
    else:
        return total_trans
# ============================================================= #

# ============================================================= #
# --------------------------- rotz ---------------------------- #
# ============================================================= #
def rotz(array, angle):
    """ Rotate the 3D array by an angle along z-axis
    """
    ang = np.radians(angle)
    cosa = np.cos(ang)
    sina = np.sin(ang)
    rot = np.array([
            [cosa, -sina, 0],
            [sina,  cosa, 0],
            [   0,     0, 1]
        ])
    return np.dot(array, rot)
# ============================================================= #

# ============================================================= #
# ------------------------ wavelength ------------------------- #
# ============================================================= #
def wavelength(freq):
    """ Convert between MHz and wavelength in meters

        Returns
        -------
        wavelength : `np.ndarray`
            Wavelength in meters
    """
    if not hasattr(freq, '__len__'):
        freq = [freq]
    if not isinstance(freq, np.ndarray):
        freq = np.array(freq)
    freq *= u.MHz
    freq = freq.to(u.Hz)
    wavel = lspeed.value / freq.value
    return wavel
# ============================================================= #

# ============================================================= #
# ------------------------ ref_location ----------------------- #
# ============================================================= #
def ref_location():
    """
    """
    return EarthLocation(
        lat=0*u.deg,
        lon=-90*u.deg,
        height=0*u.m
    )
# ============================================================= #

# ============================================================= #
# ----------------------- radio_sources ----------------------- #
# ============================================================= #
def radio_sources(time):
    """
    """
    if not isinstance(time, Time):
        time = Time(time)

    def solarsyst_eq(src, time):
        src = get_body(
            src,
            time,
            nenufar_loc()
        )
        return src.ra.deg, src.dec.deg

    with solar_system_ephemeris.set('builtin'):
        src_radec = {
            'vir a': (187.70593075, +12.39112331),
            'cyg a': (299.86815263, +40.73391583),
            'cas a': (350.850000, +58.815000),
            'her a': (252.783433, +04.993031),
            'hyd a': (139.523546, -12.095553),
            'tau a': (83.63308, +22.01450),
            '3c 380': (277.3824220006990, +48.7461552266057),
            'sun': solarsyst_eq('sun', time),
            'moon': solarsyst_eq('moon', time),
            'jupiter': solarsyst_eq('jupiter', time),
        }
    return src_radec
# ============================================================= #

# ============================================================= #
# --------------------------- ateam --------------------------- #
# ============================================================= #
def ateam():
    """
    """
    src_radec = {
        'vir a': (187.70593075, +12.39112331),
        'cyg a': (299.86815263, +40.73391583),
        'cas a': (350.850000, +58.815000),
        'her a': (252.783433, +04.993031),
        'hyd a': (139.523546, -12.095553),
        'tau a': (83.63308, +22.01450),
        '3c 380' : (277.3824220006990, +48.7461552266057)
    }
    return src_radec
# ============================================================= #

# ============================================================= #
# ------------------------- AstroPlot ------------------------- #
# ============================================================= #
class AstroPlot(object):
    """
    """
    def __init__(self, image, center, resol):
        self.npix_x = None
        self.npix_y = None
        self.image = image
        self.center = center
        self.resol = resol

        self.wcs = WCS(naxis=2)
        self.wcs.wcs.crpix = [self.npix_x/2, self.npix_y/2]
        self.wcs.wcs.cdelt = np.array([self.resol, self.resol])
        self.wcs.wcs.crval = [self.center[0], self.center[1]]
        #self.wcs.wcs.ctype = ['RA---AIR', 'DEC--AIR']
        self.wcs.wcs.ctype = ['RA---TAN', 'DEC--TAN']

    @property
    def image(self):
        return self._image
    @image.setter
    def image(self, i):
        if not len(i.shape) == 2:
            raise ValueError(
                'Weird image format.'
            )
        self.npix_x = i.shape[0]
        self.npix_y = i.shape[1]
        self._image = i
        return

    @property
    def center(self):
        return self._center
    @center.setter
    def center(self, c):
        if not isinstance(c, tuple):
            raise TypeError(
                'center should be a tuple'
            )
        if not len(c) == 2:
            raise ValueError(
                'center should be a length-2 tuple'
            )
        self._center = c
        return

    def plot(self, title='', cblabel='', sources=True, time=None, filename=None, circle=None, **kwargs):
        """
        """
        if 'cmap' not in kwargs.keys():
            kwargs['cmap'] = 'YlGnBu_r'
        if 'vmin' not in kwargs.keys():
            kwargs['vmin'] = self.image.min()
        if 'vmax' not in kwargs.keys():
            kwargs['vmax'] = self.image.max()

        cbar = True
        if 'cbar' in kwargs.keys():
            cbar = kwargs['cbar']
            del kwargs['cbar']

        import matplotlib.pyplot as plt
        from matplotlib.colorbar import ColorbarBase
        from matplotlib.ticker import LinearLocator
        from matplotlib.colors import Normalize
        from matplotlib.cm import get_cmap
        from mpl_toolkits.axes_grid1.inset_locator import inset_axes
        from astropy.visualization.wcsaxes import SphericalCircle

        fig = plt.figure(figsize=(7, 7))
        ax = plt.subplot(projection=self.wcs)
        im = ax.imshow(
            self.image,
            origin='lower',
            interpolation='nearest',
            **kwargs
        )

        ax.coords.grid(True, color='white', ls='solid', alpha=0.5)
        axra = ax.coords[0]
        axdec = ax.coords[1]
        axra.set_axislabel('RA')
        axra.set_major_formatter('d')
        axra.set_ticks(number=10)
        axdec.set_axislabel('Dec')
        axdec.set_major_formatter('d')
        axdec.set_ticks(number=10)

        if cbar:
            cax = inset_axes(ax,
               width='5%',
               height='100%',
               loc='lower left',
               bbox_to_anchor=(1.05, 0., 1, 1),
               bbox_transform=ax.transAxes,
               borderpad=0,
               )
            cb = ColorbarBase(
                cax,
                cmap=get_cmap(name='YlGnBu_r'),
                orientation='vertical',
                norm=Normalize(vmin=kwargs['vmin'], vmax=kwargs['vmax']),
                ticks=LinearLocator()
            )
            cb.solids.set_edgecolor('face')
            cb.set_label(cblabel)
            cb.formatter.set_powerlimits((0, 0))

        if sources:
            phase_center = radec(
                ra=self.center[0],
                dec=self.center[1]
            )
            if time is None:
                time = Time.now()
            srcs = radio_sources(time=time)
            for k in srcs.keys():
                src = radec(
                    ra=srcs[k][0],
                    dec=srcs[k][1]
                )
                if phase_center.separation(src).deg > self.npix_x/2 * self.resol:
                    # Source not in FoV
                    continue
                ax.text(
                    src.ra.deg,
                    src.dec.deg,
                    k.title(),
                    transform=ax.get_transform('icrs'),
                    color='white'
                )

        if circle is not None:
            r1 = SphericalCircle(
               center=(circle[0] * u.deg, circle[1] * u.deg),
               radius=circle[2]/2 * u.deg,
               resolution=100,
               edgecolor='white',
               facecolor='none',
               transform=ax.get_transform('icrs')
            )
            ax.add_patch(r1)
            r2 = SphericalCircle(
               center=(circle[0] * u.deg, circle[1] * u.deg),
               radius=circle[2] * u.deg,
               resolution=100,
               edgecolor='white',
               linestyle=':',
               facecolor='none',
               transform=ax.get_transform('icrs')
            )
            ax.add_patch(r2)

        ax.set_title(title)
        # plt.tight_layout()

        if filename is None:
            plt.show()
        else:
            plt.savefig(filename, dpi=300)
        return

    def savefits(self, fitsname, time=None, freq=None):
        """
        """
        header = self.wcs.to_header()
        hdu = fits.PrimaryHDU(self.image.data, header=header)
        hdu.writeto(fitsname, overwrite=True)
        fits.setval(fitsname, 'INSTRU', value='NenuFAR')
        fits.setval(fitsname, 'SOFTWARE', value='nenupytv')
        fits.setval(fitsname, 'VERSION', value=nenupytv.__version__)
        fits.setval(fitsname, 'TIME', value=time.isot, comment='Time in UTC')
        fits.setval(fitsname, 'FREQUENC', value=freq, comment='Mean frequency in MHz')
        fits.setval(fitsname, 'STOKES', value='I')
        fits.setval(fitsname, 'CONTACT', value='alan.loh@obspm.fr')
        return

# ============================================================= #

# ============================================================= #
# ------------------------ astro_image ------------------------ #
# ============================================================= #
def astro_image(
        image,
        center,
        npix,
        resol,
        time,
        freq=None,
        pngfile=None,
        fitsfile=None,
        show_sources=False,
        colorbar=False,
        gal_plane=False,
        **kwargs):
    """
    """
    import matplotlib.pyplot as plt
    # if 'vmin' not in kwargs.keys():
    #     kwargs['vmin'] = np.percentile(specdata.amp, 5)
    # if 'vmax' not in kwargs.keys():
    #     kwargs['vmax'] = np.percentile(specdata.amp, 95)
    if 'cmap' not in kwargs.keys():
        kwargs['cmap'] = 'YlGnBu_r'

    w = WCS(naxis=2)
    w.wcs.crpix = [npix/2, npix/2]
    w.wcs.cdelt = np.array([resol, resol])
    w.wcs.crval = [center[0], center[1]]
    w.wcs.ctype = ['RA---AIR', 'DEC--AIR']

    fig = plt.figure(figsize=(10, 10))
    ax = plt.subplot(projection=w)
    im = ax.imshow(
        image,
        origin='lower',
        aspect='equal',
        interpolation='none',
        **kwargs
    )
    if colorbar:
        plt.colorbar(im, ax=ax)

    ax.coords.grid(True, color='white', ls='solid', alpha=0.5)
    axra = ax.coords[0]
    axdec = ax.coords[1]
    axra.set_axislabel('RA')
    axra.set_major_formatter('d')
    axra.set_ticks(number=10)
    axdec.set_axislabel('Dec')
    axdec.set_major_formatter('d')
    axdec.set_ticks(number=10)

    # if sources:
    #     from astroquery.vizier import Vizier
    #     from astropy.coordinates import SkyCoord
    #     import astropy.units as un
    #     Vizier.ROW_LIMIT = -1
    #     catalog_list = Vizier.find_catalogs('VIII/1A')
    #     catalogs = Vizier.get_catalogs(catalog_list.keys())
    #     cat_3c = catalogs[0]
    #     ra_zen = center[0]
    #     dec_zen = center[1]
    #     zenith = SkyCoord(ra_zen * un.deg, dec_zen * un.deg)
    #     maxjy = np.max(np.log10(cat_3c['S159MHz']))
    #     for i in range(len(cat_3c)):
    #         src_3c = SkyCoord(
    #             cat_3c['RA1950'][i],
    #             cat_3c['DE1950'][i],
    #             unit=(un.hourangle, un.deg),
    #             equinox='B1950'
    #         )
    #         if zenith.separation(src_3c).deg < 32:
    #             ax.scatter(
    #                 src_3c.ra.deg,
    #                 src_3c.dec.deg,
    #                 transform=ax.get_transform('icrs'),
    #                 s=np.log10(cat_3c['S159MHz'][i])/maxjy * 500,
    #                 edgecolor='white',
    #                 facecolor='none',
    #                 #alpha=np.log10(cat_3c['S159MHz'][i])/maxjy
    #             )
    if show_sources:
        phase_center = radec(
            ra=center[0],
            dec=center[1]
        )

        srcs = radio_sources(time=time)

        for k in srcs.keys():
            src = radec(
                ra=srcs[k][0],
                dec=srcs[k][1]
            )
            if phase_center.separation(src).deg > npix/2 * resol:
                # Source not in FoV
                continue
            # ax.scatter(
            #     src.ra.deg,
            #     src.dec.deg,
            #     s=100,
            #     transform=ax.get_transform('icrs'),
            #     edgecolor='white',
            #     facecolor='none',
            # )
            ax.text(
                src.ra.deg,
                src.dec.deg,
                k.title(),
                transform=ax.get_transform('icrs'),
                color='white'
            )

    if pngfile is None:
        plt.show()
    else:
        plt.title('{}'.format(time.iso))
        #plt.tight_layout()
        plt.savefig(pngfile, **kwargs)

    if fitsfile is not None:
        fitsfile = path.join(fitsfile, 'nenufartv_{}.fits'.format(time.isot.split('.')[0].replace(':', '-')))
        header = w.to_header()
        hdu = fits.PrimaryHDU(image, header=header)
        hdu.writeto(fitsfile, overwrite=True)
        fits.setval(fitsfile, 'INSTRU', value='NenuFAR')
        fits.setval(fitsfile, 'SOFTWARE', value='nenupytv')
        fits.setval(fitsfile, 'VERSION', value=nenupytv.__version__)
        fits.setval(fitsfile, 'TIME', value=time.isot, comment='Time in UTC')
        fits.setval(fitsfile, 'FREQUENC', value=freq, comment='Mean frequency in MHz')
        fits.setval(fitsfile, 'STOKES', value='I')
        fits.setval(fitsfile, 'CONTACT', value='alan.loh@obspm.fr')
# ============================================================= #

import numpy,import astropy,from astropy


#!/usr/bin/env python
#
# correlate.py - The CorrelateAction class.
#
# Author: Paul McCarthy <pauldmccarthy@gmail.com>
#
"""This module provides the :class:`.PearsonCorrelateAction` class, which is
an :class:`.Action` that calculates seed-based correlation on 4D
:class:`.Image` overlays.
"""

import threading
import logging

import numpy                  as np
import scipy.spatial.distance as spd

import fsl.data.image               as fslimage
import fsl.utils.idle               as idle
import fsleyes_props                as props
import fsleyes.views.orthopanel     as orthopanel
import fsleyes_widgets.utils.status as fslstatus
import fsleyes.strings              as strings
import fsleyes.actions.base         as base

log = logging.getLogger(__name__)

class CorrelateAction(base.Action):
    """The ``CorrelateAction`` is a base class for the
    :class:`PearsonCorrelateAction` and :class:`PCACorrelateAction` classes,
    which manages adding/removing correlation overlays to/from the
    :class:`.OverlayList`, and manages execution of the correlation.

    When a 4D :class:`.Image` is selected and the ``CorrelateAction`` is
    invoked, a new 3D :class:`.Image` is created and added to the
    :class:`.OverlayList` - this image is referred to as a *correlate overlay*,
    and is used to store and display the correlation values.
    """

    @staticmethod
    def supportedViews():
        """The ``CorrelateAction`` is restricted for use with
        :class:`.OrthoPanel` views.
        """
        return [orthopanel.OrthoPanel]

    @staticmethod
    def ignoreTool():
        """Tells the FSLeyes plugin system not to add the ``CorrelateAction``
        class as an option to the FSLeyes toolsmenu.  Instead, the
        :class:`PearsonCorrelateAction` action (and other potential future
        sub-classes) is added.
        """
        return True

    def __init__(self, overlayList, displayCtx, panel):
        """Create a ``CorrelateAction``.

        :arg overlayList: The :class:`.OverlayList`.
        :arg displayCtx:  The :class:`.DisplayContext`.
        :arg panel:       The :class:`.OrthoPanel` that owns this action.
        """

        base.Action.__init__(
            self, overlayList, displayCtx, self.__runCorrelateAction)

        self.__panel = panel
        self.__name  = '{}_{}'.format(type(self).__name__, id(self))

        displayCtx .addListener('selectedOverlay',
                                self.__name,
                                self.__selectedOverlayChanged)
        overlayList.addListener('overlays',
                                self.__name,
                                self.__overlayListChanged)

        # TODO Use a single data structure -
        #      using two dicts is fragile
        self.__correlateOverlays = {}
        self.__overlayCorrelates = {}

        # The runCorrelateAction cannot be called
        # more than once at a time - this is used
        # as a semaphore to ensure that this is
        # enforced.
        self.__correlateFlag = threading.Event()

        self.__selectedOverlayChanged()

    def destroy(self):
        """Removes listeners from the :class:`.DisplayContext` and
        :class:`.OverlayList`, and calls :meth:`.Action.destroy`.
        """
        if self.destroyed:
            return

        self.displayCtx .removeListener('selectedOverlay', self.__name)
        self.overlayList.removeListener('overlays',        self.__name)
        base.Action.destroy(self)

        self.__correlateOverlays = None
        self.__overlayCorrelates = None

    def __selectedOverlayChanged(self, *a):
        """Called when the selected overlay, or overlay list, changes.

        Enables/disables this action depending on the nature of the selected
        overlay.
        """

        ovl          = self.displayCtx.getSelectedOverlay()
        isCorrOvl    = ovl in self.__overlayCorrelates

        self.enabled = isCorrOvl or  \
                       ((ovl is not None)               and
                        isinstance(ovl, fslimage.Image) and
                        ovl.ndim > 3)

    def __overlayListChanged(self, *a):
        """Called when the :class:`.OverlayList` changes. Makes sure that
        there are no obsolete correlate overlays in the list, and calls
        :meth:`__selectedOverlayChanged`.
        """
        self.__clearCorrelateOverlays()
        self.__selectedOverlayChanged()

    def __clearCorrelateOverlays(self):
        """Called by :meth:`__overlayListChanged`. Clears internal references
        to any obsolete correlate overlays.
        """

        for overlay, corrOvl in list(self.__correlateOverlays.items()):
            if overlay not in self.overlayList or \
               corrOvl not in self.overlayList:
                self.__correlateOverlays.pop(overlay)
                self.__overlayCorrelates.pop(corrOvl)

    def __createCorrelateOverlay(self, overlay, data):
        """Creates a *correlate* overlay for the given ``overlay``, adds
        it to the :class:`.OverlayList`, and initialises some display
        properties.
        """

        display = self.displayCtx.getDisplay(overlay)
        name    = '{}/correlation'.format(display.name)
        corrOvl = fslimage.Image(data, name=name, header=overlay.header)

        self.overlayList.append(corrOvl, overlayType='volume')
        self.__correlateOverlays[overlay] = corrOvl
        self.__overlayCorrelates[corrOvl] = overlay

        corrOpts = self.displayCtx.getOpts(corrOvl)

        with props.suppressAll(corrOpts), \
             props.suppressAll(display):
            corrOpts.cmap              = 'red-yellow'
            corrOpts.negativeCmap      = 'blue-lightblue'
            corrOpts.useNegativeCmap   = True
            corrOpts.displayRange      = [0.05, 1]
            corrOpts.clippingRange.xlo = 0.05

        return corrOvl

    def __runCorrelateAction(self):
        """Called when this :class:`.Action` is invoked. Calculates correlation
        values from the voxel at the current :attr:`.DisplayContext.location`
        (relative to the currently selected overlay) to all other voxels, and
        updates the correlate overlay.

        The correlation calculation and overlay update is performed on a
        separate thread (via :meth:`.idle.run`), with a call to
        :meth:`calculateCorrelation`.
        """

        # Because of the multi-threaded/asynchronous
        # way that this function does its job,
        # allowing it to be called multiple times
        # before prior calls have completed would be
        # very dangerous indeed.
        if self.__correlateFlag.is_set():
            log.debug('Correlate action is already '
                      'running - ignoring request')
            return

        # See if a correlate overlay already exists
        # for the currently selected overlay
        ovl     = self.displayCtx.getSelectedOverlay()
        corrOvl = self.__correlateOverlays.get(ovl, None)

        # If not, check to see if it is a correlate
        # overlay that is selected and, if it is,
        # look up the corresponding source overlay.
        if corrOvl is None:
            if ovl in self.__overlayCorrelates:
                corrOvl = ovl
                ovl     = self.__overlayCorrelates[corrOvl]

        # If corrOvl is still None, it means that
        # there is no correlate overlay for the
        # currently selected overlay. In this case,
        # we'll create a new correlate overlay and
        # add it to the overlay list after the
        # correlation values have been calculated.

        opts = self.displayCtx.getOpts(ovl)
        xyz  = opts.getVoxel(vround=True)

        if xyz is None:
            return

        data = ovl.data[opts.index(atVolume=False)]

        # The correlation calculation is performed
        # on a separate thread. This thread then
        # schedules a function on idle.idle to
        # update the correlation overlay back on the
        # main thread.
        def calcCorr():

            correlations = self.calculateCorrelation(xyz, data)

            # The correlation overlay is updated/
            # created on the main thread.
            def update():

                try:

                    # A correlation overlay already
                    # exists for the source overlay
                    # - update its data
                    if corrOvl is not None:
                        corrOvl[:] = correlations

                    # The correlation overlay hasn't
                    # been created yet - create a
                    # new overlay with the correlation
                    # values.
                    else:
                        self.__createCorrelateOverlay(ovl, correlations)

                finally:
                    fslstatus.clearStatus()
                    self.__correlateFlag.clear()

            idle.idle(update)

        # Protect against more calls
        # while this job is running.
        self.__correlateFlag.set()
        fslstatus.update(strings.messages[self, 'calculating'].format(*xyz))
        idle.run(calcCorr)

    def calculateCorrelation(self, seed, data):
        """Calculates correlation values between the given ``seed`` voxel (an
        ``(x, y, z)`` tuple) and all other voxels. This method must be
        implemented by sub-classes.

        :arg seed: An ``(x, y, z)`` tuple specifying the seed voxel

        :arg data: A 4D ``numpy`` array containing all of the data.

        :returns:  A 3D ``numpy`` array containing the correlation values.
        """
        raise NotImplementedError('calculateCorrelation must be '
                                  'implemented by sub-classes')

class PearsonCorrelateAction(CorrelateAction):
    """The ``PearsonCorrelateAction`` is a :class:`CorrelateAction` which
    calculates Pearson correlation coefficient values between the seed voxel
    and all other voxels.
    """

    def calculateCorrelation(self, seed, data):
        """Calculates Pearson correlation between the data at the specified
        seed voxel, and all other voxels.
        """
        return pearsonCorrelation(seed, data)

def pearsonCorrelation(seed, data):
    """Calculates Pearson correlation between the data at the specified
    seed voxel, and all other voxels.
    """

    x, y, z = seed
    npoints = data.shape[3]

    # the scipy.spatial.distance.cdist
    # function can be used to calculate
    # one-to-many correlation values.
    with np.errstate(invalid='ignore'):
        correlations = 1 - spd.cdist(
            data[x, y, z, :].reshape( 1, npoints),
            data            .reshape(-1, npoints),
            metric='correlation')

    # Set any nans to 0
    correlations[np.isnan(correlations)] = 0

    return correlations.reshape(data.shape[:3])


#Copyright (c) 2009,2010 George Dahl

import numpy as num

import cudamat as cm
from cudamat import reformat
from singleSoftmax import maskedSingleSoftmax

def getFilteringDist(net, data, index, preSigmoid = False):
    """
    We use this name to correspond more closely to Graham's matlab
    code.  This function sends the visible data stored in data through
    net to produce hidden unit activations for every valid position of
    net.  The valid positions are given by index.
    """
    assert(len(index.shape)==1)
    pred = []

    numcases = index.shape[0]
    num_mini_batches = numcases / net.mbsz
    excess = numcases - num_mini_batches*net.mbsz

    for mb in range(num_mini_batches):
        mbIdx = index[ mb*net.mbsz:(mb+1)*net.mbsz ]
        net.vis = cm.CUDAMatrix(reformat(data[:,mbIdx]))
        net.past = [ cm.CUDAMatrix(reformat(data[:,mbIdx-i-1])) for i in range(net.numPrev) ]

        if preSigmoid:
            net.hidNetInpts()
        else:
            net.hidActProbs()
        net.hActProbs.copy_to_host()
        pred.append(net.hActProbs.numpy_array.copy())
    if excess > 0:
        batch = num.zeros(net.vis.shape)
        mbIdx = index[ num_mini_batches*net.mbsz:]
        batch[:,:excess] = data[:,mbIdx]
        net.vis = cm.CUDAMatrix(reformat(batch))
        net.past = []
        for i in range(net.numPrev):
            batch[:,:excess] = data[:,mbIdx-i-1]
            net.past.append(cm.CUDAMatrix(reformat(batch)))
        if preSigmoid:
            net.hidNetInpts()
        else:
            net.hidActProbs()
        net.hActProbs.copy_to_host()
        pred.append(net.hActProbs.numpy_array.copy()[:,:excess])

    return num.hstack(pred)

class GaussianCRBM(object):
    def __init__(self, numVis, numHid, prevFrames, initHidBias = 0.0):
        self.numVis, self.numHid, self.numPrev = numVis, numHid, prevFrames
        self._mbsz = 256

        self.visToHid = 0.1*num.random.randn(numVis, numHid)
        self.visBias = num.zeros((numVis, 1))
        self.hidBias = num.zeros((numHid, 1)) + initHidBias

        #self.A[0] and self.B[0] are the weights from the most recent frame
        self.A = [0.01*num.random.randn(numVis, numVis) for i in range(self.numPrev)]
        self.B = [0.01*num.random.randn(numVis, numHid)  for i in range(self.numPrev)]

        self.init_weight_storage()

        self.initTemporary()

        #will be used for L1 reg and allocated at that point
        self.signVisToHid = None
        self.signA = None
        self.signB = None

        #set default learning parameters:
        self.setLearningParams()

        #total GPU storage costs excluding input data (self.past and self.vis):
        # 2W + 2*numHid*mbsz+2*numVis*mbsz
        #where W is the total space cost of all the weights of the model

    def setLearningParams(self, learnRate = 0.001, momentum = 0.9, weightCost = 0, regType = "L2", cdSteps = 1, \
                          pastNoise = 0, arWeightCost = None):
        self.learnRate = learnRate
        self.momentum = momentum
        self.weightCost = weightCost
        self.regType = regType
        self.cdSteps = cdSteps
        self.pastNoise = pastNoise
        self.arWeightCost = arWeightCost

    def getMBSZ(self):
        return self._mbsz

    def setMBSZ(self, newMBSZ):
        self._mbsz = newMBSZ
        self.initTemporary()
    mbsz = property(getMBSZ,setMBSZ)

    def initTemporary(self):
        self.hActs = cm.CUDAMatrix(reformat(num.zeros((self.numHid, self.mbsz))))
        self.hActProbs = cm.CUDAMatrix(reformat(num.zeros((self.numHid, self.mbsz))))
        self.negVis = cm.CUDAMatrix(reformat(num.zeros((self.numVis, self.mbsz))))
        self.tempVisMB = cm.CUDAMatrix(reformat(num.zeros((self.numVis, self.mbsz))))
        self.dynamicHidBias = cm.CUDAMatrix(reformat(num.zeros((self.numHid, self.mbsz))))
        self.dynamicVisBias = cm.CUDAMatrix(reformat(num.zeros((self.numVis, self.mbsz))))

    def init_weight_storage(self):
        """
        Initialize storage for gradients and gradient steps and build a list of
        weight/gradient/energy gradient/step tuples.
        """
        for name in self.weightVariableNames():
            w = self.__dict__[name]
            if not isinstance(w, list):
                self.__dict__[name] = cm.CUDAMatrix(reformat(w))
                self.__dict__["d"+name] = cm.CUDAMatrix(reformat(0.0 * w))
            else:
                self.__dict__[name] = [cm.CUDAMatrix(reformat(x)) for x in w]
                self.__dict__["d"+name] = [cm.CUDAMatrix(reformat(0.0*part)) for part in w]

    def scaleDerivs(self, factor):
        """
        Scales all weight derivatives by factor (used to apply
        momentum or clear the weight derivatives).
        """
        for name in self.weightVariableNames():
            w = self.__dict__[name]
            if not isinstance(w, list):
                self.__dict__["d"+name].mult_by_scalar(factor)
            else:
                for i in range(len(w)):
                    self.__dict__["d"+name][i].mult_by_scalar(factor)

    def pack_weights(self):
        w_dict = {}
        for w_name in self.weightVariableNames():
            w = self.__dict__[w_name]
            if isinstance(w, list):
                for part in w:
                    part.copy_to_host()
                w_dict[w_name] = [part.numpy_array for part in w]
            else:
                w.copy_to_host()
                w_dict[w_name] = w.numpy_array
        return w_dict

    def loadWeights(self, wDict):
        """
        This code is terrible.
        """
        assert(all(wName in wDict for wName in self.weightVariableNames()))
        for w_name in wDict:
            if w_name in self.weightVariableNames():
                w = wDict[w_name]
                if isinstance(w, list) or w_name in ["A","B"]:
                    assert( all(self.__dict__[w_name][i].numpy_array.shape == wDict[w_name][i].shape for i in range(len(wDict[w_name])) ) )
                    self.__dict__[w_name] = [cm.CUDAMatrix(reformat(part)) for part in w]
                else:
                    assert( self.__dict__[w_name].numpy_array.shape == wDict[w_name].shape )
                    self.__dict__[w_name] = cm.CUDAMatrix(reformat(w))

    def curRecErr(self):
        self.vis.subtract(self.negVis, target = self.tempVisMB)
        return self.tempVisMB.euclid_norm()**2

    def allWeightsMatlabFormat(self):
        weights = self.pack_weights()
        d = {}
        d["w"] = weights["visToHid"].transpose()
        d["bi"] = weights["visBias"]
        d["bj"] = weights["hidBias"]

        #this chunk of code depends on scipy version >= 8 so savemat works right
        d["A"] = num.empty( (self.numVis, self.numVis, self.numPrev) )
        d["B"] = num.empty( (self.numHid, self.numVis, self.numPrev) )
        for i in range(self.numPrev):
            d["A"][:,:,i] = weights["A"][i].transpose()
            d["B"][:,:,i] = weights["B"][i].transpose()

        #for i in range(self.numPrev):
        #    d["A%d" % i] = weights["A"][i].transpose()
        #    d["B%d" % i] = weights["B"][i].transpose()

        return d

    def sampleHiddens(self, hActProbsOnGPU = None):
        if hActProbsOnGPU == None:
            hActProbsOnGPU = self.hActProbs
        self.hActs.fill_with_rand()
        self.hActs.less_than(hActProbsOnGPU, target = self.hActs)

    def hidNetInpts(self, recomputeDynamicBias = True, targ = None, vis = None):
        """
        targ had better be on the gpu or None
        """
        if recomputeDynamicBias:
            self.updateDynamicHidBias()

        if targ == None:
            targ = self.hActProbs
        if vis == None:
            vis = self.vis

        cm.dot( self.visToHid.T, vis, target = targ)
        targ.add(self.dynamicHidBias)
        targ.add_col_vec(self.hidBias)

    def hidActProbs(self, recomputeDynamicBias = True, targ = None, vis = None):
        """
        targ had better be on the gpu or None
        """
        if targ == None:
            targ = self.hActProbs
        self.hidNetInpts(recomputeDynamicBias, targ, vis)
        targ.apply_sigmoid()

    def updateDynamicHidBias(self):
        self.dynamicHidBias.mult_by_scalar(0.0)
        for i in range(len(self.B)):
            #self.past[i] is the (i+1)-steps delayed frame of history
            #self.past[i] is numVis by mbsz
            #self.B[i] is numVis by numHid
            self.dynamicHidBias.add_dot(self.B[i].T, self.past[i])

    def updateDynamicVisBias(self):
        self.dynamicVisBias.mult_by_scalar(0.0)
        for i in range(len(self.A)):
            self.dynamicVisBias.add_dot(self.A[i].T, self.past[i])

    def visActProbs(self, recomputeDynamicBias):

        if recomputeDynamicBias:
            self.updateDynamicVisBias()

        cm.dot( self.visToHid, self.hActs, target = self.negVis)
        self.negVis.add(self.dynamicVisBias)
        self.negVis.add_col_vec(self.visBias)

    def weightVariableNames(self):
        """
        Returns the names of the variables for the weights that define
        this model in a cannonical order.  The order must match up
        with the way weight derivatives get returned from CDn.
        """
        return "visToHid", "hidBias", "visBias", "A", "B"

    def CDStats(self, vis, past, hid, posPhase):
        """
        hid should be self.numHid by mbsz and exist on the GPU
        vis should be self.numVis by mbsz and exist on the GPU
        past should be a length self.numPrev list of variables like vis

        This function depends on self.dynamicVisBias being up to date!!

        We modify self.d$WEIGHT_NAME as a side effect and clobber self.tempVisMB.
        """
        vis.subtract(self.dynamicVisBias, target = self.tempVisMB)
        self.tempVisMB.add_col_mult(self.visBias, -1.0) #so we are subtracting, not adding

        multiplier = 1.0 if posPhase else -1.0

        self.dhidBias.add_sums(hid, 1, mult = multiplier)
        self.dvisBias.add_sums(vis, 1, mult = multiplier)

        if posPhase:
            self.dvisToHid.add_dot(vis, hid.T)

            for i in range(self.numPrev):
                self.dA[i].add_dot( past[i], self.tempVisMB.T )
                self.dB[i].add_dot( past[i], hid.T )
        else:
            self.dvisToHid.subtract_dot(vis, hid.T)

            for i in range(self.numPrev):
                self.dA[i].subtract_dot( past[i], self.tempVisMB.T )
                self.dB[i].subtract_dot( past[i], hid.T )

    def CDn(self):
        """
        After this function runs we will have the negative data in
        self.negVis and self.hActProbs will hold the final hidden
        activation probabilities conditioned on the negative data.

        This function updates the weight derivative variables.
        """
        #we depend on the following two learning parameters
        n = self.cdSteps
        momentum = self.momentum

        #apply momentum
        self.scaleDerivs(momentum)

        #stores hidden activation probabilities in self.hActProbs and sets dynamic hidden biases
        self.hidActProbs()

        self.updateDynamicVisBias() #CDStats depends on self.dynamicVisBias being correct

        #compute positive phase statistics and add them to gradient variables
        self.CDStats(self.vis, self.past, self.hActProbs, True)

        for i in range(n):
            #updates self.hActs
            self.sampleHiddens(self.hActProbs)

            #updates self.negVis and if i == 0 computes self.dynamicVisBias
            self.visActProbs(False) #no need to recompute self.dynamicVisBias

            #stores recomputed (based on self.negVis) hidden act probs in self.hActProbs
            self.hidActProbs(False, vis = self.negVis)

        #compute negative phase statistics and subtract them from gradient variables
        self.CDStats(self.negVis, self.past, self.hActProbs, False)

    def reformatLearningRates(self, learnRate):
        if isinstance(learnRate, dict):
            assert( all(name in learnRate for name in self.weightVariableNames() ) )
            return learnRate
        rates = {}
        assert( type(learnRate) == float or type(learnRate) == int )
        for name in self.weightVariableNames():
            rates[name] = learnRate
        return rates

    def updateSignOfWeights(self):
        """
        We need the sign of the weights for L1 regularization.  Since
        we work on the GPU it is convenient to just allocate storage
        for these things once and periodically update the sign
        variables when the weights they depend on have changed and we
        need to know the signs.
        """
        if self.signVisToHid == None or self.signA == None or self.signB == None:
            self.signVisToHid = cm.CUDAMatrix(reformat(num.zeros((self.numVis, self.numHid))))
            self.signA = [cm.CUDAMatrix(reformat(num.zeros((self.numVis, self.numVis)))) for i in range(self.numPrev)]
            self.signB = [cm.CUDAMatrix(reformat(num.zeros((self.numVis, self.numHid)))) for i in range(self.numPrev)]
        self.visToHid.sign(target = self.signVisToHid)
        for i in range(self.numPrev):
            self.A[i].sign(target = self.signA[i])
            self.B[i].sign(target = self.signB[i])

    def decay(self):
        """
        Weight decay during pretraining.  LearningRates should be a
        dictionary with keys self.weightVariableNames() that holds the
        learning rate for each weight.
        """
        #here are the learning parameters this method depends on
        decayRate = self.weightCost
        arDecayRate = self.arWeightCost if self.arWeightCost != None else decayRate
        learningRates = self.reformatLearningRates(self.learnRate) # we reformat in case self.learnRate isn't a dict
        regType = self.regType

        if decayRate > 0: #hopefully this saves time when decayRate == 0
            #really for L1+bias mode we should allow different weight costs for the L1 part and the bias sparsity
            assert( regType in ["L2","L1","bias","L1+bias", "L2+bias"] )
            if "L2" in regType:
                self.visToHid.mult_by_scalar( 1-decayRate*learningRates['visToHid'] )
                for i in range(self.numPrev):
                    self.A[i].mult_by_scalar( 1-arDecayRate*learningRates['A'] )
                    self.B[i].mult_by_scalar( 1-arDecayRate*learningRates['B'] )
            if "L1" in regType:
                self.updateSignOfWeights()
                self.visToHid.subtract_mult(self.signVisToHid, decayRate*learningRates['visToHid'])
                for i in range(self.numPrev):
                    self.A[i].subtract_mult(self.signA[i], arDecayRate*learningRates['A'])
                    self.B[i].subtract_mult(self.signB[i], arDecayRate*learningRates['B'])
            if "bias" in regType:
                self.hidBias.add_scalar( -decayRate*learningRates['hidBias'] )

    def step(self, data, past):
        """
        This function sets references in self.vis and self.past to
        point to data and past.
        """
        self.vis = data
        self.past = past
        self.CDn()
        rates = self.reformatLearningRates(self.learnRate)
        self.decay() #needs dictionary of learning rates, but it will reformat the rates again on its own
        for j, wname in enumerate(self.weightVariableNames()):
            if type(self.__dict__[wname]) == list:
                for i in range(self.numPrev):
                    self.__dict__[wname][i].add_mult( self.__dict__["d"+wname][i], rates[wname]/self.mbsz )
            else: #we assume it is a numpy array
                self.__dict__[wname].add_mult( self.__dict__["d"+wname], rates[wname]/self.mbsz )

    def trainLowMemory(self, data, index, numEpochs, reportMB = False):
        assert(data.dtype == num.dtype('float32'))
        numcases = len(index)

        num_mini_batches = numcases / self.mbsz
        indexPerm = num.random.permutation(range(numcases))

        noise = cm.CUDAMatrix(reformat(num.zeros((self.numVis, self.mbsz))))
        for ep in range(numEpochs):
            recErr = 0
            for mb in range(num_mini_batches):
                mbIndex = index[ indexPerm[mb*self.mbsz:(mb+1)*self.mbsz] ]

                curInputsMB_CPU = data[:, mbIndex]
                curPastMB_CPU = [data[:, mbIndex-i-1] for i in range(self.numPrev)]
                curInputsMB = cm.CUDAMatrix(reformat(curInputsMB_CPU))
                curPastMB = [cm.CUDAMatrix(reformat(p)) for p in curPastMB_CPU]
                if self.pastNoise > 0:
                    for i in range(self.numPrev):
                        noise.fill_with_randn()
                        curPastMB[i].add_mult(noise, self.pastNoise)

                self.step(curInputsMB, curPastMB)
                recErr += self.curRecErr()
                if reportMB:
                    yield (mb, num_mini_batches)
            yield recErr

    def oneLayerReconstructions(self, inputs, past, sample = False):
        #inputs and past should be on the CPU
        hiddens = self.predictions(inputs, past, sample)
        recons = self.reconstructions(past, hiddens, False)
        return recons[:,:inputs.shape[1]]

    def predictions(self, inp, past, sample = False):

        """
        This function assumes inp and past reside on the cpu.  It
        returns a numpy array.

        We assume an integer number of minibatches and any cases
        beyond mbsz*floor(numcases/mbsz) are ignored.
        """
        #we return an array numHid by floor(numcases/mbsz)
        pred = []

        numcases = inp.shape[1]
        num_mini_batches = numcases / self.mbsz

        for i in range(num_mini_batches):
            idx = i*self.mbsz
            self.vis = cm.CUDAMatrix(reformat(inp[:,idx:idx+self.mbsz]))
            self.past = [ cm.CUDAMatrix(reformat(p[:,idx:idx+self.mbsz])) for p in past ]

            self.hidActProbs()
            if sample:
                self.sampleHiddens(self.hActProbs)
                self.hActs.copy_to_host()
                pred.append(self.hActs.numpy_array.copy())
            else:
                self.hActProbs.copy_to_host()
                pred.append(self.hActProbs.numpy_array.copy())
        return num.hstack(pred)

    def reconstructions(self, past, hiddens, onGPU = False):
        """
        We assume we have an integer number of
        minibatches.
        """
        #we return an array numVis by floor(numcases/mbsz)
        if onGPU:
            pastGPU = past
            hiddensGPU = hiddens
        else:
            pastGPU = [cm.CUDAMatrix(reformat(p)) for p in past]
            hiddensGPU = cm.CUDAMatrix(reformat(hiddens))

        numcases = hiddensGPU.numpy_array.shape[1]
        num_mini_batches = numcases / self.mbsz

        recons = []
        for i in range(num_mini_batches):
            self.past = [p.slice(i*self.mbsz, (i+1)*self.mbsz) for p in pastGPU]
            self.hActs = hiddensGPU.slice(i*self.mbsz, (i+1)*self.mbsz)
            self.visActProbs(True)
            self.negVis.copy_to_host()
            recons.append(self.negVis.numpy_array.copy())

        return num.hstack(recons)

def padToMinibatch(matrixOnCPU, mbsz):
    if matrixOnCPU.shape[1] % mbsz == 0:
        return matrixOnCPU, 0
    pad_num = mbsz - matrixOnCPU.shape[1] % mbsz
    return num.hstack( (matrixOnCPU, num.zeros((matrixOnCPU.shape[0], pad_num))) ), pad_num

class BinaryCRBM(GaussianCRBM):
    def visActProbs(self, recomputeDynamicBias):
        GaussianCRBM.visActProbs(self, recomputeDynamicBias)
        self.negVis.apply_sigmoid()

    def setLearningParams(self, learnRate = 0.04, momentum = 0.9, weightCost = 0, regType = "L2", cdSteps = 1, \
                          pastNoise = 0, arWeightCost = None, samplePast = False):
        self.learnRate = learnRate
        self.momentum = momentum
        self.weightCost = weightCost
        self.regType = regType
        self.cdSteps = cdSteps
        self.pastNoise = pastNoise
        self.samplePast = samplePast
        self.arWeightCost = arWeightCost

    def trainLowMemory(self, data, index, numEpochs, reportMB = False):
        assert(data.dtype == num.dtype('float32'))
        numcases = len(index)

        num_mini_batches = numcases / self.mbsz
        indexPerm = num.random.permutation(range(numcases))

        noise = cm.CUDAMatrix(reformat(num.zeros((self.numVis, self.mbsz))))
        noiseThresh = cm.CUDAMatrix(reformat(num.zeros((self.numVis, self.mbsz))))
        noiseThresh.assign_scalar(1.0-self.pastNoise)
        for ep in range(numEpochs):
            recErr = 0
            for mb in range(num_mini_batches):
                mbIndex = index[ indexPerm[mb*self.mbsz:(mb+1)*self.mbsz] ]

                curInputsMB_CPU = data[:, mbIndex]
                curPastMB_CPU = [data[:, mbIndex-i-1] for i in range(self.numPrev)]
                curInputsMB = cm.CUDAMatrix(reformat(curInputsMB_CPU))
                curPastMB = [cm.CUDAMatrix(reformat(p)) for p in curPastMB_CPU]
                for i in range(self.numPrev):
                    if self.pastNoise > 0 and not self.samplePast:
                        noise.fill_with_rand()
                        noise.less_than(noiseThresh, target = noise)
                        curPastMB[i].mult(noise)
                    if self.samplePast:
                        noise.fill_with_rand()
                        noise.less_than(curPastMB[i], target = curPastMB[i])

                self.step(curInputsMB, curPastMB)
                recErr += self.curRecErr()
                if reportMB:
                    yield (mb, num_mini_batches)
            yield recErr

class HybridCRBM(GaussianCRBM):
    """
    This class implements a hybrid crbm with a single softmax unit and
    some gaussian units for the visible units.
    """
    def __init__(self, numVis, numHid, prevFrames, smsz, initHidBias = 0.0):
        assert(0 <= smsz <= numVis)
        self.smsz = smsz
        GaussianCRBM.__init__(self, numVis, numHid, prevFrames, initHidBias)

    def initTemporary(self):
        self.hActs = cm.CUDAMatrix(reformat(num.zeros((self.numHid, self.mbsz))))
        self.hActProbs = cm.CUDAMatrix(reformat(num.zeros((self.numHid, self.mbsz))))
        self.negVis = cm.CUDAMatrix(reformat(num.zeros((self.numVis, self.mbsz))))
        self.tempVisMB = cm.CUDAMatrix(reformat(num.zeros((self.numVis, self.mbsz))))
        self.dynamicHidBias = cm.CUDAMatrix(reformat(num.zeros((self.numHid, self.mbsz))))
        self.dynamicVisBias = cm.CUDAMatrix(reformat(num.zeros((self.numVis, self.mbsz))))

        self.sMask = num.zeros((self.numVis, self.mbsz))
        self.sMask[:self.smsz,:] = 1
        self.gaussMask = 1-self.sMask

        self.onesCol = cm.CUDAMatrix(reformat(num.ones((self.numVis,1))))
        self.sMask = cm.CUDAMatrix(reformat(self.sMask))
        self.gaussMask = cm.CUDAMatrix(reformat(self.gaussMask))
        self.tempRow = cm.CUDAMatrix(reformat(num.zeros((1, self.mbsz))))
        #self.tempBinVisMB = cm.CUDAMatrix(reformat(num.zeros((self.numVis, self.mbsz))))

    def setLearningParams(self, learnRate = 0.001, momentum = 0.9, weightCost = 0, regType = "L2", cdSteps = 1, \
                          pastNoise = 0, arWeightCost = None, pastNoiseSM = 0):
        GaussianCRBM.setLearningParams(self, learnRate, momentum, weightCost, regType, cdSteps, pastNoise, arWeightCost)
        self.pastNoiseSM = pastNoiseSM

    def visActProbs(self, recomputeDynamicBias):
        GaussianCRBM.visActProbs(self, recomputeDynamicBias)
        maskedSingleSoftmax(self.negVis, self.tempVisMB, self.sMask, self.gaussMask, self.onesCol, self.tempRow)

    def trainLowMemory(self, data, index, numEpochs, reportMB = False):
        assert(data.dtype == num.dtype('float32'))
        numcases = len(index)

        num_mini_batches = numcases / self.mbsz
        indexPerm = num.random.permutation(range(numcases))

        noise = cm.CUDAMatrix(reformat(num.zeros((self.numVis, self.mbsz))))
        for ep in range(numEpochs):
            recErr = 0
            for mb in range(num_mini_batches):
                mbIndex = index[ indexPerm[mb*self.mbsz:(mb+1)*self.mbsz] ]

                curInputsMB_CPU = data[:, mbIndex]
                curPastMB_CPU = [data[:, mbIndex-i-1] for i in range(self.numPrev)]
                if self.pastNoiseSM > 0:
                    for i in range(self.numPrev):
                        smNoise = (self.pastNoiseSM/self.smsz)*num.random.rand(self.smsz, self.mbsz)
                        #smNoise[0,:] = 0
                        #smNoise /= self.smsz-1
                        curPastMB_CPU[i][:self.smsz,:] = (curPastMB_CPU[i][:self.smsz,:] + smNoise)/(1+self.pastNoiseSM)

                curInputsMB = cm.CUDAMatrix(reformat(curInputsMB_CPU))
                curPastMB = [cm.CUDAMatrix(reformat(p)) for p in curPastMB_CPU]
                if self.pastNoise > 0:
                    for i in range(self.numPrev):
                        noise.fill_with_randn()
                        noise.mult(self.gaussMask)
                        curPastMB[i].add_mult(noise, self.pastNoise)

                self.step(curInputsMB, curPastMB)
                recErr += self.curRecErr()
                if reportMB:
                    yield (mb, num_mini_batches)
            yield recErr

def gpu_batches(data, past, mbs, transpose = True):
    """
    We assume that the first dimension of the data is the number of cases.

    We generate minibatches of data and delayed data of the appropriate size transposed for use on the GPU.

    If we can't fill the last minibatch, we discard that data.
    """
    numCases, numDims = data.shape
    numBatches = numCases/mbs
    for i in range(numBatches):
        if transpose:
            yield (data[i*mbs:(i+1)*mbs,:].transpose(), [p[i*mbs:(i+1)*mbs,:].transpose() for p in past])
        else:
            yield (data[i*mbs:(i+1)*mbs,:], [p[i*mbs:(i+1)*mbs,:] for p in past])

def main1():
    net = BinaryCRBM(10,16,2)
    data = loadmat("brazilRainfall.mat")["batchdata"]
    chunks = [(data[i*90+2:(i+1)*90,:], [data[i*90+1:(i+1)*90-1,:], data[i*90:(i+1)*90-2,:]]) for i in range(24)]

    data = num.vstack( [c[0] for c in chunks] )
    past = [ num.vstack( [c[1][i] for c in chunks] ) for i in range(2)]

    data = data.transpose()
    past = [p.transpose() for p in past]

    print data.shape
    print data.shape[1]/64
    for p in past:
        print p.shape

    net.learnRate = 0.002
    net.momentum = 0.9
    net.weightCost = 0
    for j,err in enumerate(net.trainXFerEnMasse(data, past, 100)):
        print j+1, err

    ex = cm.CUDAMatrix(reformat(num.array([[1,1],[2,3]])))
    print ex.euclid_norm()

def main2():
    pass

from scipy.io import loadmat

if __name__ == "__main__":
    print "export LD_LIBRARY_PATH=/u/gdahl/cudaLearn/"
    print "export CUDAMATDIR=/u/gdahl/cudaLearn"

    devId = cm.cuda_get_free_device()
    cm.cuda_set_device(devId)

    cm.cublas_init()
    cm.CUDAMatrix.init_random(1)
    main1()
    cm.cublas_shutdown()


#!/usr/bin/env python
# Copyright 2014-2020 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: Qiming Sun <osirpt.sun@gmail.com>
#

import numpy
import h5py
from pyscf import lib
from pyscf import gto
from pyscf.ao2mo.outcore import balance_segs
from pyscf.pbc.lib.kpts_helper import gamma_point, unique, KPT_DIFF_TOL
from pyscf.pbc.df.incore import wrap_int3c

libpbc = lib.load_library('libpbc')

def aux_e1(cell, auxcell, erifile, intor='int3c2e', aosym='s2ij', comp=None,
           kptij_lst=None, dataname='eri_mo', shls_slice=None, max_memory=2000,
           verbose=0):
    r'''3-center AO integrals (L|ij) with double lattice sum:
    \sum_{lm} (L[0]|i[l]j[m]), where L is the auxiliary basis.
    Three-index integral tensor (kptij_idx, naux, nao_pair) or four-index
    integral tensor (kptij_idx, comp, naux, nao_pair) are stored on disk.

    Args:
        kptij_lst : (*,2,3) array
            A list of (kpti, kptj)
    '''
    intor, comp = gto.moleintor._get_intor_and_comp(cell._add_suffix(intor), comp)

    if isinstance(erifile, h5py.Group):
        feri = erifile
    elif h5py.is_hdf5(erifile):
        feri = h5py.File(erifile, 'a')
    else:
        feri = h5py.File(erifile, 'w')
    if dataname in feri:
        del(feri[dataname])
    if dataname+'-kptij' in feri:
        del(feri[dataname+'-kptij'])

    if kptij_lst is None:
        kptij_lst = numpy.zeros((1,2,3))
    feri[dataname+'-kptij'] = kptij_lst

    if shls_slice is None:
        shls_slice = (0, cell.nbas, 0, cell.nbas, 0, auxcell.nbas)

    ao_loc = cell.ao_loc_nr()
    aux_loc = auxcell.ao_loc_nr(auxcell.cart or 'ssc' in intor)[:shls_slice[5]+1]
    ni = ao_loc[shls_slice[1]] - ao_loc[shls_slice[0]]
    nj = ao_loc[shls_slice[3]] - ao_loc[shls_slice[2]]
    naux = aux_loc[shls_slice[5]] - aux_loc[shls_slice[4]]
    nkptij = len(kptij_lst)

    nii = (ao_loc[shls_slice[1]]*(ao_loc[shls_slice[1]]+1)//2 -
           ao_loc[shls_slice[0]]*(ao_loc[shls_slice[0]]+1)//2)
    nij = ni * nj

    kpti = kptij_lst[:,0]
    kptj = kptij_lst[:,1]
    aosym_ks2 = abs(kpti-kptj).sum(axis=1) < KPT_DIFF_TOL
    j_only = numpy.all(aosym_ks2)
    #aosym_ks2 &= (aosym[:2] == 's2' and shls_slice[:2] == shls_slice[2:4])
    aosym_ks2 &= aosym[:2] == 's2'
    for k, kptij in enumerate(kptij_lst):
        key = '%s/%d' % (dataname, k)
        if gamma_point(kptij):
            dtype = 'f8'
        else:
            dtype = 'c16'
        if aosym_ks2[k]:
            nao_pair = nii
        else:
            nao_pair = nij
        if comp == 1:
            shape = (naux,nao_pair)
        else:
            shape = (comp,naux,nao_pair)
        feri.create_dataset(key, shape, dtype)
    if naux == 0:
        feri.close()
        return erifile

    if j_only and aosym[:2] == 's2':
        assert(shls_slice[2] == 0)
        nao_pair = nii
    else:
        nao_pair = nij

    if gamma_point(kptij_lst):
        dtype = numpy.double
    else:
        dtype = numpy.complex128

    buflen = max(8, int(max_memory*1e6/16/(nkptij*ni*nj*comp)))
    auxdims = aux_loc[shls_slice[4]+1:shls_slice[5]+1] - aux_loc[shls_slice[4]:shls_slice[5]]
    auxranges = balance_segs(auxdims, buflen)
    buflen = max([x[2] for x in auxranges])
    buf = numpy.empty(nkptij*comp*ni*nj*buflen, dtype=dtype)
    buf1 = numpy.empty(ni*nj*buflen, dtype=dtype)

    int3c = wrap_int3c(cell, auxcell, intor, aosym, comp, kptij_lst)

    naux0 = 0
    for istep, auxrange in enumerate(auxranges):
        sh0, sh1, nrow = auxrange
        sub_slice = (shls_slice[0], shls_slice[1],
                     shls_slice[2], shls_slice[3],
                     shls_slice[4]+sh0, shls_slice[4]+sh1)
        mat = numpy.ndarray((nkptij,comp,nao_pair,nrow), dtype=dtype, buffer=buf)
        mat = int3c(sub_slice, mat)

        for k, kptij in enumerate(kptij_lst):
            h5dat = feri['%s/%d'%(dataname,k)]
            for icomp, v in enumerate(mat[k]):
                v = lib.transpose(v, out=buf1)
                if gamma_point(kptij):
                    v = v.real
                if aosym_ks2[k] and v.shape[1] == ni**2:
                    v = lib.pack_tril(v.reshape(-1,ni,ni))
                if comp == 1:
                    h5dat[naux0:naux0+nrow] = v
                else:
                    h5dat[icomp,naux0:naux0+nrow] = v
        naux0 += nrow

    if not isinstance(erifile, h5py.Group):
        feri.close()
    return erifile

def _aux_e2(cell, auxcell, erifile, intor='int3c2e', aosym='s2ij', comp=None,
            kptij_lst=None, dataname='eri_mo', shls_slice=None, max_memory=2000,
            verbose=0):
    r'''3-center AO integrals (ij|L) with double lattice sum:
    \sum_{lm} (i[l]j[m]|L[0]), where L is the auxiliary basis.
    Three-index integral tensor (kptij_idx, nao_pair, naux) or four-index
    integral tensor (kptij_idx, comp, nao_pair, naux) are stored on disk.

    **This function should be only used by df and mdf initialization function
    _make_j3c**

    Args:
        kptij_lst : (*,2,3) array
            A list of (kpti, kptj)
    '''
    intor, comp = gto.moleintor._get_intor_and_comp(cell._add_suffix(intor), comp)

    if isinstance(erifile, h5py.Group):
        feri = erifile
    elif h5py.is_hdf5(erifile):
        feri = h5py.File(erifile, 'a')
    else:
        feri = h5py.File(erifile, 'w')
    if dataname in feri:
        del(feri[dataname])
    if dataname+'-kptij' in feri:
        del(feri[dataname+'-kptij'])

    if kptij_lst is None:
        kptij_lst = numpy.zeros((1,2,3))
    feri[dataname+'-kptij'] = kptij_lst

    if shls_slice is None:
        shls_slice = (0, cell.nbas, 0, cell.nbas, 0, auxcell.nbas)

    ao_loc = cell.ao_loc_nr()
    aux_loc = auxcell.ao_loc_nr(auxcell.cart or 'ssc' in intor)[:shls_slice[5]+1]
    ni = ao_loc[shls_slice[1]] - ao_loc[shls_slice[0]]
    nj = ao_loc[shls_slice[3]] - ao_loc[shls_slice[2]]
    nkptij = len(kptij_lst)

    nii = (ao_loc[shls_slice[1]]*(ao_loc[shls_slice[1]]+1)//2 -
           ao_loc[shls_slice[0]]*(ao_loc[shls_slice[0]]+1)//2)
    nij = ni * nj

    kpti = kptij_lst[:,0]
    kptj = kptij_lst[:,1]
    aosym_ks2 = abs(kpti-kptj).sum(axis=1) < KPT_DIFF_TOL
    j_only = numpy.all(aosym_ks2)
    #aosym_ks2 &= (aosym[:2] == 's2' and shls_slice[:2] == shls_slice[2:4])
    aosym_ks2 &= aosym[:2] == 's2'

    if j_only and aosym[:2] == 's2':
        assert(shls_slice[2] == 0)
        nao_pair = nii
    else:
        nao_pair = nij

    if gamma_point(kptij_lst):
        dtype = numpy.double
    else:
        dtype = numpy.complex128

    buflen = max(8, int(max_memory*.47e6/16/(nkptij*ni*nj*comp)))
    auxdims = aux_loc[shls_slice[4]+1:shls_slice[5]+1] - aux_loc[shls_slice[4]:shls_slice[5]]
    auxranges = balance_segs(auxdims, buflen)
    buflen = max([x[2] for x in auxranges])
    buf = numpy.empty(nkptij*comp*ni*nj*buflen, dtype=dtype)
    buf1 = numpy.empty_like(buf)

    int3c = wrap_int3c(cell, auxcell, intor, aosym, comp, kptij_lst)

    kptis = kptij_lst[:,0]
    kptjs = kptij_lst[:,1]
    kpt_ji = kptjs - kptis
    uniq_kpts, uniq_index, uniq_inverse = unique(kpt_ji)
# sorted_ij_idx: Sort and group the kptij_lst according to the ordering in
# df._make_j3c to reduce the data fragment in the hdf5 file.  When datasets
# are written to hdf5, they are saved sequentially. If the integral data are
# saved as the order of kptij_lst, removing the datasets in df._make_j3c will
# lead to holes that can not be reused.
    sorted_ij_idx = numpy.hstack([numpy.where(uniq_inverse == k)[0]
                                  for k, kpt in enumerate(uniq_kpts)])
    tril_idx = numpy.tril_indices(ni)
    tril_idx = tril_idx[0] * ni + tril_idx[1]
    def save(istep, mat):
        for k in sorted_ij_idx:
            v = mat[k]
            if gamma_point(kptij_lst[k]):
                v = v.real
            if aosym_ks2[k] and nao_pair == ni**2:
                v = v[:,tril_idx]
            feri['%s/%d/%d' % (dataname,k,istep)] = v

    with lib.call_in_background(save) as bsave:
        for istep, auxrange in enumerate(auxranges):
            sh0, sh1, nrow = auxrange
            sub_slice = (shls_slice[0], shls_slice[1],
                         shls_slice[2], shls_slice[3],
                         shls_slice[4]+sh0, shls_slice[4]+sh1)
            mat = numpy.ndarray((nkptij,comp,nao_pair,nrow), dtype=dtype, buffer=buf)
            bsave(istep, int3c(sub_slice, mat))
            buf, buf1 = buf1, buf

    if not isinstance(erifile, h5py.Group):
        feri.close()
    return erifile