"""Lattice thermal conductivity calculation with RTA."""
# Copyright (C) 2020 Atsushi Togo
# All rights reserved.
#
# This file is part of phono3py.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
# * Redistributions of source code must retain the above copyright
#   notice, this list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright
#   notice, this list of conditions and the following disclaimer in
#   the documentation and/or other materials provided with the
#   distribution.
#
# * Neither the name of the phonopy project nor the names of its
#   contributors may be used to endorse or promote products derived
#   from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
# COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

import sys

import numpy as np
from phonopy.phonon.group_velocity import GroupVelocity
from phonopy.structure.tetrahedron_method import TetrahedronMethod

from phono3py.file_IO import (
    read_gamma_from_hdf5,
    read_pp_from_hdf5,
    write_gamma_detail_to_hdf5,
    write_kappa_to_hdf5,
)
from phono3py.phonon3.conductivity import Conductivity, all_bands_exist, unit_to_WmK
from phono3py.phonon3.conductivity import write_pp as _write_pp
from phono3py.phonon3.imag_self_energy import ImagSelfEnergy, average_by_degeneracy
from phono3py.phonon3.interaction import Interaction
from phono3py.phonon3.triplets import get_all_triplets
from phono3py.phonon.grid import get_grid_points_by_rotations

class Conductivity_RTA(Conductivity):
    """Lattice thermal conductivity calculation with RTA."""

    def __init__(
        self,
        interaction: Interaction,
        grid_points=None,
        temperatures=None,
        sigmas=None,
        sigma_cutoff=None,
        is_isotope=False,
        mass_variances=None,
        boundary_mfp=None,  # in micrometre
        use_ave_pp=False,
        is_kappa_star=True,
        gv_delta_q=None,
        is_full_pp=False,
        read_pp=False,
        store_pp=False,
        pp_filename=None,
        is_N_U=False,
        is_gamma_detail=False,
        is_frequency_shift_by_bubble=False,
        log_level=0,
    ):
        """Init method."""
        self._pp: Interaction
        self._grid_point_count: int
        self._num_sampling_grid_points: int
        self._gv_obj: GroupVelocity
        self._temperatures = None
        self._sigmas = None
        self._sigma_cutoff = None
        self._is_kappa_star = None
        self._is_full_pp = None
        self._is_N_U = is_N_U
        self._is_gamma_detail = is_gamma_detail
        self._is_frequency_shift_by_bubble = is_frequency_shift_by_bubble
        self._log_level = None
        self._boundary_mfp = None

        self._point_operations = None
        self._rotations_cartesian = None

        self._grid_points = None
        self._grid_weights = None

        self._read_gamma = False
        self._read_gamma_iso = False

        self._frequencies = None
        self._cv = None
        self._gv = None
        self._gv_sum2 = None
        self._gamma = None
        self._gamma_iso = None
        self._gamma_N = None
        self._gamma_U = None
        self._gamma_detail_at_q = None
        self._use_ave_pp = use_ave_pp
        self._use_const_ave_pp = None
        self._averaged_pp_interaction = None
        self._num_ignored_phonon_modes = None

        self._conversion_factor = None

        self._is_isotope = None
        self._isotope = None
        self._mass_variances = None

        super().__init__(
            interaction,
            grid_points=grid_points,
            temperatures=temperatures,
            sigmas=sigmas,
            sigma_cutoff=sigma_cutoff,
            is_isotope=is_isotope,
            mass_variances=mass_variances,
            boundary_mfp=boundary_mfp,
            is_kappa_star=is_kappa_star,
            gv_delta_q=gv_delta_q,
            is_full_pp=is_full_pp,
            log_level=log_level,
        )

        self._use_const_ave_pp = self._pp.get_constant_averaged_interaction()
        self._read_pp = read_pp
        self._store_pp = store_pp
        self._pp_filename = pp_filename

        if self._temperatures is not None:
            self._allocate_values()

    def set_kappa_at_sigmas(self):
        """Calculate kappa from ph-ph interaction results."""
        num_band = len(self._pp.primitive) * 3
        for i, grid_point in enumerate(self._grid_points):
            cv = self._cv[:, i, :]
            gp = self._grid_points[i]
            frequencies = self._frequencies[gp]

            # Kappa
            for j in range(len(self._sigmas)):
                for k in range(len(self._temperatures)):
                    g_sum = self._get_main_diagonal(i, j, k)
                    for ll in range(num_band):
                        if frequencies[ll] < self._pp.cutoff_frequency:
                            self._num_ignored_phonon_modes[j, k] += 1
                            continue

                        old_settings = np.seterr(all="raise")
                        try:
                            self._mode_kappa[j, k, i, ll] = (
                                self._gv_sum2[i, ll]
                                * cv[k, ll]
                                / (g_sum[ll] * 2)
                                * self._conversion_factor
                            )
                        except FloatingPointError:
                            # supposed that g is almost 0 and |gv|=0
                            pass
                        except Exception:
                            print("=" * 26 + " Warning " + "=" * 26)
                            print(
                                " Unexpected physical condition of ph-ph "
                                "interaction calculation was found."
                            )
                            print(
                                " g=%f at gp=%d, band=%d, freq=%f"
                                % (g_sum[ll], gp, ll + 1, frequencies[ll])
                            )
                            print("=" * 61)
                        np.seterr(**old_settings)

        N = self._num_sampling_grid_points
        self._kappa = self._mode_kappa.sum(axis=2).sum(axis=2) / N

    def get_gamma_N_U(self):
        """Return N and U parts of gamma."""
        return (self._gamma_N, self._gamma_U)

    def set_gamma_N_U(self, gamma_N, gamma_U):
        """Set N and U parts of gamma."""
        self._gamma_N = gamma_N
        self._gamma_U = gamma_U

    def get_gamma_detail_at_q(self):
        """Return contribution of each triplet to gamma at current q-point."""
        return self._gamma_detail_at_q

    def get_number_of_ignored_phonon_modes(self):
        """Return number of ignored phonon modes."""
        return self._num_ignored_phonon_modes

    def set_averaged_pp_interaction(self, ave_pp):
        """Set averaged ph-ph interaction."""
        self._averaged_pp_interaction = ave_pp

    def _run_at_grid_point(self):
        i_gp = self._grid_point_count
        self._show_log_header(i_gp)
        grid_point = self._grid_points[i_gp]
        self._set_cv(i_gp, i_gp)
        self._set_gv(i_gp, i_gp)
        self._set_gv_by_gv(i_gp, i_gp)

        if self._read_gamma:
            if self._use_ave_pp:
                self._collision.set_grid_point(grid_point)
                self._set_gamma_at_sigmas(i_gp)
        else:
            self._collision.set_grid_point(grid_point)
            num_triplets = len(self._pp.get_triplets_at_q()[0])
            if self._log_level:
                print("Number of triplets: %d" % num_triplets)

            if (
                self._is_full_pp
                or self._read_pp
                or self._store_pp
                or self._use_ave_pp
                or self._use_const_ave_pp
                or self._is_gamma_detail
            ):
                self._set_gamma_at_sigmas(i_gp)
            else:  # can save memory space
                self._set_gamma_at_sigmas_lowmem(i_gp)

        if self._isotope is not None and not self._read_gamma_iso:
            gamma_iso = self._get_gamma_isotope_at_sigmas(i_gp)
            self._gamma_iso[:, i_gp, :] = gamma_iso[:, self._pp.band_indices]

        if self._log_level:
            self._show_log(self._qpoints[i_gp], i_gp)

    def _allocate_values(self):
        num_band0 = len(self._pp.band_indices)
        num_grid_points = len(self._grid_points)
        num_temp = len(self._temperatures)
        self._kappa = np.zeros(
            (len(self._sigmas), num_temp, 6), order="C", dtype="double"
        )
        self._mode_kappa = np.zeros(
            (len(self._sigmas), num_temp, num_grid_points, num_band0, 6),
            order="C",
            dtype="double",
        )
        if not self._read_gamma:
            self._gamma = np.zeros(
                (len(self._sigmas), num_temp, num_grid_points, num_band0),
                order="C",
                dtype="double",
            )
            if self._is_gamma_detail or self._is_N_U:
                self._gamma_N = np.zeros_like(self._gamma)
                self._gamma_U = np.zeros_like(self._gamma)
        self._gv = np.zeros((num_grid_points, num_band0, 3), order="C", dtype="double")
        self._gv_sum2 = np.zeros(
            (num_grid_points, num_band0, 6), order="C", dtype="double"
        )
        self._cv = np.zeros(
            (num_temp, num_grid_points, num_band0), order="C", dtype="double"
        )
        if self._isotope is not None:
            self._gamma_iso = np.zeros(
                (len(self._sigmas), num_grid_points, num_band0),
                order="C",
                dtype="double",
            )
        if self._is_full_pp or self._use_ave_pp or self._use_const_ave_pp:
            self._averaged_pp_interaction = np.zeros(
                (num_grid_points, num_band0), order="C", dtype="double"
            )
        self._num_ignored_phonon_modes = np.zeros(
            (len(self._sigmas), num_temp), order="C", dtype="intc"
        )
        self._collision = ImagSelfEnergy(
            self._pp, with_detail=(self._is_gamma_detail or self._is_N_U)
        )

    def _set_gamma_at_sigmas(self, i):
        for j, sigma in enumerate(self._sigmas):
            self._collision.set_sigma(sigma, sigma_cutoff=self._sigma_cutoff)
            self._collision.set_integration_weights()

            if self._log_level:
                text = "Collisions will be calculated with "
                if sigma is None:
                    text += "tetrahedron method."
                else:
                    text += "sigma=%s" % sigma
                    if self._sigma_cutoff is None:
                        text += "."
                    else:
                        text += "(%4.2f SD)." % self._sigma_cutoff
                print(text)

            if self._read_pp:
                pp, _g_zero = read_pp_from_hdf5(
                    self._pp.mesh_numbers,
                    grid_point=self._grid_points[i],
                    sigma=sigma,
                    sigma_cutoff=self._sigma_cutoff,
                    filename=self._pp_filename,
                    verbose=(self._log_level > 0),
                )
                _, g_zero = self._collision.get_integration_weights()
                if self._log_level:
                    if len(self._sigmas) > 1:
                        print(
                            "Multiple sigmas or mixing smearing and "
                            "tetrahedron method is not supported."
                        )
                if _g_zero is not None and (_g_zero != g_zero).any():
                    raise ValueError("Inconsistency found in g_zero.")
                self._collision.set_interaction_strength(pp)
            elif self._use_ave_pp:
                self._collision.set_averaged_pp_interaction(
                    self._averaged_pp_interaction[i]
                )
            elif self._use_const_ave_pp:
                if self._log_level:
                    print(
                        "Constant ph-ph interaction of %6.3e is used."
                        % self._pp.get_constant_averaged_interaction()
                    )
                self._collision.run_interaction()
                self._averaged_pp_interaction[i] = self._pp.get_averaged_interaction()
            elif j != 0 and (self._is_full_pp or self._sigma_cutoff is None):
                if self._log_level:
                    print("Existing ph-ph interaction is used.")
            else:
                if self._log_level:
                    print("Calculating ph-ph interaction...")
                self._collision.run_interaction(is_full_pp=self._is_full_pp)
                if self._is_full_pp:
                    self._averaged_pp_interaction[
                        i
                    ] = self._pp.get_averaged_interaction()

            # Number of triplets depends on q-point.
            # So this is allocated each time.
            if self._is_gamma_detail:
                num_temp = len(self._temperatures)
                self._gamma_detail_at_q = np.empty(
                    ((num_temp,) + self._pp.get_interaction_strength().shape),
                    dtype="double",
                    order="C",
                )
                self._gamma_detail_at_q[:] = 0

            if self._log_level:
                print("Calculating collisions at temperatures...")
            for k, t in enumerate(self._temperatures):
                self._collision.set_temperature(t)
                self._collision.run()
                self._gamma[j, k, i] = self._collision.get_imag_self_energy()
                if self._is_N_U or self._is_gamma_detail:
                    g_N, g_U = self._collision.get_imag_self_energy_N_and_U()
                    self._gamma_N[j, k, i] = g_N
                    self._gamma_U[j, k, i] = g_U
                if self._is_gamma_detail:
                    self._gamma_detail_at_q[
                        k
                    ] = self._collision.get_detailed_imag_self_energy()

    def _set_gamma_at_sigmas_lowmem(self, i):
        """Calculate gamma without storing ph-ph interaction strength.

        `svecs` and `multi` below must not be simply replaced by
        `self._pp.primitive.get_smallest_vectors()` because they must be in
        dense format as always so in Interaction class instance.
        `p2s`, `s2p`, and `masses` have to be also given from Interaction
        class instance.

        """
        band_indices = self._pp.band_indices
        (
            svecs,
            multi,
            p2s,
            s2p,
            masses,
        ) = self._pp.get_primitive_and_supercell_correspondence()
        fc3 = self._pp.fc3
        triplets_at_q, weights_at_q, _, _ = self._pp.get_triplets_at_q()
        symmetrize_fc3_q = 0

        if None in self._sigmas:
            thm = TetrahedronMethod(self._pp.bz_grid.microzone_lattice)

        # It is assumed that self._sigmas = [None].
        for j, sigma in enumerate(self._sigmas):
            self._collision.set_sigma(sigma)
            if self._is_N_U:
                collisions = np.zeros(
                    (2, len(self._temperatures), len(band_indices)),
                    dtype="double",
                    order="C",
                )
            else:
                collisions = np.zeros(
                    (len(self._temperatures), len(band_indices)),
                    dtype="double",
                    order="C",
                )
            import phono3py._phono3py as phono3c

            if sigma is None:
                phono3c.pp_collision(
                    collisions,
                    np.array(
                        np.dot(thm.get_tetrahedra(), self._pp.bz_grid.P.T),
                        dtype="int_",
                        order="C",
                    ),
                    self._frequencies,
                    self._eigenvectors,
                    triplets_at_q,
                    weights_at_q,
                    self._pp.bz_grid.addresses,
                    self._pp.bz_grid.gp_map,
                    self._pp.bz_grid.store_dense_gp_map * 1 + 1,
                    self._pp.bz_grid.D_diag,
                    self._pp.bz_grid.Q,
                    fc3,
                    svecs,
                    multi,
                    masses,
                    p2s,
                    s2p,
                    band_indices,
                    self._temperatures,
                    self._is_N_U * 1,
                    symmetrize_fc3_q,
                    self._pp.cutoff_frequency,
                )
            else:
                if self._sigma_cutoff is None:
                    sigma_cutoff = -1
                else:
                    sigma_cutoff = float(self._sigma_cutoff)
                phono3c.pp_collision_with_sigma(
                    collisions,
                    sigma,
                    sigma_cutoff,
                    self._frequencies,
                    self._eigenvectors,
                    triplets_at_q,
                    weights_at_q,
                    self._pp.bz_grid.addresses,
                    self._pp.bz_grid.D_diag,
                    self._pp.bz_grid.Q,
                    fc3,
                    svecs,
                    multi,
                    masses,
                    p2s,
                    s2p,
                    band_indices,
                    self._temperatures,
                    self._is_N_U * 1,
                    symmetrize_fc3_q,
                    self._pp.cutoff_frequency,
                )
            col_unit_conv = self._collision.get_unit_conversion_factor()
            pp_unit_conv = self._pp.get_unit_conversion_factor()
            if self._is_N_U:
                col = collisions.sum(axis=0)
                col_N = collisions[0]
                col_U = collisions[1]
            else:
                col = collisions
            for k in range(len(self._temperatures)):
                self._gamma[j, k, i, :] = average_by_degeneracy(
                    col[k] * col_unit_conv * pp_unit_conv,
                    band_indices,
                    self._frequencies[self._grid_points[i]],
                )
                if self._is_N_U:
                    self._gamma_N[j, k, i, :] = average_by_degeneracy(
                        col_N[k] * col_unit_conv * pp_unit_conv,
                        band_indices,
                        self._frequencies[self._grid_points[i]],
                    )
                    self._gamma_U[j, k, i, :] = average_by_degeneracy(
                        col_U[k] * col_unit_conv * pp_unit_conv,
                        band_indices,
                        self._frequencies[self._grid_points[i]],
                    )

    def _show_log(self, q, i):
        gp = self._grid_points[i]
        frequencies = self._frequencies[gp][self._pp.band_indices]
        gv = self._gv[i]
        if self._averaged_pp_interaction is not None:
            ave_pp = self._averaged_pp_interaction[i]
        else:
            ave_pp = None
        self._show_log_value_names()

        if self._log_level > 2:
            self._show_log_values_on_kstar(frequencies, gv, ave_pp, gp, q)
        else:
            self._show_log_values(frequencies, gv, ave_pp)

        sys.stdout.flush()

    def _show_log_values(self, frequencies, gv, ave_pp):
        if self._is_full_pp or self._use_ave_pp or self._use_const_ave_pp:
            for f, v, pp in zip(frequencies, gv, ave_pp):
                print(
                    "%8.3f   (%8.3f %8.3f %8.3f) %8.3f %11.3e"
                    % (f, v[0], v[1], v[2], np.linalg.norm(v), pp)
                )
        else:
            for f, v in zip(frequencies, gv):
                print(
                    "%8.3f   (%8.3f %8.3f %8.3f) %8.3f"
                    % (f, v[0], v[1], v[2], np.linalg.norm(v))
                )

    def _show_log_values_on_kstar(self, frequencies, gv, ave_pp, gp, q):
        rotation_map = get_grid_points_by_rotations(gp, self._pp.bz_grid)
        for i, j in enumerate(np.unique(rotation_map)):
            for k, (rot, rot_c) in enumerate(
                zip(self._point_operations, self._rotations_cartesian)
            ):
                if rotation_map[k] != j:
                    continue

                print(
                    " k*%-2d (%5.2f %5.2f %5.2f)" % ((i + 1,) + tuple(np.dot(rot, q)))
                )
                if self._is_full_pp or self._use_ave_pp or self._use_const_ave_pp:
                    for f, v, pp in zip(frequencies, np.dot(rot_c, gv.T).T, ave_pp):
                        print(
                            "%8.3f   (%8.3f %8.3f %8.3f) %8.3f %11.3e"
                            % (f, v[0], v[1], v[2], np.linalg.norm(v), pp)
                        )
                else:
                    for f, v in zip(frequencies, np.dot(rot_c, gv.T).T):
                        print(
                            "%8.3f   (%8.3f %8.3f %8.3f) %8.3f"
                            % (f, v[0], v[1], v[2], np.linalg.norm(v))
                        )
        print("")

    def _show_log_value_names(self):
        if self._is_full_pp or self._use_ave_pp or self._use_const_ave_pp:
            text = "Frequency     group velocity (x, y, z)     |gv|       Pqj"
        else:
            text = "Frequency     group velocity (x, y, z)     |gv|"
        if self._gv_obj.q_length is None:
            pass
        else:
            text += "  (dq=%3.1e)" % self._gv_obj.q_length
        print(text)

def get_thermal_conductivity_RTA(
    interaction: Interaction,
    temperatures=None,
    sigmas=None,
    sigma_cutoff=None,
    mass_variances=None,
    grid_points=None,
    is_isotope=False,
    boundary_mfp=None,  # in micrometre
    use_ave_pp=False,
    is_kappa_star=True,
    gv_delta_q=None,
    is_full_pp=False,
    write_gamma=False,
    read_gamma=False,
    is_N_U=False,
    write_kappa=False,
    write_pp=False,
    read_pp=False,
    write_gamma_detail=False,
    compression="gzip",
    input_filename=None,
    output_filename=None,
    log_level=0,
):
    """Run RTA thermal conductivity calculation."""
    if temperatures is None:
        _temperatures = np.arange(0, 1001, 10, dtype="double")
    else:
        _temperatures = temperatures

    if log_level:
        print(
            "-------------------- Lattice thermal conducitivity (RTA) "
            "--------------------"
        )
    br = Conductivity_RTA(
        interaction,
        grid_points=grid_points,
        temperatures=_temperatures,
        sigmas=sigmas,
        sigma_cutoff=sigma_cutoff,
        is_isotope=is_isotope,
        mass_variances=mass_variances,
        boundary_mfp=boundary_mfp,
        use_ave_pp=use_ave_pp,
        is_kappa_star=is_kappa_star,
        gv_delta_q=gv_delta_q,
        is_full_pp=is_full_pp,
        read_pp=read_pp,
        store_pp=write_pp,
        pp_filename=input_filename,
        is_N_U=is_N_U,
        is_gamma_detail=write_gamma_detail,
        log_level=log_level,
    )

    if read_gamma:
        if not _set_gamma_from_file(br, filename=input_filename):
            print("Reading collisions failed.")
            return False

    for i in br:
        if write_pp:
            _write_pp(
                br, interaction, i, compression=compression, filename=output_filename
            )
        if write_gamma:
            _write_gamma(
                br,
                interaction,
                i,
                compression=compression,
                filename=output_filename,
                verbose=log_level,
            )
        if write_gamma_detail:
            _write_gamma_detail(
                br,
                interaction,
                i,
                compression=compression,
                filename=output_filename,
                verbose=log_level,
            )

    if grid_points is None and all_bands_exist(interaction):
        br.set_kappa_at_sigmas()
        if log_level:
            _show_kappa(br, log_level)
        if write_kappa:
            _write_kappa(
                br,
                interaction.primitive.volume,
                compression=compression,
                filename=output_filename,
                log_level=log_level,
            )

    return br

def _write_gamma_detail(
    br, interaction, i, compression="gzip", filename=None, verbose=True
):
    gamma_detail = br.get_gamma_detail_at_q()
    temperatures = br.get_temperatures()
    mesh = br.get_mesh_numbers()
    grid_points = br.get_grid_points()
    gp = grid_points[i]
    sigmas = br.get_sigmas()
    sigma_cutoff = br.get_sigma_cutoff_width()
    triplets, weights, _, _ = interaction.get_triplets_at_q()
    all_triplets = get_all_triplets(gp, interaction.bz_grid)

    if all_bands_exist(interaction):
        for j, sigma in enumerate(sigmas):
            write_gamma_detail_to_hdf5(
                temperatures,
                mesh,
                gamma_detail=gamma_detail,
                grid_point=gp,
                triplet=triplets,
                weight=weights,
                triplet_all=all_triplets,
                sigma=sigma,
                sigma_cutoff=sigma_cutoff,
                compression=compression,
                filename=filename,
                verbose=verbose,
            )
    else:
        for j, sigma in enumerate(sigmas):
            for k, bi in enumerate(interaction.get_band_indices()):
                write_gamma_detail_to_hdf5(
                    temperatures,
                    mesh,
                    gamma_detail=gamma_detail[:, :, k, :, :],
                    grid_point=gp,
                    triplet=triplets,
                    weight=weights,
                    band_index=bi,
                    sigma=sigma,
                    sigma_cutoff=sigma_cutoff,
                    compression=compression,
                    filename=filename,
                    verbose=verbose,
                )

def _write_gamma(br, interaction, i, compression="gzip", filename=None, verbose=True):
    """Write mode kappa related properties into a hdf5 file."""
    grid_points = br.grid_points
    group_velocities = br.group_velocities
    gv_by_gv = br.gv_by_gv
    mode_heat_capacities = br.mode_heat_capacities
    ave_pp = br.averaged_pp_interaction
    mesh = br.mesh_numbers
    temperatures = br.temperatures
    gamma = br.gamma
    gamma_isotope = br.gamma_isotope
    sigmas = br.sigmas
    sigma_cutoff = br.sigma_cutoff_width
    volume = interaction.primitive.volume
    gamma_N, gamma_U = br.get_gamma_N_U()

    gp = grid_points[i]
    if all_bands_exist(interaction):
        if ave_pp is None:
            ave_pp_i = None
        else:
            ave_pp_i = ave_pp[i]
        frequencies = interaction.get_phonons()[0][gp]
        for j, sigma in enumerate(sigmas):
            if gamma_isotope is not None:
                gamma_isotope_at_sigma = gamma_isotope[j, i]
            else:
                gamma_isotope_at_sigma = None
            if gamma_N is None:
                gamma_N_at_sigma = None
            else:
                gamma_N_at_sigma = gamma_N[j, :, i]
            if gamma_U is None:
                gamma_U_at_sigma = None
            else:
                gamma_U_at_sigma = gamma_U[j, :, i]

            write_kappa_to_hdf5(
                temperatures,
                mesh,
                frequency=frequencies,
                group_velocity=group_velocities[i],
                gv_by_gv=gv_by_gv[i],
                heat_capacity=mode_heat_capacities[:, i],
                gamma=gamma[j, :, i],
                gamma_isotope=gamma_isotope_at_sigma,
                gamma_N=gamma_N_at_sigma,
                gamma_U=gamma_U_at_sigma,
                averaged_pp_interaction=ave_pp_i,
                grid_point=gp,
                sigma=sigma,
                sigma_cutoff=sigma_cutoff,
                kappa_unit_conversion=unit_to_WmK / volume,
                compression=compression,
                filename=filename,
                verbose=verbose,
            )
    else:
        for j, sigma in enumerate(sigmas):
            for k, bi in enumerate(interaction.band_indices):
                if ave_pp is None:
                    ave_pp_ik = None
                else:
                    ave_pp_ik = ave_pp[i, k]
                frequencies = interaction.get_phonons()[0][gp, bi]
                if gamma_isotope is not None:
                    gamma_isotope_at_sigma = gamma_isotope[j, i, k]
                else:
                    gamma_isotope_at_sigma = None
                if gamma_N is None:
                    gamma_N_at_sigma = None
                else:
                    gamma_N_at_sigma = gamma_N[j, :, i, k]
                if gamma_U is None:
                    gamma_U_at_sigma = None
                else:
                    gamma_U_at_sigma = gamma_U[j, :, i, k]
                write_kappa_to_hdf5(
                    temperatures,
                    mesh,
                    frequency=frequencies,
                    group_velocity=group_velocities[i, k],
                    gv_by_gv=gv_by_gv[i, k],
                    heat_capacity=mode_heat_capacities[:, i, k],
                    gamma=gamma[j, :, i, k],
                    gamma_isotope=gamma_isotope_at_sigma,
                    gamma_N=gamma_N_at_sigma,
                    gamma_U=gamma_U_at_sigma,
                    averaged_pp_interaction=ave_pp_ik,
                    grid_point=gp,
                    band_index=bi,
                    sigma=sigma,
                    sigma_cutoff=sigma_cutoff,
                    kappa_unit_conversion=unit_to_WmK / volume,
                    compression=compression,
                    filename=filename,
                    verbose=verbose,
                )

def _show_kappa(br, log_level):
    temperatures = br.temperatures
    sigmas = br.sigmas
    kappa = br.kappa
    num_ignored_phonon_modes = br.get_number_of_ignored_phonon_modes()
    num_band = br.frequencies.shape[1]
    num_phonon_modes = br.get_number_of_sampling_grid_points() * num_band
    for i, sigma in enumerate(sigmas):
        text = "----------- Thermal conductivity (W/m-k) "
        if sigma:
            text += "for sigma=%s -----------" % sigma
        else:
            text += "with tetrahedron method -----------"
        print(text)
        if log_level > 1:
            print(
                ("#%6s       " + " %-10s" * 6 + "#ipm")
                % ("T(K)", "xx", "yy", "zz", "yz", "xz", "xy")
            )
            for j, (t, k) in enumerate(zip(temperatures, kappa[i])):
                print(
                    ("%7.1f" + " %10.3f" * 6 + " %d/%d")
                    % (
                        (t,)
                        + tuple(k)
                        + (num_ignored_phonon_modes[i, j], num_phonon_modes)
                    )
                )
        else:
            print(
                ("#%6s       " + " %-10s" * 6)
                % ("T(K)", "xx", "yy", "zz", "yz", "xz", "xy")
            )
            for j, (t, k) in enumerate(zip(temperatures, kappa[i])):
                print(("%7.1f " + " %10.3f" * 6) % ((t,) + tuple(k)))
        print("")

def _write_kappa(br, volume, compression="gzip", filename=None, log_level=0):
    temperatures = br.temperatures
    sigmas = br.sigmas
    sigma_cutoff = br.sigma_cutoff_width
    gamma = br.gamma
    gamma_isotope = br.gamma_isotope
    gamma_N, gamma_U = br.get_gamma_N_U()
    mesh = br.mesh_numbers
    frequencies = br.frequencies
    gv = br.group_velocities
    gv_by_gv = br.gv_by_gv
    mode_cv = br.mode_heat_capacities
    ave_pp = br.averaged_pp_interaction
    qpoints = br.qpoints
    weights = br.grid_weights
    kappa = br.kappa
    mode_kappa = br.mode_kappa

    for i, sigma in enumerate(sigmas):
        kappa_at_sigma = kappa[i]
        if gamma_isotope is not None:
            gamma_isotope_at_sigma = gamma_isotope[i]
        else:
            gamma_isotope_at_sigma = None
        if gamma_N is None:
            gamma_N_at_sigma = None
        else:
            gamma_N_at_sigma = gamma_N[i]
        if gamma_U is None:
            gamma_U_at_sigma = None
        else:
            gamma_U_at_sigma = gamma_U[i]

        write_kappa_to_hdf5(
            temperatures,
            mesh,
            frequency=frequencies,
            group_velocity=gv,
            gv_by_gv=gv_by_gv,
            heat_capacity=mode_cv,
            kappa=kappa_at_sigma,
            mode_kappa=mode_kappa[i],
            gamma=gamma[i],
            gamma_isotope=gamma_isotope_at_sigma,
            gamma_N=gamma_N_at_sigma,
            gamma_U=gamma_U_at_sigma,
            averaged_pp_interaction=ave_pp,
            qpoint=qpoints,
            weight=weights,
            sigma=sigma,
            sigma_cutoff=sigma_cutoff,
            kappa_unit_conversion=unit_to_WmK / volume,
            compression=compression,
            filename=filename,
            verbose=log_level,
        )

def _set_gamma_from_file(br, filename=None, verbose=True):
    """Read kappa-*.hdf5 files for thermal conductivity calculation."""
    sigmas = br.get_sigmas()
    sigma_cutoff = br.get_sigma_cutoff_width()
    mesh = br.get_mesh_numbers()
    grid_points = br.get_grid_points()
    temperatures = br.get_temperatures()
    num_band = br.get_frequencies().shape[1]

    gamma = np.zeros(
        (len(sigmas), len(temperatures), len(grid_points), num_band), dtype="double"
    )
    gamma_N = np.zeros_like(gamma)
    gamma_U = np.zeros_like(gamma)
    gamma_iso = np.zeros((len(sigmas), len(grid_points), num_band), dtype="double")
    ave_pp = np.zeros((len(grid_points), num_band), dtype="double")

    is_gamma_N_U_in = False
    is_ave_pp_in = False
    read_succeeded = True

    for j, sigma in enumerate(sigmas):
        data = read_gamma_from_hdf5(
            mesh,
            sigma=sigma,
            sigma_cutoff=sigma_cutoff,
            filename=filename,
            verbose=verbose,
        )
        if data:
            gamma[j] = data["gamma"]
            if "gamma_isotope" in data:
                gamma_iso[j] = data["gamma_isotope"]
            if "gamma_N" in data:
                is_gamma_N_U_in = True
                gamma_N[j] = data["gamma_N"]
                gamma_U[j] = data["gamma_U"]
            if "ave_pp" in data:
                is_ave_pp_in = True
                ave_pp[:] = data["ave_pp"]
        else:
            for i, gp in enumerate(grid_points):
                data_gp = read_gamma_from_hdf5(
                    mesh,
                    grid_point=gp,
                    sigma=sigma,
                    sigma_cutoff=sigma_cutoff,
                    filename=filename,
                    verbose=verbose,
                )
                if data_gp:
                    gamma[j, :, i] = data_gp["gamma"]
                    if "gamma_iso" in data_gp:
                        gamma_iso[j, i] = data_gp["gamma_iso"]
                    if "gamma_N" in data_gp:
                        is_gamma_N_U_in = True
                        gamma_N[j, :, i] = data_gp["gamma_N"]
                        gamma_U[j, :, i] = data_gp["gamma_U"]
                    if "ave_pp" in data_gp:
                        is_ave_pp_in = True
                        ave_pp[i] = data_gp["ave_pp"]
                else:
                    for bi in range(num_band):
                        data_band = read_gamma_from_hdf5(
                            mesh,
                            grid_point=gp,
                            band_index=bi,
                            sigma=sigma,
                            sigma_cutoff=sigma_cutoff,
                            filename=filename,
                            verbose=verbose,
                        )
                        if data_band:
                            gamma[j, :, i, bi] = data_band["gamma"]
                            if "gamma_iso" in data_band:
                                gamma_iso[j, i, bi] = data_band["gamma_iso"]
                            if "gamma_N" in data_band:
                                is_gamma_N_U_in = True
                                gamma_N[j, :, i, bi] = data_band["gamma_N"]
                                gamma_U[j, :, i, bi] = data_band["gamma_U"]
                            if "ave_pp" in data_band:
                                is_ave_pp_in = True
                                ave_pp[i, bi] = data_band["ave_pp"]
                        else:
                            read_succeeded = False

    if read_succeeded:
        br.set_gamma(gamma)
        if is_ave_pp_in:
            br.set_averaged_pp_interaction(ave_pp)
        if is_gamma_N_U_in:
            br.set_gamma_N_U(gamma_N, gamma_U)
        return True
    else:
        return False


#!/usr/bin/env python
# Copyright 2014-2018 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>
#

'''
Restricted open-shell Hartree-Fock for periodic systems at a single k-point

See Also:
    pyscf/pbc/scf/khf.py : Hartree-Fock for periodic systems with k-point sampling
'''

import numpy as np
from pyscf import lib
from pyscf.lib import logger
from pyscf.scf import rohf as mol_rohf
from pyscf.pbc.scf import hf as pbchf
from pyscf.pbc.scf import uhf as pbcuhf
from pyscf import __config__

get_fock = mol_rohf.get_fock
get_occ = mol_rohf.get_occ
get_grad = mol_rohf.get_grad
make_rdm1 = mol_rohf.make_rdm1
energy_elec = mol_rohf.energy_elec

class ROHF(mol_rohf.ROHF, pbchf.RHF):
    '''ROHF class for PBCs.
    '''

    direct_scf = getattr(__config__, 'pbc_scf_SCF_direct_scf', False)

    def __init__(self, cell, kpt=np.zeros(3),
                 exxdiv=getattr(__config__, 'pbc_scf_SCF_exxdiv', 'ewald')):
        pbchf.SCF.__init__(self, cell, kpt, exxdiv)
        self.nelec = None
        self._keys = self._keys.union(['nelec'])

    def dump_flags(self):
        pbchf.SCF.dump_flags(self)
        if self.nelec is None:
            nelec = self.cell.nelec
        else:
            nelec = self.nelec
        logger.info(self, 'number of electrons per unit cell  '
                    'alpha = %d beta = %d', *nelec)
        return self

    build = pbchf.SCF.build
    check_sanity = pbchf.SCF.check_sanity
    get_hcore = pbchf.SCF.get_hcore
    get_ovlp = pbchf.SCF.get_ovlp
    get_jk = pbchf.SCF.get_jk
    get_j = pbchf.SCF.get_j
    get_k = pbchf.SCF.get_k
    get_jk_incore = pbchf.SCF.get_jk_incore
    energy_tot = pbchf.SCF.energy_tot

    def get_veff(self, cell=None, dm=None, dm_last=0, vhf_last=0, hermi=1,
                 kpt=None, kpts_band=None):
        if cell is None: cell = self.cell
        if dm is None: dm = self.make_rdm1()
        if kpt is None: kpt = self.kpt
        if isinstance(dm, np.ndarray) and dm.ndim == 2:
            dm = np.asarray((dm*.5,dm*.5))
        if hasattr(dm, 'mo_coeff'):
            mo_coeff = dm.mo_coeff
            mo_occ_a = (dm.mo_occ > 0).astype(np.double)
            mo_occ_b = (dm.mo_occ ==2).astype(np.double)
            dm = lib.tag_array(dm, mo_coeff=(mo_coeff,mo_coeff),
                               mo_occ=(mo_occ_a,mo_occ_b))
        vj, vk = self.get_jk(cell, dm, hermi, kpt, kpts_band)
        vhf = vj[0] + vj[1] - vk
        return vhf

    def get_bands(self, kpts_band, cell=None, dm=None, kpt=None):
        '''Get energy bands at the given (arbitrary) 'band' k-points.

        Returns:
            mo_energy : (nmo,) ndarray or a list of (nmo,) ndarray
                Bands energies E_n(k)
            mo_coeff : (nao, nmo) ndarray or a list of (nao,nmo) ndarray
                Band orbitals psi_n(k)
        '''
        raise NotImplementedError

    def dip_moment(self, mol=None, dm=None, unit='Debye', verbose=logger.NOTE,
                   **kwargs):
        # skip dipole memont for crystal
        return

    def get_init_guess(self, cell=None, key='minao'):
        if cell is None: cell = self.cell
        dm = mol_rohf.ROHF.get_init_guess(self, cell, key)
        if cell.dimension < 3:
            if isinstance(dm, np.ndarray) and dm.ndim == 2:
                ne = np.einsum('ij,ji->', dm, self.get_ovlp(cell))
            else:
                ne = np.einsum('xij,ji->', dm, self.get_ovlp(cell))
            if abs(ne - cell.nelectron).max() > 1e-7:
                logger.warn(self, 'Big error detected in the electron number '
                            'of initial guess density matrix (Ne/cell = %g)!\n'
                            '  This can cause huge error in Fock matrix and '
                            'lead to instability in SCF for low-dimensional '
                            'systems.\n  DM is normalized to correct number '
                            'of electrons', ne)
                dm *= cell.nelectron / ne
        return dm

    def init_guess_by_1e(self, cell=None):
        if cell is None: cell = self.cell
        if cell.dimension < 3:
            logger.warn(self, 'Hcore initial guess is not recommended in '
                        'the SCF of low-dimensional systems.')
        return mol_uhf.UHF.init_guess_by_1e(cell)

    def init_guess_by_chkfile(self, chk=None, project=True, kpt=None):
        if chk is None: chk = self.chkfile
        if kpt is None: kpt = self.kpt
        return pbcuhf.init_guess_by_chkfile(self.cell, chk, project, kpt)

    dump_chk = pbchf.SCF.dump_chk
    _is_mem_enough = pbchf.SCF._is_mem_enough

    density_fit = pbchf.SCF.density_fit
    # mix_density_fit inherits from hf.SCF.mix_density_fit

    x2c = x2c1e = sfx2c1e = pbchf.SCF.sfx2c1e

    def convert_from_(self, mf):
        '''Convert given mean-field object to RHF/ROHF'''
        from pyscf.pbc.scf import addons
        addons.convert_to_rhf(mf, self)
        return self

    stability = None
    nuc_grad_method = None


#!/usr/bin/env python3
"""create metamers

"""
from . import utils
import re
import torch
import torchvision
import plenoptic as po
import pyrtools as pt
import time
import numpy as np
import os.path as op
import imageio
import matplotlib as mpl
from skimage import color
import warnings
import sys
import yaml
with open(op.join(op.dirname(op.realpath(__file__)), '..', 'config.yml')) as f:
    fov_path = yaml.safe_load(f)['FOVEATED_METAMERS_PATH']
sys.path.append(op.join(fov_path, 'extra_packages'))
import plenoptic_part as pop

def setup_initial_image(initial_image_type, image):
    r"""Set up the initial image.

    Parameters
    ----------
    initial_image_type : {'white', 'pink', 'gray', 'blue'} or path to file
        What to use for the initial image. If 'white', we use white
        noise. If 'pink', we use pink noise
        (``pyrtools.synthetic_images.pink_noise(fract_dim=1)``). If
        'blue', we use blue noise
        (``pyrtools.synthetic_images.blue_noise(fract_dim=1)``). If
        'gray', we use a flat image with values of .5 everywhere. If
        path to a file, that's what we use as our initial image (and so
        the seed will have no effect on this).
    image : torch.Tensor
        The reference image tensor

    Returns
    -------
    initial_image : torch.Tensor
        The initial image to pass to metamer.synthesize

    """
    if initial_image_type == 'white':
        initial_image = torch.rand_like(image, dtype=torch.float32)
    elif initial_image_type == 'gray':
        initial_image = .5 * torch.ones_like(image, dtype=torch.float32)
    elif initial_image_type == 'pink':
        # this `.astype` probably isn't necessary, but just in case
        initial_image = pt.synthetic_images.pink_noise(image.shape[-2:]).astype(np.float32)
        # need to rescale this so it lies between 0 and 1
        initial_image += np.abs(initial_image.min())
        initial_image /= initial_image.max()
        initial_image = torch.Tensor(initial_image).unsqueeze(0).unsqueeze(0)
    elif initial_image_type == 'blue':
        # this `.astype` probably isn't necessary, but just in case
        initial_image = pt.synthetic_images.blue_noise(image.shape[-2:]).astype(np.float32)
        # need to rescale this so it lies between 0 and 1
        initial_image += np.abs(initial_image.min())
        initial_image /= initial_image.max()
        initial_image = torch.Tensor(initial_image).unsqueeze(0).unsqueeze(0)
    elif op.isfile(initial_image_type):
        warnings.warn("Using image %s as initial image!" % initial_image_type)
        initial_image = imageio.imread(initial_image_type)
        initial_image = convert_im_to_float(initial_image)
        initial_image = torch.tensor(initial_image, dtype=torch.float32).unsqueeze(0).unsqueeze(0)
    else:
        raise Exception("Don't know how to handle initial_image_type %s! Must be one of {'white',"
                        " 'gray', 'pink', 'blue'}" % initial_image_type)
    return torch.nn.Parameter(initial_image)

def setup_image(image, n_channels=1):
    r"""Set up the image.

    We load in the image, if it's not already done so (converting it to
    gray-scale in the process), make sure it lies between 0 and 1, and
    make sure it's a tensor of the correct type and specified device

    Parameters
    ----------
    image : str or array_like
        Either the path to the file to load in or the loaded-in
        image. If array_like, we assume it's already 2d (i.e.,
        grayscale)
    n_channels : int, optional
        How many channels the image should have. Will always be grayscale, but
        we may duplicate along one of the channels so we can feed this into
        VGG16

    Returns
    -------
    image : torch.Tensor
        The image tensor, ready to go

    """
    if isinstance(image, str):
        print("Loading in reference image from %s" % image)
        image = imageio.imread(image)
    if image.dtype == np.uint8:
        warnings.warn("Image is int8, with range (0, 255)")
        image = utils.convert_im_to_float(image)
    elif image.dtype == np.uint16:
        warnings.warn("Image is int16 , with range (0, 65535)")
        image = utils.convert_im_to_float(image)
    else:
        warnings.warn("Image is float 32, so we assume image range is (0, 1)")
        if image.max() > 1:
            raise Exception("Image is neither int8 nor int16, but its max is greater than 1!")
    # we use skimage.color.rgb2gray in order to handle rgb
    # correctly. this uses the ITU-R 601-2 luma transform, same as
    # matlab. we do this after the above, because it changes the image
    # dtype to float32
    if image.ndim == 3:
        # then it's a color image, and we need to make it grayscale
        image = color.rgb2gray(image)
    image = torch.tensor(image, dtype=torch.float32)
    while image.ndimension() < 4:
        image = image.unsqueeze(0)
    return image.repeat(1, n_channels, 1, 1)

def setup_model(model_name, image, min_ecc, max_ecc, cache_dir,
                normalize_dict=None):
    r"""Set up the model.

    We initialize the model, with the specified parameters, and return it.

    `model_name` must be 'VGG16_poolN' (where N is an int between 1 and 5),
    'PSTexture' or one of our foveated models. If a foveated model, it must be
    constructed of several parts, for which you have several chocies:
    `'{visual_area}{options}_{window_type}_scaling-{scaling}'`:
    - `visual_area`: which visual area we're modeling.`'RGC'` (retinal
      ganglion cells, `plenoptic.simul.PooledRGC` class) or
      `'V1'` (primary visual cortex,
      `plenoptic.simul.PrimaryVisualCortex` class)
    - `options`: you can additionally include the following strs,
      separated by `_`:
      - `'norm'`: if included, we normalize the models' `cone_responses`
        and (if V1) `complex_cell_responses` attributes. In this case,
        `normalize_dict` must also be set (and include those two
        keys). If not included, the model is not normalized
        (normalization makes the optimization easier because the
        different scales of the steerable pyramid have different
        magnitudes).
      - `s#` (V1 only), where `#` is an integer. The number of scales to
        inlude in the steerable pyramid that forms the basis fo the `V1`
        models. If not included, will use 4.
    - `window_type`: `'gaussian'` or `'cosine'`. whether to build the
      model with gaussian or raised-cosine windows. Regardless, scaling
      will always give the ratio between the FWHM and eccentricity of
      the windows, but the gaussian windows are much tighter packed, and
      so require more windows (and thus more memory), but also seem to
      have fewer aliasing issues.
    - `scaling`: float giving the scaling values of these models

    The recommended model_name values that correspond to our foveated models
    are: `RGC_norm_gaussian_scaling-{scaling}`,
    `V1_norm_s6_gaussian_scaling-{scaling}` (pick whatever scaling value you
    like).

    For the other model_name choices:

    - PSTexture: the Portilla-Simoncelli texture stats with n_scales=4,
      n_orientations=4, spatial_corr_width=9, use_true_correlations=True

    - VGG16_poolN: pretrained VGG16 from torchvision, through Nth max pooling
      layer (where N is an int from 1 to 5)

    Parameters
    ----------
    model_name : str
        str specifying which of the models we should initialize. See above for
        more details.
    image : torch.tensor or np.array
        The image we will call the model on. This is only necessary
        because we need to know how big it is; we just use its shape
    min_ecc : float
        The minimum eccentricity for the pooling windows (see
        plenoptic.simul.VentralStream for more details)
    max_ecc : float
        The maximum eccentricity for the pooling windows (see
        plenoptic.simul.VentralStream for more details)
    cache_dir : str or None, optional
        The directory to cache the windows tensor in. If set, we'll look
        there for cached versions of the windows we create, load them if
        they exist and create and cache them if they don't. If None, we
        don't check for or cache the windows.
    normalize_dict : dict or None, optional
        If a dict, should contain the stats to use for normalization. If
        None, we don't normalize. This can only be set (and must be set)
        if the model is "V1_norm". In any other case, we'll throw an
        Exception.

    Returns
    -------
    model : torch.nn.Module
        A ventral stream model, ready to use

    """
    if 'gaussian' in model_name:
        window_type = 'gaussian'
        t_width = None
        std_dev = 1
    elif 'cosine' in model_name:
        window_type = 'cosine'
        t_width = 1
        std_dev = None
    if model_name.startswith('RGC') or model_name.startswith('V1'):
        model_name, scaling = re.findall('([a-zA-z_0-9-.]+)_scaling-([0-9.]+)', model_name)[0]
        scaling = float(scaling)
        if 'norm' not in model_name:
            if normalize_dict:
                raise Exception(f"Cannot normalize model {model_name} (norm must be part of model_name to do so)!")
            normalize_dict = {}
        if not normalize_dict and 'norm' in model_name:
            raise Exception(f"If model_name is {model_name}, normalize_dict must be set!")
        if model_name.startswith('RGC'):
            model = pop.PooledRGC(scaling, image.shape[-2:],
                                  min_eccentricity=min_ecc,
                                  max_eccentricity=max_ecc,
                                  window_type=window_type,
                                  transition_region_width=t_width,
                                  cache_dir=cache_dir,
                                  std_dev=std_dev,
                                  normalize_dict=normalize_dict)
        elif model_name.startswith('V1'):
            try:
                num_scales = int(re.findall('_s([0-9]+)_', model_name)[0])
            except (IndexError, ValueError):
                num_scales = 4
            try:
                moments = int(re.findall('_m([0-9]+)_', model_name)[0])
                moments = list(range(2, moments+1))
            except (IndexError, ValueError):
                moments = []
            model = pop.PooledV1(scaling, image.shape[-2:],
                                 min_eccentricity=min_ecc,
                                 max_eccentricity=max_ecc,
                                 std_dev=std_dev,
                                 transition_region_width=t_width,
                                 cache_dir=cache_dir,
                                 normalize_dict=normalize_dict,
                                 num_scales=num_scales,
                                 window_type=window_type,
                                 moments=moments)
    elif model_name.startswith('OnOff'):
        pretrained, kernel_size = re.findall('OnOff_pretrained-([a-zA-z]+)_size-([0-9]+)', model_name)[0]
        if pretrained == 'True':
            pretrained = True
        else:
            pretrained = False
        if pretrained:
            assert int(kernel_size) == 31
        model = po.simul.OnOff(int(kernel_size), pretrained=pretrained)
    elif model_name == 'PSTexture':
        model = po.simul.PortillaSimoncelli(image.shape[-2:], n_scales=4,
                                            n_orientations=4,
                                            spatial_corr_width=9,
                                            use_true_correlations=True)
    elif 'VGG16' in model_name:
        model = torchvision.models.vgg16(pretrained=True).eval()
        # through the first max pooling layer
        if 'pool1' in model_name:
            model = torch.nn.Sequential(*list(model.children())[0][:5])
        # through the second max pooling layer
        elif 'pool2' in model_name:
            model = torch.nn.Sequential(*list(model.children())[0][:10])
        # etc
        elif 'pool3' in model_name:
            model = torch.nn.Sequential(*list(model.children())[0][:17])
        elif 'pool4' in model_name:
            model = torch.nn.Sequential(*list(model.children())[0][:24])
        elif 'pool5' in model_name:
            model = torch.nn.Sequential(*list(model.children())[0][:31])
        else:
            raise Exception(f"Don't know what to do with model_name {model_name}!")
    else:
        raise Exception("Don't know how to handle model_name %s" % model_name)
    return model

def setup_device(*args, gpu_id=None):
    r"""Setup device and get everything onto it

    This simple function checks whether ``torch.cuda.is_available()``
    and ``gpu_id`` is not None. If not, we use the cpu as the device

    We then call a.to(device) for every a in args (so this can be called
    with an arbitrary number of objects, each of which just needs to have
    .to method).

    Note that we always return a list (even if you only pass one item),
    so if you pass a single object, you'll need to either grab it
    specifically, either by doing ``im = setup_device(im,
    gpu_id=0)[0]`` or ``im, = setup_device(im)`` (notice the
    comma).

    Parameters
    ----------
    args :
        Some number of torch objects that we want to get on the proper
        device
    gpu_id : int or None, optional
        If not None, the GPU we will use. If None, we run on CPU. We
        don't do anything clever to handle that here, but the
        contextmanager utils.get_gpu_id does, so you should use that to
        make sure you're using a GPU that exists and is available (see
        Snakefile for example). Note that, to set this,
        you must set it as a keyword, i.e., ``setup_device(im, 0)``
        won't work but ``setup_device(im, gpu_id=True)`` will (this is
        because the ``*args`` in our function signature will greedily
        grab every non-keyword argument).

    Returns
    -------
    args : list
        Every item we were passed in arg, now on the proper device

    """
    if gpu_id is not None:
        if not torch.cuda.is_available():
            raise Exception("CUDA is not available but gpu_id is not None!")
        device = torch.device("cuda:%s" % gpu_id)
        dtype = torch.float32
    else:
        device = torch.device("cpu")
        dtype = torch.float32
    print("On device %s" % device)
    if dtype is not None:
        print("Changing dtype to %s" % dtype)
        args = [a.to(dtype) for a in args]
    return [a.to(device) for a in args]

def add_center_to_image(model, image, reference_image):
    r"""Add the reference image center to an image

    The VentralStream class of models will do nothing to the center of
    the image (they don't see the fovea), so we add the fovea to the
    image before synthesis.

    Parameters
    ----------
    model : plenoptic.simul.VentralStream
        The model used to create the metamer. Specifically, we need its
        windows attribute
    image : torch.Tensor
        The image to add the center back to
    reference_image : torch.Tensor
        The reference/target image for synthesis
        (``metamer.base_signal``); the center comes from this image.

    Returns
    -------
    recentered_image : torch.Tensor
        ``image`` with the reference image center added back in

    """
    model(image)
    rep = model.representation['mean_luminance']
    dummy_ones = torch.ones_like(rep)
    windows = model.PoolingWindows.project(dummy_ones).squeeze().to(image.device)
    # these aren't exactly zero, so we can't convert it to boolean
    anti_windows = 1 - windows
    return ((windows * image) + (anti_windows * reference_image))

def save(save_path, metamer):
    """Save Metamer object and its outputs.

    We save the object itself, plus:
    - The finished metamer in its original float32 format (with
      values between 0 and 1, as a numpy array), at
      ``os.path.splitext(save_path)[0] + "_metamer.npy"``.
    - The finished metamer 8-bit image, at
      ``os.path.splitext(save_path)[0] + "_metamer.png"``.
    - Picture showing synthesis progress summary at
      ``os.path.splitext(save_path)[0] + "_synthesis.png"``.

    Parameters
    ----------
    save_path : str
        The path to save the MADCompetition object at, which we use as a
        starting-point for the other save paths
    metamer : plenoptic.synth.MADCompetition
        The Metamer object after synthesis

    """
    if metamer.model(metamer.synthesized_signal).ndimension() == 4:
        # these VGG representations have many channels, plotting them all takes
        # too much time
        plot_model_response_error = False
    else:
        plot_model_response_error = True
    print("Saving at %s" % save_path)
    if hasattr(metamer.model, 'PoolingWindows'):
        # If we're using one of our foveated models, we add the center back at
        # the end because our gradients are not exactly zero in the center, and
        # thus those pixels end up getting moved around a little bit. Not
        # entirely sure why, but probably not worth tracing down, since we're
        # interested in the periphery
        metamer.synthesized_signal = torch.nn.Parameter(add_center_to_image(metamer.model,
                                                                            metamer.synthesized_signal,
                                                                            metamer.target_signal))
    metamer.save(save_path)
    # save png of mad
    metamer_path = op.splitext(save_path)[0] + "_metamer.png"
    metamer_image = po.to_numpy(metamer.synthesized_signal).squeeze()
    print("Saving metamer float32 array at %s" % metamer_path.replace('.png', '.npy'))
    np.save(metamer_path.replace('.png', '.npy'), metamer_image)
    print("Saving metamer image at %s" % metamer_path)
    if metamer_image.ndim == 3:
        # then this is an RGB, from the VGG16 models, and we want to move the
        # channels dim to the last dimension
        metamer_image = metamer_image.transpose(1, 2, 0)
    imageio.imwrite(metamer_path, utils.convert_im_to_int(metamer_image))
    synthesis_path = op.splitext(save_path)[0] + "_synthesis.png"
    print(f"Saving synthesis image at {synthesis_path}")
    fig, _ = po.synth.metamer.plot_synthesis_status(metamer,
                                                    model_response_error=plot_model_response_error)
    fig.savefig(synthesis_path)

def main(model_name, image, seed=0, min_ecc=.5, max_ecc=15, learning_rate=1,
         max_iter=100, stop_criterion=1e-4, stop_iters_to_check=50,
         save_path=None, initial_image='white', gpu_id=None, cache_dir=None,
         normalize_dict=None, optimizer='Adam', loss_func='mse',
         range_penalty_lambda=.1, coarse_to_fine=False,
         coarse_to_fine_kwargs={}, continue_path=None, num_threads=None):
    r"""Create metamer images.

    Given a model_name, model parameters, a target image, and some
    optimization parameters, we do our best to synthesize a metamer,
    saving the outputs after it finishes.

    `model_name` must either a model, 'VGG16_poolN' (where N is an int between
    1 and 5), 'PSTexture' or one of our foveated models. If a foveated model,
    it must be constructed of several parts, for which you have several
    chocies:
    `'{visual_area}{options}_{window_type}_scaling-{scaling}'`:
    - `visual_area`: which visual area we're modeling.`'RGC'` (retinal
      ganglion cells, `plenoptic.simul.PooledRGC` class) or
      `'V1'` (primary visual cortex,
      `plenoptic.simul.PrimaryVisualCortex` class)
    - `options`: you can additionally include the following strs,
      separated by `_`:
      - `'norm'`: if included, we normalize the models' `cone_responses`
        and (if V1) `complex_cell_responses` attributes. In this case,
        `normalize_dict` must also be set (and include those two
        keys). If not included, the model is not normalized
        (normalization makes the optimization easier because the
        different scales of the steerable pyramid have different
        magnitudes).
      - `s#` (V1 only), where `#` is an integer. The number of scales to
        inlude in the steerable pyramid that forms the basis fo the `V1`
        models. If not included, will use 4.
    - `window_type`: `'gaussian'` or `'cosine'`. whether to build the
      model with gaussian or raised-cosine windows. Regardless, scaling
      will always give the ratio between the FWHM and eccentricity of
      the windows, but the gaussian windows are much tighter packed, and
      so require more windows (and thus more memory), but also seem to
      have fewer aliasing issues.
    - `scaling`: float giving the scaling values of these models

    The recommended model_name values that correspond to our foveated models
    are: `RGC_norm_gaussian_scaling-{scaling}`,
    `V1_norm_s6_gaussian_scaling-{scaling}` (pick whatever scaling value you
    like).

    For the other model_name choices:

    - PSTexture: the Portilla-Simoncelli texture stats with n_scales=4,
      n_orientations=4, spatial_corr_width=9, use_true_correlations=True

    - VGG16_poolN: pretrained VGG16 from torchvision, through Nth max pooling
      layer (where N is an int from 1 to 5)

    If you want to resume synthesis from an earlier run that didn't
    finish, set `continue_path` to the path of the `.pt` file created by
    that earlier run. We will then load it in and continue. For right
    now, we don't do anything to make sure that the arguments you pass
    to the function are the same as the first time, we just use the ones
    passed in. Generally, they should be identical, with the exception
    of learning_rate (which can be None to resume where you left off)
    and max_iter (which gives the number of extra iterations you want to
    do). Specifically, I think things might get weird if you do this
    initially on a GPU and then try to resume on a CPU (or vice versa),
    for example. When resuming, there's always a slight increase in the
    loss that, as far as I can tell, is unavoidable; it goes away
    quickly (and the loss continues its earlier trend) and so I don't
    think is an issue.

    Parameters
    ----------
    model_name : str
        str specifying which of the model we should use. See above for more
        details.
    image : str or array_like
        Either the path to the file to load in or the loaded-in
        image. If array_like, we assume it's already 2d (i.e.,
        grayscale)
    seed : int, optional
        The number to use for initializing numpy and torch's random
        number generators
    min_ecc : float, optional
        The minimum eccentricity for the pooling windows (see
        plenoptic.simul.VentralStream for more details)
    max_ecc : float, optional
        The maximum eccentricity for the pooling windows (see
        plenoptic.simul.VentralStream for more details)
    learning_rate : float, optional
        The learning rate to pass to metamer.synthesize's optimizer
    max_iter : int, optional
        The maximum number of iterations we allow the synthesis
        optimization to run for
    stop_criterion : float, optional
        The stop criterion. If the loss has changed by less than this over the
        past stop_iters_to_check iterations, we quit out.
    stop_iters_to_check : int, optional
        How many iterations back to check in order to see if the loss has
        stopped decreasing.
    save_path : str or None, optional
        If a str, the path to the file to save the metamer object to. If
        None, we don't save the synthesis output (that's probably a bad
        idea)
    initial_image : {'white', 'pink', 'gray', 'blue'} or path to a file
        What to use for the initial image. If 'white', we use white
        noise. If 'pink', we use pink noise
        (``pyrtools.synthetic_images.pink_noise(fract_dim=1)``). If
        'blue', we use blue noise
        (``pyrtools.synthetic_images.blue_noise(fract_dim=1)``). If
        'gray', we use a flat image with values of .5 everywhere. If
        path to a file, that's what we use as our initial image (and so
        the seed will have no effect on this).
        std dev of Gaussian noise added to image to initialize synthesis.
    gpu_id : int or None, optional
        If not None, the GPU we will use. If None, we run on CPU. We
        don't do anything clever to handle that here, but the
        contextmanager utils.get_gpu_id does, so you should use that to
        make sure you're using a GPU that exists and is available (see
        Snakefile for example)
    cache_dir : str or None, optional
        The directory to cache the windows tensor in. If set, we'll look
        there for cached versions of the windows we create, load them if
        they exist and create and cache them if they don't. If None, we
        don't check for or cache the windows.
    normalize_dict : str or None, optional
        If a str, the path to the dictionary containing the statistics to use
        for normalization for the model. If None, we don't normalize anything
    optimizer: {'Adam', 'SGD'}
        The choice of optimization algorithm
    loss_func : {'mse', 'l2'}, optional
        Which loss function to use.
    range_penalty_lambda :
        Lambda to multiply by range penalty and add to loss.
    coarse_to_fine : { 'together', 'separate', False}, optional
        If False, don't do coarse-to-fine optimization. Else, there
        are two options for how to do it:
        - 'together': start with the coarsest scale, then gradually
          add each finer scale. this is like blurring the objective
          function and then gradually adding details and is probably
          what you want.
        - 'separate': compute the gradient with respect to each
          scale separately (ignoring the others), then with respect
          to all of them at the end.
    coarse_to_fine_kwargs : dict, optional
        Dictionary of args for coarse to fine optimization. See
        Metamer.synthesize() docstring for details.
    continue_path : str or None, optional
        If None, we synthesize a new metamer. If str, this should be the
        path to a previous synthesis run, which we are resuming. In that
        case, you may set learning_rate to None (in which case we resume
        where we left off) and set max_iter to a different value (the
        number of extra iterations to run) otherwise the rest of the
        arguments should be the same as the first run.
    num_threads : int or None, optional
        If int, the number of CPU threads to use. If None, we don't restrict it
        and so we'll use all available resources. If using the GPU, this won't
        matter (all costly computations are done on the GPU). If one the CPU,
        we seem to only improve performance up to ~12 threads (at least with
        RGC model), and actively start to harm performance as we get above 40.

    """
    print("Using seed %s" % seed)
    if num_threads is not None:
        print(f"Using {num_threads} threads")
        torch.set_num_threads(num_threads)
    else:
        print("Not restricting number of threads, will probably use max "
              f"available ({torch.get_num_threads()})")
    po.tools.set_seed(seed)
    image = setup_image(image, 3 if 'VGG16' in model_name else 1)
    print(f"Using initial image {initial_image}")
    initial_image = setup_initial_image(initial_image, image)
    # this will be false if normalize_dict is None or an empty list
    if normalize_dict:
        normalize_dict = torch.load(normalize_dict)
    model = setup_model(model_name, image, min_ecc, max_ecc, cache_dir,
                        normalize_dict)
    model_str = f"Using model {model_name}"
    if model_name.startswith('RGC') or model_name.startswith("V1"):
        model_str += f" from {min_ecc} degrees to {max_ecc} degrees"
    print(model_str)
    image, initial_image, model = setup_device(image, initial_image, model, gpu_id=gpu_id)
    store_progress = max(10, max_iter//100)
    if loss_func == 'mse':
        loss_func = po.tools.optim.mse
    elif loss_func == 'l2':
        loss_func = po.tools.optim.l2_norm
    else:
        raise Exception(f"Don't know how to handle loss_func {loss_func}!")
    metamer = po.synth.Metamer(image, model, loss_func, range_penalty_lambda,
                               initial_image=initial_image)
    print(f"Using optimizer {optimizer}")
    if optimizer == 'Adam':
        opt = torch.optim.Adam([metamer.synthesized_signal], lr=learning_rate, amsgrad=True)
    elif optimizer == 'SGD':
        opt = torch.optim.SGD([metamer.synthesized_signal], lr=learning_rate)
    if continue_path is not None:
        print("Resuming synthesis saved at %s" % continue_path)
        metamer = metamer.load(continue_path)
        opt = None
    print(f"Using learning rate {learning_rate}, stop_criterion {stop_criterion} (stop_iters_to_check "
          f"{stop_iters_to_check}), and max_iter {max_iter}")
    print(f"Using coarse-to-fine {coarse_to_fine} with kwargs {coarse_to_fine_kwargs}")
    start_time = time.time()
    metamer.synthesize(max_iter=max_iter, optimizer=opt,
                       store_progress=store_progress,
                       stop_criterion=stop_criterion,
                       stop_iters_to_check=stop_iters_to_check,
                       coarse_to_fine=coarse_to_fine,
                       coarse_to_fine_kwargs=coarse_to_fine_kwargs)
    duration = time.time() - start_time
    print(f"Synthesis took {duration} seconds")
    # make sure everything's on the cpu for saving
    metamer.to('cpu')
    if save_path is not None:
        save(save_path, metamer)


import pickle
from typing import Tuple, Callable, Sequence

import numpy as np

__all__ = ["Neuron", "Callback"]

Callback = Callable[[float], np.ndarray]

class BaseNeuron:
    @staticmethod
    def logistic(z: np.ndarray) -> np.ndarray:
        return 1 / (2 + np.expm1(-z))

    @staticmethod
    def tanh(z: np.ndarray, *args, **kwargs) -> np.ndarray:
        return np.tanh(z, *args, **kwargs)

    @staticmethod
    def softplus(z: np.ndarray) -> np.ndarray:
        return np.log(np.expm1(z) + 2)

    @staticmethod
    def relu(z: np.ndarray) -> np.ndarray:
        return np.array([max(0, zz) for zz in z])

    @staticmethod
    def relu6(z: np.ndarray) -> np.ndarray:
        return np.array([min(max(zz, 0), 6) for zz in z])

    @staticmethod
    def binary(z: np.ndarray) -> np.ndarray:
        return np.array([[-1, 1][zz > 0] for zz in z])

    @staticmethod
    def identity(z: np.ndarray) -> np.ndarray:
        return z

class Neuron(BaseNeuron):
    """
    In the framework of the Echo State Neuron, an individual neuron is vastly more complicated and capable than
    traditional neural network nodes.  In the traditional case the dot product of an input vector and a learned
    weight vector plus a bias are passed through some non linearity (or not) to produce the neuron's activation.  The
    formulation is thus:

    $$ a_{out} = \sigma(\vec{a_{in}} \cdot \vec{w} + b) $$

    The basic problem with this is that it is static. It completely misses the dynamics possible in a real neuron and
    does now allow for spatiotemporal receptive field mapping without playing with approximations of time,
    etc. And even then, the output really does not vary based on the input to the system.

    I propose an ESNeuron.  The ESNeuron gets its name from echo state networks.  In an ESNeuron multiple inputs
    still exist and there is a single output.  The distinction comes when everything else is considered.  In an
    ESNeuron the inputs are fully connected to a dynamical pool.  If we assume learning or processing occurs in the
    dendritic arbor of a neuron, then the dynamical pool is an analog to the dendritic arbor.  We can likewise assume
    that the axon hillock is the integration point in a neuron. In this case, our ESNeuron performs a linear
    transformation with the reservoir resulting in a value taken to be the ESNeuron activation.

    The final step in the system is to connect the output via fixed weights to the reservoir.  Because we're passing
    through an echo state network, our output is not entirely driven by the input to the system but works with inputs
    received, feedback from prior outputs and the internal dynamical nature of the ESN.

    In theory any output pattern can be trained given any input.  One could produce, given some random fixed pattern
    virtually any wave form.  This wave form can act as the input to another ESN.  In that case, the ESN receives an
    input pattern over time and this drives another potentially more complex output pattern.

    Given the above there are several parameters that decide the ESNeuron:
    """
    __version__ = "0.9"

    def __init__(self,
                 numInputs: int,
                 numReservoir: int,
                 numOutputs: int,
                 pct: float,
                 alpha: float,
                 *,
                 f: Callable[[np.ndarray], np.ndarray] = BaseNeuron.tanh,
                 g: Callable[[np.ndarray], np.ndarray] = BaseNeuron.identity,
                 feedback: bool = True
                 ):
        """
        Initialize the neuron with the architecture defined by the parameters

        :param numInputs: The number of connections from precursor neurons to this neuron
        :param numReservoir: The number of nodes in the internal echo state network
        :param numOutputs: The number of connections to the internal integrator from the reservoir
        :param pct: The probability of a connection between nodes of the echo state network
        :param alpha: float: measure of speed of reservoir lower values faster 0 < alpha <= 1
        :param f: callable: default Neuron.tanh
        :param g: callable: default Neuron.identity
        :param feedback: bool: default True - feed outputs back into the reservoir
        """
        assert 0 <= numInputs, "We restrict the number of inputs to be non-negative"
        assert 0 < numReservoir, "We restrict the number of reservoir nodes to be strictly positive"
        assert 0 < numOutputs, "We restrict the number of outputs to be strictly positive"

        assert 0 < pct <= 1, "Probability of connection must be between 0 exclusive and 1 inclusive"
        # assert 0 < alpha < 1, "alpha - must be 0 < alpha < 1"

        assert callable(f), "f must be a function(arg) or None"
        assert callable(g), "f must be a function(arg) or None"

        # These parameters are sufficient to describe the echo state network dimensions
        self._numInputs = numInputs  # I  K
        self._numReservoir = numReservoir  # H  N
        self._numOutputs = numOutputs  # O  L

        # Store the functions used to define how the neuron operates
        self.f = f
        self.g = g

        # Construct the Input Matrix that transforms the input vector from length I to length H
        self.Win = np.random.standard_normal((numReservoir, numInputs))

        # Construct the Reservoir Matrix that transforms Hidden to Hidden.
        self.W = np.eye(numReservoir)
        self.W[np.random.rand(numReservoir, numReservoir) < pct] = 1
        self.W[self.W > 0] = np.random.standard_normal((numReservoir, numReservoir))[self.W > 0]
        self.W = self.W * alpha / np.max(np.abs(np.linalg.eig(self.W)[0]))

        # Construct the feedback network
        self.Wfb = np.random.randn(numReservoir, numOutputs) * int(feedback)

        # Construct the output network
        self._Wout = np.random.standard_normal((numOutputs, numReservoir + numInputs))
        self.baseWout = self._Wout.copy()
        # self._Wout = np.zeros(shape=(numOutputs, numReservoir + numInputs))

        # Construct the basic activation vectors (neurons)
        self.x = np.random.uniform(-1, 1, size=(numReservoir, 1))
        self.y = np.random.uniform(-1, 1, size=(numOutputs, 1))
        self.z = np.random.uniform(-1, 1, size=(numReservoir + numInputs, 1))

        # number of iterations of no input (zero vector) used internally as setup and during training
        self._washout = 500

        # make sure the network has settled prior to finishing initialization
        self.forget()
        # self.settle(self.washout)

    @property
    def washout(self):
        return self._washout

    @washout.setter
    def washout(self, washout):
        self._washout = washout

    @property
    def Wout(self) -> np.ndarray:
        return self._Wout

    @Wout.setter
    def Wout(self, Wout: np.ndarray):
        assert self._Wout.shape == Wout.shape, "Cannot set weights. The dimensions do not match"
        self._Wout = Wout.copy()

    @property
    def numInputs(self):
        return self._numInputs

    @property
    def numReservoir(self):
        return self._numReservoir

    @property
    def numOutputs(self):
        return self._numOutputs

    @classmethod
    def load(cls):
        try:
            with open("esn_parameters.pkl", "rb") as inp:
                x = pickle.load(inp)
            assert isinstance(x, cls), "Pickled object is not a {}".format(cls)
            assert x.__version__ == cls.__version__, "Version mismatch - pkl file is out of date"
            return x
        except FileNotFoundError:
            return None

    def save(self):
        with open("esn_parameters.pkl", "wb") as out:
            pickle.dump(self, out, pickle.HIGHEST_PROTOCOL)

    def forget(self):
        """
        Forget any prior learning by zeroing the output weight matrix and then running the network with
        no input for a <washout> number of steps.
        :return: None
        """

        # Zero the output matrix and then let the network settle for the washout period
        # self._Wout = np.zeros(shape=self._Wout.shape)
        self._Wout = self.baseWout.copy()
        self.settle(self.washout)

    def settle(self, count: int):
        """
        Cause the network to run <count> steps with no input.  This is intended to remove prior input traces
        from the reservoir prior to training or when ever a clean slate is desired. The number of steps <count>
        really should be computed based on the value of alpha - the spectral radius of the Reservoir.
        :param count: int; The number of cycles to run the ESN with no input.
        :return: None
        """

        def fin(_: float) -> np.ndarray:
            """function returning zero vector size numInputs"""
            return np.zeros((self.numInputs, 1))

        self.cycle(1, count, [fin for _ in range(count)])

    def cycle(self,
              t0: int,
              tn: int,
              f_in: Sequence[Callback]) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
        """
        (1) x(n+1) = f(W * x(n) + Win * u(n+1) + Wfb * y(n))
            z(n) = [x(n); u(n)]
        (2) y(n) = g(Wout * z(n))

        Execute a single pass through the echo state network potentially taking in some stimulus.
        There is no learning when using this method. To train your ESN, please use force_learn.
        :param t0: int: Starting index passed to f_teach to retrieve the t0'th teaching pair
        :param tn: int: Ending index passed to f_teach to retrieve the tn'th teaching pair
        :param f_in: callable(t: float) -> np.array a function that returns the network input at time t
        :return: the current output vector
        """
        assert t0 <= tn, "Cannot observe backwards through time. tn must not be less than t0"

        nmax = tn - t0 + 1
        u = None

        for t in range(nmax):
            u = f_in[t](t0 + t)  # retrieve the next input pattern
            self.x = self.f((self.W @ self.x) + (self.Win @ u) + (self.Wfb @ self.y))
            self.z = np.concatenate((self.x, u))
            self.y = self.g(self._Wout @ self.z)

        return u, self.x, self.y, self.z

    def learn(self,
              t0: int,
              tn: int,
              f_in: Sequence[Callback],
              f_out: Sequence[Callback]) -> Tuple[np.ndarray, np.ndarray]:
        """
        Execute a single forced learning phase taking stimulus and desired response into consideration.

        (1) x(n+1) = f(W * x(n) + Win * u(n+1) + Wfb * y(n))
            z(n) = [x(n); u(n)]
        (2) y(n) = g(Wout * z(n))

        Learning equations. In the state harvesting stage of the training, the ESN is driven by an input sequence
        u(1),…,u(nmax) , which yields a sequence z(1),…,z(nmax) of extended system states. The system equations (1),
        (2) are used here. If the model includes output feedback (i.e., nonzero Wfb), then during the generation of the
        system states, the correct outputs d(n) (part of the training data) are written into the output units ("teacher
        forcing"). The obtained extended system states are filed row-wise into a state collection matrix S of size
        nmax×(N+K) . Usually some initial portion of the states thus collected are discarded to accommodate for a
        washout of the arbitrary (random or zero) initial reservoir state needed at time 1. Likewise, the desired
        outputs d(n) are sorted row-wise into a teacher output collection matrix D of size nmax×L .

        The desired output weights Wout are the linear regression weights of the desired outputs d(n) on the harvested
        extended states z(n) . A mathematically straightforward way to compute Wout is to invoke the pseudoinverse
        (denoted by ⋅†) of S :

        (3) Wout = (S†D)′
             S = z(t0:tn)
             D = y(t0:tn)
             Wout = np.matmul(np.linalg.pinv(S), D).T

        which is an offline algorithm (the prime denotes matrix transpose). Online adaptive methods known from linear
        signal processing can also be used to compute output weights (Jaeger 2003).

        A note on the function f_teach.  f_teach(t) should return a tuple of functions where the first
        of the tuple, when called returns the t'th input pattern and the second tuple, when called returns
        the t'th target pattern.  These returned functions take no parameters and each returns a vector.

        :param t0: int: Starting index passed to f_teach to retrieve the t0'th teaching pair
        :param tn: int: Ending index passed to f_teach to retrieve the tn'th teaching pair
        :param f_in: callable(t: int) -> np.array a function that returns the network input at time t
        :param f_out: callable(t: int) -> np.array a function that returns the target output at time t
        :return:
        """
        assert t0 <= tn, "Cannot train backwards through time. tn must not be less than t0"

        nmax = tn - t0 + 1

        S = np.empty(shape=(nmax, self.numReservoir + self.numInputs))
        D = np.empty(shape=(nmax, self.numOutputs))

        print("Learning sequence of {} inputs".format(nmax))
        for t in range(nmax):
            # self.cycle(t0 + t, t0 + t, f_in)

            u = f_in[t](t0 + t)  # retrieve the next input pattern
            y = f_out[t](t0 + t)  # retrieve the next output pattern

            reservoir = (self.W @ self.x)
            inputData = (self.Win @ u)
            feedBack = (self.Wfb @ self.y)

            self.x = self.f(reservoir + inputData + feedBack)
            self.z = np.concatenate((self.x, u))
            self.y = self.g(self._Wout @ self.z)

            # Teacher Forcing
            self.y = y

            # Save our state and target outputs
            S[t] = self.z.copy().reshape(-1)
            D[t] = y.copy().reshape(-1)

        self._Wout += (np.linalg.pinv(S) @ D).T
        return S, D