"""multipy: Python library for multicomponent mass transfer"""

__author__ = "James C. Sutherland, Kamila Zdybal"
__copyright__ = "Copyright (c) 2022, James C. Sutherland, Kamila Zdybal"
__license__ = "MIT"
__version__ = "1.0.0"
__maintainer__ = ["Kamila Zdybal"]
__email__ = ["kamilazdybal@gmail.com"]
__status__ = "Production"

import numpy as np
import pandas as pd
import random
import copy
import scipy
import multipy
import warnings

gas_constant = 8.31446261815324

################################################################################
################################################################################
####
####    Class: Flux
####
################################################################################
################################################################################

class Flux:
    """
    Supports computing and storing fluxes. This class assumes that the species velocities, :math:`\\mathbf{u}_i`, are known.

    Diffusive fluxes:

    - mass diffusive flux relative to a mass-averaged velocity, :math:`\mathbf{j}_i`
    - mass diffusive flux relative to a molar-averaged velocity, :math:`\mathbf{j}_i^u`
    - molar diffusive flux relative to a mass-averaged velocity, :math:`\mathbf{J}_i^v`
    - molar diffusive flux relative to a molar-averaged velocity, :math:`\mathbf{J}_i`

    :param species_velocities:
        vector ``numpy.ndarray`` specifying the species velocities :math:`\mathbf{u}_i` in :math:`[m/s]`. It should be of size ``(n_species,n_observations)``.

    **Getters:**

    - **get_species_velocities**
    - **get_diffusive_molar_molar** (is set to ``None`` at class init)
    - **get_diffusive_molar_mass** (is set to ``None`` at class init)
    - **get_diffusive_mass_molar** (is set to ``None`` at class init)
    - **get_diffusive_mass_mass** (is set to ``None`` at class init)

    **Setters:**

    - **set_species_velocities**
    - **set_diffusive_molar_molar** (is set to ``None`` at class init)
    - **set_diffusive_molar_mass** (is set to ``None`` at class init)
    - **set_diffusive_mass_molar** (is set to ``None`` at class init)
    - **set_diffusive_mass_mass** (is set to ``None`` at class init)
    """

    # --------------------------------------------------------------------------

    def __init__(self, species_velocities):

        if not isinstance(species_velocities, np.ndarray):
            raise ValueError("Parameter `species_velocities` has to be of type `numpy.ndarray`.")

        try:
            (n_species, n_observations) = np.shape(species_velocities)
        except:
            raise ValueError("Parameter `species_velocities` has to be a matrix.")

        if n_species < 2:
            raise ValueError("Parameter `species_velocities` has to have at least two species.")

        self.__species_velocities = species_velocities
        self.__velocity = multipy.Velocity(self.get_species_velocities)
        self.__diffusive_molar_molar = None
        self.__diffusive_molar_mass = None
        self.__diffusive_mass_molar = None
        self.__diffusive_mass_mass = None

    @property
    def get_species_velocities(self):
        return self.__species_velocities

    @property
    def get_diffusive_molar_molar(self):
        return self.__diffusive_molar_molar

    @property
    def get_diffusive_molar_mass(self):
        return self.__diffusive_molar_mass

    @property
    def get_diffusive_mass_molar(self):
        return self.__diffusive_mass_molar

    @property
    def get_diffusive_mass_mass(self):
        return self.__diffusive_mass_mass

    @get_species_velocities.setter
    def set_species_velocities(self, new_species_velocities):

        if new_species_velocities is not None:
            if not isinstance(new_species_velocities, np.ndarray):
                raise ValueError("Parameter `species_velocities` has to be of type `numpy.ndarray`.")

            try:
                (n_species, n_observations) = np.shape(new_species_velocities)
            except:
                raise ValueError("Parameter `species_velocities` has to be a matrix.")

        self.__species_velocities = new_species_velocities

    @get_diffusive_molar_molar.setter
    def set_diffusive_molar_molar(self, new_diffusive_molar_molar):

        if new_diffusive_molar_molar is not None:
            if not isinstance(new_diffusive_molar_molar, np.ndarray):
                raise ValueError("Parameter `diffusive_molar_molar` has to be of type `numpy.ndarray`.")

            try:
                (n_species, n_observations) = np.shape(new_diffusive_molar_molar)
            except:
                raise ValueError("Parameter `diffusive_molar_molar` has to be a matrix.")

        self.__diffusive_molar_molar = new_diffusive_molar_molar

    @get_diffusive_molar_mass.setter
    def set_diffusive_molar_mass(self, new_diffusive_molar_mass):

        if new_diffusive_molar_mass is not None:
            if not isinstance(new_diffusive_molar_mass, np.ndarray):
                raise ValueError("Parameter `diffusive_molar_mass` has to be of type `numpy.ndarray`.")

            try:
                (n_species, n_observations) = np.shape(new_diffusive_molar_mass)
            except:
                raise ValueError("Parameter `diffusive_molar_mass` has to be a matrix.")

        self.__diffusive_molar_mass = new_diffusive_molar_mass

    @get_diffusive_mass_molar.setter
    def set_diffusive_mass_molar(self, new_diffusive_mass_molar):

        if new_diffusive_mass_molar is not None:
            if not isinstance(new_diffusive_mass_molar, np.ndarray):
                raise ValueError("Parameter `diffusive_mass_molar` has to be of type `numpy.ndarray`.")

            try:
                (n_species, n_observations) = np.shape(new_diffusive_mass_molar)
            except:
                raise ValueError("Parameter `diffusive_mass_molar` has to be a matrix.")

        self.__diffusive_mass_molar = new_diffusive_mass_molar

    @get_diffusive_mass_mass.setter
    def set_diffusive_mass_mass(self, new_diffusive_mass_mass):

        if new_diffusive_mass_mass is not None:
            if not isinstance(new_diffusive_mass_mass, np.ndarray):
                raise ValueError("Parameter `diffusive_mass_mass` has to be of type `numpy.ndarray`.")

            try:
                (n_species, n_observations) = np.shape(new_diffusive_mass_mass)
            except:
                raise ValueError("Parameter `diffusive_mass_mass` has to be a matrix.")

        self.__diffusive_mass_mass = new_diffusive_mass_mass

    # --------------------------------------------------------------------------

    def plot_diffusive_flux(self, species_names=None, colors=None, figsize=(10,5), filename=None):
        """
        Plots the computed diffusive fluxes.

        **Example:**

        .. image:: ../images/stefan-tube-diffusive-flux-molar-diff-molar-avg.svg
          :width: 400

        :param species_names: (optional)
            ``list`` of ``str`` specifying the species names.
        :param colors: (optional)
            ``list`` of ``str`` specifying the plotting colors for each species. Example: ``colors=['#C7254E', '#BBBBBB', '#008CBA']``.
        :param figsize: (optional)
            ``tuple`` specifying the figure size.
        :param filename: (optional)
            ``str`` specifying the filename. If set to ``None``, plot will not be saved to a file.
        """

        if filename is not None:

            path = False

            if filename[0:2] == '..':
                __filename = filename[2::]
                path = True
            else:
                __filename = filename

            __base = __filename.split('.')[0]
            __extension = __filename.split('.')[1]

            if path:
                __filename = '..' + __base
            else:
                __filename = __base

        if self.get_diffusive_molar_molar is not None:
            if filename is not None:
                plt = multipy.plot.plot_1d_diffusive_flux(self.get_diffusive_molar_molar, flux='molar', velocity='molar', species_names=species_names, colors=colors, figsize=figsize, filename=__filename + '-molar-diff-molar-avg.' + __extension)
            else:
                plt = multipy.plot.plot_1d_diffusive_flux(self.get_diffusive_molar_molar, flux='molar', velocity='molar', species_names=species_names, colors=colors, figsize=figsize, filename=None)

        if self.get_diffusive_molar_mass is not None:
            if filename is not None:
                plt = multipy.plot.plot_1d_diffusive_flux(self.get_diffusive_molar_mass, flux='molar', velocity='mass', species_names=species_names, colors=colors, figsize=figsize, filename=__filename + '-molar-diff-mass-avg.' + __extension)
            else:
                plt = multipy.plot.plot_1d_diffusive_flux(self.get_diffusive_molar_mass, flux='molar', velocity='mass', species_names=species_names, colors=colors, figsize=figsize, filename=None)

        if self.get_diffusive_mass_molar is not None:
            if filename is not None:
                plt = multipy.plot.plot_1d_diffusive_flux(self.get_diffusive_mass_molar, flux='mass', velocity='molar', species_names=species_names, colors=colors, figsize=figsize, filename=__filename + '-mass-diff-molar-avg.' + __extension)
            else:
                plt = multipy.plot.plot_1d_diffusive_flux(self.get_diffusive_mass_molar, flux='mass', velocity='molar', species_names=species_names, colors=colors, figsize=figsize, filename=None)

        if self.get_diffusive_mass_mass is not None:
            if filename is not None:
                plt = multipy.plot.plot_1d_diffusive_flux(self.get_diffusive_mass_mass, flux='mass', velocity='mass', species_names=species_names, colors=colors, figsize=figsize, filename=__filename + '-mass-diff-mass-avg.' + __extension)
            else:
                plt = multipy.plot.plot_1d_diffusive_flux(self.get_diffusive_mass_mass, flux='mass', velocity='mass', species_names=species_names, colors=colors, figsize=figsize, filename=None)

    # --------------------------------------------------------------------------

    def diffusive_molar_molar(self, species_mole_fractions, species_molar_densities):
        """
        Computes the molar diffusive flux relative to a molar-averaged velocity:

        .. math::

            \mathbf{J}_i = c_i \mathbf{u}_i + c_i \mathbf{u}

        :param species_mole_fractions:
            scalar ``numpy.ndarray`` specifying the species mole fractions, :math:`X_i`, in :math:`[-]`. It should be of size ``(n_species,n_observations)``.
        :param species_molar_densities:
            scalar ``numpy.ndarray`` specifying the molar densities of species, :math:`c_i`, in :math:`[mole/m^3]`. It should be of size ``(n_species,n_observations)``.

        :return:
            - **diffusive_flux** - vector ``numpy.ndarray`` of molar diffusive fluxes relative to a molar-averaged velocity :math:`\mathbf{J}_i` in :math:`[mole/(m^2s)]`. It has size ``(n_species,n_observations)``.
        """

        if not isinstance(species_mole_fractions, np.ndarray):
            raise ValueError("Parameter `species_mole_fractions` has to be of type `numpy.ndarray`.")

        try:
            (n_species_1, n_observations_1) = np.shape(species_mole_fractions)
        except:
            raise ValueError("Parameter `species_mole_fractions` has to be a matrix.")

        if not isinstance(species_molar_densities, np.ndarray):
            raise ValueError("Parameter `species_molar_densities` has to be of type `numpy.ndarray`.")

        try:
            (n_species_2, n_observations_2) = np.shape(species_molar_densities)
        except:
            raise ValueError("Parameter `species_molar_densities` has to be a matrix.")

        if n_observations_1 != n_observations_2:
            raise ValueError("Parameters `species_mole_fractions` and `species_molar_densities` have different number of observations `n_observations`.")

        if n_species_1 != n_species_2:
            raise ValueError("Parameters `species_mole_fractions` and `species_molar_densities` have different number of species `n_species`.")

        (n_species, n_observations) = np.shape(self.get_species_velocities)

        if n_observations != n_observations_1:
            raise ValueError("Parameters `species_mole_fractions`, `species_molar_densities` and `species_velocities` have different number of observations `n_observations`.")

        if n_species != n_species_1:
            raise ValueError("Parameters `species_mole_fractions`, `species_molar_densities` and `species_velocities` have different number of species `n_species`.")

        molar_averaged_velocity = self.__velocity.molar_averaged(species_mole_fractions)

        diffusive_flux = np.multiply(species_molar_densities, self.get_species_velocities) - np.multiply(species_molar_densities, molar_averaged_velocity)
        self.__diffusive_molar_molar = diffusive_flux

        return diffusive_flux

    # --------------------------------------------------------------------------

    def diffusive_molar_mass(self, species_mass_fractions, species_molar_densities):
        """
        Computes the molar diffusive flux relative to a mass-averaged velocity:

        .. math::

            \mathbf{J}_i^v = c_i \mathbf{u}_i + c_i \mathbf{v}

        :param species_mass_fractions:
            scalar ``numpy.ndarray`` specifying the species mass fractions, :math:`Y_i`, in :math:`[-]`. It should be of size ``(n_species,n_observations)``.
        :param species_molar_densities:
            scalar ``numpy.ndarray`` specifying the species molar densities :math:`c_i` in :math:`[mole/m^3]`. It should be of size ``(n_species,n_observations)``.

        :return:
            - **diffusive_flux** - vector ``numpy.ndarray`` of molar diffusive fluxes relative to a mass-averaged velocity :math:`\mathbf{J}_i^v` in :math:`[mole/(m^2s)]`. It has size ``(n_species,n_observations)``.
        """

        if not isinstance(species_mass_fractions, np.ndarray):
            raise ValueError("Parameter `species_mass_fractions` has to be of type `numpy.ndarray`.")

        try:
            (n_species_1, n_observations_1) = np.shape(species_mass_fractions)
        except:
            raise ValueError("Parameter `species_mass_fractions` has to be a matrix.")

        if not isinstance(species_molar_densities, np.ndarray):
            raise ValueError("Parameter `species_molar_densities` has to be of type `numpy.ndarray`.")

        try:
            (n_species_2, n_observations_2) = np.shape(species_molar_densities)
        except:
            raise ValueError("Parameter `species_molar_densities` has to be a matrix.")

        if n_observations_1 != n_observations_2:
            raise ValueError("Parameters `species_mass_fractions` and `species_molar_densities` have different number of observations `n_observations`.")

        if n_species_1 != n_species_2:
            raise ValueError("Parameters `species_mass_fractions` and `species_molar_densities` have different number of species `n_species`.")

        (n_species, n_observations) = np.shape(self.get_species_velocities)

        if n_observations != n_observations_1:
            raise ValueError("Parameters `species_mass_fractions`, `species_molar_densities` and `species_velocities` have different number of observations `n_observations`.")

        if n_species != n_species_1:
            raise ValueError("Parameters `species_mass_fractions`, `species_molar_densities` and `species_velocities` have different number of species `n_species`.")

        mass_averaged_velocity = self.__velocity.mass_averaged(species_mass_fractions)

        diffusive_flux = np.multiply(species_molar_densities, self.get_species_velocities) - np.multiply(species_molar_densities, mass_averaged_velocity)
        self.__diffusive_molar_mass = diffusive_flux

        return diffusive_flux

    # --------------------------------------------------------------------------

    def diffusive_mass_molar(self, species_mole_fractions, species_mass_densities):
        """
        Computes the mass diffusive flux relative to a molar-averaged velocity:

        .. math::

            \mathbf{j}_i^u = \\rho_i \mathbf{u}_i + \\rho_i \mathbf{u}

        :param species_mole_fractions:
            scalar ``numpy.ndarray`` specifying the species mole fractions :math:`X_i` in :math:`[-]`. It should be of size ``(n_species,n_observations)``.
        :param species_mass_densities:
            scalar ``numpy.ndarray`` specifying the species mass densities :math:`\mathbf{\\rho}_i` in :math:`[kg/m^3]`. It should be of size ``(n_species,n_observations)``.

        :return:
            - **diffusive_flux** - vector ``numpy.ndarray`` of mass diffusive fluxes relative to a molar-averaged velocity :math:`\mathbf{j}_i^u` in :math:`[kg/(m^2s)]`. It has size ``(n_species,n_observations)``.
        """

        if not isinstance(species_mole_fractions, np.ndarray):
            raise ValueError("Parameter `species_mole_fractions` has to be of type `numpy.ndarray`.")

        try:
            (n_species_1, n_observations_1) = np.shape(species_mole_fractions)
        except:
            raise ValueError("Parameter `species_mole_fractions` has to be a matrix.")

        if not isinstance(species_mass_densities, np.ndarray):
            raise ValueError("Parameter `species_mass_densities` has to be of type `numpy.ndarray`.")

        try:
            (n_species_2, n_observations_2) = np.shape(species_mass_densities)
        except:
            raise ValueError("Parameter `species_mass_densities` has to be a matrix.")

        if n_observations_1 != n_observations_2:
            raise ValueError("Parameters `species_mole_fractions` and `species_mass_densities` have different number of observations `n_observations`.")

        if n_species_1 != n_species_2:
            raise ValueError("Parameters `species_mole_fractions` and `species_mass_densities` have different number of species `n_species`.")

        (n_species, n_observations) = np.shape(self.get_species_velocities)

        if n_observations != n_observations_1:
            raise ValueError("Parameters `species_mole_fractions`, `species_mass_densities` and `species_velocities` have different number of observations `n_observations`.")

        if n_species != n_species_1:
            raise ValueError("Parameters `species_mole_fractions`, `species_mass_densities` and `species_velocities` have different number of species `n_species`.")

        molar_averaged_velocity = self.__velocity.molar_averaged(species_mole_fractions)

        diffusive_flux = np.multiply(species_mass_densities, self.get_species_velocities) - np.multiply(species_mass_densities, molar_averaged_velocity)
        self.__diffusive_mass_molar = diffusive_flux

        return diffusive_flux

    # --------------------------------------------------------------------------

    def diffusive_mass_mass(self, species_mass_fractions, species_mass_densities):
        """
        Computes the mass diffusive flux relative to a mass-averaged velocity:

        .. math::

            \mathbf{j}_i = \\rho_i \mathbf{u}_i + \\rho_i \mathbf{v}

        :param species_mass_fractions:
            scalar ``numpy.ndarray`` specifying the species mass fractions :math:`Y_i` in :math:`[-]`. It should be of size ``(n_species, n_observations)``.
        :param species_mass_densities:
            scalar ``numpy.ndarray`` specifying the species mass densities :math:`\mathbf{\\rho}_i` in :math:`[kg/m^3]`. It should be of size ``(n_species, n_observations)``.

        :return:
            - **diffusive_flux** - vector ``numpy.ndarray`` of mass diffusive fluxes relative to a mass-averaged velocity :math:`\mathbf{j}_i` in :math:`[kg/(m^2s)]`. It has size ``(n_species, n_observations)``.
        """
        if not isinstance(species_mass_fractions, np.ndarray):
            raise ValueError("Parameter `species_mass_fractions` has to be of type `numpy.ndarray`.")

        try:
            (n_species_1, n_observations_1) = np.shape(species_mass_fractions)
        except:
            raise ValueError("Parameter `species_mass_fractions` has to be a matrix.")

        if not isinstance(species_mass_densities, np.ndarray):
            raise ValueError("Parameter `species_mass_densities` has to be of type `numpy.ndarray`.")

        try:
            (n_species_2, n_observations_2) = np.shape(species_mass_densities)
        except:
            raise ValueError("Parameter `species_mass_densities` has to be a matrix.")

        if n_observations_1 != n_observations_2:
            raise ValueError("Parameters `species_mass_fractions` and `species_mass_densities` have different number of observations `n_observations`.")

        if n_species_1 != n_species_2:
            raise ValueError("Parameters `species_mass_fractions` and `species_mass_densities` have different number of species `n_species`.")

        (n_species, n_observations) = np.shape(self.get_species_velocities)

        if n_observations != n_observations_1:
            raise ValueError("Parameters `species_mass_fractions`, `species_mass_densities` and `species_velocities` have different number of observations `n_observations`.")

        if n_species != n_species_1:
            raise ValueError("Parameters `species_mass_fractions`, `species_mass_densities` and `species_velocities` have different number of species `n_species`.")

        mass_averaged_velocity = self.__velocity.mass_averaged(species_mass_fractions)

        diffusive_flux = np.multiply(species_mass_densities, self.get_species_velocities) - np.multiply(species_mass_densities, mass_averaged_velocity)
        self.__diffusive_mass_mass = diffusive_flux

        return diffusive_flux

    # --------------------------------------------------------------------------


#!/bin/python
# -*- coding: utf-8 -*-

import numpy as np
import pandas as pd
import os
import time
import tqdm
from datetime import datetime
from .mpile import get_par

def mcmc(
    self,
    p0=None,
    nsteps=3000,
    nwalks=None,
    tune=None,
    moves=None,
    temp=False,
    seed=None,
    backend=True,
    suffix=None,
    linear=None,
    resume=False,
    append=False,
    update_freq=None,
    lprob_seed=None,
    report=None,
    maintenance_interval=10,
    verbose=False,
    debug=False,
    **samplerargs
):
    """Run the emcee ensemble MCMC sampler.

    Parameters:
    ----------

    p0 : ndarray of initial states of the walkers in the parameterspace
    moves : emcee.moves object
    lprob_seed : lprob_seed must be one of ('vec', 'rand', 'set').
    """

    import emcee
    from grgrlib.multiprocessing import serializer

    if not hasattr(self, "ndim"):
        # if it seems to be missing, lets do it.
        # but without guarantee...
        self.prep_estim(load_R=True)

    if seed is None:
        seed = self.fdict["seed"]

    self.tune = tune
    if tune is None:
        self.tune = int(nsteps * 1 / 5.0)

    if update_freq is None:
        update_freq = int(nsteps / 5.0)

    if linear is None:
        linear = self.filter.name == "KalmanFilter"

    if "description" in self.fdict.keys():
        self.description = self.fdict["description"]

    if hasattr(self, "pool"):
        from .estimation import create_pool

        create_pool(self)

    lprob_global = serializer(self.lprob)

    if isinstance(temp, bool) and not temp:
        temp = 1

    def lprob(par):
        return lprob_global(
            par,
            linear=linear,
            verbose=verbose,
            temp=temp,
            lprob_seed=lprob_seed or "set",
        )

    bnd = np.array(self.fdict["prior_bounds"])

    if self.pool:
        self.pool.clear()

    if p0 is None and not resume:
        if temp < 1:
            p0 = get_par(
                self,
                "prior_mean",
                asdict=False,
                full=False,
                nsample=nwalks,
                verbose=verbose,
            )
        else:
            p0 = get_par(
                self, "best", asdict=False, full=False, nsample=nwalks, verbose=verbose
            )
    elif not resume:
        nwalks = p0.shape[0]

    if backend:

        if isinstance(backend, str):
            # backend_file will only be loaded later if explicitely defined before
            self.fdict["backend_file"] = backend
        try:
            backend = self.fdict["backend_file"]
        except KeyError:
            # this is the default case
            suffix = str(suffix) if suffix else "_sampler.h5"
            backend = os.path.join(self.path, self.name + suffix)

            if os.path.exists(backend) and not (resume or append):
                print(
                    "[mcmc:]".ljust(15, " ")
                    + " HDF backend at %s already exists. Deleting..." % backend
                )
                os.remove(backend)

        backend = emcee.backends.HDFBackend(backend)

        if not (resume or append):
            if not nwalks:
                raise TypeError(
                    "If neither `resume`, `append` or `p0` is given I need to know the number of walkers (`nwalks`)."
                )
            try:
                backend.reset(nwalks, self.ndim)
            except KeyError as e:
                raise KeyError(str(e) + ". Your `*.h5` file is likely to be damaged...")
    else:
        backend = None

    if resume:
        nwalks = backend.get_chain().shape[1]

    if debug:
        sampler = emcee.EnsembleSampler(nwalks, self.ndim, lprob)
    else:
        sampler = emcee.EnsembleSampler(
            nwalks, self.ndim, lprob, moves=moves, pool=self.pool, backend=backend
        )

    if resume and not p0:
        p0 = sampler.get_last_sample()

    self.sampler = sampler
    self.temp = temp

    if not verbose:
        np.warnings.filterwarnings("ignore")

    if verbose > 2:
        report = report or print
    else:
        pbar = tqdm.tqdm(total=nsteps, unit="sample(s)", dynamic_ncols=True)
        report = report or pbar.write

    old_tau = np.inf
    cnt = 0

    for result in sampler.sample(p0, iterations=nsteps, **samplerargs):

        if not verbose:
            lls = list(result)[1]
            maf = np.mean(sampler.acceptance_fraction[-update_freq:]) * 100
            pbar.set_description(
                "[ll/MAF:%s(%1.0e)/%1.0f%%]" % (str(np.max(lls))[:7], np.std(lls), maf)
            )

        if cnt and update_freq and not cnt % update_freq:

            prnttup = "[mcmc:]".ljust(
                15, " "
            ) + "Summary from last %s of %s iterations" % (update_freq, cnt)

            if temp < 1:
                prnttup += " with temp of %s%%" % (np.round(temp * 100, 6))

            if self.description is not None:
                prnttup += " (%s)" % str(self.description)

            prnttup += ":"

            report(prnttup)

            sample = sampler.get_chain()

            tau = emcee.autocorr.integrated_time(sample, tol=0)
            min_tau = np.min(tau).round(2)
            max_tau = np.max(tau).round(2)
            dev_tau = np.max(np.abs(old_tau - tau) / tau)

            tau_sign = ">" if max_tau > sampler.iteration / 50 else "<"
            dev_sign = ">" if dev_tau > 0.01 else "<"

            self.mcmc_summary(
                chain=sample,
                tune=update_freq,
                calc_mdd=False,
                calc_ll_stats=True,
                out=lambda x: report(str(x)),
            )

            report(
                "Convergence stats: tau is in (%s,%s) (%s%s) and change is %s (%s0.01)."
                % (
                    min_tau,
                    max_tau,
                    tau_sign,
                    sampler.iteration / 50,
                    dev_tau.round(3),
                    dev_sign,
                )
            )

        if cnt and update_freq and not (cnt + 1) % update_freq:
            sample = sampler.get_chain()
            old_tau = emcee.autocorr.integrated_time(sample, tol=0)

        if not verbose:
            pbar.update(1)

        # avoid mem leakage
        if cnt and not cnt % maintenance_interval:
            self.pool.clear()

        cnt += 1

    pbar.close()
    if self.pool:
        self.pool.close()

    if not verbose:
        np.warnings.filterwarnings("default")

    log_probs = sampler.get_log_prob()[-self.tune :]
    chain = sampler.get_chain()[-self.tune :]
    chain = chain.reshape(-1, chain.shape[-1])

    arg_max = log_probs.argmax()
    mode_f = log_probs.flat[arg_max]
    mode_x = chain[arg_max].flatten()

    if temp == 1:

        self.fdict["mcmc_mode_x"] = mode_x
        self.fdict["mcmc_mode_f"] = mode_f

        if "mode_f" in self.fdict.keys() and mode_f < self.fdict["mode_f"]:
            print(
                "[mcmc:]".ljust(15, " ")
                + " New mode of %s is below old mode of %s. Rejecting..."
                % (mode_f, self.fdict["mode_f"])
            )
        else:
            self.fdict["mode_x"] = mode_x
            self.fdict["mode_f"] = mode_f

    self.fdict["datetime"] = str(datetime.now())

    return

def tmcmc(
    self,
    nsteps,
    nwalks,
    ntemps,
    target,
    update_freq=False,
    test_lprob=False,
    verbose=True,
    debug=False,
    **mcmc_args
):
    """Run Tempered Ensemble MCMC

    Parameters
    ----------
    ntemps : int
    target : float
    nsteps : float
    """

    from grgrlib.core import map2arr
    from .mpile import prior_sampler

    update_freq = update_freq if update_freq <= nsteps else False

    # sample pars from prior
    pars = prior_sampler(
        self, nwalks, test_lprob=test_lprob, verbose=max(verbose, 2 * debug)
    )

    x = get_par(self, "prior_mean", asdict=False, full=False, verbose=verbose > 1)

    pbar = tqdm.tqdm(total=ntemps, unit="temp(s)", dynamic_ncols=True)
    tmp = 0

    for i in range(ntemps):

        # update tmp
        ll = self.lprob(x)
        lp = self.lprior(x)

        tmp = tmp * (ntemps - i - 1) / (ntemps - i) + (target - lp) / (ntemps - i) / (
            ll - lp
        )
        aim = lp + (ll - lp) * tmp

        if tmp >= 1:
            # print only once
            pbar.write(
                "[tmcmc:]".ljust(15, " ")
                + "Increasing temperature to %s°. Too hot! I'm out..."
                % np.round(100 * tmp, 3)
            )
            pbar.update()
            self.temp = 1
            # skip for-loop to exit
            continue

        pbar.write(
            "[tmcmc:]".ljust(15, " ")
            + "Increasing temperature to %2.5f°, aiming @ %4.3f." % (100 * tmp, aim)
        )
        pbar.set_description("[tmcmc: %2.3f°" % (100 * tmp))

        self.mcmc(
            p0=pars,
            nsteps=nsteps,
            temp=tmp,
            update_freq=update_freq,
            verbose=verbose > 1,
            append=i,
            report=pbar.write,
            debug=debug,
            **mcmc_args
        )

        self.temp = tmp
        self.mcmc_summary(tune=int(nsteps / 10), calc_mdd=False, calc_ll_stats=True)

        pbar.update()

        pars = self.get_chain()[-1]
        lprobs_adj = self.get_log_prob()[-1]
        x = pars[lprobs_adj.argmax()]

    pbar.close()
    self.fdict["datetime"] = str(datetime.now())

    return pars


# -*- coding: utf-8 -*-
# Copyright (c) Vlachos Group, Jung Group
# GNU v3.0 license

__author__ = 'Geun Ho Gu'
__copyright__ = "Copyright 2019, Vlachos Group, Jung Group"
__version__ = "1.0"
__maintainer__ = "Geun Ho Gu"
__email__ = "googhgoo@hotmail.com"
__date__ = "July 31, 2019"

import tensorflow as tf
import numpy as np

__all__ = [
        'ModelBuilder',
        ]

class ModelBuilder(object):
    def __init__(self,n_occupancy,n_neighbor_sites_list,n_permutation_list,\
                 X_mean, X_std, Y_mean, Y_std,
                 n_conv=1,n_feature=150,sitewise_n_feature=25,
                 ):
        """
        Model builder.

        Parameters
        ----------
        n_occupancy: int. number of possible occupancy
        n_neighbor_sites_list: list of int. Number of neighbors of each site.
        X_mean : no. list of float. Mean of the one hot encoding matrix by column.
            no is the number of occupancy state.
        X_std : no. list of float. Standard deviation of the one hot encoding matrix by column.
        Y_mean : float. Mean of target property
        Y_std : float. Standard deviation of target property
        nconv: int. number of convolutions performed
        n_feature: int. number of feature for each site
        sitewise_n_feature: int. number of features for atoms for site-wise activation
        """
        tf.Graph()
        tf.reset_default_graph()
        # Not to be used within the graph.. Used in Trainer.
        if X_mean is None: # None is used when evaluation mode is activated
            self.X_mean = tf.get_variable('X_mean',
                shape=(n_occupancy,),trainable=False)
            self.X_std = tf.get_variable('X_std',
                shape=(n_occupancy,),trainable=False)
            self.Y_mean = tf.get_variable('Y_mean',
                shape=(),trainable=False)
            self.Y_std = tf.get_variable('Y_std',
                shape=(),trainable=False)
        else:
            self.X_mean = tf.get_variable('X_mean',
                initializer=tf.convert_to_tensor(np.array(X_mean,dtype=np.float32)),
                trainable=False)
            self.X_std = tf.get_variable('X_std',
                initializer=tf.convert_to_tensor(np.array(X_std,dtype=np.float32)),
                trainable=False)
            self.Y_mean = tf.get_variable('Y_mean',
                initializer=tf.convert_to_tensor(np.array(Y_mean,dtype=np.float32)),
                trainable=False)
            self.Y_std = tf.get_variable('Y_std',
                initializer=tf.convert_to_tensor(np.array(Y_std,dtype=np.float32)),
                trainable=False)

        self.Y = tf.placeholder(dtype=tf.float32, shape = [None])
        self.Idx_config = tf.placeholder(dtype=tf.int32, shape = [None,1])
        self.N_Sites = tf.placeholder(dtype=tf.int32, shape = ())
        self.N_Sites_per_config = tf.placeholder(dtype=tf.float32, shape = [None])
        self.X_Sites = tf.placeholder(dtype=tf.float32, shape = [None,n_occupancy])

        self.X_NSs = []
        for n_neighbor_sites,n_permutation in zip(n_neighbor_sites_list,n_permutation_list):
            self.X_NSs.append(tf.placeholder(dtype=tf.int32, shape = [None,n_permutation,n_neighbor_sites]))
        self.X_NSs = tuple(self.X_NSs)
        self.Lambda = tf.placeholder(dtype=tf.float32, shape = [])
        self.Dropout_rate = tf.placeholder(dtype=tf.float32, shape = [])
        self.is_training = tf.placeholder(dtype=tf.bool, shape = [])

        """ initiaite Model Helper """
        self.ModelHelper = CNNHelper(dropout_rate = self.Dropout_rate, UseBN=True,is_training=self.is_training)

        """ keeps tracks of site layer """
        self.SiteLayer = [self.X_Sites]

        """ Featurize one-hot-encoding """
        #self.SiteLayer.append(self.ModelHelper.Atom_Wise_Linear(self.X_Sites,n_feature))

        """ Perform Convolution """
        for i in range(0,n_conv):
            self.SiteLayer.append(self.ModelHelper.LCNNConvolution(self.SiteLayer[-1],self.X_NSs,self.N_Sites,depth=n_feature))
        self.SiteLayer.append(self.ModelHelper.Atom_Wise_Convolution(self.SiteLayer[-1],depth=sitewise_n_feature))

        """ Perform Linear Multiplication """
        self.Y_atom_wise = tf.expand_dims(tf.reduce_sum(self.ModelHelper.Atom_Wise_Linear(self.SiteLayer[-1]),axis=1),axis=1)

        """ Sum up contributions """
        self.Y_conf_wise = self.ModelHelper.Unflatten(self.Y_atom_wise,self.Idx_config,self.N_Sites_per_config)

        """ Loss + L2 norm """
        Params = self.ModelHelper.Params
        Params = tf.concat(Params,0)
        self.Loss = tf.losses.mean_squared_error(labels=self.Y, predictions=self.Y_conf_wise) + self.Lambda*tf.norm(Params)

    def GetStatistics(self):
        """Returns statistics for site layer and target property. Used in trainer in Train.py
        """
        return self.X_mean,self.X_std,self.Y_mean,self.Y_std

class CNNHelper(object):
    """This code is a helper for constructing CNN model.
    """
    def __init__(self,dropout_rate = tf.constant(1,dtype=tf.float32) ,UseBN = False, is_training = None,e = 1e-3,initializer=tf.contrib.layers.xavier_initializer()):
        """
        Initialize a CNN model builder. For the low level tensorflow models,
        keeping track of model weights can be difficult. This code automates
        this process.

        parameters
        ----------
        dropout_rate: 0D tensor. drop out rate
        UseBN: boolean . Flag for using batch normalization. Recommended to turn it on.
        is_training: 0D placeholder boolean tensor. Whether training or not. Used for batch norm
        e: float. This is a offset value for batch normalization for convergence
        initializer: tensorflow intializer class. tensorflow initializer for weights. Default(xavier) works well
        """
        if UseBN and is_training == None:
            raise NotImplemented('When Batch normalization is used, is_training boolean placeholder must be provided')
        self.e = e
        self.init = initializer
        self.UseBN = UseBN
        self._nCNNLayer = 0  # Number of CNN layers
        self.dropout_rate = dropout_rate
        self.Params = []
        self.is_training = is_training

    def LCNNConvolution(self,X_Sites, X_NSs, N_Sites, depth=1):
        """Construct convolution layer

        parameters
        ----------
        X_Sites: 2D tensor of shape (?,one hot encoding classification size)
            Site layer flattened over data index

        X_NSs: 3D tensor list of (?,number of permutation, 1 (site itself) + number of neighbors)
              given that the shape is (x,y,z), x is the site index which has been
              flattened over data, y is the permutation index, and z is the
              site itself index and neighbor site index.

        depth: int. depth/feature of convolution

        Return
        ----------
        X_Sites: 2D tensor
          Convoluted tensor
        """

        with tf.variable_scope("layer"+str(self._nCNNLayer)):
            indices = []
            updates = []
            for i in range(0,len(X_NSs)): # iterating over each site type

                with tf.variable_scope("site"+str(i)):
                    X_NSs_i = X_NSs[i]
                    # Apply permutations
                    # X_Sites_Permed : R^Number of Site x Permutation Index x Neighbor Indicex x Depth
                    X_Sites_Permed = tf.gather(X_Sites,X_NSs_i,axis=0)
                    # X_Sites_Concatenated : R^Number of Site x Permutation Index x (Neighbor Indicex x Depth)
                    X_Sites_Concatenated = tf.reshape(X_Sites_Permed,[-1,X_Sites_Permed.shape[1],X_Sites_Permed.shape[2]*X_Sites_Permed.shape[3]])
                    ## Filter shape: [filter height, in_channels, out_channels]
                    ## Stride shape: Filter height
                    w = tf.get_variable("w",[1,X_Sites_Concatenated.shape[2],depth],dtype=tf.float32,initializer=self.init)
                    self.Params.append(tf.reshape(w,[-1]))
                    b = tf.get_variable("b",[depth],dtype=tf.float32,initializer=self.init)
                    self.Params.append(tf.reshape(b,[-1]))

                    Conv1d = tf.nn.conv1d(X_Sites_Concatenated,w,1,"SAME")
                    Conv1dB = tf.nn.bias_add(Conv1d,b)
                    ## Batch Normalization
                    # See https://stackoverflow.com/questions/33949786/how-could-i-use-batch-normalization-in-tensorflow/38320613#38320613
                    if self.UseBN:
                        u = tf.get_variable("u",[1,1,depth],dtype=tf.float32,initializer=self.init)
                        v = tf.get_variable("v",[1,1,depth],dtype=tf.float32,initializer=self.init)
                        def updateuv():
                            ut,vt = tf.nn.moments(Conv1dB,[0,1],keep_dims=True)
                            return (u.assign(ut),v.assign(vt))
                        u,v = tf.cond(self.is_training,updateuv,lambda: (u,v))

                        o = tf.get_variable("o",[1,1,depth],dtype=tf.float32,initializer=self.init)
                        self.Params.append(tf.reshape(o,[-1]))
                        s = tf.get_variable("s",[1,1,depth],dtype=tf.float32,initializer=self.init)
                        self.Params.append(tf.reshape(s,[-1]))
                        Conv1dBN = tf.nn.batch_normalization(Conv1dB,u,v,o,s,self.e)
                    X_Sites1 = activation(Conv1dBN)
                    X_Sites2 = tf.nn.dropout(X_Sites1,self.dropout_rate,noise_shape=[1,1,depth]) # Dropout
                    X_Sites3 = tf.reduce_sum(X_Sites2,axis=1) # Sum over each permutations
                    indices.append(X_NSs_i[:,0,0])
                    updates.append(X_Sites3)
            # gather over multiple site types
            indices = tf.expand_dims(tf.concat(indices,axis=0),1)
            updates = tf.concat(updates,axis=0)
            NewSiteLayer = tf.scatter_nd(indices,updates,[N_Sites,depth])
            self._nCNNLayer += 1
            return NewSiteLayer

    def Atom_Wise_Linear(self,X_Sites, depth=1):
        """
        Perform Matrix Multiplication

        parameters
        ----------
        X_Sites: 2D tensor of shape (?,one hot encoding classification size)
            Site layer flattened over data index

        depth: int. depth/feature of convolution

        return
        ------
        X_Sites: 2D tensor. Convoluted tensor
        """
        with tf.variable_scope("layer"+str(self._nCNNLayer)):
            ## Filter shape: [filter height, in_channels, out_channels]
            ## Stride shape: Filter height
            w = tf.get_variable("w",[X_Sites.shape[1],depth],dtype=tf.float32,initializer=self.init)
            self.Params.append(tf.reshape(w,[-1]))
            b = tf.get_variable("b",[depth],dtype=tf.float32,initializer=self.init)
            self.Params.append(tf.reshape(b,[-1]))

            X_Sites = tf.matmul(X_Sites,w)
            X_Sites = tf.nn.bias_add(X_Sites,b)

            self._nCNNLayer += 1
            return X_Sites

    def Atom_Wise_Convolution(self,X_Sites, depth=1):
        """
        Perform self convolution to a site layer.

        parameters
        ----------
        X_Sites: 2D tensor of shape (?,one hot encoding classification size)
            Site layer flattened over data index
        depth: int. depth of convolution
        """

        with tf.variable_scope("layer"+str(self._nCNNLayer)):
            ## Filter shape: [filter height, in_channels, out_channels]
            ## Stride shape: Filter height
            w = tf.get_variable("w",[X_Sites.shape[1],depth],dtype=tf.float32,initializer=self.init)
            self.Params.append(tf.reshape(w,[-1]))
            b = tf.get_variable("b",[depth],dtype=tf.float32,initializer=self.init)
            self.Params.append(tf.reshape(b,[-1]))

            Conv = tf.matmul(X_Sites,w)
            Conv = tf.nn.bias_add(Conv,b)
            ## Batch Normalization
            if self.UseBN:
                u = tf.get_variable("u",[1,depth],dtype=tf.float32,initializer=self.init)
                v = tf.get_variable("v",[1,depth],dtype=tf.float32,initializer=self.init)
                def updateuv():
                    ut,vt = tf.nn.moments(Conv,[0],keep_dims=True)
                    return (u.assign(ut),v.assign(vt))
                u,v = tf.cond(self.is_training,updateuv,lambda: (u,v))

                o = tf.get_variable("o",[1,depth],dtype=tf.float32,initializer=self.init)
                self.Params.append(tf.reshape(o,[-1]))
                s = tf.get_variable("s",[1,depth],dtype=tf.float32,initializer=self.init)
                self.Params.append(tf.reshape(s,[-1]))
                Conv = tf.nn.batch_normalization(Conv,u,v,o,s,self.e)
            X_Sites = activation(Conv)
            X_Sites = tf.nn.dropout(X_Sites,self.dropout_rate,noise_shape=[1,depth]) # Dropout

            self._nCNNLayer += 1
            return X_Sites

    def SoftMax(self,X_Sites,depth = 1):
        """Perform Convolution to a site layer.

        parameters
        ----------
        X_Sites: 2D tensor of shape (?,one hot encoding classification size)
            Site layer flattened over data index
        depth: int. depth/features of the softmax

        return
        ------
        X_Sites: 2D tensor. softmaxed sitelayer
        """
        with tf.variable_scope("SoftMax"):
            w = tf.get_variable("w",[X_Sites.shape[1],depth],dtype=tf.float32,initializer=self.init)
            self.Params.append(tf.reshape(w,[-1]))
            b = tf.get_variable("b",[depth],dtype=tf.float32,initializer=self.init)
            self.Params.append(tf.reshape(b,[-1]))
            X_Sites = tf.matmul(X_Sites,w)
            X_Sites = tf.nn.bias_add(X_Sites,b)
            if self.UseBN:
                u = tf.get_variable("u",[1,depth],dtype=tf.float32,initializer=self.init)
                v = tf.get_variable("v",[1,depth],dtype=tf.float32,initializer=self.init)
                def updateuv():
                    ut,vt = tf.nn.moments(X_Sites,[0],keep_dims=True)
                    return (u.assign(ut),v.assign(vt))
                u,v = tf.cond(self.is_training,updateuv,lambda: (u,v))

                o = tf.get_variable("o",[1,depth],dtype=tf.float32,initializer=self.init)
                self.Params.append(tf.reshape(o,[-1]))
                s = tf.get_variable("s",[1,depth],dtype=tf.float32,initializer=self.init)
                self.Params.append(tf.reshape(s,[-1]))
                X_Sites = tf.nn.batch_normalization(X_Sites,u,v,o,s,self.e)
            X_Sites = tf.nn.softmax(X_Sites,1)
            return X_Sites

    def Unflatten(self,X_Sites,Idx_config,N_Sites):
        """This function unflattens the flattened site layer by summation. Then
        the unflattened values are dvidided by number of sites.

        parameters
        ----------
        X_Sites: 2D tensor. Site layer flattened over data index
        Idx_config: 1D tensor of integer. In the algorithm, Site layer is
            flattened over each data points. This way, we avoid using padding,
            and having an upper limit for the maximum number of sites. calculations
            are faster, too. To do this, we need data index for each site.
            This vector contains that information.

        N_Sites: 1D tensor. Each integer indicates number of sites in each
            configuration. Used to compute per site formation energy

        return
        ------
        ConfigLayer: 1D tensor. Unflattened per each datum

        """
        # All the other failed attempts to densifying X_Sites
#        ConfigLayer = tf.Variable(tf.zeros([Idx_config[-1]+1,X_Sites.shape[1]]), tf.float32)
#        ConfigLayer = tf.Variable(tf.zeros([N_Sites.get_shape()[0],X_Sites.shape[1]]), tf.float32)
#        ConfigLayer = tf.fill([Idx_config[-1]+1,X_Sites.shape[1]],0.0)
#        ConfigLayer = tf.scatter_add(ConfigLayer,Idx_config,X_Sites)
#        ConfigLayer = tf.scatter_nd(Idx_config,X_Sites,[Idx_config[-1]+1,X_Sites.shape[1]])
#        ConfigLayer = tf.scatter_nd(Idx_config,X_Sites,[N_Sites.shape[0],X_Sites.shape[1]])
#        ConfigLayer = tf.scatter_nd(Idx_config,X_Sites,[Idx_config[-1]+1,tf.constant([1])])
#        ConfigLayer = tf.scatter_nd(Idx_config,X_Sites,[Idx_config[-1]+1,[1]])
#        ConfigLayer = tf.scatter_nd(Idx_config,X_Sites,Idx_config[-1]+1)
#        ConfigLayer = tf.scatter_nd(Idx_config,X_Sites,[1,1])
        ConfigLayer = tf.scatter_nd(Idx_config,X_Sites,[tf.reshape(Idx_config[-1]+1,[]),X_Sites.shape[1]])
        ConfigLayer = tf.reshape(ConfigLayer,[-1])
        ConfigLayer = tf.div(ConfigLayer,N_Sites)
        return ConfigLayer

def activation(tensor):
    """Return activation function"""
    return tf.nn.softplus(tensor)-np.log(2.0)


# -*- coding: utf-8 -*-
"""QLearning.ipynb

Automatically generated by Colaboratory.
"""

#remove " > /dev/null 2>&1" to see what is going on under the hood
# !pip install gym pyvirtualdisplay > /dev/null 2>&1
# !apt-get install -y xvfb python-opengl ffmpeg > /dev/null 2>&1

# !apt-get update > /dev/null 2>&1
# !apt-get install cmake > /dev/null 2>&1
# !pip install --upgrade setuptools 2>&1
# !pip install ez_setup > /dev/null 2>&1
# !pip install gym[atari] > /dev/null 2>&1

import gym
from gym import logger as gymlogger
from gym.wrappers import Monitor
gymlogger.set_level(40) #error only
import numpy as np
import random
import matplotlib
import matplotlib.pyplot as plt
# %matplotlib inline
import glob
import io
import base64
from IPython.display import HTML

from IPython import display as ipythondisplay

from pyvirtualdisplay import Display
display = Display(visible=0, size=(1400, 900))
display.start()

"""
Utility functions to enable video recording of gym environment and displaying it
To enable video, just do "env = wrap_env(env)""
"""

def show_video():
  mp4list = glob.glob('video/*.mp4')
  if len(mp4list) > 0:
    mp4 = mp4list[0]
    video = io.open(mp4, 'r+b').read()
    encoded = base64.b64encode(video)
    ipythondisplay.display(HTML(data='''<video alt="test" autoplay
                loop controls style="height: 400px;">
                <source src="data:video/mp4;base64,{0}" type="video/mp4" />
             </video>'''.format(encoded.decode('ascii'))))
  else:
    print("Could not find video")

def wrap_env(env):
  env = Monitor(env, './video', force=True)
  return env

env = wrap_env(gym.make('MountainCar-v0'))
env.reset()

#check out the Mountain Car action space!
print(env.action_space)

def QLearning(env, learning, discount, epsilon, min_eps, episodes):
    # Determine size of discretized state space
    num_states = (env.observation_space.high - env.observation_space.low)*\
                    np.array([10, 100])
    num_states = np.round(num_states, 0).astype(int) + 1

    # Initialize Q table
    Q = np.random.uniform(low = -1, high = 1,
                          size = (num_states[0], num_states[1],
                                  env.action_space.n))

    # Initialize variables to track rewards
    reward_list = []
    ave_reward_list = []

    # Calculate episodic reduction in epsilon
    reduction = (epsilon - min_eps)/episodes

    # Run Q learning algorithm
    for i in range(episodes):
        # Initialize parameters
        done = False
        tot_reward, reward = 0,0
        state = env.reset()

        # Discretize state
        state_adj = (state - env.observation_space.low)*np.array([10, 100])
        state_adj = np.round(state_adj, 0).astype(int)

        while done != True:
            # Render environment for last five episodes
            if i >= (episodes - 20):
                env.render()

            # Determine next action - epsilon greedy strategy
            if np.random.random() < 1 - epsilon:
                action = np.argmax(Q[state_adj[0], state_adj[1]])
            else:
                action = np.random.randint(0, env.action_space.n)

            # Get next state and reward
            state2, reward, done, info = env.step(action)

            # Discretize state2
            state2_adj = (state2 - env.observation_space.low)*np.array([10, 100])
            state2_adj = np.round(state2_adj, 0).astype(int)

            #Allow for terminal states
            if done and state2[0] >= 0.5:
                Q[state_adj[0], state_adj[1], action] = reward

            # Adjust Q value for current state
            else:
                delta = learning*(reward +
                                 discount*np.max(Q[state2_adj[0],
                                                   state2_adj[1]]) -
                                 Q[state_adj[0], state_adj[1],action])
                Q[state_adj[0], state_adj[1],action] += delta

            # Update variables
            tot_reward += reward
            state_adj = state2_adj

        # Decay epsilon
        if epsilon > min_eps:
            epsilon -= reduction

        # Track rewards
        reward_list.append(tot_reward)

        if (i+1) % 100 == 0:
            ave_reward = np.mean(reward_list)
            ave_reward_list.append(ave_reward)
            reward_list = []

        if (i+1) % 100 == 0:
            print('Episode {} Average Reward: {}'.format(i+1, ave_reward))

    env.close()

    return ave_reward_list

# Run Q-learning algorithm
rewards = QLearning(env, 0.2, 0.9, 0.8, 0, 10000)

# Plot Rewards
plt.plot(100*(np.arange(len(rewards)) + 1), rewards)
plt.xlabel('Episodes')
plt.ylabel('Average Reward')
plt.title('Average Reward vs Episodes')
plt.savefig('rewards.jpg')
plt.show()
plt.close()
show_video()


import os

import sympy.physics.units.definitions as u
from sympy.physics.units.quantities import Quantity
from sympy.physics.units.systems.si import dimsys_SI

from mathpad.val import Val
import mathpad.dimensions as dims

HERE = os.path.dirname(__file__)

AUTOGEN_WARNING = '"WARNING: this file was automatically generated by commented code within mathpad.py, and then modified slightly"\n'

print("writing to " + f"{HERE}/units.py\nand {HERE}/constants.py")

with open(f"{HERE}/units.py", "w") as units_f, open(
    f"{HERE}/constants.py", "w"
) as constants_f:

    # top of 'units.py'
    units_f.write(AUTOGEN_WARNING)
    units_f.write("import sympy.physics.units.definitions.unit_definitions as u\n")
    units_f.write("from mathpad.physical_quantities import *\n")

    # top of 'constants.py'
    constants_f.write(AUTOGEN_WARNING)
    constants_f.write("import sympy.physics.units.definitions.unit_definitions as u\n")
    constants_f.write("from mathpad.physical_quantities import *\n")

    for name in dir(u):
        qty = getattr(u, name)

        if name.endswith("constant"):
            constants_f.write(
                f"""
{name} = OutputVal(u.{name})"""
            )

        elif isinstance(qty, Quantity):

            for quantity_cls_name in dir(dims):
                if quantity_cls_name in "du":
                    continue

                quantity_cls = getattr(dims, quantity_cls_name)

                # handle the Angle dimension specially
                if str(qty).startswith("angular") or str(qty).startswith("rad"):
                    units_f.write(
                        f"""
{name} = Angle(u.{name})"""
                    )
                    break

                elif (
                    isinstance(quantity_cls, type)
                    and issubclass(quantity_cls, Val)
                    and quantity_cls is not Val
                    and dimsys_SI.equivalent_dims(qty.dimension, quantity_cls.dimension)
                ):
                    units_f.write(
                        f"""
{name} = {quantity_cls_name}(u.{name})"""
                    )
                    break
            else:
                print(f"Warning: Quantity '{name}' could not be matched; skipping...")


"""
Core for the backend of the project.
All the computations on the Neural Networks are done here.
Author: Antonio Strippoli
"""
# General imports
import os
import sys
import PIL.Image as pil
import numpy as np

import matplotlib

# Set non interactive backend for matplolib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

# Project's imports
from flaskr import paths

# TensorFlow & Keras
import tensorflow as tf
from tensorflow.keras.models import Model

# TF-Keras-Vis
from tf_keras_vis.gradcam import GradcamPlusPlus
from tf_keras_vis.utils import normalize

# Utilities
from .utils import np_array_to_img_rgb
from utils import pickle_load
from capslayers import compute_vectors_length

def load_model(model_name):
    """
    """
    # Import model (HACK)
    sys.path.append(os.path.join(paths["trainer"], model_name))
    from capsnet import CapsuleNet

    sys.path.pop()

    # Load model
    try:
        model_params = pickle_load(
            os.path.join(paths["trainer"], model_name, "outs", "model_params.pkl")
        )
    except FileNotFoundError:
        return None, None

    dataset = model_params.pop("dataset")
    _, model = CapsuleNet(**model_params)
    model_params["dataset"] = dataset

    # Clean modules (HACK)
    del sys.modules["capsnet"]

    return model, model_params

def get_processable_layers(layers):
    """
    Returns computable layers from a list of layers, along with their types.
    We check that by searching specific keywords in layer's name,
    since we can't be sure about layer's type.

    If you wish to extend the project and add compatibility
    for more layers, you should as well update this function.
    """
    processable_layers = []

    for layer in layers:
        layer_name = layer.name.lower()

        if "conv" in layer_name:
            processable_layers.append([layer_name, "CONVOLUTIONAL"])
        elif "primary" in layer_name and "caps" in layer_name:
            processable_layers.append([layer_name, "PRIMARY_CAPS"])
        elif "caps" in layer_name:
            processable_layers.append([layer_name, "DENSE_CAPS"])
        elif "mask" in layer_name:
            processable_layers.append([layer_name, "MASK"])

    return processable_layers

def conv_out_process(act, img_mode, layer_img_dir):
    # Extract from batch
    act = act[0]

    # Prepare new image to contain convolutional's features
    new_img_width = act.shape[0] * act.shape[-1]
    new_img_height = act.shape[1]
    new_img = pil.new("RGB", (new_img_width, new_img_height))

    # Save features
    act = tf.multiply(act, 255.0)
    for i in range(act.shape[-1]):
        feature = act[:, :, i].numpy().astype("int8")
        feature_image = pil.fromarray(np.uint8(matplotlib.cm.viridis(feature) * 255))

        new_img.paste(
            feature_image, (act.shape[0] * i, 0),
        )

    # Save the new image
    new_img.save(os.path.join(layer_img_dir, f"out.jpeg"))

    # Save filters (TODO?)
    pass

    return {
        "Activations after ReLU": {
            "filename": "out.jpeg",
            "rows": 1,
            "cols": act.shape[-1],
            "chunk_width": act.shape[0],
            "chunk_height": act.shape[1],
        }
    }

def pcap_out_process(act, prep_img, img_mode, layer, layer_conf, layer_img_dir):
    # Extract from batch
    act = act[0]
    prep_img = prep_img[0]

    # Move preprocessed image to RGB domain (if not)
    prep_img = np_array_to_img_rgb(prep_img, img_mode)

    # Get new features' dimension
    feature_dim = int(
        (layer.input_shape[1] - layer_conf["kernel_size"] + 1) / layer_conf["strides"]
    )
    # Compute vectors' length and reshape
    act = tf.reshape(compute_vectors_length(act), (feature_dim, feature_dim, -1))

    # Prepare new image to contain capsules' activations
    chunk_width, chunk_height = prep_img.size
    new_img_width = chunk_width * act.shape[-1]
    new_img_height = chunk_height
    new_img = pil.new("RGB", (new_img_width, new_img_height))

    # Save as images with the same size of the inputs
    act = tf.multiply(act, 255)
    for i in range(act.shape[-1]):
        capsules_length = act[:, :, i].numpy()
        capsules_length = np.interp(
            capsules_length, (capsules_length.min(), capsules_length.max()), (0.0, 1.0),
        )
        heatmap = pil.fromarray(np.uint8(matplotlib.cm.jet(capsules_length) * 255))

        # If we stop here, we can get a visualized matrix showing capsules activations,
        # but we go further, rescaling and superimposing the heatmap to the original image
        heatmap = heatmap.resize((chunk_width, chunk_height))
        pcaps_img = pil.blend(prep_img, heatmap, alpha=0.7)

        new_img.paste(
            pcaps_img, (chunk_width * i, 0),
        )

    # Save the new image
    new_img.save(os.path.join(layer_img_dir, f"out.jpeg"))

    return {
        "Capsules' length as heatmap": {
            "filename": "out.jpeg",
            "rows": 1,
            "cols": act.shape[-1],
            "chunk_width": chunk_width,
            "chunk_height": chunk_height,
        }
    }

def dcap_out_process(prev_layer_pack, layer_pack, prep_img, img_mode, layer_img_dir):
    # --------------------------
    # The following is a compact version of routing path visualization. All the credits goes to Aman Bhullar:
    # https://atrium.lib.uoguelph.ca/xmlui/bitstream/handle/10214/17834/Bhullar_Aman_202003_Msc.pdf?sequence=1&isAllowed=y
    # --------------------------

    # TODO: Implement the complete routing path visualization.
    # Also, atm we assume that the previous layer is always a primary caps, which is not always be the case

    # Unpack variables
    ((pl, _), prev_act) = prev_layer_pack
    ((cl, _), curr_act) = layer_pack

    pl_conf = pl.get_config()
    cl_conf = cl.get_config()

    # Get new features' dimension
    feature_dim = int(
        (pl.input_shape[1] - pl_conf["kernel_size"] + 1) / pl_conf["strides"]
    )

    # HACK: Assume that every input image to the network is squared,
    # which would produce squared features as well
    dims = [
        feature_dim,
        feature_dim,
        pl_conf["n_caps"],
        cl_conf["n_caps"],
    ]

    # Prepare output of previous caps layer
    tmp_dims = dims[0:3] + [1] * (len(dims) - 3)
    prev_caps_lengths = tf.reshape(compute_vectors_length(prev_act), tmp_dims)
    tmp_dims = [1] * (3) + dims[3 : len(dims)]
    prev_caps_lengths_tiled = tf.tile(prev_caps_lengths, tmp_dims)

    # Prepare routing weights
    tmp_dims = dims[0:4] + [1] * (len(dims) - 4)
    routing_weights_reshape = tf.reshape(curr_act[1], tmp_dims)
    tmp_dims = [1] * (4) + dims[4 : len(dims)]
    routing_weights_reshape_tiled = tf.tile(routing_weights_reshape, tmp_dims)

    # Prepare output of current caps layer
    tmp_dims = [1, 1, 1, dims[3]] + [1] * (len(dims) - 4)
    curr_caps_lengths = tf.reshape(compute_vectors_length(curr_act[0]), tmp_dims)
    tmp_dims = dims[0:3] + [1] + dims[4 : len(dims)]
    curr_caps_lengths_tiled = tf.tile(curr_caps_lengths, tmp_dims)

    # Calculate routing path visualization
    tmp = tf.multiply(routing_weights_reshape_tiled, curr_caps_lengths_tiled)
    all_paths = tf.multiply(prev_caps_lengths_tiled, tmp)
    all_paths_average = tf.reduce_sum(all_paths, axis=2)

    # Preprocess input image
    input_img_width, input_img_height = prep_img.shape[1:3]
    prep_img = prep_img[0]

    # Move preprocessed image to RGB domain (if not)
    prep_img = np_array_to_img_rgb(prep_img, img_mode)

    # Prepare new image to contain capsules' activations
    new_img = pil.new("RGB", (input_img_width * dims[-1], input_img_height))

    # Save outputs as a single image
    all_paths_average = all_paths_average.numpy()
    all_paths_average = np.interp(
        all_paths_average,
        (all_paths_average.min(), all_paths_average.max()),
        (0.0, 1.0),
    )
    for i in range(all_paths_average.shape[-1]):
        rpv = all_paths_average[:, :, i]
        heatmap = pil.fromarray(np.uint8(matplotlib.cm.jet(rpv) * 255))

        # Rescale and blend
        heatmap = heatmap.resize((input_img_width, input_img_height))
        dcaps_img = pil.blend(prep_img, heatmap, alpha=0.7)

        new_img.paste(dcaps_img, (i * input_img_width, 0))

    # Save the new image
    new_img.save(os.path.join(layer_img_dir, f"out.jpeg"))

    return {
        "Routing Path Visualization": {
            "filename": "out.jpeg",
            "rows": 1,
            "cols": dims[-1],
            "chunk_width": input_img_width,
            "chunk_height": input_img_height,
        }
    }

def dcap_gradcam(model, layer, prep_img, img_mode, layer_img_dir):
    # Edit model to stop at the current layer
    def model_modifier(model):
        return tf.keras.Model(inputs=model.inputs, outputs=layer.output)

    # Instantiate gradcam
    gradcam = GradcamPlusPlus(model, model_modifier, clone=False)

    def loss(output):
        """Maximize the prediction with greater probability"""
        output = compute_vectors_length(output)
        i = np.argmax(output, 1)[0]
        return output[0, i]

    # Compute gradcam
    cam = normalize(
        gradcam(
            [loss, lambda x: tf.convert_to_tensor(0.0)],
            prep_img,
            penultimate_layer=layer,
        )
    )
    heatmap = pil.fromarray(np.uint8(matplotlib.cm.jet(cam[0]) * 255))

    # Move preprocessed image to RGB domain (if not)
    prep_img = prep_img[0]
    prep_img = np_array_to_img_rgb(prep_img, img_mode)

    # Superimpose heatmap
    gradcam_final = pil.blend(prep_img, heatmap, alpha=0.7)

    # Save the new image
    gradcam_final.convert("RGB").save(os.path.join(layer_img_dir, f"out.jpeg"))

    return {"filename": "out.jpeg"}

def mask_out_process(act, img_mode, layer, model_params, layer_img_dir):
    # Extract from batch and convert to numpy
    act = act[0]
    act_numpy = act.numpy()

    # Get next layer ( suppose decoder (HACK) )
    next_layer = layer._outbound_nodes[1].outbound_layer

    # Prepare dictionary for outputs
    out = {}

    # === MANIPULATIONS ON MAGNITUDE ===
    n_manip = 11

    # Prepare new output image
    new_img_width = model_params["input_shape"][0] * n_manip
    new_img_height = model_params["input_shape"][1]
    new_img = pil.new(img_mode, (new_img_width, new_img_height))

    # Edit tensor values and feed the activation forward to get reconstruction
    for i_r, r in enumerate(np.linspace(0, 1, n_manip)):
        r = round(r, 1)
        # Rescale magnitude value
        act_to_feed = np.copy(act_numpy)
        act_to_feed *= r / compute_vectors_length(act_to_feed)

        # Convert back to tensor and feed-forward
        act_to_feed = tf.expand_dims(tf.convert_to_tensor(act_to_feed), 0)
        reconstructed_image = next_layer(act_to_feed)

        # Pre-process reconstruction
        reconstructed_image = tf.squeeze(reconstructed_image)
        reconstructed_image = tf.multiply(reconstructed_image, 255.0)
        reconstructed_image = reconstructed_image.numpy().astype("int8")

        # Convert to image and paste to the new image
        reconstructed_image = pil.fromarray(reconstructed_image, mode=img_mode)
        new_img.paste(
            reconstructed_image, (model_params["input_shape"][0] * i_r, 0,),
        )

    # Save the new image
    new_img.save(os.path.join(layer_img_dir, f"magnitude.jpeg"))

    # Prepare partial output
    out["1) Manipulations on magnitude"] = {
        "filename": "magnitude.jpeg",
        "rows": 1,
        "cols": n_manip,
        "chunk_width": model_params["input_shape"][0],
        "chunk_height": model_params["input_shape"][1],
    }

    # === MANIPULATIONS ON DIMENSIONS ===
    # Prepare variables for the iterations
    start = np.argmax(act != 0)
    n_dims = act.shape[0] // model_params["n_class"]
    n_manip = 11

    # Prepare new output image
    new_img_width = model_params["input_shape"][0] * n_manip
    new_img_height = model_params["input_shape"][1] * n_dims
    new_img = pil.new(img_mode, (new_img_width, new_img_height))

    # Edit tensor values and feed the activation forward to get reconstruction
    for i_dim in range(start, start + n_dims):
        for i_r, r in enumerate(np.linspace(-0.25, 0.25, n_manip)):
            r = round(r, 2)
            # Edit dimension value
            act_to_feed = np.copy(act_numpy)
            act_to_feed[i_dim] += r

            # Convert back to tensor and feed-forward
            act_to_feed = tf.expand_dims(tf.convert_to_tensor(act_to_feed), 0)
            reconstructed_image = next_layer(act_to_feed)

            # Pre-process reconstruction
            reconstructed_image = tf.squeeze(reconstructed_image)
            reconstructed_image = tf.multiply(reconstructed_image, 255.0)
            reconstructed_image = reconstructed_image.numpy().astype("int8")

            # Convert to image and paste to the new image
            reconstructed_image = pil.fromarray(reconstructed_image, mode=img_mode)
            new_img.paste(
                reconstructed_image,
                (
                    model_params["input_shape"][0] * i_r,
                    model_params["input_shape"][1] * (i_dim - start),
                ),
            )

    # Save the new image
    new_img.save(os.path.join(layer_img_dir, f"dimensions.jpeg"))

    out["2) Manipulations on dimensions"] = {
        "filename": "dimensions.jpeg",
        "rows": n_dims,
        "cols": n_manip,
        "chunk_width": model_params["input_shape"][0],
        "chunk_height": model_params["input_shape"][1],
    }

    return out

def model_out_process(predictions, model_out_dir):
    """
    Visualize model's predictions in an histogram.
    """
    plt.rcdefaults()
    fig, ax = plt.subplots()

    classes = np.arange(len(predictions))

    ax.barh(classes, predictions, align="center")
    ax.set_yticks(classes)
    ax.set_yticklabels(classes)
    ax.set_xlim(0, 1.0)

    plt.title(
        f"Prediction: {np.argmax(predictions)}",
        fontsize=30.0,
        color="blue",
        fontweight="bold",
    )
    ax.invert_yaxis()
    ax.spines["right"].set_visible(False)
    ax.spines["top"].set_visible(False)
    ax.tick_params(axis="both", colors="blue", labelsize=20.0)

    for i, v in enumerate(predictions):
        ax.text(
            v + 0.01,
            i + 0.06,
            str(round(v, 4)),
            color="blue",
            fontweight="bold",
            verticalalignment="center",
            size=15.0,
        )

    plt.savefig(os.path.join(model_out_dir, "out.png"), transparent=True)
    plt.close("all")
    return {"filename": "out.png"}

def compute_step(model, model_params, prep_img, img_mode, req_out_dir):
    """
    Computes and saves the outputs of the processable layers in the given model.
    """
    # Create a temporary model to retrieve all the layers' activations
    processable_layers = get_processable_layers(model.layers)
    processable_layers = [(model.get_layer(pl[0]), pl[1]) for pl in processable_layers]

    activation_model = Model(model.input, [pl[0].output for pl in processable_layers])
    activations = activation_model(prep_img)

    # Prepare dictionary to contain outputs
    out_dir = req_out_dir.replace("\\", "/")
    out_dir = out_dir[out_dir.find("/static") :] + "/"

    out_info = {"out_dir": out_dir}
    outs = {}

    layer_pack = list(zip(processable_layers, activations))
    for i, ((layer, layer_type), act) in enumerate(layer_pack):
        # Get layer config
        layer_conf = layer.get_config()

        # Prepare directory for the layer
        layer_name = layer_conf["name"]
        layer_img_dir = os.path.join(req_out_dir, layer_name)
        os.mkdir(layer_img_dir)

        # === PROCESS LAYER ACTIVATION BASED ON TYPE ===
        if layer_type == "CONVOLUTIONAL":
            outs[layer_name] = conv_out_process(act, img_mode, layer_img_dir)
        elif layer_type == "PRIMARY_CAPS":
            outs[layer_name] = pcap_out_process(
                act, prep_img, img_mode, layer, layer_conf, layer_img_dir
            )
        elif layer_type == "DENSE_CAPS":
            curr_lp = layer_pack[i]
            prev_lp = layer_pack[i - 1]
            outs[layer_name] = dcap_out_process(
                prev_lp, curr_lp, prep_img, img_mode, layer_img_dir
            )

            # Compute GradCAM++
            gradcam_dir = os.path.join(req_out_dir, "gradcam")
            os.mkdir(gradcam_dir)
            out_info["gradcam"] = dcap_gradcam(
                model, layer, prep_img, img_mode, gradcam_dir
            )
        elif layer_type == "MASK":
            outs[layer_name] = mask_out_process(
                act, img_mode, layer, model_params, layer_img_dir
            )
        else:
            raise TypeError(f"Layer type '{layer_type}' cannot be computed.")

    # Add layers outputs
    out_info["layers_outs"] = outs

    # Process model's output
    model_out_dir = os.path.join(req_out_dir, "model_out")
    os.mkdir(model_out_dir)

    model_outs = model(prep_img)
    if type(model_outs) == list:
        predictions = model_outs[0][0].numpy()
    else:
        predictions = model_outs[0].numpy()

    out_info["model_out"] = model_out_process(predictions, model_out_dir)

    return out_info


'''
Personalised argumentation paper: are user features are required or latent variables sufficient/correlated with observed features?
Is prior stance a useful user feature for predicting belief change? It should be, since a user can only be convinced by
an argument if they did not previously believe in it.
Topic-specific nature means predictions based on linguistic features are likely to be weak?

Created on 19 Jun 2017

@author: simpson
'''
import os
import sys
import logging
import matplotlib.pyplot as plt # do this here so we don't get the debugging crap later from the logger

from sklearn.gaussian_process.gpr import GaussianProcessRegressor
from sklearn.gaussian_process.kernels import Matern
from sklearn.metrics import accuracy_score

logging.basicConfig(level=logging.DEBUG)

sys.path.append("./python/analysis/habernal_comparison")

from tests import TestRunner
from gp_regressor_svi import GPRegressorSVI
from collab_pref_learning_svi import CollabPrefLearningSVI
import numpy as np

nfactors = 50
max_Kw_size = 2000

rate_s = 200

rate_sy0 = 10

delay = 0 # some default will be used depending on the function unless this is set to > 0

class PersonalisedTestRunner(TestRunner):

    def run_crowd_bt(self):

        nitems = self.items_feat.shape[0]
        workers = np.unique(self.person_train)
        nworkers = np.max(workers) + 1

        scales = [1]  # [0.01, 0.1, 1, 10, 100]  # 10 was found to be optimal on the random selection tests.

        tr_acc_best = 0

        for scale in scales:

            # initialise variational parameters
            Es = np.zeros(nitems)
            Eeta = np.ones(nworkers) * 0.9
            sigma = np.ones(nitems) * scale
            alpha = np.ones(nworkers) * 9
            beta = np.ones(nworkers)

            balance = 0#1e-6 # tiny amount to ensure numerical stability

            for pair_idx in range(len(self.a1_train)):

                # get the indices
                a1 = self.a1_train[pair_idx]
                a2 = self.a2_train[pair_idx]

                if self.prefs_train[pair_idx] == 1:
                    continue

                if self.prefs_train[pair_idx] != 2:  # swap so a2 is the preferred one
                    tmp = a1
                    a1 = a2
                    a2 = tmp

                k = self.person_train[pair_idx]

                # update the means
                prob_incr = alpha[k] * np.exp(Es[a1]) / (alpha[k] * np.exp(Es[a1]) + beta[k] * np.exp(Es[a2]) + balance) \
                             - np.exp(Es[a1]) / (np.exp(Es[a1]) + np.exp(Es[a2]) + balance)
                Es[a1] = Es[a1] + sigma[a1] ** 2 * prob_incr
                Es[a2] = Es[a2] - sigma[a2] ** 2 * prob_incr

                var_diff = alpha[k]  *np.exp(Es[a1]) * beta[k] * np.exp(Es[a2]) / \
                           ((alpha[k] * np.exp(Es[a1]) + beta[k] * np.exp(Es[a2]))**2 + balance) \
                           - np.exp(Es[a1]) * np.exp(Es[a2]) / ((np.exp(Es[a1]) + np.exp(Es[a2]))**2 + balance)
                sigma[a1] = np.sqrt(sigma[a1] ** 2 * np.max([1 + sigma[a1] ** 2 * (var_diff), 10e-4]))
                sigma[a2] = np.sqrt(sigma[a2] ** 2 * np.max([1 + sigma[a2] ** 2 * (var_diff), 10e-4]))

                C1 = np.exp(Es[a1]) / (np.exp(Es[a1]) + np.exp(Es[a2]) + balance) \
                     + 0.5 * (sigma[a1] ** 2 + sigma[a2] ** 2) \
                     * np.exp(Es[a1]) * np.exp(Es[a2]) * (np.exp(Es[a2]) - np.exp(Es[a1])) \
                     / (np.exp(Es[a1]) + np.exp(Es[a2]) + balance) ** 3

                C2 = 1 - C1

                C = (C1 * alpha[k] + C2 * beta[k]) / (alpha[k] + beta[k] + balance)  # normalisation constant for p( 1 > 2 | worker k)

                Eeta[k] = (C1 * (alpha[k] + 1) * alpha[k] + C2 * alpha[k] * beta[k]) / (
                            C * (alpha[k] + beta[k] + 1) * (alpha[k] + beta[k]) + balance)
                Eeta_sq_k = (C1 * (alpha[k] + 2) * (alpha[k] + 1) * alpha[k] + C2 * (alpha[k] + 1) * alpha[k] * beta[k]) / \
                            (C * (alpha[k] + beta[k] + 2) * (alpha[k] + beta[k] + 1) * (alpha[k] + beta[k]) + balance)

                alpha[k] = (Eeta[k] - Eeta_sq_k) * Eeta[k] / (Eeta_sq_k - Eeta[k] ** 2 + balance)
                beta[k] = (Eeta[k] - Eeta_sq_k) * (1 - Eeta[k]) / (Eeta_sq_k - Eeta[k] ** 2 + balance)

                if np.mod(pair_idx, 1000) == 0:
                    print('Learning crowdBT, iteration %i' % pair_idx)

            if np.any(np.isnan(Es)):
                continue

            tr_proba = np.exp(Es[self.a1_train]) / (np.exp(Es[self.a1_train]) + np.exp(Es[self.a2_train]) + balance)
            tr_acc = accuracy_score(self.prefs_train[self.prefs_train != 1]==2, np.round(tr_proba[self.prefs_train != 1]))
            print('training set accuracy = %f with scale %f' % (tr_acc, scale) )
            if tr_acc > tr_acc_best:
                Es_best = Es
                scale_best = scale
                tr_acc_best = tr_acc

        Es = Es_best

        print('Completed online learning of crowd BT. Found best scale is %f' % scale_best)

        proba = np.exp(Es[self.a1_test]) / (np.exp(Es[self.a1_test]) + np.exp(Es[self.a2_test]) + balance)

        self.crowdBT_sigma = sigma
        self.crowdBT_s = Es

        scores = Es[self.a_rank_test]

        tr_proba = np.exp(Es[self.a1_unseen]) / (
                    np.exp(Es[self.a1_unseen]) + np.exp(Es[self.a2_unseen]) + balance)

        return proba, scores, tr_proba, Es[self.a_rank_train]

    def run_crowd_bt_gpr(self):

        # we first train crowd_bt as above. Then, we use the scores for items that were compared in training
        # to train a GP regression model. The GP then predicts the scores of all items. This means we can generalise
        # from the training items to all items, plus the GP will do some smoothing over the training items in case they
        # had sparse noisy data.

        proba, predicted_f, tr_proba, tr_f = self.run_crowd_bt()

        if 'additive' in self.method:
            kernel_combination = '+'
        else:
            kernel_combination = '*'

        if 'shrunk' in self.method:
            ls_initial = self.ls_initial / float(len(self.ls_initial))
        else:
            ls_initial = self.ls_initial

        if 'weaksprior' in self.method:
            shape_s0 = 2.0
            rate_s0 = 200.0
        elif 'lowsprior' in self.method:
            shape_s0 = 1.0
            rate_s0 = 1.0
        elif 'weakersprior' in self.method:
            shape_s0 = 2.0
            rate_s0 = 2000.0
        else:
            shape_s0 = 200.0
            rate_s0 = 20000.0

        if '_M' in self.method:
            validx = self.method.find('_M') + 2
            M = int(self.method[validx:])
        else:
            M = 500

        if '_SS' in self.method:
            validx = self.method.find('_SS') + 3
            SS = int(self.method[validx:])
        else:
            SS = 200

        self.model = GPRegressorSVI(ninput_features=self.ndims, ls_initial=ls_initial, verbose=self.verbose,
                                    shape_s0=2, rate_s0=200, rate_ls=1.0 / np.mean(ls_initial),
                                    use_svi=True,
                                    ninducing=M, max_update_size=SS, kernel_combination=kernel_combination,
                                    forgetting_rate=0.7,
                                    delay=1.0)
        self.model.max_iter_VB = 2  #00
        new_items_feat = self.items_feat  # pass only when initialising

        print("no. features: %i" % new_items_feat.shape[1])
        self.model.fit(self.items_feat, self.crowdBT_s, obs_noise=self.crowdBT_sigma ** 2)

        predicted_f, _ = self.model.predict_f()  #self.model.obs_f

        balance = 0
        proba = np.exp(predicted_f[self.a1_test]) / (
                    np.exp(predicted_f[self.a1_test]) + np.exp(predicted_f[self.a2_test]) + balance)

        f = predicted_f[self.a_rank_test]

        tr_f = predicted_f[self.a_rank_train]

        tr_proba = np.exp(predicted_f[self.a1_unseen]) / (
                    np.exp(predicted_f[self.a1_unseen]) + np.exp(predicted_f[self.a2_unseen]) + balance)

        return proba, f, tr_proba, tr_f

    def _train_persgppl(self, delay):
        common_mean = False

        if '_commonmean' in self.method:
            common_mean = True

        if 'weaksprior' in self.method:
            shape_s0 = 2.0
            rate_s0 = rate_s #200.0
        elif 'lowsprior' in self.method:
            shape_s0 = 1.0
            rate_s0 = 1.0
        elif 'weakersprior' in self.method:
            shape_s0 = 2.0
            rate_s0 = 2000.0
        else:
            shape_s0 = 200.0
            rate_s0 = 20000.0

        if '_M' in self.method:
            validx = self.method.find('_M') + 2
            M = int(self.method[validx:].split('_')[0])
        else:
            M = 500

        if M == 0:
            M = self.items_feat.shape[0]

        if '_F' in self.method:
            valididx = self.method.find('_F') + 2
            F = int(self.method[valididx:].split('_')[0])
        else:
            F = nfactors

        if '_SS' in self.method:
            validx = self.method.find('_SS') + 3
            SS = int(self.method[validx:])

            niter = 200 #* (200.0 / float(SS))
        else:
            SS = 200
            niter = 200

        self.model = CollabPrefLearningSVI(nitem_features=self.ndims, ls=self.ls_initial, verbose=self.verbose,
                                           nfactors=F, rate_ls=1.0 / np.mean(self.ls_initial),
                                           use_common_mean_t=common_mean, max_update_size=SS, use_lb=True,
                                           shape_s0=shape_s0, rate_s0=rate_s0,
                                           shape_st0=shape_s0, rate_st0=rate_s0,
                                           shape_sy0=1, rate_sy0=rate_sy0,
                                           ninducing=M, forgetting_rate=0.9,
                                           delay=delay,
                                           exhaustive_train_count=1)

        self.model.max_iter = niter # same as for single user GPPL
        self.model.max_Kw_size = max_Kw_size

        zero_centered_prefs = np.array(self.prefs_train, dtype=float) - 1

        # subsample for debugging!!!
        # self.chosen_people = np.unique(self.person_test)[:50]
        # tridxs = np.in1d(self.person_train, self.chosen_people)

        #self.model.uselowerbound = False
        self.model.use_local_obs_posterior_y = False

        self.model.fit(self.person_train, self.a1_train, self.a2_train, self.items_feat, zero_centered_prefs,
                       optimize=self.optimize_hyper, nrestarts=1, input_type='zero-centered')

    def run_persgppl(self):
        '''
        Make personalised predictions
        :return:
        '''
        global delay
        if delay == 0:
            delay = 10
        self._train_persgppl(delay=delay)

        if self.vscales is not None:
            self.vscales.append(np.sort((self.model.rate_sw / self.model.shape_sw) *
                                        (self.model.rate_sw / self.model.shape_sw))[::-1])

        proba = self.model.predict(self.person_test, self.a1_test, self.a2_test)
        tr_proba = self.model.predict(self.person_unseen, self.a1_unseen, self.a2_unseen)

        # what did we change?
        # - more iterations ( 200 --> 500 )
        # - smaller rates ( 200 --> 20 ) because there are multiple factors whose scales all add up
        # next: try increasing delay so that y doesn't disappear to zero so easily when the person is not seen until
        # later batch of training data
        # Also: consider that when users are independent, then the stochastic updates are also updating only some of
        # the variables. But the others are getting set back to zero... this is only a problem in this model because
        # if they are initialised to zero, they don't move much because w and y are scaled by each other's current
        # estimates. It might be okay if variance of w and y is large because this is added to the scale factor

        if self.a_rank_test is not None:
            predicted_f = self.model.predict_f_item_person(self.a_rank_test, self.person_rank_test)
        else:
            predicted_f = None

        if self.a_rank_train is not None:
            tr_f = self.model.predict_f_item_person(self.a_rank_train, self.person_rank_train)
        else:
            tr_f = None

        return proba, predicted_f, tr_proba, tr_f

    def run_persgppl_consensus(self):
        '''
        Predict the consensus from multiple people's opinions.
        '''

        # look for a file that was trained on the same data but with the personalised predictions instead of MACE consensus.
        # pretrainedmodelfile = self.modelfile.replace('_evalMACE', '')
        # pretrainedmodelfile = pretrainedmodelfile.replace('Consensus', '')
        #
        # logging.info('Looking for a pretrained model at %s' % pretrainedmodelfile)
        #
        # if os.path.exists(pretrainedmodelfile):
        #     with open(pretrainedmodelfile, 'rb') as fh:
        #         self.model = pickle.load(fh)
        #         logging.info('Reloaded a pretrained model :)')
        # else:
        #     logging.info('I didnae find any pretrained model :(')
        #     self._train_persgppl()

        print('Training crowdGPPL to predict consensus...')
        global delay
        if delay == 0:
            delay = 10

        self._train_persgppl(delay=delay)

        if self.vscales is not None:
            self.vscales.append(np.sort(self.model.rate_sw / self.model.shape_sw)[::-1])

        print('Testing crowdGPPL on consensus -- making predictions now.')
        proba = self.model.predict_common(None, self.a1_test, self.a2_test)
        tr_proba = self.model.predict_common(None, self.a1_unseen, self.a2_unseen)

        if self.a_rank_test is not None:
            predicted_f = self.model.predict_t()[self.a_rank_test]
        else:
            predicted_f = None

        if self.a_rank_train is not None:
            tr_f = self.model.predict_t()[self.a_rank_train]
        else:
            predicted_f = None

        print('Max probability = %f, min = %f' % (np.max(proba), np.min(proba)))

        return proba, predicted_f, tr_proba, tr_f

    def _choose_method_fun(self, feature_type):
        if 'PersPrefGP' in self.method:
            method_runner_fun = self.run_persgppl
        elif 'PersConsensusPrefGP' in self.method:
            method_runner_fun = self.run_persgppl_consensus
        elif 'IndPrefGP' in self.method:
            method_runner_fun = self.run_persgppl # switches to correct class inside the method
        elif 'crowdBT' in self.method:
            method_runner_fun = self.run_crowd_bt
        elif 'cBT_GP' in self.method:
            method_runner_fun = self.run_crowd_bt_gpr
        else:
            method_runner_fun = super(PersonalisedTestRunner, self)._choose_method_fun(feature_type)

        return method_runner_fun

if __name__ == '__main__':

    test_to_run = int(sys.argv[1])

    if len(sys.argv) > 2:
        npairs = int(sys.argv[2])
    else:
        npairs = 0 #5000

    if len(sys.argv) > 3:
        lsm = int(sys.argv[3])
    else:
        lsm = 1

    test_dir = 'D05-%i_P%i' % (lsm, npairs)  #'rate_s_tests_single'

    methods = ['SinglePrefGP_noOpt_weaksprior']
    datasets = ['UKPConvArgCrowdSample_evalMACE']
    dataset_increment = 0
    # UKPConvArgCrowdSample tests prediction of personal data.
    # UKPConvArgCrowdSample_evalMACE uses the personal data as input, but predicts the global labels/rankings.
    feature_types = ['both']  # can be 'embeddings' or 'ling' or 'both' or 'debug'
    embeddings_types = ['word_mean']

    datasets = ['UKPConvArgCrowdSample']
    methods = ['PersPrefGP_commonmean_noOpt_weaksprior']

    runner = PersonalisedTestRunner(test_dir, datasets, feature_types, embeddings_types, methods,
                                    dataset_increment)

    max_fold = 32

    rate_sy0 = 10

    # PERSONALISED PREDICTION
    if test_to_run == 0:
        runner.run_test_set(min_no_folds=0, max_no_folds=max_fold, npairs=npairs, ls_factor=lsm)

    elif test_to_run == 13:
        # tune up on training set accuracy

        rateyvals = [1, 10, 100]
        delays = [1, 10]

        for ratey in rateyvals:
            for delay in delays:
                rate_sy0 = ratey

                test_dir = 'D05-%i_P%i-ratey%i-delay%i' % (lsm, npairs, ratey, delay)  # 'rate_s_tests_single'
                runner = PersonalisedTestRunner(test_dir, datasets, feature_types, embeddings_types, methods,
                                            dataset_increment)
                runner.run_test_set(min_no_folds=0, max_no_folds=10, npairs=npairs, ls_factor=lsm)

    elif test_to_run == 14:

        rateyvals = [1, 10, 100]
        delays = [1, 10]

        for ratey in rateyvals:
            for delay in delays:
                rate_sy0 = ratey

                test_dir = 'D05-%i_P%i-ratey%i-delay%i' % (lsm, npairs, ratey, delay)  # 'rate_s_tests_single'
                runner = PersonalisedTestRunner(test_dir, datasets, feature_types, embeddings_types, methods,
                                            dataset_increment)
                runner.datasets = ['UKPConvArgCrowdSample_evalMACE']
                runner.methods = ['PersConsensusPrefGP_commonmean_noOpt_weaksprior']
                runner.run_test_set(min_no_folds=0, max_no_folds=10, npairs=npairs, ls_factor=lsm)

    elif test_to_run == 12:
        runner.datasets = ['UKPConvArgCrowdSample']
        runner.methods = ['PersPrefGP_noOpt_weaksprior']
        runner.run_test_set(min_no_folds=0, max_no_folds=max_fold, npairs=npairs, ls_factor=lsm)

    # CONSENSUS PREDICTION
    elif test_to_run == 1:
        runner.datasets = ['UKPConvArgCrowdSample_evalMACE']
        runner.methods = ['PersConsensusPrefGP_commonmean_noOpt_weaksprior']
        runner.run_test_set(min_no_folds=0, max_no_folds=max_fold, npairs=npairs, ls_factor=lsm)

    # Plot the scales of the latent factors ----------------------------------------------------------------------
    if test_to_run < 4 and len(runner.vscales):
        vscales = np.mean(runner.vscales, axis=0)

        logging.getLogger().setLevel(logging.WARNING) # matplotlib prints loads of crap to the debug and info outputs

        import matplotlib
        matplotlib.use('Agg')
        import matplotlib.pyplot as plt

        fig = plt.figure(figsize=(5, 4))

        markers = ['x', 'o', '+', '>', '<', '*']

        plt.plot(np.arange(vscales.shape[0]), vscales, marker=markers[0], label='UKPConvArgCrowdSample',
                 linewidth=2, markersize=8)

        plt.ylabel('Inverse scale 1/s')
        plt.xlabel('Factor ID')

        plt.grid('on', axis='y')
        plt.legend(loc='best')
        plt.tight_layout()

        figure_root_path = './results/conv_factors'
        if not os.path.exists(figure_root_path):
            os.mkdir(figure_root_path)

        plt.savefig(figure_root_path + '/UKPConvArgCrowdSample_factor_scales.pdf')

        np.savetxt(figure_root_path + '/UKPConvArgCrowdSample_factor_scales.csv', vscales, delimiter=',', fmt='%f')

        logging.getLogger().setLevel(logging.DEBUG)

    # PERSONALISED PREDICTION for other methods -----------------------------------------------------------------
    elif test_to_run == 6:
        methods = [
               'SinglePrefGP_noOpt_weaksprior' # 'SinglePrefGP_noOpt_weaksprior',
            ]
        runner.datasets = ['UKPConvArgCrowdSample']
        runner.methods = methods
        runner.run_test_set(min_no_folds=0, max_no_folds=max_fold, npairs=npairs, ls_factor=lsm)

    elif test_to_run == 7:
        methods = [
               'SinglePrefGP_noOpt_weaksprior' # 'SinglePrefGP_noOpt_weaksprior',
            ]
        runner.datasets = ['UKPConvArgCrowdSample_evalMACE']
        runner.methods = methods
        runner.run_test_set(min_no_folds=0, max_no_folds=max_fold, npairs=npairs, ls_factor=lsm)

    elif test_to_run == 8:
        methods = [
               #'crowdBT', # no point running this because it cannot predict on the test instances, for aggregation only
               'cBT_GP',
            ]
        runner.datasets = ['UKPConvArgCrowdSample']
        runner.methods = methods
        runner.run_test_set(min_no_folds=0, max_no_folds=max_fold, npairs=npairs, ls_factor=lsm)

    elif test_to_run == 9: # commented so we run both tests with cBT
        methods = [
               #'crowdBT', # no point running this because it cannot predict on the test instances, for aggregation only
               'cBT_GP',
        ]
        runner.datasets = ['UKPConvArgCrowdSample_evalMACE']
        runner.methods = methods
        runner.run_test_set(min_no_folds=0, max_no_folds=max_fold, npairs=npairs, ls_factor=lsm)

    elif test_to_run == 11: # commented so we run both tests with cBT
        methods = [
               'crowdBT', # no point running this because it cannot predict on the test instances, for aggregation only
        ]
        runner.datasets = ['UKPConvArgCrowdSample_evalMACE']
        runner.methods = methods
        runner.run_test_set(min_no_folds=0, max_no_folds=max_fold, npairs=npairs, ls_factor=lsm)