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

"""
   This file belong to https://github.com/snolfi/evorobotpy
   and has been written by Stefano Nolfi and Paolo Pagliuca, stefano.nolfi@istc.cnr.it, paolo.pagliuca@istc.cnr.it
   salimans.py include an implementation of the OpenAI-ES algorithm described in
   Salimans T., Ho J., Chen X., Sidor S & Sutskever I. (2017). Evolution strategies as a scalable alternative to reinforcement learning. arXiv:1703.03864v2
   requires es.py, policy.py, and evoalgo.py
"""

import numpy as np
from numpy import zeros, ones, dot, sqrt
import math
import time
from evoalgo import EvoAlgo
from utils import ascendent_sort
import sys
import os
import configparser

# Parallel implementation of Open-AI-ES algorithm developed by Salimans et al. (2017)
# the workers evaluate a fraction of the population in parallel
# the master post-evaluate the best sample of the last generation and eventually update the input normalization vector

class Algo(EvoAlgo):
    def __init__(self, env, policy, seed, fileini, filedir):
        EvoAlgo.__init__(self, env, policy, seed, fileini, filedir)

    def loadhyperparameters(self):

        if os.path.isfile(self.fileini):

            config = configparser.ConfigParser()
            config.read(self.fileini)
            self.maxsteps = 1000000
            self.stepsize = 0.01
            self.batchSize = 20
            self.noiseStdDev = 0.02
            self.wdecay = 0
            self.symseed = 1
            self.saveeach = 60
            options = config.options("ALGO")
            for o in options:
                found = 0
                if o == "maxmsteps":
                    self.maxsteps = config.getint("ALGO","maxmsteps") * 1000000
                    found = 1
                if o == "stepsize":
                    self.stepsize = config.getfloat("ALGO","stepsize")
                    found = 1
                if o == "noisestddev":
                    self.noiseStdDev = config.getfloat("ALGO","noiseStdDev")
                    found = 1
                if o == "samplesize":
                    self.batchSize = config.getint("ALGO","sampleSize")
                    found = 1
                if o == "wdecay":
                    self.wdecay = config.getint("ALGO","wdecay")
                    found = 1
                if o == "symseed":
                    self.symseed = config.getint("ALGO","symseed")
                    found = 1
                if o == "saveeach":
                    self.saveeach = config.getint("ALGO","saveeach")
                    found = 1

                if found == 0:
                    print("\033[1mOption %s in section [ALGO] of %s file is unknown\033[0m" % (o, filename))
                    print("available hyperparameters are: ")
                    print("maxmsteps [integer]       : max number of (million) steps (default 1)")
                    print("stepsize [float]          : learning stepsize (default 0.01)")
                    print("samplesize [int]          : popsize/2 (default 20)")
                    print("noiseStdDev [float]       : samples noise (default 0.02)")
                    print("wdecay [0/2]              : weight decay (default 0), 1 = L1, 2 = L2")
                    print("symseed [0/1]             : same environmental seed to evaluate symmetrical samples [default 1]")
                    print("saveeach [integer]        : save file every N minutes (default 60)")

                    sys.exit()
        else:
            print("\033[1mERROR: configuration file %s does not exist\033[0m" % (self.fileini))

    def setProcess(self):
        self.loadhyperparameters()               # load hyperparameters
        self.center = np.copy(self.policy.get_trainable_flat())  # the initial centroid
        self.nparams = len(self.center)          # number of adaptive parameters
        self.cgen = 0                            # currrent generation
        self.samplefitness = zeros(self.batchSize * 2) # the fitness of the samples
        self.samples = None                      # the random samples
        self.m = zeros(self.nparams)             # Adam: momentum vector
        self.v = zeros(self.nparams)             # Adam: second momentum vector (adam)
        self.epsilon = 1e-08                     # Adam: To avoid numerical issues with division by zero...
        self.beta1 = 0.9                         # Adam: beta1
        self.beta2 = 0.999                       # Adam: beta2
        self.bestgfit = -99999999                # the best generalization fitness
        self.bfit = 0                            # the fitness of the best sample
        self.gfit = 0                            # the postevaluation fitness of the best sample of last generation
        self.rs = None                           # random number generator
        self.inormepisodes = self.batchSize * 2 * self.policy.ntrials / 100.0 # number of normalization episode for generation (1% of generation episodes)
        self.tnormepisodes = 0.0                 # total epsidoes in which normalization data should be collected so far
        self.normepisodes = 0                    # numer of episodes in which normalization data has been actually collected so far
        self.normalizationdatacollected = False  # whether we collected data for updating the normalization vector

    def savedata(self):
        self.save()             # save the best agent so far, the best postevaluated agent so far, and progress data across generations
        fname = self.filedir + "/S" + str(self.seed) + ".fit"
        fp = open(fname, "w")   # save summary
        fp.write('Seed %d (%.1f%%) gen %d msteps %d bestfit %.2f bestgfit %.2f bestsam %.2f avgfit %.2f paramsize %.2f \n' %
             (self.seed, self.steps / float(self.maxsteps) * 100, self.cgen, self.steps / 1000000, self.bestfit, self.bestgfit, self.bfit, self.avgfit, self.avecenter))
        fp.close()

    def evaluate(self):
        cseed = self.seed + self.cgen * self.batchSize  # Set the seed for current generation (master and workers have the same seed)
        self.rs = np.random.RandomState(cseed)
        self.samples = self.rs.randn(self.batchSize, self.nparams)
        self.cgen += 1

        # evaluate samples
        candidate = np.arange(self.nparams, dtype=np.float64)
        for b in range(self.batchSize):
            for bb in range(2):
                if (bb == 0):
                    candidate = self.center + self.samples[b,:] * self.noiseStdDev
                else:
                    candidate = self.center - self.samples[b,:] * self.noiseStdDev
                self.policy.set_trainable_flat(candidate)
                self.policy.nn.normphase(0) # normalization data is collected during the post-evaluation of the best sample of he previous generation
                eval_rews, eval_length = self.policy.rollout(self.policy.ntrials, seed=(self.seed + (self.cgen * self.batchSize) + b))
                self.samplefitness[b*2+bb] = eval_rews
                self.steps += eval_length

        fitness, self.index = ascendent_sort(self.samplefitness)       # sort the fitness
        self.avgfit = np.average(fitness)                         # compute the average fitness

        self.bfit = fitness[(self.batchSize * 2) - 1]
        bidx = self.index[(self.batchSize * 2) - 1]
        if ((bidx % 2) == 0):                                     # regenerate the genotype of the best samples
            bestid = int(bidx / 2)
            self.bestsol = self.center + self.samples[bestid] * self.noiseStdDev
        else:
            bestid = int(bidx / 2)
            self.bestsol = self.center - self.samples[bestid] * self.noiseStdDev

        self.updateBest(self.bfit, self.bestsol)                  # Stored if it is the best obtained so far

        # postevaluate best sample of the last generation
        # in openaiesp.py this is done the next generation, move this section before the section "evaluate samples" to produce identical results
        gfit = 0
        if self.bestsol is not None:
            self.policy.set_trainable_flat(self.bestsol)
            self.tnormepisodes += self.inormepisodes
            for t in range(self.policy.nttrials):
                if self.policy.normalize == 1 and self.normepisodes < self.tnormepisodes:
                    self.policy.nn.normphase(1)
                    self.normepisodes += 1  # we collect normalization data
                    self.normalizationdatacollected = True
                else:
                    self.policy.nn.normphase(0)
                eval_rews, eval_length = self.policy.rollout(1, seed=(self.seed + 100000 + t))
                gfit += eval_rews
                self.steps += eval_length
            gfit /= self.policy.nttrials
            self.updateBestg(gfit, self.bestsol)

    def optimize(self):

        popsize = self.batchSize * 2                              # compute a vector of utilities [-0.5,0.5]
        utilities = zeros(popsize)
        for i in range(popsize):
            utilities[self.index[i]] = i
        utilities /= (popsize - 1)
        utilities -= 0.5

        weights = zeros(self.batchSize)                           # Assign the weights (utility) to samples on the basis of their fitness rank
        for i in range(self.batchSize):
            idx = 2 * i
            weights[i] = (utilities[idx] - utilities[idx + 1])    # merge the utility of symmetric samples

        g = 0.0
        i = 0
        while i < self.batchSize:                                 # Compute the gradient (the dot product of the samples for their utilities)
            gsize = -1
            if self.batchSize - i < 500:                          # if the popsize is larger than 500, compute the gradient for multiple sub-populations
                gsize = self.batchSize - i
            else:
                gsize = 500
            g += dot(weights[i:i + gsize], self.samples[i:i + gsize,:])
            i += gsize
        g /= popsize                                              # normalize the gradient for the popsize

        if self.wdecay == 1:
            globalg = -g + 0.005 * self.center                    # apply weight decay
        else:
            globalg = -g

        # adam stochastic optimizer
        a = self.stepsize * sqrt(1.0 - self.beta2 ** self.cgen) / (1.0 - self.beta1 ** self.cgen)
        self.m = self.beta1 * self.m + (1.0 - self.beta1) * globalg
        self.v = self.beta2 * self.v + (1.0 - self.beta2) * (globalg * globalg)
        dCenter = -a * self.m / (sqrt(self.v) + self.epsilon)

        self.center += dCenter                                    # move the center in the direction of the momentum vectors
        self.avecenter = np.average(np.absolute(self.center))

    def run(self):

        self.setProcess()                           # initialize class variables
        start_time = time.time()
        last_save_time = start_time
        elapsed = 0
        self.steps = 0
        print("Salimans: seed %d maxmsteps %d batchSize %d stepsize %lf noiseStdDev %lf wdecay %d symseed %d nparams %d" % (self.seed, self.maxsteps / 1000000, self.batchSize, self.stepsize, self.noiseStdDev, self.wdecay, self.symseed, self.nparams))

        while (self.steps < self.maxsteps):

            self.evaluate()                           # evaluate samples

            self.optimize()                           # estimate the gradient and move the centroid in the gradient direction

            self.stat = np.append(self.stat, [self.steps, self.bestfit, self.bestgfit, self.bfit, self.avgfit, self.avecenter])  # store performance across generations

            if ((time.time() - last_save_time) > (self.saveeach * 60)):
                self.savedata()                       # save data on files
                last_save_time = time.time()

            if self.normalizationdatacollected:
                self.policy.nn.updateNormalizationVectors()  # update the normalization vectors with the new data collected
                self.normalizationdatacollected = False

            print('Seed %d (%.1f%%) gen %d msteps %d bestfit %.2f bestgfit %.2f bestsam %.2f avg %.2f weightsize %.2f' %
                      (self.seed, self.steps / float(self.maxsteps) * 100, self.cgen, self.steps / 1000000, self.bestfit, self.bestgfit, self.bfit, self.avgfit, self.avecenter))

        self.savedata()                           # save data at the end of evolution

        # print simulation time
        end_time = time.time()
        print('Simulation time: %dm%ds ' % (divmod(end_time - start_time, 60)))


import copy
import numpy as np
import scipy as sp
from scipy import sparse
from mesohops.dynamics.eom_functions import operator_expectation
from mesohops.util.physical_constants import hbar
from mesohops.util.exceptions import UnsupportedRequest

__title__ = "Basis Class"
__author__ = "D. I. G. Bennett"
__version__ = "1.0"

class HopsBasis(object):
    """
    Every HOPS calculation is defines by the HopsSystem, HopsHierarchy, and HopsEOM
    classes (and their associated parameters). These form the basis set for the
    calculation. HopsBasis is the class that contains all of these sub-classes and
    mediates the way the HopsTrajectory will interact with them.
    """

    def __init__(self, system, hierarchy, eom):
        """
        INPUTS:
        -------
        1. system: dictionary of user inputs
            [see hops_system.py]
            a. HAMILTONIAN
            b. GW_SYSBATH
            c. CORRELATION_FUNCTION_TYPE
            d. LOPERATORS
            e. CORRELATION_FUNCTION
        2. hierarchy_parameters: dictionary of user inputs
            [see hops_hierarchy.py]
            f. MAXHIER
            g. TERMINATOR
            h. STATIC_FILTERS
        3. eom_parameters: dictionary of user inputs
            [see hops_eom.py]
            i. TIME_DEPENDENCE
            j. EQUATION_OF_MOTION
            k. ADAPTIVE_H
            l. ADAPTIVE_S
            m. DELTA_H
            n. DELTA_S

        RETURNS
        -------
        None
        """
        self.system = system
        self.hierarchy = hierarchy
        self.eom = eom

    def initialize(self, psi_0):
        """
        This function initializes the hierarchy and equations of motion classes
        so that everything is prepared for integration. It returns the
        dsystem_dt function to be used in the integrator.

        PARAMETERS
        ----------
        1. psi_0 : np.array
                  the initial wave function

        RETURNS
        -------
        1. dsystem_dt : function
                       this is the core function for calculating the time-evolution of
                       the wave function
        """
        self.hierarchy.initialize(self.adaptive_h)
        self.system.initialize(self.adaptive_s, psi_0)
        dsystem_dt = self.eom._prepare_derivative(self.system, self.hierarchy)

        return dsystem_dt

    def define_basis(self, Φ, delta_t, z_step):
        """
        This is the function that determines what basis is needed for a given
        full hierarchy (Φ) in order to construct an approximate derivative with
        error below the specified threshold.

        PARAMETERS
        ----------
        1. Φ : np.array
               the current full hierarchy
        2. delta_t : float
                     the timestep for the calculation
        3. z_step : list
                    the list of noise terms (compressed) for the next timestep

        RETURNS
        -------
        1. state_update :
            a. list_state_new : list
                               list of states in the new basis (S_1)
            b. list_state_stable : list
                                   list of stable states in the new basis (S_S)
            c. list_state_bound : list
                                  list of boundary states in the new basis (S_B)
        2. hierarchy_update :
            d. list_aux_new : list
                              list of auxiliaries in new basis (H_1)
            e. list_stable_aux : list
                                 list of stable auxiliaries in the new basis (H_S)
            f. list_aux_bound : list
                                list of boundary auxiliaries in the new basis (H_B)

        """
        # ==========================================
        # =======      Calculate Updates      ======
        # ==========================================

        # Calculate New Hierarchy List
        # ----------------------------
        if self.adaptive_h:
            list_aux_stable, list_aux_bound = self._check_hierarchy_list(
                Φ, delta_t, z_step
            )
            list_aux_new = list(set(list_aux_stable) | set(list_aux_bound))
            list_index_stable_aux = [
                self.hierarchy._aux_index(aux) for aux in list_aux_stable
            ]
            list_index_stable_aux.sort()
        else:
            list_aux_new = self.hierarchy.auxiliary_list
            list_aux_stable = self.hierarchy.auxiliary_list
            list_aux_bound = []
            list_index_stable_aux = np.arange(len(self.hierarchy.auxiliary_list))
            E2_flux = None

        # Calculate New State List
        # ------------------------
        if self.adaptive_s:
            list_state_stable, list_state_bound = self._check_state_list(
                Φ, delta_t, z_step, list_index_stable_aux
            )
            list_state_new = list(set(list_state_stable) | set(list_state_bound))
            list_state_stable.sort()
        else:
            list_state_new = list(set(self.system.state_list))
            list_state_bound = []
            list_state_stable = list_state_new

        return [
            (list_state_new, list_state_stable, list_state_bound),
            (list_aux_new, list_aux_stable, list_aux_bound),
        ]

    def update_basis(self, Φ, state_update, aux_update):
        """
        This function updates the derivative function and full hierarchy vector (Φ) for the
        new basis (hierarchy and/or system).

        PARAMETERS
        ----------
        1. Φ : np.array
               the current full hierarchy
        2. state_update : list
                          list of list containing list_state_new, list_stable_state, and list_add_state
        3. aux_update : list
                        list of list containing list_aux_new, list_stable_aux, and list_add_aux

        RETURNS
        -------
        1. Φ_new : np.array
                   the updated full hierarchy
        2. dsystem_dt : function
                        the updated derivative function
        """
        # Unpack input values
        # ===================
        (list_state_new, list_stable_state, list_add_state) = state_update
        (list_aux_new, list_stable_aux, list_add_aux) = aux_update

        # Update State List
        # =================
        flag_update = False
        if set(list_state_new) != set(self.system.state_list):
            flag_update = True
            list_state_previous = copy.deepcopy(self.system.state_list)
            list_absindex_l2_old = copy.deepcopy(self.system.list_absindex_L2)
            self.system.state_list = np.array(list_state_new)
        else:
            list_state_previous = self.system.state_list
            list_absindex_l2_old = self.system.list_absindex_L2

        # Update Hierarchy List
        # =====================
        if set(list_aux_new) != set(self.hierarchy.auxiliary_list):
            flag_update = True
            # Update Auxiliary List
            list_old_aux = self.hierarchy.auxiliary_list
            self.hierarchy.auxiliary_list = list_aux_new
        else:
            list_old_aux = self.hierarchy.auxiliary_list

        # Update state of calculation for new basis
        # =========================================
        if flag_update:

            # Define permutation matrix from old basis --> new basis
            # ------------------------------------------------------
            permute_aux_row = []
            permute_aux_col = []
            nstate_old = len(list_state_previous)
            list_index_old_stable_state = np.array(
                [
                    i_rel
                    for (i_rel, i_abs) in enumerate(list_state_previous)
                    if i_abs in list_stable_state
                ]
            )
            list_index_new_stable_state = np.array(
                [
                    i_rel
                    for (i_rel, i_abs) in enumerate(self.system.state_list)
                    if i_abs in list_stable_state
                ]
            )

            for aux in list_stable_aux:
                permute_aux_row.extend(
                    self.hierarchy._aux_index(aux) * self.n_state
                    + list_index_new_stable_state
                )
                permute_aux_col.extend(
                    list_old_aux.index(aux) * nstate_old + list_index_old_stable_state
                )

            # Update phi
            # ----------
            Φ_new = np.zeros(self.n_hier * self.n_state, dtype=np.complex128)
            Φ_new[permute_aux_row] = Φ[permute_aux_col]

            # Update dsystem_dt
            # -----------------
            dsystem_dt = self.eom._prepare_derivative(
                self.system,
                self.hierarchy,
                list_stable_aux,
                list_add_aux,
                list_stable_state,
                list_absindex_l2_old,
                len(list_old_aux) * nstate_old,
                [permute_aux_row, permute_aux_col],
                update=True,
            )

            return (Φ_new, dsystem_dt)
        else:
            return (Φ, self.eom.dsystem_dt)

    def _check_state_list(self, Φ, delta_t, z_step, list_index_aux_stable):
        """
        This is a function that determines the states which should be
        included in the adaptive integration for the next time point.

        PARAMETERS
        ----------
        1. Φ : np.array
               the current full hierarchy
        2. delta_t : float
                     the timestep for the calculation
        3. z_step : list
                    the list of noise terms (compressed) for the next timestep
        4. list_index_aux_stable : list
                                   a list of relative indices for the stable auxiliaries

        RETURNS
        -------
        1. list_state_stable : list
                               list of stable states (absolute state index, S_S)
        2. list_state_boundary : list
                                 a list of the boundary states (absolute state index, S_B)
        """
        # Define Constants
        # ----------------
        delta_state = self.delta_s  # * np.linalg.norm(Φ)

        # CONSTRUCT STABLE STATE (S_S)
        # ============================

        # Construct Error For Excluding Member of S0
        # ------------------------------------------
        error_by_state = self.error_stable_state(
            Φ, delta_t, z_step, list_index_aux_stable
        )

        # Determine the Stable States (S_S = S_0 & S_1)
        # ---------------------------------------------
        list_index_stable, list_state_stable = self._determine_basis_from_list(
            error_by_state, delta_state / 2, self.system.state_list
        )

        # CONSTRUCT BOUNDARY STATE (S_B)
        # ==============================

        # Establish the error available for the boundary states
        # -----------------------------------------------------
        stable_error = np.sqrt(
            np.max([
                np.sum(error_by_state ** 2) - np.sum(error_by_state[list_index_stable] ** 2),
                0])
        )
        bound_error = delta_state - stable_error

        # Construct Error for Excluding Member of S0^C
        # --------------------------------------------
        list_index_nonzero, list_error_nonzero = self.error_boundary_state(
            Φ, list_index_stable, list_index_aux_stable
        )

        # Determine Boundary States
        # -------------------------
        if len(list_error_nonzero) > 0:
            _, list_state_boundary = self._determine_basis_from_list(
                list_error_nonzero, bound_error, list_index_nonzero
            )
        else:
            list_state_boundary = []

        # Check for overlap with populated states
        # ---------------------------------------
        list_state_boundary = list(
            set(list_state_boundary) - set(self.system.state_list)
        )
        list_state_boundary.sort()

        return (
            np.array(list_state_stable, dtype=np.int),
            np.array(list_state_boundary, dtype=np.int),
        )

    def _check_hierarchy_list(self, Φ, delta_t, z_step):
        """
        This is a function that determines the auxiliaries which should be
        included in the adaptive integration for the next time point.

        PARAMETERS
        ----------
        1. Φ : np.array
               the current full hierarchy
        2. delta_t : float
                     the timestep for the calculation
        3. z_step : list
                    the list of noise terms (compressed) for the next timestep

        RETURNS
        -------
        1. list_aux_stable : list
                             a list of the stable auxiliaries (H_S)
        2. list_aux_boundary : list
                               a list of auxiliaries that share a boundary with stable
                               auxiliaries (H_B)
        """
        # Define Constants
        # ----------------
        delta_hier = self.delta_h  # * np.linalg.norm(Φ)

        # CONSTRUCT STABLE HIERARCHY
        # ==========================

        # Construct Error For Excluding Member of H0
        # ------------------------------------------
        error_by_aux, list_e2_kflux = self.hier_stable_error(Φ, delta_t, z_step)

        # Determine the Stable Auxiliaries (H_S = H_0 & H_1)
        # --------------------------------------------------
        list_index_stable, list_aux_stable = self._determine_basis_from_list(
            error_by_aux, delta_hier / 2, self.hierarchy.auxiliary_list
        )

        # CONSTRUCT BOUNDARY HIERARCHY
        # ============================

        # Establish the error available for the boundary auxiliaries
        # ----------------------------------------------------------
        stable_error = np.sqrt(
            np.max([
                np.sum(error_by_aux ** 2) - np.sum(error_by_aux[list_index_stable] ** 2),
                0])
        )
        bound_error = delta_hier - stable_error

        # Construct Error For Excluding Members of H0^C
        # ---------------------------------------------
        E2_flux_up = list_e2_kflux[0][:, list_index_stable]
        E2_flux_down = list_e2_kflux[1][:, list_index_stable]

        # Determine the Boundary Auxiliaries (H_B = H0^C & H_1)
        # -----------------------------------------------------
        list_aux_up, list_aux_down = self._determine_boundary_hier(
            [E2_flux_up, E2_flux_down], list_index_stable, bound_error
        )

        # Check Boundary Set For Duplication
        # ----------------------------------
        # NOTE: Theoretically, the boundary should not contain any members of H0, but
        #       when we implemented that in practice it resulted in a huge explosion
        #       of the number of inch wormsteps that were required.
        list_aux_boundary = list(
            (set(list_aux_up) | set(list_aux_down)) - set(self.hierarchy.auxiliary_list)
        )

        # Filter Boundary Set for Auxiliaries That Are Not Part of H_T
        # ------------------------------------------------------------
        if len(list_aux_boundary) > 0:
            list_aux_boundary = self.hierarchy.filter_aux_list(list_aux_boundary)
        return list_aux_stable, list_aux_boundary

    def _determine_boundary_hier(
        self, list_e2_kflux, list_index_aux_stable, bound_error
    ):
        """
        This function determines the set of boundary auxiliaries for the next time step

        PARAMETERS
        ----------
        1. list_e2_kflux : list
                           a list of list containing the error values for the flux up
                           and flux down terms
        2. list_index_aux_stable : list
                                   a list of the indices for stable auxiliaries
        3. bound_error : float
                         the boundary error value

        RETURNS
        -------
        1. list_aux_up : list
                         a list of the flux up auxiliaries
        2. list_aux_down : list
                           a list of the flux down auxiliaries
        """
        # Construct constants
        # -------------------
        E2_flux_up = list_e2_kflux[0]
        E2_flux_down = list_e2_kflux[1]

        # Find the error threshold for edge auxiliaries
        # ---------------------------------------------
        E1_nonzero_flux = E2_flux_up[E2_flux_up != 0]
        sorted_error = np.sort(
            np.append(E1_nonzero_flux, E2_flux_down[E2_flux_down != 0])
        )
        error_thresh = self._determine_error_thresh(sorted_error, bound_error)

        # Loop over residual fluxes to identify boundary auxiliaries
        # ----------------------------------------------------------
        list_aux_up = [
            self.hierarchy.auxiliary_list[list_index_aux_stable[i_aux]].e_step(
                self.system.list_absindex_mode[i_mode_rel], 1
            )
            for (i_mode_rel, i_aux) in zip(*np.where(E2_flux_up > error_thresh))
        ]

        list_aux_down = [
            self.hierarchy.auxiliary_list[list_index_aux_stable[i_aux]].e_step(
                self.system.list_absindex_mode[i_mode_rel], -1
            )
            for (i_mode_rel, i_aux) in zip(*np.where(E2_flux_down > error_thresh))
        ]
        return (list_aux_up, list_aux_down)

    def _determine_basis_from_list(self, error_by_member, max_error, list_member):
        """
        This function determines the members of a list that must be kept in order
        for the total error (terms that are dropped) to be below the max_error value.

        PARAMETERS
        ----------
        1. error_by_member : np.array
                             a list of error values
        2. max_error : float
                       the maximum error value
        3. list_member : np.array
                         a list of members

        RETURNS
        -------
        1. list_index : np.array
                        a list of indices for the members
        2. list_new_member : list
                             a list of the members
        """
        error_thresh = self._determine_error_thresh(np.sort(error_by_member), max_error)
        list_index = np.where(error_by_member > error_thresh)[0]
        list_new_member = [list_member[i_aux] for i_aux in list_index]
        return (list_index, list_new_member)

    @staticmethod
    def _determine_error_thresh(sorted_error, max_error):
        """
        This function determines which error value becomes the error threshold such
        that the sum of all errors below the threshold remains less then max_error.

        PARAMETERS
        ----------
        1. sorted_error : np.array
                          a list of error values
        2. max_error : float
                       the maximum error value

        RETURNS
        -------
        3. error_thresh : float
                          the error value at which the threshold is established

        """
        index_thresh = np.argmax(np.sqrt(np.cumsum(sorted_error ** 2)) > max_error)

        if index_thresh > 0:
            error_thresh = sorted_error[index_thresh - 1]
        else:
            error_thresh = 0.0

        return error_thresh

    def error_boundary_state(self, Φ, list_index_stable, list_index_aux_stable):
        """
        This function determines the error associated with neglecting flux into n not
        a member of S_t. This corresponds to equations 43-45 in arXiv:2008.06496

        PARAMETERS
        ----------
        1. Φ : np.array
               the current full hierarchy vector
        2. list_index_stable : list
                               a list of the stable states
        3. list_index_aux_stable : list
                                   a list relative indices for the stable auxiliaries

        RETURNS
        -------
        1. list_index_nonzero : list
                                a list of the nonzero state boundary indices
        2. list_error_nonzero : list
                                a list of the nonzero state boundary error values
        """
        if len(list_index_stable) < self.system.param["NSTATES"]:
            # Remove aux components from H0\H1
            # -------------------------------------
            C2_phi = np.zeros([self.n_state, self.n_hier], dtype=np.complex128)
            C2_phi[np.ix_(list_index_stable, list_index_aux_stable)] = np.array(
                Φ
            ).reshape([self.n_state, self.n_hier], order="F")[
                np.ix_(list_index_stable, list_index_aux_stable)
            ]

            # Construct Hamiltonian
            # ---------------------------------------------------------
            list_s0 = np.array(self.system.state_list)

            # First construct ST<--S0 Hamiltonian
            H2_sparse_hamiltonian = self.system.param["SPARSE_HAMILTONIAN"][:, list_s0]

            # Remove components that map S0<--S0
            H2_sparse_subset = sparse.coo_matrix(
                H2_sparse_hamiltonian[np.ix_(list_s0, range(len(list_s0)))]
            )
            H2_removal = sparse.csc_matrix(
                (
                    H2_sparse_subset.data,
                    (
                        list_s0[H2_sparse_subset.row],
                        np.arange(len(list_s0))[H2_sparse_subset.col],
                    ),
                ),
                shape=H2_sparse_hamiltonian.shape,
            )
            H2_sparse_hamiltonian = H2_sparse_hamiltonian - H2_removal

            # Determine Boundary States
            # -------------------------
            C2_phi_deriv = abs(
                H2_sparse_hamiltonian @ sparse.csc_matrix(C2_phi / hbar)
            ).power(2)
            C1_sum_deriv = np.sum(C2_phi_deriv, axis=1)
            return (
                C1_sum_deriv.nonzero()[0],
                np.sqrt(np.array(C1_sum_deriv[C1_sum_deriv.nonzero()])[0]),
            )

        else:
            return [], []

    def error_stable_state(self, Φ, delta_t, z_step, list_index_aux_stable):
        """
        This function finds the total error associated with neglecting n in the state
        basis S_t. This includes the error associated with deleting a state from the basis,
        the error associated with losing all flux into a state, and the error associated
        with losing all flux out of a state. This corresponds to equations 37-42 in
        arXiv:2008.06496

        PARAMETERS
        ----------
        1. Φ : np.array
               the current full hierarchy vector
        2. delta_t : float
                     the timestep for the calculation
        3. z_step : list
                    the list of noise terms (compressed) for the next timestep
        4. list_index_aux_stable : list
                                   a list relative indices for the stable auxiliaries (A_p)

        RETURNS
        -------
        1. error : np.array
                   list of error associated with removing each state
        """
        M2_state_from_mode = np.zeros([self.n_state, self.system.n_hmodes])
        M2_state_from_mode[
            self.system.list_state_indices_by_hmode[:, 0],
            np.arange(np.shape(self.system.list_state_indices_by_hmode)[0]),
        ] = 1

        # Construct the error terms
        # -------------------------
        E2_deriv_state = self.error_deriv(Φ, z_step, list_index_aux_stable)[
            :, list_index_aux_stable
        ]
        E2_deletion = self.error_deletion(Φ, delta_t)[:, list_index_aux_stable]
        E1_state_flux = self.error_sflux_state(
            Φ, list_index_aux_stable, self.system.state_list
        )
        E2_flux_down = self.error_flux_down(Φ, "S")[:, list_index_aux_stable]
        E2_flux_up = self.error_flux_up(Φ)[:, list_index_aux_stable]
        # Map flux_up error from mode to state space
        E2_flux_up = np.sqrt(M2_state_from_mode @ E2_flux_up ** 2)

        # Compress the error onto the state/mode axis
        # -------------------------------------------
        return np.sqrt(
            np.sum(E2_deriv_state ** 2, axis=1)
            + np.sum(E2_deletion ** 2, axis=1)
            + np.sum(E2_flux_up ** 2, axis=1)
            + np.sum(E2_flux_down ** 2, axis=1)
            + E1_state_flux ** 2
        )

    def hier_stable_error(self, Φ, delta_t, z_step):
        """
        This function finds the total error associated with removing k in A_t.
        This corresponds to the sum of equations 29,30,31,33, and 34 in arXiv:2008.06496

        PARAMETERS
        ----------
        1. Φ : np.array
               the current full hierarchy vector
        2. delta_t : float
                     the timestep for the calculation
        3. z_step : list
                    the list of noise terms (compressed) for the next timestep

        RETURNS
        -------
        1. error : np.array
                   list of error associated with removing each auxiliary in A_t
        2. E2_flux_up : np.array
                        the error induced by neglecting flux from A_t (or A_p)
                                to auxiliaries with lower summed index in A_t^C.
        3. E2_flux_down : np.array
                          the error induced by neglecting flux from A_t (or A_p)
                                to auxiliaries with higher summed index in A_t^C.
        """
        # Construct the error terms
        # -------------------------
        E2_deletion = self.error_deletion(Φ, delta_t)
        E2_deriv_self = self.error_deriv(Φ, z_step)
        E1_flux_state = self.error_sflux_hier(Φ, self.system.state_list)
        E2_flux_up = self.error_flux_up(Φ)
        E2_flux_down = self.error_flux_down(Φ, "H")

        # Compress the error onto the aux axis
        # ------------------------------------
        return (
            np.sqrt(
                np.sum(E2_deriv_self ** 2, axis=0)
                + np.sum(E2_deletion ** 2, axis=0)
                + np.sum(E2_flux_down ** 2, axis=0)
                + np.sum(E2_flux_up ** 2, axis=0)
                + E1_flux_state
            ),
            [E2_flux_up, E2_flux_down],
        )

    def error_sflux_state(self, Φ, list_index_aux_stable, list_states):
        """
        The error associated with losing all flux out of n in S_t. This flux always involves
        changing the state index and, as a result, can be rewritten in terms of the -iH
        component of the self-interaction. This corresponds to equation 38 and 39 in
        arXiv:2008.06496

        PARAMETERS
        ----------
        1. Φ : np.array
               the current full hierarchy vector
        2. list_index_aux_stable : list
                                   a list relative indices for the stable auxiliaries (A_p)
        3. list_states : list
                         the list of current states (absolute index)

        RETURNS
        -------
        1. E1_state_flux : array
                           the error associated with flux out of each state in S_t
        """
        C2_phi = np.asarray(Φ).reshape([self.n_state, self.n_hier], order="F")[
            :, list_index_aux_stable
        ]
        H2_sparse_hamiltonian = self.system.param["SPARSE_HAMILTONIAN"][:, list_states]
        H2_sparse_hamiltonian = H2_sparse_hamiltonian - sp.sparse.diags(
            H2_sparse_hamiltonian.diagonal(0),
            format="csc",
            shape=H2_sparse_hamiltonian.shape,
        )
        V1_norm_squared = np.array(
            np.sum(np.abs(H2_sparse_hamiltonian).power(2), axis=0))[:, 0]
        C1_norm_squared_by_state = np.sum(np.abs(C2_phi) ** 2, axis=1)
        return np.sqrt(V1_norm_squared * C1_norm_squared_by_state) / hbar

    def error_sflux_hier(self, Φ, list_s0):
        """
        The error associated with losing all flux terms inside the kth auxiliary to
        states not contained in S_t. This corresponds to equation 30 in arXiv:2008.06496

        PARAMETERS
        ----------
        1. Φ : np.array
               the current full hierarchy vector
        2. list_s0 : list
                     a list of the current states (absolute index)

        RETURNS
        -------
        1. E2_flux_state : array
                           the error introduced by losing flux within k from S_t to S_t^C
                           for each k in A_t

        """
        # Construct the 2D phi and sparse Hamiltonian
        # -------------------------------------------
        list_s0 = np.array(list_s0)
        C2_phi = np.asarray(Φ).reshape([self.n_state, self.n_hier], order="F")

        H2_sparse_hamiltonian = self.system.param["SPARSE_HAMILTONIAN"][:, list_s0]

        # Remove the components of the Hamiltonian that map S0-->S0
        # ---------------------------------------------------------
        H2_sparse_subset = sparse.coo_matrix(
            H2_sparse_hamiltonian[np.ix_(list_s0, range(len(list_s0)))]
        )
        H2_removal = sparse.csc_matrix(
            (
                H2_sparse_subset.data,
                (
                    list_s0[H2_sparse_subset.row],
                    np.arange(len(list_s0))[H2_sparse_subset.col],
                ),
            ),
            shape=H2_sparse_hamiltonian.shape,
        )
        H2_sparse_hamiltonian = H2_sparse_hamiltonian - H2_removal
        D2_derivative_abs_sq = np.abs(
            H2_sparse_hamiltonian @ sparse.csc_matrix(C2_phi) / hbar).power(2)
        return np.sqrt(np.array(np.sum(D2_derivative_abs_sq, axis=0))[0])

    def error_deriv(self, Φ, z_step, list_index_aux_stable=None):
        """
        The error associated with losing all flux terms into the k auxiliary and n state,
        where k is in A_t and n is in S_t. This corresponds to equation 29 in arXiv:2008.06496

        PARAMETERS
        ----------
        1. Φ : np.array
               The current full hierarchy vector
        2. z_step : list
                    the list of noise terms (compressed) for the next timestep
        3. list_index_aux_stable : list
                                   a list relative indices for the stable auxiliaries

        RETURNS
        -------
        1. E2_del_phi : np.array
                        the error associated with losing flux to a component (either
                        hierarchy or state basis element) in H_t direct sum S_t
        """
        if list_index_aux_stable is not None:
            # Error arises for flux only out of the stable auxiliaries
            # --------------------------------------------------------
            Φ_stab = np.zeros(self.n_state * self.n_hier, dtype=np.complex128)
            Φ_stab_v = Φ_stab.view().reshape([self.n_state, self.n_hier], order="F")
            Φ_stab_v[:, list_index_aux_stable] = Φ.view().reshape(
                [self.n_state, self.n_hier], order="F"
            )[:, list_index_aux_stable]
        else:
            Φ_stab = Φ

        list_avg_L2 = [
            operator_expectation(L, Φ[: self.n_state]) for L in self.system.list_L2_coo
        ]

        P1_del_phi = (
            self.eom.K2_k @ Φ_stab + self.eom.K2_kp1 @ Φ_stab + self.eom.K2_km1 @ Φ_stab
        )

        for j in range(len(self.system.list_absindex_L2)):
            P1_del_phi += z_step[j] * self.eom.Z2_k[j] @ Φ_stab
            P1_del_phi += np.conj(list_avg_L2[j]) * self.eom.Z2_kp1[j] @ Φ_stab

        E2_del_phi = np.abs(
            P1_del_phi.reshape([self.n_state, self.n_hier], order="F") / hbar
        )

        return E2_del_phi

    def error_deletion(self, Φ, delta_t):
        """
        The error associated with setting the corresponding component of Phi to 0,
        corresponding to equation 34 and 42 in arXiv:2008.06496

        PARAMETERS
        ----------
        1. Φ : np.array
               The current position of the full hierarchy vector
        2. delta_t : float
                     the timestep for the calculation

        RETURNS
        -------
        1. E2_site_aux : np.array
                         the error induced by removing components of Φ in A_t+S_t
        """

        # Error arising from removing the auxiliary directly
        # --------------------------------------------------
        E2_site_aux = np.abs(
            np.asarray(Φ).reshape([self.n_state, self.n_hier], order="F") / delta_t
        )

        return E2_site_aux

    def error_flux_down(self, Φ, type):
        """
        A function that returns the error associated with neglecting flux from members of
        A_t to auxiliaries in A_t^C that arise due to flux from higher auxiliaries to
        lower auxiliaries. This corresponds to equation 33 and 41 in arXiv:2008.06496

        PARAMETERS
        ----------
        1. Φ : np.array
               The current state of the hierarchy
        2. type: string
                 'H' - Hierarchy type calculation
                 'S' - State type calculation

        RETURNS
        -------
        1. E2_flux_down_error : np.array
                                the error induced by neglecting flux from A_t (or A_S)
                                to auxiliaries with higher summed index in A_t^C.
        """
        # Constants
        # ---------
        list_new_states = [
            self.system.list_state_indices_by_hmode[:, 0] + i * self.n_state
            for i in range(self.n_hier)
        ]
        list_modes_from_site_index = [
            item for sublist in list_new_states for item in sublist
        ]

        # Reshape hierarchy (to matrix)
        # ------------------------------
        P2_pop_site = (
            np.abs(np.asarray(Φ).reshape([self.n_state, self.n_hier], order="F")) ** 2
        )
        P1_aux_norm = np.sqrt(np.sum(P2_pop_site, axis=0))
        P2_modes_from0 = np.asarray(Φ)[
            np.tile(self.system.list_state_indices_by_hmode[:, 0], self.n_hier)
        ]
        P2_pop_modes_down_1 = (np.abs(P2_modes_from0) ** 2).reshape(
            [self.n_hmodes, self.n_hier], order="F"
        )
        P1_modes = np.asarray(Φ)[list_modes_from_site_index]
        P2_pop_modes = np.abs(P1_modes).reshape([self.n_hmodes, self.n_hier], order="F")

        # Get flux factors
        # ----------------
        G2_bymode = np.array(self.system.g)[
            np.tile(list(range(self.n_hmodes)), self.n_hier)
        ].reshape([self.n_hmodes, self.n_hier], order="F")
        W2_bymode = np.array(self.system.w)[
            np.tile(list(range(self.n_hmodes)), self.n_hier)
        ].reshape([self.n_hmodes, self.n_hier], order="F")

        F2_filter_aux = np.array(
            [
                [
                    1 if aux[self.system.list_absindex_mode[i_mode_rel]] - 1 >= 0 else 0
                    for aux in self.hierarchy.auxiliary_list
                ]
                for i_mode_rel in range(self.n_hmodes)
            ]
        )

        if type == "H":
            # Hierarchy Type Downward Flux
            # ============================
            E2_flux_down_error = (
                np.real(
                    F2_filter_aux
                    * np.abs(G2_bymode / W2_bymode)
                    * (P2_pop_modes_down_1 * P1_aux_norm[None, :] + P2_pop_modes)
                )
                / hbar
            )
        elif type == "S":
            # State Type Downward Flux
            # ========================
            # Construct <L_m> term
            # --------------------
            E2_lm = np.tile(
                np.sum(
                    F2_filter_aux * np.abs(G2_bymode / W2_bymode) * P2_pop_modes_down_1,
                    axis=0,
                ),
                [self.n_state, 1],
            )

            # Map Error to States
            # -------------------
            M2_state_from_mode = np.zeros([self.n_state, self.system.n_hmodes])
            M2_state_from_mode[
                self.system.list_state_indices_by_hmode[:, 0],
                np.arange(np.shape(self.system.list_state_indices_by_hmode)[0]),
            ] = 1
            E2_flux_down_error = (
                M2_state_from_mode
                @ np.real(F2_filter_aux * np.abs(G2_bymode / W2_bymode) * P2_pop_modes)
                / hbar
            )
            E2_flux_down_error += E2_lm * P2_pop_site / hbar
        else:
            E2_flux_down_error = 0
            raise UnsupportedRequest(type, "error_flux_down")

        return E2_flux_down_error

    def error_flux_up(self, Φ):
        """
        A function that returns the error associated with neglecting flux from members of
        A_t to auxiliaries in A_t^C that arise due to flux from lower auxiliaries to
        higher auxiliaries. This corresponds to equation 31 and 40 in arXiv:2008.06496

        PARAMETERS
        ----------
        1. Φ : np.array
               The current state of the hierarchy

        RETURNS
        -------
        1. E2_flux_up_error : np. array
                              the error induced by neglecting flux from A_t (or A_S)
                              to auxiliaries with lower summed index in A_t^C.
        """
        # Constants
        # ---------
        list_new_states = [
            self.system.list_state_indices_by_hmode[:, 0] + i * self.n_state
            for i in range(self.n_hier)
        ]
        list_modes_from_site_index = [
            item for sublist in list_new_states for item in sublist
        ]

        # Reshape hierarchy (to matrix)
        # ------------------------------
        P1_modes = np.asarray(Φ)[list_modes_from_site_index]
        P2_pop_modes = np.sqrt(np.abs(P1_modes) ** 2).reshape(
            [self.n_hmodes, self.n_hier], order="F"
        )

        # Get flux factors
        # ----------------
        W2_bymode = np.array(self.system.w)[
            np.tile(list(range(self.n_hmodes)), self.n_hier)
        ].reshape([self.n_hmodes, self.n_hier], order="F")
        K2aux_bymode = np.transpose(
            np.array(
                [
                    aux.get_values(self.system.list_absindex_mode)
                    for aux in self.hierarchy.auxiliary_list
                ]
            )
        )

        # Filter out fluxes beyond the hierarchy depth
        # --------------------------------------------
        filter_aux = np.array(
            [
                i_aux
                for (i_aux, aux) in enumerate(self.hierarchy.auxiliary_list)
                if np.sum(aux) + 1 > self.hierarchy.param["MAXHIER"]
            ]
        )
        F2_filter = np.ones([self.n_hmodes, self.n_hier])
        if filter_aux.size > 0:
            F2_filter[:, filter_aux] = 0

        # Filter out Markovian Modes
        # --------------------------
        array2D_mark_param = np.array(
            [
                np.array(param)[self.system.list_absindex_mode]
                for (name, param) in self.hierarchy.param["STATIC_FILTERS"]
                if name == "Markovian"
            ]
        )
        if len(array2D_mark_param) > 0:
            array_mark_param = np.any(array2D_mark_param, axis=0)
            F2_filter_markov = np.ones([self.n_hmodes, self.n_hier])
            F2_filter_markov[array_mark_param, 1:] = 0
            F2_filter = F2_filter * F2_filter_markov

        # Test upward fluxes
        # ------------------
        return F2_filter * np.abs(W2_bymode) * (1 + K2aux_bymode) * P2_pop_modes / hbar

    @property
    def n_hmodes(self):
        return self.system.n_hmodes

    @property
    def n_state(self):
        return self.system.size

    @property
    def n_hier(self):
        return self.hierarchy.size

    @property
    def adaptive(self):
        return self.eom.param["ADAPTIVE"]

    @property
    def adaptive_h(self):
        return self.eom.param["ADAPTIVE_H"]

    @property
    def adaptive_s(self):
        return self.eom.param["ADAPTIVE_S"]

    @property
    def delta_h(self):
        return self.eom.param["DELTA_H"]

    @property
    def delta_s(self):
        return self.eom.param["DELTA_S"]


# -*- coding: utf-8 -*-
import re
import os
import gc
import math
import glob
import nltk
import shutil
import numpy as np
from os import listdir
from zipfile import ZipFile
from collections import Counter
from os.path import isfile, join
from collections import defaultdict
from nltk.stem import PorterStemmer
from nltk.corpus import stopwords as stopWords
from nltk.tokenize import word_tokenize as wordTokenize

nltk.download('punkt')
nltk.download('stopwords')

stopWordsSet = set(stopWords.words('english'))
ps = PorterStemmer()

class Node:
    def __init__(self, docId, freq):
        self.freq = freq
        self.doc = docId
        self.next = None

    def __str__(self):
        return 'doc:' + str(self.doc) + ', freq:' + str(self.freq)

class LinkedList:
    def __init__(self):
        self.head = None
        self.tail = None
        self.n_docs = 0

    def print_list(self):
        aux = self.head
        while aux:
            print(aux)
            aux = aux.next

    def get_doclist(self):
        l = []
        aux = self.head
        while aux:
            l.append([aux.doc, aux.freq])
            aux = aux.next
        return l

    def add_doc(self, doc, freq):
        node = Node(doc, freq)
        if self.head == None:
            self.head = node
        else:
            self.tail.next = node
        self.tail = node
        self.n_docs += 1

def removeSpecialCharacters(text):
    regex = re.compile('[^a-zA-Z0-9\s]')
    return re.sub(regex, '', text)

def tokenizeAndRemoveSpecialCharacters(text):
    return wordTokenize(removeSpecialCharacters(text))

def listsIntersection(lst1, lst2):
    return list(set(lst1) & set(lst2))

def computeNorm(word, maxFreq, linkedListData, wordsInDocument, itf, docsIds):
    sum = 0
    docs = linkedListData[word].get_doclist()

    for doc in wordsInDocument.keys():
        if doc in docs:
            freq = docs[1]
        else:
            freq = 0
        a = pow((0.5 + 0.5 * (freq / maxFreq)), 2)
        b = pow(itf[docsIds[doc]], 2)
        sum = sum + a * b

    return math.sqrt(sum)

def getMaxFreq(word, linkedListData):
    maxFreq = 0
    docs = linkedListData[word].get_doclist()

    for doc in docs:
        freq = doc[1]
        if freq > maxFreq:
            maxFreq = freq
    return maxFreq

def getFileText(text):
    try:
        text = re.search('<TEXT>(.+?)</TEXT>', text, flags=re.DOTALL).group(1)
    except AttributeError: # <TEXT></TEXT> not found in the original string
        text = ''
    return text

def getFileName(text):
    try:
        name = re.search('<DOCNO>(.+?)</DOCNO>', text, flags=re.DOTALL).group(1)
    except AttributeError: # <DOCNO></DOCNO> not found in the original string
        name = ''
    return name

def filterText(removeStopWords, removeStemmingWords, value):
    words = []
    for w in value:
        if removeStopWords == True and w not in stopWordsSet: # stop words
            if removeStemmingWords == True:
                w = ps.stem(w)  # stemming
        elif removeStopWords == False:
            if removeStemmingWords == True:
                w = ps.stem(w)  # stemming
        words.append(w)
    return words

def readDocs(path):
    numDocs = 0
    docsIds = {}
    docsNames = []
    wordsInDocument = {}

    for folder in glob.glob(path):
        for subfolder in glob.glob(folder + '/*'):
            for directory in glob.glob(subfolder + '/*'):
                for subdir in glob.glob(directory):
                    with open(subdir, "r") as doc:
                        text = doc.read()

                        fileName = getFileName(text)

                        wordsInDocument[fileName] = getFileText(text)

                        wordsInDocument[fileName] = tokenizeAndRemoveSpecialCharacters(wordsInDocument[fileName].lower())

                        # atribuindo um id pra cada documento
                        docsIds[fileName] = numDocs

                        # salvando nomes dos arquivos
                        docsNames.append(fileName)

                        numDocs = numDocs + 1
                doc.close()

    return numDocs, docsIds, docsNames, wordsInDocument

def readQueriesDoc(path):
    """Lê o documento en.topics e armazena as queries, e seus respectivos ids, que serão usados"""

    queries = defaultdict()
    with open(path) as file:
        text = file.read()
        splitted = text.split("\n\n")

    for line in splitted:
        if 'topics' in line: continue
        id = re.search(r'<num>(.*?)</num>', line).group(1)
        query = re.search(r'<title>(.*?)</title>.', line, flags=re.DOTALL).group(1)
        query += re.search(r'<narr>(.*?)</narr>', line, flags=re.DOTALL).group(1)
        query = removeSpecialCharacters(query)
        queries[id] = wordTokenize(query)
        queries[id] = filterText(True, True, queries[id])

    return queries

def prob_model(wordsInDocument, query, linkedListData, numDocs, uniqueWords):
    """
    Modelo probabilistico
    """

    print("\nRealizando seleção inicial no Modelo probabilístico...\n")
    # Seleção inicial
    initial_sel_answer = {}

    for doc in wordsInDocument.keys():
        words = wordsInDocument[doc]
        commonWordsBetweenQueryAndDoc = listsIntersection(query, words)
        relevance = 0
        for ki in commonWordsBetweenQueryAndDoc:
            numDocsWithKi = linkedListData[ki].n_docs
            relevance += math.log10((numDocs + 0.5)/(numDocsWithKi + 0.5))
            if relevance > 0:
                initial_sel_answer[doc] = relevance

    # verifica qual a maior relevancia encontrada
    # na seleção inicial dentre os documentos
    max = 0
    doc = ''
    for item in initial_sel_answer.keys():
        if initial_sel_answer[item] > max:
            max = initial_sel_answer[item]
            doc = item

    print("Documento mais relevante na seleção inicial:", doc)

    # Ranqueamento final
    print("\nRealizando Rnaqueamento final no Modelo probabilístico...\n")

    relevantDocs = []
    numRelevantDocs = {}

    sorted_answer = sorted(initial_sel_answer.items(), key = lambda x: x[1], reverse = True)

    relevantDocs.append(sorted_answer[0][0])

    print("Documentos mais relevantes:", relevantDocs)

    final_rank_answer = {}

    for word in query:
        if word not in uniqueWords:
            continue
        postings = linkedListData[word].head
        numRelevantDocs[word] = 0
        while postings:
            if postings.doc in relevantDocs:
                numRelevantDocs[word] += 1
            postings = postings.next

    numRelevantDocsToQuery = len(relevantDocs)

    for doc in wordsInDocument.keys():
        words = wordsInDocument[doc]
        commonWordsBetweenQueryAndDoc = listsIntersection(query, words)
        relevance = 0

        for ki in commonWordsBetweenQueryAndDoc:
            numDocsWithKi = linkedListData[ki].n_docs
            numRelevantDocsWithKi = numRelevantDocs[ki]
            numerator = ((numRelevantDocsWithKi + 0.5) * (numDocs - numDocsWithKi - numRelevantDocsToQuery + numRelevantDocsWithKi + 0.5))
            denominator = ((numRelevantDocsToQuery - numRelevantDocsWithKi + 0.5) * (numDocsWithKi - numRelevantDocsWithKi + 0.5))
            relevance += math.log10(numerator / denominator)
        final_rank_answer[doc] = relevance

    # final_rank_answer

    # Se os dois prints tiverem o mesmo valor, significa que
    # o documento com mais relevância na seleção inicial
    # é o mesmo na seleção final

    max = 0
    for item in final_rank_answer.keys():
        if final_rank_answer[item] > max:
            max = final_rank_answer[item]

    if final_rank_answer[relevantDocs[0]] == max:
        print("\nResultados iguais = ", final_rank_answer[relevantDocs[0]])

    print("\nFim do Modelo probabilístico...\n\n")

    return initial_sel_answer, final_rank_answer

def vect_model(wordsIds, uniqueWords, numDocs, linkedListData, query, docsNames, docsIds):
    """
    Modelo Vetorial

    Iremos utilizar a ponderação de termos TF-IDF, que considera a frequência de cada termo e a frequência inversa de documento.

    Alguns termos possuem maior importância semântica para representar determinado elemento. Por exemplo, um termo que aparece em todos os documentos indexados não nos auxilia em nada no rankeamento dos mesmos, da mesma forma que um termo raro pode ser de extrema importância para identificar um determinado assunto.

    """
    # Ponderação TF-IDF
    print("Realizando ponderação TF-IDF no Modelo Vetorial...\n")

    m = np.zeros((len(uniqueWords), numDocs))

    for word in uniqueWords:
        postings = linkedListData[word].get_doclist()

        idf = math.log2(numDocs/ linkedListData[word].n_docs)

        for node in postings:

            docName = node[0]
            freq = node[1]

            # tf-idf = (1+log2(freq))*idf, se freq > 0
            if freq > 0:
              m[wordsIds[word], docsIds[docName]] = (1 + math.log2(freq)) * idf
            # tf-idf = 0, caso contrario
            #   m[wordsIds[word], docsIds[docName]] = 0

    """Normalização pelo tamanho dos documentos"""
    print("Realizando Normalização pelo tamanho dos documentos no Modelo Vetorial...\n")

    m = m**2
    norm = np.sum(m, axis=0)
    norm = [math.sqrt(norm[i]) for i in range(len(norm))]

    """Realizando a ponderação TF-IDF na query"""
    print("Realizando ponderação TF-IDF na query no Modelo Vetorial...\n")

    i = 0
    queryVector = np.zeros(len(uniqueWords))

    for word in uniqueWords:
        if word in query:
            idf = math.log2(numDocs/ linkedListData[word].n_docs)
            queryVector[i] = (1 + math.log2(query.count(word))) * idf
            i += 1

    """Ranqueamento"""
    print("Realizando Ranqueamento no Modelo Vetorial...\n")

    ranking = {}

    for docName in docsNames:
        ranking[docName] = np.dot(m[:,docsIds[docName]], queryVector) / norm[docsIds[docName]]

    sorted_ranking = {k: v for k, v in sorted(ranking.items(), key = lambda item: item[1], reverse=True)}

    del m
    gc.collect()

    print("\nFim do Modelo Vetorial...\n\n")

    return sorted_ranking

def query_expansion(numDocs, wordsInDocument, docsIds, uniqueWords, linkedListData, query, wordsIds):
    """
    Expansão de consultas

    """

    """Frequência inversa de termo (itf)"""
    print("Realizando Frequência inversa de termo (itf) na Expansao\n")

    # Frequencia inversa de termo (itf) pra cada documento j
    # Calculo dele é dado por:
    # numerador = nro de termos distintos em tds os docs
    # denominador = nro de termos distintos no doc j
    # itf(doc j) = numerador / denominador
    itf = [0] * numDocs

    for doc in wordsInDocument.keys():
        uniqueWordsInDoc = set(wordsInDocument[doc])
        if(len(uniqueWordsInDoc) == 0):
            continue
        docId = docsIds[doc]
        itf[docId] = math.log2(len(uniqueWords) / len(uniqueWordsInDoc))

    """Pesagem dos documentos"""
    print("Realizando Pesagem dos documentos na Expansao\n")

    w = np.zeros((len(uniqueWords), numDocs))

    for word in uniqueWords:
        maxFreq = getMaxFreq(word, linkedListData)
        norm = computeNorm(word, maxFreq, linkedListData, wordsInDocument, itf, docsIds)

        for docs in linkedListData[word].get_doclist():
            fileName = docs[0]
            freq = docs[1]
            a = (0.5 + 0.5 * (freq / maxFreq))
            b = itf[docsIds[fileName]]
            w[wordsIds[word]][docsIds[fileName]] = (a * b) / norm

    """Similaridade"""
    print("Calculando similaridade na Expansao\n")

    # Similaridade
    c = np.dot(w, w.transpose())

    del w
    gc.collect()

    """Pesagem da consulta"""
    print("Realizando Pesagem da consulta na Expansao\n")

    # Pesagem da consulta
    q = [0] * len(query)

    uniqueWordsInQuery = set(query)

    i = 0
    wordsInQueryIds = {}
    for word in query:
        wordsInQueryIds[word] = i
        i = i + 1

    sum = 0
    for i in range(len(itf)):
        sum = sum + math.pow(itf[i], 2)
    norm = math.sqrt(sum)

    for word in query:
        q[wordsInQueryIds[word]] = math.log2(len(uniqueWords)/ len(uniqueWordsInQuery)) / norm

    """Combinação da similaridade com pesagem da consulta"""
    print("Combinando similaridade com pesagem da consulta na Expansao\n")

    result = {}

    # Combinação da similaridade com a pesagem da consulta
    for word in uniqueWords:
        similaridade = 0
        result[word] = similaridade
        for item in query:
            a = q[wordsInQueryIds[item]]
            if item not in uniqueWords:
                b = 0
            else:
                b = c[wordsIds[item]][wordsIds[word]]
            similaridade = similaridade + a * b
            result[word] = similaridade

    """Resultado da expansão"""

    sorted_result = {k: v for k, v in sorted(result.items(), key = lambda item: item[1], reverse=True)}

    keys = list(sorted_result.keys())

    print("Query expandida: ", keys[:4])

    print("\nFim da Expansao...\n\n")

    return sorted_result

def getRelDocs(path):
    rel_docs = defaultdict()

    with open(path) as file:
        for line in file:
            token = wordTokenize(line)
            if token[3] == '1':
                if int(token[0]) in rel_docs.keys():
                    rel_docs[int(token[0])].append(token[2])
                else:
                    rel_docs[int(token[0])] = [token[2]]

    return rel_docs

def evaluation(queryId, rel_docs, answer):
    """
    Avaliação

    A avaliação serve para julgar o quão bem o sistema atende à necessidade de
    informação do usuário. Estamos considerando precisão e revocação.
    """

    intersec = 0

    for rel_doc in answer.keys():
        if rel_doc in rel_docs[int(queryId)]:
            intersec += 1

    precisao = intersec/len(rel_docs[int(queryId)])
    revocacao = intersec/len(answer)

    return precisao, revocacao

def main():

    docsPath = 'a/*'
    topicsPath = 'en.topics.76-125.2010.txt'
    queriesPath = 'en.qrels.76-125.2010.txt'

    numDocs, docsIds, docsNames, wordsInDocument = readDocs(docsPath)

    # Remoção de stopwords, radicalização e indexação

    dict = defaultdict(int)

    for key, value in wordsInDocument.items():
        wordsInDocument[key] = filterText(True, True, value)

    for key, value in wordsInDocument.items():
        for w in set(value):
            dict[w] += 1

    i = 0
    wordsIds = {}
    linkedListData = {}
    uniqueWords = sorted(set(dict.keys()))

    # atribuindo um id pra cada palavra
    for word in uniqueWords:
        wordsIds[word] = i
        linkedListData[word] = LinkedList()
        i = i + 1

    print("\nFound %d unique words" % len(uniqueWords))

    for doc in wordsInDocument.keys():
        words = wordsInDocument[doc]
        for word in set(words):
            linkedListData[word].add_doc(doc, words.count(word))

    queries = readQueriesDoc(topicsPath)

    # Define qual query iremos trabalhar em cima
    queryId = '76'
    query = queries[queryId]

    initial_sel_answer, final_rank_answer = prob_model(wordsInDocument, query, linkedListData, numDocs, uniqueWords)

    ranked_vect_answer = vect_model(wordsIds, uniqueWords, numDocs, linkedListData, query, docsNames, docsIds)

    expansion_answer = query_expansion(numDocs, wordsInDocument, docsIds, uniqueWords, linkedListData, query, wordsIds)

    rel_docs = getRelDocs(queriesPath)

    precisao = defaultdict()
    revocacao = defaultdict()

    precisao[queryId], revocacao[queryId] = evaluation(queryId, rel_docs, initial_sel_answer)

    print("precisao ", precisao[queryId], "\nrevocacao ", revocacao[queryId])

main()


""" Implementation of the Constrained Policy Optimization (CPO) algorithm.

    see:
    Joshua Achiam, David Held, Aviv Tamar, Peter Abbeel
    Constrained Policy Optimization
    ICML 2017

Author:     Sven Gronauer (sven.gronauer@tum.de)
Created:    12.10.2020
Updated:    27.10.2020
"""

import numpy as np
import torch
from rl_safety_algorithms.common import utils
from rl_safety_algorithms.algs.trpo.trpo import TRPOAlgorithm
from rl_safety_algorithms.algs.npg.npg import NaturalPolicyGradientAlgorithm
import rl_safety_algorithms.algs.utils as U
import rl_safety_algorithms.common.mpi_tools as mpi_tools

class CPOAlgorithm(TRPOAlgorithm):
    """Constrained Policy Optimization Algorithm.

    This implementation does not use cost shaping, but relies on exploration
    noise annealing.
    Large parts are inspired and taken from:
    https://github.com/openai/safety-starter-agents/blob/master/safe_rl/pg/agents.py#L209
    (MIT License)
    """
    def __init__(
            self,
            alg: str = 'cpo',
            cost_limit: float = 25.,
            **kwargs
    ):
        super().__init__(
            alg=alg,
            cost_limit=cost_limit,
            use_cost_value_function=True,
            **kwargs
        )
        self.cost_limit = cost_limit
        self.loss_pi_cost_before = 0.

    def adjust_cpo_step_direction(
            self,
            step_dir,
            g_flat,
            c,
            optim_case,
            p_dist,
            data,
            total_steps: int = 25,
            decay: float = 0.8
    ):
        """CPO algorithm performs line-search to ensure constraint satisfaction
        for rewards and costs."""
        step_frac = 1.0
        _theta_old = U.get_flat_params_from(self.ac.pi.net)
        _, old_log_p = self.ac.pi(data['obs'], data['act'])
        expected_rew_improve = g_flat.dot(step_dir)

        # while not within_trust_region:
        for j in range(total_steps):
            new_theta = _theta_old + step_frac * step_dir
            U.set_param_values_to_model(self.ac.pi.net, new_theta)
            acceptance_step = j + 1

            with torch.no_grad():
                loss_pi_rew, _ = self.compute_loss_pi(data=data)
                loss_pi_cost, _ = self.compute_loss_cost_performance(data=data)
                # determine KL div between new and old policy
                q_dist = self.ac.pi.dist(data['obs'])
                torch_kl = torch.distributions.kl.kl_divergence(
                    p_dist, q_dist).mean().item()
            loss_rew_improve = self.loss_pi_before - loss_pi_rew.item()
            cost_diff = loss_pi_cost.item() - self.loss_pi_cost_before

            # Average across MPI processes...
            torch_kl = mpi_tools.mpi_avg(torch_kl)
            loss_rew_improve = mpi_tools.mpi_avg(loss_rew_improve)
            cost_diff = mpi_tools.mpi_avg(cost_diff)

            self.logger.log("Expected Improvement: %.3f Actual: %.3f" % (
                expected_rew_improve, loss_rew_improve))

            if not torch.isfinite(loss_pi_rew) and not torch.isfinite(
                    loss_pi_cost):
                self.logger.log('WARNING: loss_pi not finite')
            elif loss_rew_improve < 0 if optim_case > 1 else False:
                self.logger.log('INFO: did not improve improve <0')

            elif cost_diff > max(-c, 0):
                self.logger.log(f'INFO: no improve {cost_diff} > {max(-c, 0)}')
            elif torch_kl > self.target_kl * 1.5:
                self.logger.log(
                    f'INFO: violated KL constraint {torch_kl} at step {j + 1}.')
            else:
                # step only if surrogate is improved and we are
                # within the trust region
                self.logger.log(f'Accept step at i={j + 1}')
                break
            step_frac *= decay
        else:
            self.logger.log('INFO: no suitable step found...')
            step_dir = torch.zeros_like(step_dir)
            acceptance_step = 0

        U.set_param_values_to_model(self.ac.pi.net, _theta_old)
        return step_frac * step_dir, acceptance_step

    def algorithm_specific_logs(self):
        NaturalPolicyGradientAlgorithm.algorithm_specific_logs(self)
        self.logger.log_tabular('Misc/cost_gradient_norm')
        self.logger.log_tabular('Misc/A')
        self.logger.log_tabular('Misc/B')
        self.logger.log_tabular('Misc/q')
        self.logger.log_tabular('Misc/r')
        self.logger.log_tabular('Misc/s')
        self.logger.log_tabular('Misc/Lambda_star')
        self.logger.log_tabular('Misc/Nu_star')
        self.logger.log_tabular('Misc/OptimCase')

    def compute_loss_cost_performance(self, data):
        dist, _log_p = self.ac.pi(data['obs'], data['act'])
        ratio = torch.exp(_log_p - data['log_p'])
        cost_loss = (ratio * data['cost_adv']).mean()
        # ent = dist.entropy().mean().item()
        info = {}
        return cost_loss, info

    def update_policy_net(self, data):
        # Get loss and info values before update
        theta_old = U.get_flat_params_from(self.ac.pi.net)
        self.pi_optimizer.zero_grad()
        loss_pi, pi_info = self.compute_loss_pi(data=data)
        self.loss_pi_before = loss_pi.item()
        self.loss_v_before = self.compute_loss_v(data['obs'],
                                                 data['target_v']).item()
        self.loss_c_before = self.compute_loss_c(data['obs'],
                                                 data['target_c']).item()
        # get prob. distribution before updates
        p_dist = self.ac.pi.dist(data['obs'])
        # Train policy with multiple steps of gradient descent
        loss_pi.backward()
        # average grads across MPI processes
        mpi_tools.mpi_avg_grads(self.ac.pi.net)
        g_flat = U.get_flat_gradients_from(self.ac.pi.net)

        # flip sign since policy_loss = -(ration * adv)
        g_flat *= -1

        x = U.conjugate_gradients(self.Fvp, g_flat, self.cg_iters)
        assert torch.isfinite(x).all()
        eps = 1.0e-8
        # Note that xHx = g^T x, but calculating xHx is faster than g^T x
        xHx = torch.dot(x, self.Fvp(x))  # equivalent to : g^T x
        alpha = torch.sqrt(2 * self.target_kl / (xHx + eps))
        assert xHx.item() >= 0, 'No negative values'

        # get the policy cost performance gradient b (flat as vector)
        self.pi_optimizer.zero_grad()
        loss_cost, _ = self.compute_loss_cost_performance(data=data)
        loss_cost.backward()
        # average grads across MPI processes
        mpi_tools.mpi_avg_grads(self.ac.pi.net)
        self.loss_pi_cost_before = loss_cost.item()
        b_flat = U.get_flat_gradients_from(self.ac.pi.net)

        ep_costs = self.logger.get_stats('EpCosts')[0]
        c = ep_costs - self.cost_limit
        c /= (self.logger.get_stats('EpLen')[0] + eps)  # rescale
        self.logger.log(f'c = {c}')
        self.logger.log(f'b^T b = {b_flat.dot(b_flat).item()}')

        # set variable names as used in the paper
        p = U.conjugate_gradients(self.Fvp, b_flat, self.cg_iters)
        q = xHx
        r = g_flat.dot(p)  # g^T H^{-1} b
        s = b_flat.dot(p)  # b^T H^{-1} b

        # print('b^T g', g_flat.dot(b_flat))

        if b_flat.dot(b_flat) <= 1e-6 and c < 0:
            # feasible step and cost grad is zero: use plain TRPO update...
            A = torch.zeros(1)
            B = torch.zeros(1)
            optim_case = 4
        else:

            self.logger.log(f'q={q.item()}')
            self.logger.log(f'r={r.item()}')
            self.logger.log(f's={s.item()}')
            self.logger.log(f'r/c={(r / c).item()}')
            assert torch.isfinite(r).all()
            assert torch.isfinite(s).all()

            A = q - r ** 2 / s  # must be always >= 0 (Cauchy-Schwarz inequality)
            B = 2 * self.target_kl - c ** 2 / s  # safety line intersects trust-region if B > 0

            if c < 0 and B < 0:
                # point in trust region is feasible and safety boundary doesn't intersect
                # ==> entire trust region is feasible
                optim_case = 3
            elif c < 0 and B >= 0:
                # x = 0 is feasible and safety boundary intersects
                # ==> most of trust region is feasible
                optim_case = 2
            elif c >= 0 and B >= 0:
                # x = 0 is infeasible and safety boundary intersects
                # ==> part of trust region is feasible, recovery possible
                optim_case = 1
                self.logger.log('Alert! Attempting feasible recovery!',
                                'yellow')
            else:
                # x = 0 infeasible, and safety halfspace is outside trust region
                # ==> whole trust region is infeasible, try to fail gracefully
                optim_case = 0
                self.logger.log('Alert! Attempting infeasible recovery!', 'red')

        if optim_case in [3, 4]:
            alpha = torch.sqrt(2 * self.target_kl / (xHx + 1e-8))
            nu_star = torch.zeros(1)
            lambda_star = 1 / alpha
            step_dir = alpha * x

        elif optim_case in [1, 2]:
            def project_on_set(t: torch.Tensor,
                               low: float,
                               high: float
                               ) -> torch.Tensor:
                return torch.Tensor([max(low, min(t, high))])

            lambda_a = torch.sqrt(A / B)
            lambda_b = torch.sqrt(q / (2 * self.target_kl))
            if c < 0:
                lambda_a_star = project_on_set(lambda_a, 0., r / c)
                lambda_b_star = project_on_set(lambda_b, r / c, np.inf)
            else:
                lambda_a_star = project_on_set(lambda_a, r / c, np.inf)
                lambda_b_star = project_on_set(lambda_b, 0., r / c)

            def f_a(lam):
                return -0.5 * (A / (lam + eps) + B * lam) - r * c / (s + eps)

            def f_b(lam):
                return -0.5 * (q / (lam + eps) + 2 * self.target_kl * lam)

            lambda_star = lambda_a_star \
                if f_a(lambda_a_star) >= f_b(lambda_b_star) else lambda_b_star

            # Discard all negative values with torch.clamp(x, min=0)
            nu_star = torch.clamp(lambda_star * c - r, min=0) / (s + eps)
            step_dir = 1. / (lambda_star + eps) * (x - nu_star * p)

        else:  # case == 0
            # purely decrease costs
            lambda_star = torch.zeros(1)
            nu_star = np.sqrt(2 * self.target_kl / (s + eps))
            step_dir = -nu_star * p

        final_step_dir, accept_step = self.adjust_cpo_step_direction(
            step_dir,
            g_flat,
            c=c,
            optim_case=optim_case,
            p_dist=p_dist,
            data=data,
            total_steps=20
        )
        # update actor network parameters
        new_theta = theta_old + final_step_dir
        U.set_param_values_to_model(self.ac.pi.net, new_theta)

        q_dist = self.ac.pi.dist(data['obs'])
        torch_kl = torch.distributions.kl.kl_divergence(
            p_dist, q_dist).mean().item()

        self.logger.store(**{
            'Values/Adv': data['act'].numpy(),
            'Entropy': pi_info['ent'],
            'KL': torch_kl,
            'PolicyRatio': pi_info['ratio'],
            'Loss/Pi': self.loss_pi_before,
            'Loss/DeltaPi': loss_pi.item() - self.loss_pi_before,
            'Misc/StopIter': 1,
            'Misc/AcceptanceStep': accept_step,
            'Misc/Alpha': alpha.item(),
            'Misc/FinalStepNorm': final_step_dir.norm().numpy(),
            'Misc/xHx': xHx.numpy(),
            'Misc/H_inv_g': x.norm().item(),  # H^-1 g
            'Misc/gradient_norm': torch.norm(g_flat).numpy(),
            'Misc/cost_gradient_norm': torch.norm(b_flat).numpy(),
            'Misc/Lambda_star': lambda_star.item(),
            'Misc/Nu_star': nu_star.item(),
            'Misc/OptimCase': int(optim_case),
            'Misc/A': A.item(),
            'Misc/B': B.item(),
            'Misc/q': q.item(),
            'Misc/r': r.item(),
            'Misc/s': s.item(),
        })

def get_alg(env_id, **kwargs) -> CPOAlgorithm:
    defaults = utils.get_defaults_kwargs(alg='cpo', env_id=env_id)
    defaults.update(**kwargs)
    return CPOAlgorithm(
        env_id=env_id,
        **defaults
    )

def learn(
        env_id,
        **kwargs
) -> tuple:
    defaults = utils.get_defaults_kwargs(alg='cpo', env_id=env_id)
    defaults.update(**kwargs)
    alg = CPOAlgorithm(
        env_id=env_id,
        **defaults
    )
    ac, env = alg.learn()
    return ac, env