"""
Base classes

.. codeauthor:: David Zwicker <david.zwicker@ds.mpg.de>
"""

from abc import ABCMeta, abstractmethod
import logging
from typing import Callable, Dict, Optional, TYPE_CHECKING  # @UnusedImport

import numpy as np

from ..fields.base import FieldBase
from ..trackers.base import TrackerCollectionDataType
from ..tools.numba import jit

if TYPE_CHECKING:
    from ..solvers.controller import TRangeType  # @UnusedImport

class PDEBase(metaclass=ABCMeta):
    """ base class for solving partial differential equations """

    explicit_time_dependence: Optional[bool] = None

    def __init__(self, noise: float = 0):
        """
        Args:
            noise (float):
                Magnitude of the additive Gaussian white noise that is supported
                by default. If set to zero, a determinitics partial differential
                equation will be solved. If another noise structure is required
                the respective methods need to be overwritten.
        """
        self._logger = logging.getLogger(self.__class__.__name__)
        self.noise = noise

    @property
    def is_sde(self) -> bool:
        """ flag indicating whether this is a stochastic differential equation

        The :class:`BasePDF` class supports additive Gaussian white noise, whose
        magnitude is controlled by the `noise` property. In this case, `is_sde`
        is `True` if `self.noise != 0`.
        """
        # check for self.noise, in case __init__ is not called in a subclass
        return hasattr(self, 'noise') and self.noise != 0

    @abstractmethod
    def evolution_rate(self, field: FieldBase, t: float = 0) \
        -> FieldBase: pass

    def _make_pde_rhs_numba(self, state: FieldBase) -> Callable:
        """ create a compiled function for evaluating the right hand side """
        raise NotImplementedError

    def make_pde_rhs(self, state: FieldBase, backend: str = 'auto') -> Callable:
        """ return a function for evaluating the right hand side of the PDE

        Args:
            state (:class:`~pde.fields.FieldBase`):
                An example for the state from which the grid and other
                information can be extracted
            backend (str): Determines how the function is created. Accepted
                values are 'python` and 'numba'. Alternatively, 'auto' lets the
                code decide for the most optimal backend.

        Returns:
            Function determining the right hand side of the PDE
        """
        if backend == 'auto':
            try:
                result = self._make_pde_rhs_numba(state)
            except NotImplementedError:
                backend = 'numpy'
            else:
                result._backend = 'numba'  # type: ignore
                return result

        if backend == 'numba':
            result = self._make_pde_rhs_numba(state)
            result._backend = 'numba'  # type: ignore

        elif backend == 'numpy':
            state = state.copy()

            def evolution_rate_numpy(state_data, t: float):
                """ evaluate the rhs given only a state without the grid """
                state.data = state_data
                return self.evolution_rate(state, t).data

            result = evolution_rate_numpy
            result._backend = 'numpy'  # type: ignore

        else:
            raise ValueError(f'Unknown backend `{backend}`')

        return result

    def noise_realization(self, state: FieldBase, t: float = 0) -> FieldBase:
        """ returns a realization for the noise

        Args:
            state (:class:`~pde.fields.ScalarField`):
                The scalar field describing the concentration distribution
            t (float): The current time point

        Returns:
            :class:`~pde.fields.ScalarField`:
            Scalar field describing the evolution rate of the PDE
        """
        if self.noise:
            data = np.random.normal(scale=self.noise, size=state.data.shape)
            return state.copy(data=data, label='Noise realization')
        else:
            return state.copy(data=0, label='Noise realization')

    def _make_noise_realization_numba(self, state: FieldBase) -> Callable:
        """ return a function for evaluating the noise term of the PDE

        Args:
            state (:class:`~pde.fields.FieldBase`):
                An example for the state from which the grid and other
                information can be extracted

        Returns:
            Function determining the right hand side of the PDE
        """
        if self.noise:
            noise_strength = float(self.noise)
            data_shape = state.data.shape

            @jit
            def noise_realization(state_data: np.ndarray, t: float):
                """ compiled helper function returning a noise realization """
                return noise_strength * np.random.randn(*data_shape)

        else:
            @jit
            def noise_realization(state_data: np.ndarray, t: float):
                """ compiled helper function returning a noise realization """
                return None

        return noise_realization  # type: ignore

    def _make_sde_rhs_numba(self, state: FieldBase) -> Callable:
        """ return a function for evaluating the noise term of the PDE

        Args:
            state (:class:`~pde.fields.FieldBase`):
                An example for the state from which the grid and other
                information can be extracted

        Returns:
            Function determining the right hand side of the PDE
        """
        evolution_rate = self._make_pde_rhs_numba(state)
        noise_realization = self._make_noise_realization_numba(state)

        @jit
        def sde_rhs(state_data: np.ndarray, t: float):
            """ compiled helper function returning a noise realization """
            return (evolution_rate(state_data, t),
                    noise_realization(state_data, t))

        return sde_rhs  # type: ignore

    def make_sde_rhs(self, state: FieldBase, backend: str = 'auto') \
            -> Callable:
        """ return a function for evaluating the right hand side of the SDE

        Args:
            state (:class:`~pde.fields.FieldBase`):
                An example for the state from which the grid and other
                information can be extracted
            backend (str): Determines how the function is created. Accepted
                values are 'python` and 'numba'. Alternatively, 'auto' lets the
                code decide for the most optimal backend.

        Returns:
            Function determining the deterministic part of the right hand side
            of the PDE together with a noise realization.
        """
        if backend == 'auto':
            try:
                sde_rhs = self._make_sde_rhs_numba(state)
            except NotImplementedError:
                backend = 'numpy'
            else:
                sde_rhs._backend = 'numba'  # type: ignore
                return sde_rhs

        if backend == 'numba':
            sde_rhs = self._make_sde_rhs_numba(state)
            sde_rhs._backend = 'numba'  # type: ignore

        elif backend == 'numpy':
            state = state.copy()

            def sde_rhs(state_data, t: float):
                """ evaluate the rhs given only a state without the grid """
                state.data = state_data
                return (self.evolution_rate(state, t).data,
                        self.noise_realization(state, t).data)

            sde_rhs._backend = 'numpy'  # type: ignore

        else:
            raise ValueError(f'Unknown backend `{backend}`')

        return sde_rhs

    def solve(self, state: FieldBase,
              t_range: "TRangeType",
              dt: float = None,
              tracker: TrackerCollectionDataType = ['progress', 'consistency'],
              method: str = 'auto',
              **kwargs):
        """ convenience method for solving the partial differential equation

        The method constructs a suitable solver
        (:class:`~pde.solvers.base.SolverBase`) and controller
        (:class:`~pde.controller.Controller`) to advance the state over the
        temporal range specified by `t_range`. To obtain full flexibility, it is
        advisable to construct these classes explicitly.

        Args:
            state (:class:`~pde.fields.base.FieldBase`):
                The initial state (which also defines the grid)
            t_range (float or tuple):
                Sets the time range for which the PDE is solved. If only a
                single value `t_end` is given, the time range is assumed to be
                `[0, t_end]`.
            dt (float):
                Time step of the chosen stepping scheme. If `None`, a default
                value based on the stepper will be chosen.
            tracker:
                Defines a tracker that process the state of the simulation at
                fixed time intervals. Multiple trackers can be specified as a
                list. The default value is ['progress', 'consistency'], which
                displays a progress bar and checks the state for consistency,
                aborting the simulation when not-a-number values appear.
            method (:class:`~pde.solvers.base.SolverBase` or str):
                Specifies a method for solving the differential equation. This
                can either be an instance of
                :class:`~pde.solvers.base.SolverBase` or a descriptive name
                like 'explicit' or 'scipy'. The valid names are given by
                :meth:`pde.solvers.base.SolverBase.registered_solvers`.
            **kwargs:
                Additional keyword arguments are forwarded to the solver class

        Returns:
            :class:`~pde.fields.base.FieldBase`:
            The state at the final time point.
        """
        from ..solvers.base import SolverBase

        if method == 'auto':
            method = 'scipy' if dt is None else 'explicit'

        # create solver
        if callable(method):
            solver = method(pde=self, **kwargs)
            if not isinstance(solver, SolverBase):
                self._logger.warn('Solver is not an instance of `SolverBase`. '
                                  'Specified wrong method?')
        else:
            solver = SolverBase.from_name(method, pde=self, **kwargs)

        # create controller
        from ..solvers import Controller
        controller = Controller(solver, t_range=t_range, tracker=tracker)

        # run the simulation
        return controller.run(state, dt)


##############################################################################
# Copyright 2016-2017 Rigetti Computing
#
#    Licensed under the Apache License, Version 2.0 (the "License");
#    you may not use this file except in compliance with the License.
#    You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS,
#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#    See the License for the specific language governing permissions and
#    limitations under the License.
##############################################################################

import pyquil.api as api
import pyquil.quil as pq
import numpy as np
from collections import Counter
from pyquil.gates import STANDARD_GATES, RX, RY
from pyquil.paulis import PauliTerm, PauliSum
import funcsigs

class OptResults(dict):
    """
    Object for holding optimization results from VQE.
    """
    def __getattr__(self, name):
        try:
            return self[name]
        except KeyError:
            raise AttributeError(name)

    __setattr__ = dict.__setitem__
    __delattr__ = dict.__delitem__

class VQE(object):
    """
    The Variational-Quantum-Eigensolver algorithm

    VQE is an object that encapsulates the VQE algorithm (functional
    minimization). The main components of the VQE algorithm are a minimizer
    function for performing the functional minimization, a function that takes a
    vector of parameters and returns a pyQuil program, and a
    Hamiltonian of which to calculate the expectation value.

    Using this object:

        1) initialize with `inst = VQE(minimizer)` where `minimizer` is a
        function that performs a gradient free minization--i.e
        scipy.optimize.minimize(. , ., method='Nelder-Mead')

        2) call `inst.vqe_run(variational_state_evolve, hamiltonian,
        initial_parameters)`. Returns the optimal parameters and minimum
        expecation

    :param minimizer: function that minimizes objective f(obj, param). For
                      example the function scipy.optimize.minimize() needs
                      at least two parameters, the objective and an initial
                      point for the optimization.  The args for minimizer
                      are the cost function (provided by this class),
                      initial parameters (passed to vqe_run() method, and
                      jacobian (defaulted to None).  kwargs can be passed
                      in below.
    :param minimizer_args: (list) arguments for minimizer function. Default=None
    :param minimizer_kwargs: (dict) arguments for keyword args.
                              Default=None

    """

    def __init__(self, minimizer, minimizer_args=[], minimizer_kwargs={}):
        self.minimizer = minimizer
        self.minimizer_args = minimizer_args
        self.minimizer_kwargs = minimizer_kwargs
        self.n_qubits = None

    def vqe_run(self, variational_state_evolve, hamiltonian, initial_params,
                gate_noise=None, measurement_noise=None,
                jacobian=None, qvm=None, disp=None, samples=None, return_all=False):
        """
        functional minimization loop.

        :param variational_state_evolve: function that takes a set of parameters
                                        and returns a pyQuil program.
        :param hamiltonian: (PauliSum) object representing the hamiltonian of
                            which to take the expectation value.
        :param initial_params: (ndarray) vector of initial parameters for the
                               optimization
        :param gate_noise: list of Px, Py, Pz probabilities of gate being
                           applied to every gate after each get application
        :param measurement_noise: list of Px', Py', Pz' probabilities of a X, Y
                                  or Z being applied before a measurement.
        :param jacobian: (optional) method of generating jacobian for parameters
                         (Default=None).
        :param qvm: (optional, QVM) forest connection object.
        :param disp: (optional, bool) display level. If True then each iteration
                     expectation and parameters are printed at each
                     optimization iteration.
        :param samples: (int) Number of samples for calculating the expectation
                        value of the operators.  If `None` then faster method
                        ,dotting the wave function with the operator, is used.
                        Default=None.
        :param return_all: (optional, bool) request to return all intermediate
                           parameters determined during the optimization.
        :return: (vqe.OptResult()) object :func:`OptResult <vqe.OptResult>`.
                 The following fields are initialized in OptResult:
                 -x: set of w.f. ansatz parameters
                 -fun: scalar value of the objective function

                 -iteration_params: a list of all intermediate parameter vectors. Only
                                    returned if 'return_all=True' is set as a vqe_run()
                                    option.

                 -expectation_vals: a list of all intermediate expectation values. Only
                                    returned if 'return_all=True' is set as a
                                    vqe_run() option.
        """
        self._disp_fun = disp if disp is not None else lambda x: None
        iteration_params = []
        expectation_vals = []
        self._current_expectation = None
        if samples is None:
            print("""WARNING: Fast method for expectation will be used. Noise
                     models will be ineffective""")

        if qvm is None:
            qvm = api.QVMConnection(
                    gate_noise=gate_noise,
                    measurement_noise=measurement_noise)
        else:
            self.qvm = qvm

        def objective_function(params):
            """
            closure representing the functional

            :param params: (ndarray) vector of parameters for generating the
                           the function of the functional.
            :return: (float) expectation value
            """
            pyquil_prog = variational_state_evolve(params)
            mean_value = self.expectation(pyquil_prog, hamiltonian, samples, qvm)
            self._current_expectation = mean_value  # store for printing
            return mean_value

        def print_current_iter(iter_vars):
            self._disp_fun("\tParameters: {} ".format(iter_vars))
            if jacobian is not None:
                grad = jacobian(iter_vars)
                self._disp_fun("\tGrad-L1-Norm: {}".format(np.max(np.abs(grad))))
                self._disp_fun("\tGrad-L2-Norm: {} ".format(np.linalg.norm(grad)))

            self._disp_fun("\tE => {}".format(self._current_expectation))
            if return_all:
                iteration_params.append(iter_vars)
                expectation_vals.append(self._current_expectation)

        # using self.minimizer
        arguments = funcsigs.signature(self.minimizer).parameters.keys()

        if disp is not None and 'callback' in arguments:
            self.minimizer_kwargs['callback'] = print_current_iter

        args = [objective_function, initial_params]
        args.extend(self.minimizer_args)
        if 'jac' in arguments:
            self.minimizer_kwargs['jac'] = jacobian

        result = self.minimizer(*args, **self.minimizer_kwargs)

        if hasattr(result, 'status'):
            if result.status != 0:
                self._disp_fun("Classical optimization exited with an error index: %i"
                               % result.status)

        results = OptResults()
        if hasattr(result, 'x'):
            results.x = result.x
            results.fun = result.fun
        else:
            results.x = result

        if return_all:
            results.iteration_params = iteration_params
            results.expectation_vals = expectation_vals
        return results

    @staticmethod
    def expectation(pyquil_prog, pauli_sum, samples, qvm):
        """
        Computes the expectation value of pauli_sum over the distribution
        generated from pyquil_prog.

        :param pyquil_prog: (pyQuil program)
        :param pauli_sum: (PauliSum, ndarray) PauliSum representing the
                          operator of which to calculate the expectation value
                          or a numpy matrix representing the Hamiltonian
                          tensored up to the appropriate size.
        :param samples: (int) number of samples used to calculate the
                        expectation value.  If samples is None then the expectation
                        value is calculated by calculating <psi|O|psi> on the
                        QVM.  Error models will not work if samples is None.

        :param qvm: (qvm connection)

        :returns: (float) representing the expectation value of pauli_sum given
                  given the distribution generated from quil_prog.
        """
        if isinstance(pauli_sum, np.ndarray):
            # debug mode by passing an array
            wf = qvm.wavefunction(pyquil_prog)
            wf = np.reshape(wf.amplitudes, (-1, 1))
            average_exp = np.conj(wf).T.dot(pauli_sum.dot(wf)).real
            return average_exp
        else:
            if not isinstance(pauli_sum, (PauliTerm, PauliSum)):
                raise TypeError("pauli_sum variable must be a PauliTerm or"
                                "PauliSum object")

            if isinstance(pauli_sum, PauliTerm):
                pauli_sum = PauliSum([pauli_sum])

            if samples is None:
                operator_progs = []
                operator_coeffs = []
                for p_term in pauli_sum.terms:
                    op_prog = pq.Program()
                    for qindex, op in p_term:
                        op_prog.inst(STANDARD_GATES[op](qindex))
                    operator_progs.append(op_prog)
                    operator_coeffs.append(p_term.coefficient)

                result_overlaps = qvm.expectation(pyquil_prog,
                                                  operator_programs=operator_progs)
                result_overlaps = list(result_overlaps)
                assert len(result_overlaps) == len(operator_progs), """Somehow we
                didn't get the correct number of results back from the QVM"""
                expectation = sum(list(map(lambda x: x[0]*x[1],
                                           zip(result_overlaps, operator_coeffs))))
                return expectation.real
            else:
                if not isinstance(samples, int):
                    raise TypeError("samples variable must be an integer")
                if samples <= 0:
                    raise ValueError("samples variable must be a postive integer")

                # normal execution via fake sampling
                # stores the sum of contributions to the energy from each operator term
                expectation = 0.0
                for j, term in enumerate(pauli_sum.terms):
                    meas_basis_change = pq.Program()
                    qubits_to_measure = []
                    if term.id() == "":
                        meas_outcome = 1.0
                    else:
                        for index, gate in term:
                            qubits_to_measure.append(index)
                            if gate == 'X':
                                meas_basis_change.inst(RY(-np.pi / 2, index))
                            elif gate == 'Y':
                                meas_basis_change.inst(RX(np.pi / 2, index))

                            meas_outcome = \
                                expectation_from_sampling(pyquil_prog + meas_basis_change,
                                                          qubits_to_measure,
                                                          qvm,
                                                          samples)

                    expectation += term.coefficient * meas_outcome

                return expectation.real

def parity_even_p(state, marked_qubits):
    """
    Calculates the parity of elements at indexes in marked_qubits

    Parity is relative to the binary representation of the integer state.

    :param state: The wavefunction index that corresponds to this state.
    :param marked_qubits: The indexes to be considered in the parity sum.
    :returns: A boolean corresponding to the parity.
    """
    assert isinstance(state, int), "{} is not an integer. Must call " \
                                   "parity_even_p with an integer " \
                                   "state.".format(state)
    mask = 0
    for q in marked_qubits:
        mask |= 1 << q
    return bin(mask & state).count("1") % 2 == 0

def expectation_from_sampling(pyquil_program, marked_qubits, qvm, samples):
    """
    Calculation of Z_{i} at marked_qubits

    Given a wavefunctions, this calculates the expectation value of the Zi
    operator where i ranges over all the qubits given in marked_qubits.

    :param pyquil_program: pyQuil program generating some state
    :param marked_qubits: The qubits within the support of the Z pauli
                          operator whose expectation value is being calculated
    :param qvm: A QVM connection.
    :param samples: Number of bitstrings collected to calculate expectation
                    from sampling.
    :returns: The expectation value as a float.
    """
    # construct program to measure
    for qindex in marked_qubits:
        pyquil_program.measure(qindex, qindex)

    bitstring_samples = qvm.run(pyquil_program, range(max(marked_qubits) + 1), trials=samples)
    bitstring_tuples = list(map(tuple, bitstring_samples))

    freq = Counter(bitstring_tuples)

    # perform weighted average
    expectation = 0
    for bitstring, count in freq.items():
        bitstring_int = int("".join([str(x) for x in bitstring[::-1]]), 2)
        if parity_even_p(bitstring_int, marked_qubits):
            expectation += float(count)/samples
        else:
            expectation -= float(count)/samples
    return expectation


import collections
import logging
import random
from dataclasses import dataclass
from typing import Any, List, Tuple, cast

import numpy as np
import tensorflow as tf
import tensorflow.keras as keras
from srl.base.define import RLObservationType
from srl.base.env.base import EnvRun
from srl.base.rl.algorithms.discrete_action import DiscreteActionConfig, DiscreteActionWorker
from srl.base.rl.base import RLParameter, RLTrainer
from srl.base.rl.registration import register
from srl.base.rl.remote_memory import PriorityExperienceReplay
from srl.rl.functions.common import (
    create_beta_list,
    create_epsilon_list,
    create_gamma_list,
    inverse_rescaling,
    render_discrete_action,
    rescaling,
)
from srl.rl.functions.dueling_network import create_dueling_network_layers
from srl.rl.functions.model import ImageLayerType, create_input_layers
from tensorflow.keras import layers as kl

logger = logging.getLogger(__name__)

"""
DQN
    window_length       : o (config selection)
    Target Network      : o
    Huber loss function : o
    Delay update Target Network: o
    Experience Replay   : o
    Frame skip          : -
    Annealing e-greedy  : x
    Reward clip         : x
    Image preprocessor  : -
Rainbow
    Double DQN               : o (config selection)
    Priority Experience Reply: o (config selection)
    Dueling Network          : o (config selection)
    Multi-Step learning      : x
    Noisy Network            : x
    Categorical DQN          : x
Recurrent Replay Distributed DQN(R2D2)
    LSTM                     : x
    Value function rescaling : o
Never Give Up(NGU)
    Intrinsic Reward : o
    UVFA             : TODO
    Retrace          : x
Agent57
    Meta controller(sliding-window UCB) : o
    Intrinsic Reward split              : o
Other
    invalid_actions : TODO
"""

# ------------------------------------------------------
# config
# ------------------------------------------------------
@dataclass
class Config(DiscreteActionConfig):

    # test
    test_epsilon: float = 0
    test_beta: float = 0

    # model
    window_length: int = 1
    hidden_layer_sizes: Tuple[int, ...] = (512,)
    activation: str = "relu"
    image_layer_type: ImageLayerType = ImageLayerType.DQN
    batch_size: int = 32
    q_ext_lr: float = 0.001
    q_int_lr: float = 0.001
    target_model_update_interval: int = 100

    # double dqn
    enable_double_dqn: bool = True

    # DuelingNetwork
    enable_dueling_network: bool = True
    dueling_network_type: str = "average"

    # Priority Experience Replay
    capacity: int = 100_000
    memory_name: str = "RankBaseMemory"
    memory_warmup_size: int = 1000
    memory_alpha: float = 0.6
    memory_beta_initial: float = 0.4
    memory_beta_steps: int = 1_000_000

    # ucb(160,0.5 or 3600,0.01)
    actor_num: int = 32
    ucb_window_size: int = 160  # UCB上限
    ucb_epsilon: float = 0.5  # UCBを使う確率
    ucb_beta: float = 1  # UCBのβ

    # episodic
    episodic_lr: float = 0.0005
    episodic_count_max: int = 10  # k
    episodic_epsilon: float = 0.001
    episodic_cluster_distance: float = 0.008
    episodic_memory_capacity: int = 30000
    episodic_pseudo_counts: float = 0.1  # 疑似カウント定数

    # lifelong
    lifelong_lr: float = 0.00001
    lifelong_max: float = 5.0  # L

    dummy_state_val: float = 0.0

    def __post_init__(self):
        super().__init__()

    @property
    def observation_type(self) -> RLObservationType:
        return RLObservationType.CONTINUOUS

    @staticmethod
    def getName() -> str:
        return "Agent57_light"

    def assert_params(self) -> None:
        super().assert_params()
        assert self.window_length > 0
        assert self.memory_warmup_size < self.capacity
        assert self.batch_size < self.memory_warmup_size
        assert len(self.hidden_layer_sizes) > 0

register(
    Config,
    __name__ + ":RemoteMemory",
    __name__ + ":Parameter",
    __name__ + ":Trainer",
    __name__ + ":Worker",
)

# ------------------------------------------------------
# RemoteMemory
# ------------------------------------------------------
class RemoteMemory(PriorityExperienceReplay):
    def __init__(self, *args):
        super().__init__(*args)
        self.config = cast(Config, self.config)

        self.init(
            self.config.memory_name,
            self.config.capacity,
            self.config.memory_alpha,
            self.config.memory_beta_initial,
            self.config.memory_beta_steps,
        )

# ------------------------------------------------------
# network
# ------------------------------------------------------
class _QNetwork(keras.Model):
    def __init__(self, config: Config):
        super().__init__()

        in_state, c = create_input_layers(
            config.window_length,
            config.observation_shape,
            config.env_observation_type,
            config.image_layer_type,
        )

        for i in range(len(config.hidden_layer_sizes) - 1):
            c = kl.Dense(
                config.hidden_layer_sizes[i],
                activation=config.activation,
                kernel_initializer="he_normal",
            )(c)

        if config.enable_dueling_network:
            c = create_dueling_network_layers(
                c,
                config.nb_actions,
                config.hidden_layer_sizes[-1],
                config.dueling_network_type,
                activation=config.activation,
            )
        else:
            c = kl.Dense(config.hidden_layer_sizes[-1], activation=config.activation, kernel_initializer="he_normal")(
                c
            )
            c = kl.Dense(
                config.nb_actions, kernel_initializer="truncated_normal", bias_initializer="truncated_normal"
            )(c)

        self.model = keras.Model(in_state, c)

        # 重みを初期化
        dummy_state = np.zeros(shape=(1, config.window_length) + config.observation_shape, dtype=np.float32)
        val = self(dummy_state)
        assert val.shape == (1, config.nb_actions)

    def call(self, state):
        return self.model(state)

# ------------------------------------------------------
# エピソード記憶部(episodic_reward)
# ------------------------------------------------------
class _EmbeddingNetwork(keras.Model):
    def __init__(self, config: Config):
        super().__init__()

        in_state, c = create_input_layers(
            config.window_length,
            config.observation_shape,
            config.env_observation_type,
            config.image_layer_type,
        )

        c = kl.Dense(
            32, activation="relu", kernel_initializer="he_normal", bias_initializer=keras.initializers.constant(0.001)
        )(c)
        self.model = keras.Model(in_state, c)

        # out layer
        self.concatenate = kl.Concatenate()
        self.d1 = kl.Dense(128, activation="relu", kernel_initializer="he_normal")
        c = kl.LayerNormalization()(c)
        self.out = kl.Dense(config.nb_actions, activation="softmax")

        # 重みを初期化
        dummy_state = np.zeros(shape=(1, config.window_length) + config.observation_shape, dtype=np.float32)
        val = self(dummy_state, dummy_state)
        assert val.shape == (1, config.nb_actions)

    def call(self, state1, state2):
        c1 = self.model(state1)
        c2 = self.model(state2)
        c = self.concatenate([c1, c2])
        c = self.d1(c)
        c = self.out(c)
        return c

    def predict(self, state):
        return self.model(state)

# ------------------------------------------------------
# 生涯記憶部(life long novelty module)
# ------------------------------------------------------
class _LifelongNetwork(keras.Model):
    def __init__(self, config: Config):
        super().__init__()

        in_state, c = create_input_layers(
            config.window_length,
            config.observation_shape,
            config.env_observation_type,
            config.image_layer_type,
        )

        c = kl.Dense(
            128,
            activation="relu",
            kernel_initializer="he_normal",
            bias_initializer="he_normal",
        )(c)
        c = kl.LayerNormalization()(c)
        self.model = keras.Model(in_state, c)

        # 重みを初期化
        dummy_state = np.zeros(shape=(1, config.window_length) + config.observation_shape, dtype=np.float32)
        val = self(dummy_state)
        assert val.shape == (1, 128)

    def call(self, state):
        return self.model(state)

# ------------------------------------------------------
# Parameter
# ------------------------------------------------------
class Parameter(RLParameter):
    def __init__(self, *args):
        super().__init__(*args)
        self.config = cast(Config, self.config)

        self.q_ext_online = _QNetwork(self.config)
        self.q_ext_target = _QNetwork(self.config)
        self.q_int_online = _QNetwork(self.config)
        self.q_int_target = _QNetwork(self.config)
        self.emb_network = _EmbeddingNetwork(self.config)
        self.lifelong_target = _LifelongNetwork(self.config)
        self.lifelong_train = _LifelongNetwork(self.config)

    def restore(self, data: Any) -> None:
        self.q_ext_online.set_weights(data[0])
        self.q_ext_target.set_weights(data[0])
        self.q_int_online.set_weights(data[1])
        self.q_int_target.set_weights(data[1])
        self.emb_network.set_weights(data[2])
        self.lifelong_target.set_weights(data[3])
        self.lifelong_train.set_weights(data[4])

    def backup(self):
        d = [
            self.q_ext_online.get_weights(),
            self.q_int_online.get_weights(),
            self.emb_network.get_weights(),
            self.lifelong_target.get_weights(),
            self.lifelong_train.get_weights(),
        ]
        return d

    def summary(self):
        self.q_ext_online.model.summary()
        self.emb_network.summary()
        self.lifelong_target.model.summary()

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

    def calc_target_q(
        self,
        q_online,
        q_target,
        rewards,
        n_states,
        done_list,
        gamma_list,
    ):
        n_q = q_online(n_states).numpy()
        n_q_target = q_target(n_states).numpy()

        target_q = np.zeros(len(rewards))
        for i in range(len(rewards)):
            if done_list[i]:
                gain = rewards[i]
            else:
                # DoubleDQN: indexはonlineQから選び、値はtargetQを選ぶ
                if self.config.enable_double_dqn:
                    n_act_idx = np.argmax(n_q[i])
                else:
                    n_act_idx = np.argmax(n_q_target[i])
                maxq = n_q_target[i][n_act_idx]
                gain = rewards[i] + gamma_list[i] * inverse_rescaling(maxq)
            gain = rescaling(gain)
            target_q[i] = gain

        return target_q

# ------------------------------------------------------
# Trainer
# ------------------------------------------------------
class Trainer(RLTrainer):
    def __init__(self, *args):
        super().__init__(*args)
        self.config = cast(Config, self.config)
        self.parameter = cast(Parameter, self.parameter)
        self.remote_memory = cast(RemoteMemory, self.remote_memory)

        self.q_ext_optimizer = keras.optimizers.Adam(learning_rate=self.config.q_ext_lr)
        self.q_int_optimizer = keras.optimizers.Adam(learning_rate=self.config.q_int_lr)
        self.q_loss = keras.losses.Huber()

        self.emb_optimizer = keras.optimizers.Adam(learning_rate=self.config.episodic_lr)
        self.emb_loss = keras.losses.MeanSquaredError()

        self.lifelong_optimizer = keras.optimizers.Adam(learning_rate=self.config.lifelong_lr)
        self.lifelong_loss = keras.losses.MeanSquaredError()

        self.beta_list = create_beta_list(self.config.actor_num)
        self.gamma_list = create_gamma_list(self.config.actor_num)
        self.epsilon_list = create_epsilon_list(self.config.actor_num)

        self.train_count = 0

    def get_train_count(self):
        return self.train_count

    def train(self):

        if self.remote_memory.length() < self.config.memory_warmup_size:
            return {}

        indices, batchs, weights = self.remote_memory.sample(self.train_count, self.config.batch_size)
        td_error, info = self._train_on_batchs(batchs, weights)
        priorities = abs(td_error) + 0.0001
        self.remote_memory.update(indices, batchs, priorities)

        # targetと同期
        if self.train_count % self.config.target_model_update_interval == 0:
            self.parameter.q_ext_target.set_weights(self.parameter.q_ext_online.get_weights())
            self.parameter.q_int_target.set_weights(self.parameter.q_int_online.get_weights())

        self.train_count += 1
        info["priority"] = np.mean(priorities)
        return info

    def _train_on_batchs(self, batchs, weights):

        # データ形式を変形
        states = []
        actions = []
        n_states = []
        rewards_ext = []
        rewards_int = []
        done_list = []
        gamma_list = []
        beta_list = []
        for b in batchs:
            states.append(b["states"][:-1])
            actions.append(b["action"])
            n_states.append(b["states"][1:])
            rewards_ext.append(b["reward_ext"])
            rewards_int.append(b["reward_int"])
            done_list.append(b["done"])
            gamma_list.append(self.gamma_list[b["actor"]])
            beta_list.append(self.beta_list[b["actor"]])

        states = np.asarray(states)
        n_states = np.asarray(n_states)
        rewards_ext = np.asarray(rewards_ext)
        rewards_int = np.asarray(rewards_int)

        actions_onehot = tf.one_hot(actions, self.config.nb_actions)

        # ----------------------------------------
        # Q network
        # ----------------------------------------
        _params = [
            states,
            n_states,
            actions_onehot,
            done_list,
            weights,
            gamma_list,
        ]
        td_error_ext, loss_ext = self._update_q(
            self.parameter.q_ext_online,
            self.parameter.q_ext_target,
            self.q_ext_optimizer,
            rewards_ext,
            *_params,
        )
        td_error_int, loss_int = self._update_q(
            self.parameter.q_int_online,
            self.parameter.q_int_target,
            self.q_int_optimizer,
            rewards_int,
            *_params,
        )
        td_errors = td_error_ext
        # td_errors = td_error_ext + beta_list * td_error_int

        # ----------------------------------------
        # embedding network
        # ----------------------------------------
        with tf.GradientTape() as tape:
            actions_probs = self.parameter.emb_network(states, n_states)
            emb_loss = self.emb_loss(actions_probs, actions_onehot)

        grads = tape.gradient(emb_loss, self.parameter.emb_network.trainable_variables)
        self.emb_optimizer.apply_gradients(zip(grads, self.parameter.emb_network.trainable_variables))

        # ----------------------------------------
        # lifelong network
        # ----------------------------------------
        lifelong_target_val = self.parameter.lifelong_target(states)
        with tf.GradientTape() as tape:
            lifelong_train_val = self.parameter.lifelong_train(states)
            lifelong_loss = self.lifelong_loss(lifelong_target_val, lifelong_train_val)

        grads = tape.gradient(lifelong_loss, self.parameter.lifelong_train.trainable_variables)
        self.lifelong_optimizer.apply_gradients(zip(grads, self.parameter.lifelong_train.trainable_variables))

        return td_errors, {
            "loss_ext": loss_ext,
            "loss_int": loss_int,
            "emb_loss": emb_loss.numpy(),
            "lifelong_loss": lifelong_loss.numpy(),
        }

    def _update_q(
        self,
        q_online,
        q_target,
        optimizer,
        rewards,
        #
        states,
        n_states,
        actions_onehot,
        done_list,
        weights,
        gamma_list,
    ):
        target_q = self.parameter.calc_target_q(
            q_online,
            q_target,
            rewards,
            n_states,
            done_list,
            gamma_list,
        )

        with tf.GradientTape() as tape:
            q = q_online(states)
            q = tf.reduce_sum(q * actions_onehot, axis=1)
            loss = self.q_loss(target_q * weights, q * weights)

        grads = tape.gradient(loss, q_online.trainable_variables)
        optimizer.apply_gradients(zip(grads, q_online.trainable_variables))

        td_error = (target_q - q).numpy()
        loss = loss.numpy()

        return td_error, loss

# ------------------------------------------------------
# Worker
# ------------------------------------------------------
class Worker(DiscreteActionWorker):
    def __init__(self, *args):
        super().__init__(*args)
        self.config = cast(Config, self.config)
        self.parameter = cast(Parameter, self.parameter)
        self.remote_memory = cast(RemoteMemory, self.remote_memory)

        self.dummy_state = np.full(self.config.observation_shape, self.config.dummy_state_val, dtype=np.float32)

        # actor
        self.beta_list = create_beta_list(self.config.actor_num)
        self.epsilon_list = create_epsilon_list(self.config.actor_num)
        self.gamma_list = create_gamma_list(self.config.actor_num)

        # ucb
        self.actor_index = -1
        self.ucb_recent = []
        self.ucb_actors_count = [1 for _ in range(self.config.actor_num)]  # 1回は保証
        self.ucb_actors_reward = [0.0 for _ in range(self.config.actor_num)]

    def call_on_reset(self, state: np.ndarray, invalid_actions: List[int]) -> None:
        if self.training:
            # エピソード毎に actor を決める
            self.actor_index = self._calc_actor_index()
            self.beta = self.beta_list[self.actor_index]
            self.epsilon = self.epsilon_list[self.actor_index]
            self.gamma = self.gamma_list[self.actor_index]

        else:
            self.epsilon = self.config.test_epsilon
            self.beta = self.config.test_beta

        self.recent_states = [self.dummy_state for _ in range(self.config.window_length + 1)]
        self.recent_states.pop(0)
        self.recent_states.append(state)
        self.invalid_actions = invalid_actions

        # Q値取得用
        self.onehot_actor_idx = tf.one_hot(np.array(self.actor_index), self.config.actor_num)[np.newaxis, ...]

        # sliding-window UCB 用に報酬を保存
        self.episode_reward = 0.0

        # エピソードメモリ(エピソード毎に初期化)
        self.episodic_memory = collections.deque(maxlen=self.config.episodic_memory_capacity)

    # (sliding-window UCB)
    def _calc_actor_index(self) -> int:

        # UCB計算用に保存
        if self.actor_index != -1:
            self.ucb_recent.append(
                (
                    self.actor_index,
                    self.episode_reward,
                )
            )
            self.ucb_actors_count[self.actor_index] += 1
            self.ucb_actors_reward[self.actor_index] += self.episode_reward
            if len(self.ucb_recent) >= self.config.ucb_window_size:
                d = self.ucb_recent.pop(0)
                self.ucb_actors_count[d[0]] -= 1
                self.ucb_actors_reward[d[0]] -= d[1]

        N = len(self.ucb_recent)

        # 全て１回は実行
        if N < self.config.actor_num:
            return N

        # ランダムでactorを決定
        if random.random() < self.config.ucb_epsilon:
            return random.randint(0, self.config.actor_num - 1)

        # UCB値を計算
        ucbs = []
        for i in range(self.config.actor_num):
            n = self.ucb_actors_count[i]
            u = self.ucb_actors_reward[i] / n
            ucb = u + self.config.ucb_beta * np.sqrt(np.log(N) / n)
            ucbs.append(ucb)

        # UCB値最大のポリシー（複数あればランダム）
        return random.choice(np.where(ucbs == np.max(ucbs))[0])

    def call_policy(self, _state: np.ndarray, invalid_actions: List[int]) -> int:
        state = np.asarray([self.recent_states[1:]])

        q_ext = self.parameter.q_ext_online(state)[0].numpy()
        q_int = self.parameter.q_int_online(state)[0].numpy()
        q = q_ext + self.beta * q_int

        if random.random() < self.epsilon:
            self.action = random.choice([a for a in range(self.config.nb_actions) if a not in invalid_actions])
        else:
            # valid actions以外は -inf にする
            q = [(-np.inf if a in invalid_actions else v) for a, v in enumerate(q)]

            # 最大値を選ぶ（複数はほぼないので無視）
            self.action = int(np.argmax(q))

        self.q_ext = q_ext[self.action]
        self.q = q[self.action]
        return self.action

    def call_on_step(
        self,
        next_state: np.ndarray,
        reward_ext: float,
        done: bool,
        next_invalid_actions: List[int],
    ):
        self.episode_reward += reward_ext

        self.recent_states.pop(0)
        self.recent_states.append(next_state)
        self.invalid_actions = next_invalid_actions

        # 内部報酬
        n_s = np.asarray([self.recent_states[1:]])
        episodic_reward = self._calc_episodic_reward(n_s)
        lifelong_reward = self._calc_lifelong_reward(n_s)
        reward_int = episodic_reward * lifelong_reward

        _info = {
            "episodic": episodic_reward,
            "lifelong": lifelong_reward,
            "reward_int": reward_int,
        }
        if not self.training:
            return _info

        batch = {
            "states": self.recent_states[:],
            "action": self.action,
            "reward_ext": reward_ext,
            "reward_int": reward_int,
            "done": done,
            "actor": self.actor_index,
            "next_invalid_actions": next_invalid_actions,
        }

        # priority
        if self.config.memory_name == "ReplayMemory":
            priority = 0
        elif not self.distributed:
            priority = 0
        else:
            _params = [
                n_s,
                [done],
                [self.gamma],
            ]
            target_q_ext = self.parameter.calc_target_q(
                self.parameter.q_ext_online,
                self.parameter.q_ext_target,
                [reward_ext],
                *_params,
            )
            if True:
                priority = abs(target_q_ext - self.q_ext) + 0.0001
            else:
                target_q_int = self.parameter.calc_target_q(
                    self.parameter.q_int_online,
                    self.parameter.q_int_target,
                    [reward_int],
                    *_params,
                )
                priority = abs((target_q_ext + self.beta * target_q_int) - self.q) + 0.0001

        self.remote_memory.add(batch, priority)

        _info["priority"] = priority
        return _info

    def _calc_episodic_reward(self, state):
        k = self.config.episodic_count_max
        epsilon = self.config.episodic_epsilon
        cluster_distance = self.config.episodic_cluster_distance
        c = self.config.episodic_pseudo_counts

        # 埋め込み関数から制御可能状態を取得
        cont_state = self.parameter.emb_network.predict(state)[0].numpy()

        # 初回
        if len(self.episodic_memory) == 0:
            self.episodic_memory.append(cont_state)
            return 1 / c

        # エピソードメモリ内の全要素とユークリッド距離を求める
        euclidean_list = [np.linalg.norm(m - cont_state, ord=2) for m in self.episodic_memory]

        # エピソードメモリに制御可能状態を追加
        self.episodic_memory.append(cont_state)

        # 近いk個を対象
        euclidean_list = np.sort(euclidean_list)[:k]

        # 上位k個の移動平均を出す
        mode_ave = np.mean(euclidean_list)
        if mode_ave == 0.0:
            # ユークリッド距離は正なので平均0は全要素0のみ
            dn = euclidean_list
        else:
            dn = euclidean_list / mode_ave  # 正規化

        # 一定距離以下を同じ状態とする
        dn = np.maximum(dn - cluster_distance, 0)

        # 訪問回数を計算(Dirac delta function の近似)
        dn = epsilon / (dn + epsilon)
        N = np.sum(dn)

        # 報酬の計算
        reward = 1 / (np.sqrt(N) + c)
        return reward

    def _calc_lifelong_reward(self, state):
        # RND取得
        rnd_target_val = self.parameter.lifelong_target(state)[0]
        rnd_train_val = self.parameter.lifelong_train(state)[0]

        # MSE
        error = np.square(rnd_target_val - rnd_train_val).mean()
        reward = 1 + error

        if reward < 1:
            reward = 1
        if reward > self.config.lifelong_max:
            reward = self.config.lifelong_max

        return reward

    def call_render(self, env: EnvRun) -> None:
        state = np.asarray([self.recent_states[1:]])
        q_ext = self.parameter.q_ext_online(state)[0].numpy()
        q_int = self.parameter.q_int_online(state)[0].numpy()
        q = q_ext + self.beta * q_int
        invalid_actions = self.get_invalid_actions(env)

        maxa = np.argmax(q)

        def _render_sub(a: int) -> str:
            return f"{q[a]:5.3f} = {q_ext[a]:5.3f} + {self.beta} * {q_int[a]:5.3f}"

        render_discrete_action(invalid_actions, maxa, env, _render_sub)


#!/usr/bin/env python

__author__ = "Geoffroy Hautier, Bharat Medasani,  Danny Broberg"
__copyright__ = "Copyright 2014, The Materials Project"
__version__ = "1.0"
__maintainer__ = "Geoffroy Hautier, Bharat Medasani"
__email__ = "geoffroy@uclouvain.be, mbkumar@gmail.com"
__status__ = "Development"
__date__ = "November 4, 2012"

from math import sqrt, pi, exp
from collections import defaultdict
from itertools import combinations

import os
import numpy as np

from pymatgen.core import Element
from pymatgen.core.structure import PeriodicSite, Structure
from pymatgen.entries.computed_entries import ComputedStructureEntry
from pymatgen.symmetry.analyzer import SpacegroupAnalyzer

from pycdt.corrections.finite_size_charge_correction import get_correction_freysoldt, get_correction_kumagai
from pycdt.utils.parse_calculations import SingleDefectParser
from pycdt.utils.units import kb, conv, hbar

import warnings
warnings.simplefilter('default')

def freysoldt_correction_from_paths( defect_file_path, bulk_file_path, dielectric,
                                     defect_charge, plot=False):
    """
    A function for performing the Freysoldt correction with a set of file paths.
    If this correction is used, please reference Freysoldt's original paper.
    doi: 10.1103/PhysRevLett.102.016402

    Does not require transformation.json file to exist in file path.

    :param defect_file_path (str): file path to defect folder of interest
    :param bulk_file_path (str): file path to bulk folder of interest
    :param dielectric (float or 3x3 matrix): Dielectric constant (or tensor) for the structure
    :param defect_charge (int): charge of defect structure of interest
    :param plot (bool): allow for plotting electrostatic potential
    :return:
        Dictionary of Freysoldt Correction for defect
    """
    sdp = SingleDefectParser.from_paths( defect_file_path, bulk_file_path, dielectric, defect_charge)
    _ = sdp.freysoldt_loader()
    plt_title = os.path.join( defect_file_path,
                             "{}_chg_{}".format(sdp.defect_entry.name, defect_charge)) if plot else None
    correction = get_correction_freysoldt(sdp.defect_entry,
                                          dielectric,
                                          title=plt_title)

    return correction

def kumagai_correction_from_paths( defect_file_path, bulk_file_path, dielectric,
                                   defect_charge, plot=False):
    """
    A function for performing the Kumagai correction with a set of file paths.
    If this correction is used, please reference Kumagai and Oba's original paper
    (doi: 10.1103/PhysRevB.89.195205) as well as Freysoldt's original
    paper (doi: 10.1103/PhysRevLett.102.016402

    Does not require transformation.json file to exist in file path.

    :param defect_file_path (str): file path to defect folder of interest
    :param bulk_file_path (str): file path to bulk folder of interest
    :param dielectric (float or 3x3 matrix): Dielectric constant (or tensor) for the structure
    :param defect_charge (int): charge of defect structure of interest
    :param plot (bool): allow for plotting electrostatic potential
    :return:
        Dictionary of Kumagai Correction for defect
    """
    sdp = SingleDefectParser.from_paths( defect_file_path, bulk_file_path, dielectric, defect_charge)
    _ = sdp.kumagai_loader()
    plt_title = os.path.join( defect_file_path,
                             "{}_chg_{}".format(sdp.defect_entry.name, defect_charge)) if plot else None
    correction = get_correction_kumagai(sdp.defect_entry,
                                          dielectric,
                                          title=plt_title)

    return correction

warnings.warn("Replaced PyCDT usage of ComputedDefect objects with "
              "DefectEntry objects from pymatgen.analysis.defects.core\n"
              "Will remove ComputedDefect with Version 2.5 of PyCDT.",
              DeprecationWarning)
class ComputedDefect(object):
    """
    Holds all the info concerning a defect computation:
    composition+structure, energy, correction on energy and name
    """
    def __init__(self, entry_defect, site_in_bulk, multiplicity=None,
                 supercell_size=(1, 1, 1), charge=0.0,
                 charge_correction=0.0, other_correction=0.0, name=None):
        """
        Args:
            entry_defect:
                An ComputedStructureEntry object corresponding to the
                defect supercell
            site_in_bulk:
                Site of the defect in bulk supercell. Defect positions
                are often required to perform posteriori corrections
            multiplicity:
                Multiplicity of defect site in a cell. Useful to
                evaluate defect concentrations
            supercell_size:
                Size of the defect supercell in terms of unit cell
            charge:
                The charge of the defect
            charge_correction:
                Correction to the energy due to charge
            other_correction:
                Correction to the energy due to other factors
            name:
                The name of the defect
        """

        self.entry = entry_defect
        self.site = site_in_bulk
        self.multiplicity = multiplicity
        self.supercell_size = supercell_size
        self.charge = charge
        self.charge_correction = charge_correction # Can be added after initialization
        self.other_correction = other_correction
        self.name = name
        if self.name:
            self.full_name = self.name + "_" + str(charge)
        else:
            self.full_name = "defect_" + str(charge)

    def as_dict(self):
        return {'entry': self.entry.as_dict(),
                'site': self.site.as_dict(),
                'multiplicity': self.multiplicity,
                'supercell_size': self.supercell_size,
                'charge': self.charge,
                'charge_correction': self.charge_correction,
                'other_correction': self.other_correction,
                'name': self.name,
                'full_name': self.full_name,
                '@module': self.__class__.__module__,
                '@class': self.__class__.__name__}

    @classmethod
    def from_dict(cls, d):
        return cls(
                ComputedStructureEntry.from_dict(d['entry']),
                PeriodicSite.from_dict(d['site']),
                multiplicity=d.get('multiplicity', None),
                supercell_size=d.get('supercell_size', [1,1,1]),
                charge=d.get('charge', 0.0),
                charge_correction=d.get('charge_correction', 0.0),
                other_correction=d.get('other_correction', 0.0),
                name=d.get('name', None))

warnings.warn("Replaced PyCDT usage of DefectsAnalyzer objects with "
              "DefectPhaseDiagram objects from pymatgen.analysis.defects.thermodynamics\n"
              "Will remove DefectsAnalyzer with Version 2.5 of PyCDT.",
              DeprecationWarning)
class DefectsAnalyzer(object):
    """
    a class aimed at performing standard analysis of defects
    """
    def __init__(self, entry_bulk, e_vbm, mu_elts, band_gap):
        """
        Args:
            entry_bulk:
                the bulk data as an Entry
            e_vbm:
                the energy of the vbm (in eV)
            mu_elts:
                a dictionnary of {Element:value} giving the chemical
                potential of each element
            band_gap:
                the band gap (in eV)
        """
        self._entry_bulk = entry_bulk
        self._e_vbm = e_vbm
        self._mu_elts = mu_elts
        self._band_gap = band_gap
        self._defects = []
        self._formation_energies = []

    def as_dict(self):
        d = {'entry_bulk': self._entry_bulk.as_dict(),
             'e_vbm': self._e_vbm,
             'mu_elts': {k.symbol:v for k,v in self._mu_elts.items()},
             'band_gap': self._band_gap,
             'defects': [d.as_dict() for d in self._defects],
             'formation_energies': self._formation_energies,
             "@module": self.__class__.__module__,
             "@class": self.__class__.__name__}
        return d

    @classmethod
    def from_dict(cls, d):
        struct = d['entry_bulk']['structure']
        struct = struct if isinstance(struct, Structure) \
            else Structure.from_dict(struct)
        entry_bulk = ComputedStructureEntry(struct, d['entry_bulk']['energy'])
        analyzer = DefectsAnalyzer(
            entry_bulk, d['e_vbm'],
            {Element(el): d['mu_elts'][el] for el in d['mu_elts']}, d['band_gap'])
        for ddict in d['defects']:
            analyzer.add_computed_defect(ComputedDefect.from_dict(ddict))
        return analyzer

    def add_computed_defect(self, defect):
        """
        add a parsed defect to the analyzer
        Args:
            defect:
                a ComputedDefect object
        """
        self._defects.append(defect)
        self._compute_form_en()

    def change_charge_correction(self, i, correction):
        """
        Change the charge correction for defect at index i
        Args:
            i:
                Index of defects
            correction:
                New correction to be applied for defect
        """
        self._defects[i].charge_correction = correction
        self._compute_form_en()

    def change_other_correction(self, i, correction):
        """
        Change the charge correction for defect at index i
        Args:
            i:
                Index of defects
            correction:
                New correction to be applied for defect
        """
        self._defects[i].other_correction = correction
        self._compute_form_en()

    def _get_all_defect_types(self):
        to_return = []
        for d in self._defects:
            if d.name not in to_return: to_return.append(d.name)
        return to_return

    def _compute_form_en(self):
        """
        compute the formation energies for all defects in the analyzer
        """
        self._formation_energies = []
        for d in self._defects:
            #compensate each element in defect with the chemical potential
            mu_needed_coeffs = {}
            for elt in d.entry.composition.elements:
                el_def_comp = d.entry.composition[elt]
                el_blk_comp = self._entry_bulk.composition[elt]
                mu_needed_coeffs[Element(elt)] = el_blk_comp - el_def_comp

            sum_mus = 0.0
            for elt in mu_needed_coeffs:
                sum_mus += mu_needed_coeffs[elt] * self._mu_elts[elt]

            self._formation_energies.append(
                    d.entry.energy - self._entry_bulk.energy + \
                            sum_mus + d.charge*self._e_vbm + \
                            d.charge_correction + d.other_correction)

    def correct_bg_simple(self, vbm_correct, cbm_correct):
        """
        correct the band gap in the analyzer.
        We assume the defects level remain the same when moving the
        band edges
        Args:
            vbm_correct:
                The correction on the vbm as a positive number. e.g.,
                if the VBM goes 0.1 eV down vbm_correct=0.1
            cbm_correct:
                The correction on the cbm as a positive number. e.g.,
                if the CBM goes 0.1 eV up cbm_correct=0.1

        """
        self._band_gap = self._band_gap + cbm_correct + vbm_correct
        self._e_vbm = self._e_vbm - vbm_correct
        self._compute_form_en()

    def get_transition_levels(self):
        """
        Charge transition levels are computed
        :return: Transition levels for each pair of the defects.
        If any pair is missing, the transition level for that pair is
        not within the band gap.
        """
        xlim = (-0.5, self._band_gap+1.5)
        nb_steps = 1000
        x = np.arange(xlim[0], xlim[1], (xlim[1]-xlim[0])/nb_steps)

        y = defaultdict(defaultdict)
        for i, dfct in enumerate(self._defects):
            yval = self._formation_energies[i] + dfct.charge*x
            y[dfct.name][dfct.charge] = yval

        transit_levels = defaultdict(defaultdict)
        for dfct_name in y:
            q_ys = y[dfct_name]
            for qpair in combinations(q_ys.keys(), 2):
                qpair_s = tuple(sorted(list(qpair)))
                y_absdiff = abs(q_ys[qpair_s[1]] - q_ys[qpair_s[0]])
                if y_absdiff.min() < 0.4:
                    transit_levels[dfct_name][qpair_s] = x[np.argmin(y_absdiff)]
        return transit_levels

    def _get_form_energy(self, ef, i):
        return self._formation_energies[i] + self._defects[i].charge*ef

    def get_formation_energies(self, ef=0.0):
        """
        Get the defect formation energies for a given Fermi level
        Args:
            ef:
                the fermi level in eV (with respect to the VBM)
        Returns:
            a list of dict of {'name': defect name, 'charge': defect charge
                               'energy': defect formation energy in eV}
        """
        energies = []
        i = 0
        for i, d in enumerate(self._defects):
            energies.append({
                'name': d.name,
                'charge': d.charge,
                'energy': self._get_form_energy(ef, i)
                })
        return energies

    def get_defects_concentration(self, temp=300, ef=0.0):
        """
        Get the defect concentration for a temperature and Fermi level.
        Note: This method has an approximation and can fail
        [Ref: PRB 86, 144109 (2012)]
        Args:
            temp:
                the temperature in K
            Ef:
                the fermi level in eV (with respect to the VBM)
        Returns:
            A list of dict of {'name': defect name, 'charge': defect charge
                               'conc': defects concentration in m-3}
        """
        conc = []
        struct = self._entry_bulk.structure
        for i, d in enumerate(self._defects):
            cell_multiplier = np.prod(d.supercell_size)
            n = d.multiplicity * cell_multiplier * 1e30 / struct.volume
            conc.append({'name': d.name, 'charge': d.charge,
                         'conc': n*exp(
                             -self._get_form_energy(ef, i)/(kb*temp))})

        return conc

    def get_defects_concentration_old(self, temp=300, ef=0.0):
        """
        get the defect concentration for a temperature and Fermi level
        Written when the site multiplicity is not supplied with ComputedDefect
        Args:
            temp:
                the temperature in K
            Ef:
                the fermi level in eV (with respect to the VBM)
        Returns:
            a list of dict of {'name': defect name, 'charge': defect charge
                               'conc': defects concentration in m-3}
        """
        conc=[]
        spga = SpacegroupAnalyzer(self._entry_bulk.structure, symprec=1e-1)
        struct = spga.get_symmetrized_structure()
        i = 0
        for d in self._defects:
            df_coords = d.site.frac_coords
            target_site=None
            for s in struct.sites:
                sf_coords = s.frac_coords
                if abs(s.frac_coords[0]-df_coords[0]) < 0.1 \
                        and abs(s.frac_coords[1]-df_coords[1]) < 0.1 \
                        and abs(s.frac_coords[2]-df_coords[2]) < 0.1:
                    target_site=s
                    break
            equiv_site_no = len(struct.find_equivalent_sites(target_site))
            n = equiv_site_no * 1e30 / struct.volume
            conc.append({'name': d.name, 'charge': d.charge,
                         'conc': n*exp(
                             -self._get_form_energy(ef, i)/(kb*temp))})
            i += 1
        return conc

    def _get_dos(self, e, m1, m2, m3, e_ext):
        return sqrt(2) / (pi**2*hbar**3) * sqrt(m1*m2*m3) * sqrt(e-e_ext)

    def _get_dos_fd_elec(self, e, ef, t, m1, m2, m3):
        return conv * (2.0/(exp((e-ef)/(kb*t))+1)) * \
               (sqrt(2)/(pi**2)) * sqrt(m1*m2*m3) * \
               sqrt(e-self._band_gap)

    def _get_dos_fd_hole(self, e, ef, t, m1, m2, m3):
        return conv * (exp((e-ef)/(kb*t))/(exp((e-ef)/(kb*t))+1)) * \
               (2.0 * sqrt(2)/(pi**2)) * sqrt(m1*m2*m3) * \
               sqrt(-e)

    def _get_qd(self, ef, t):
        summation = 0.0
        for d in self.get_defects_concentration(t, ef):
            summation += d['charge'] * d['conc']
        return summation

    def get_qi(self, ef, t, m_elec, m_hole):
        from scipy import integrate as intgrl

        elec_den_fn = lambda e: self._get_dos_fd_elec(
                e, ef, t, m_elec[0], m_elec[1], m_elec[2])
        hole_den_fn = lambda e: self._get_dos_fd_hole(
                e, ef, t, m_hole[0], m_hole[1], m_hole[2])

        bg = self._band_gap
        elec_count = -intgrl.quad(elec_den_fn, bg, bg+5)[0]
        hole_count = intgrl.quad(hole_den_fn, -5, 0.0)[0]

        return elec_count + hole_count

    def _get_qtot(self, ef, t, m_elec, m_hole):
        return self._get_qd(ef, t) + self.get_qi(ef, t, m_elec, m_hole)

    def get_eq_ef(self, t, m_elec, m_hole):
        """
        access to equilibrium values of Fermi level and concentrations
        in defects and carriers obtained by self-consistent solution of
        charge balance + defect and carriers concentrations
        Args:
            t: temperature in K
            m_elec: electron effective mass as a 3 value list
                    (3 eigenvalues for the tensor)
            m_hole:: hole effective mass as a 3 value list
                    (3 eigenvalues for the tensor)
        Returns:
            a dict with {
                'ef':eq fermi level,
                'Qi': the concentration of carriers
                      (positive for holes, negative for e-) in m^-3,
                'conc': the concentration of defects as a list of dicts
                }
        """
        from scipy.optimize import bisect
        e_vbm = self._e_vbm
        e_cbm = self._e_vbm+self._band_gap
        ef = bisect(lambda e:self._get_qtot(e,t,m_elec,m_hole), 0,
                self._band_gap)
        return {'ef': ef, 'Qi': self.get_qi(ef, t, m_elec, m_hole),
                'QD': self._get_qd(ef,t),
                'conc': self.get_defects_concentration(t, ef)}

    def get_non_eq_ef(self, tsyn, teq, m_elec, m_hole):
        """
        access to the non-equilibrium values of Fermi level and
        concentrations in defects and carriers obtained by
        self-consistent solution of charge balance + defect and
        carriers concentrations

        Implemented following Sun, R., Chan, M. K. Y., Kang, S.,
        and Ceder, G. (2011). doi:10.1103/PhysRevB.84.035212

        Args:
            tsyn: the synthesis temperature in K
            teq: the temperature of use in K
            m_elec: electron effective mass as a 3 value list
                    (3 eigenvalues for the tensor)
            m_hole: hole effective mass as a 3 value list
                    (3 eigenvalues for the tensor)
        Returns:
            a dict with {
                'ef':eq fermi level,
                'Qi': the concentration of carriers
                      (positive for holes, negative for e-) in m^-3,
                'conc': the concentration of defects as a list of dict
                }
        """
        from scipy.optimize import bisect
        eqsyn = self.get_eq_ef(tsyn, m_elec, m_hole)
        cd = {}
        for c in eqsyn['conc']:
            if c['name'] in cd:
                cd[c['name']] += c['conc']
            else:
                cd[c['name']] = c['conc']
        ef = bisect(lambda e:self._get_non_eq_qtot(cd, e, teq, m_elec, m_hole),
                    -1.0, self._band_gap+1.0)
        return {'ef':ef, 'Qi':self.get_qi(ef, teq, m_elec, m_hole),
                'conc_syn':eqsyn['conc'],
                'conc':self._get_non_eq_conc(cd, ef, teq)}

    def _get_non_eq_qd(self, cd, ef, t):
        sum_tot = 0.0
        for n in cd:
            sum_d = 0.0
            sum_q = 0.0
            i = 0
            for d in self._defects:
                if d.name == n:
                    sum_d += exp(-self._get_form_energy(ef, i)/(kb*t))
                    sum_q += d.charge * exp(
                            -self._get_form_energy(ef, i)/(kb*t))
                i += 1
            sum_tot += cd[n]*sum_q/sum_d
        return sum_tot

    def _get_non_eq_conc(self, cd, ef, t):
        sum_tot = 0.0
        res=[]
        for n in cd:
            sum_tot = 0
            i = 0
            for d in self._defects:
                if d.name == n:
                    sum_tot += exp(-self._get_form_energy(ef,i)/(kb*t))
                i += 1
            i=0
            for d in self._defects:
                if d.name == n:
                    res.append({'name':d.name,'charge':d.charge,
                                'conc':cd[n]*exp(-self._get_form_energy(
                                    ef,i)/(kb*t))/sum_tot})
                i += 1
        return res

    def _get_non_eq_qtot(self, cd, ef, t, m_elec, m_hole):
        return self._get_non_eq_qd(cd, ef, t) + \
               self.get_qi(ef, t, m_elec, m_hole)

    def correct_bg(self, dict_levels, vbm_correct, cbm_correct):
        """
        NOTE from developers: This code uses deprecated concepts and
            will not be used or maintained going forward (as of 12/15/2017).
            However, we are keeping function here to allow for
            current users to make use of it...

        correct the band gap in the analyzer and make sure the levels move
        accordingly.
        There are two types of defects vbm_like and cbm_like and we need
        to provide a formal oxidation state
        The vbm-like will follow the vbm and the cbm_like the cbm. If nothing
        is specified the defect transition level does not move
        Args:
            dict_levels: a dictionary of type {defect_name:
            {'type':type_of_defect,'q*':formal_ox_state}}
            Where type_of_defect is a string: 'vbm_like' or 'cbm_like'
        """
        self._band_gap = self._band_gap + cbm_correct + vbm_correct
        self._e_vbm = self._e_vbm - vbm_correct
        self._compute_form_en()
        for i in range(len(self._defects)):
            name = self._defects[i].name
            if not name in dict_levels:
                continue

            if dict_levels[name]['type'] == 'vbm_like':
                z = self._defects[i].charge - dict_levels[name]['q*']
                self._formation_energies[i] += z * vbm_correct
            if dict_levels[name]['type'] == 'cbm_like':
                z = dict_levels[name]['q*'] - self._defects[i].charge
                self._formation_energies[i] +=  z * cbm_correct

    def get_defect_occupancies(self):
        """
        NOTE from developers: This code uses deprecated concepts and
            will not be used or maintained going forward (as of 12/15/2017).
            However, we are keeping function here to allow for
            current users to make use of it...

        Defect occupancies with respect to defect chargest.
        The assumption is that the highest charge of defect numerically
        has zero occupancy:
        Ex: In Cr2O3, V_O has 0 occupancy for +2 q and 2 occupancy for 0 q.
            V_{Cr} has 0 occupancy for 0 q and 3 occupancy for -3 q
        Caution: Didn't check for semiconductor with large # of defect q's.
        Returns:
            Defect occupancies as a nested dict
        """
        charges = defaultdict(list)
        for dfct in self._defects:
            charges[dfct.name].append(dfct.charge)

        occupancies = defaultdict(lambda: defaultdict(int))
        for dfct_name in charges:
            for i, q in enumerate(sorted(charges[dfct_name], reverse=True)):
                occupancies[dfct_name][q] = i
            occupancies[dfct_name]['0_occupancy'] = \
                    sorted(charges[dfct_name], reverse=True)[0]
        return occupancies


# -*- coding: utf-8 -*-
"""
Created on Tue Mar 29 12:31:30 2022

@author: LENOVO
"""

import biosteam as bst
import thermosteam as tmo
import numpy as np
from chaospy import distributions as shape
import biosteam as bst
import numpy as np
# All units are explicitly defined here for transparency and easy reference
# Naming conventions:
#     D = Distillation column
#     M = Mixer
#     E = Multi effect evaporator
#     C = Crystalliser
#     H = Heat exchange
#     L = Liquid-liquid extraction unit (Multi stage Mixer settlers)
#     P = Pump (including conveying belt)
#     R = Reactor
#     S = Splitter (including solid/liquid separator)
#     T = Tank or bin for storage
#     U = Other units
#     PS = Process specificiation, not physical units, but for adjusting streams
#     MT = Mixing Tank

# Processes:
#     100: Conversion
#     200: PrimarySeparation
#     300: SecondarySeparation
#     400: Wastewater
#     500: Facilities

#Defining chemicals
#ref:lactic acid code
#Chemicals used
Water = tmo.Chemical('Water')
#New chemical
#FFA = C18H34O2
#C18H34O2 + CH3OH → C19H36O2 + H2O
#FAME (C19H36O2)
#WCO
#C27H54O2 + 3CH3OH → C3H8O3 + FAME
#Triglyceride considered is triolein

#all_chemicals.set_synonym('Tripalmitin' ,'Oil1')
# Oil2 = tmo.Chemical()
Oil2 = tmo.Chemical('TAG', search_ID='122-32-7', phase = 'l')

Oil2.copy_models_from(tmo.Chemical('Water'), ['mu'])
Oil2.Hf = -1607.83*1000

Glycerol = tmo.Chemical('Glycerol', search_ID='56-81-5')
Glycerol_missing_properties = Glycerol.get_missing_properties()
Furfural = tmo.Chemical('Furfural')
Furfural_missing_properties = Furfural.get_missing_properties()
for i in Glycerol_missing_properties:
    try:
        Glycerol.copy_models_from(Furfural,[i])
    except: pass

CaO = tmo.Chemical('Calcium_oxide', phase = 's')
Calcium_Sulphate = tmo.Chemical('Calcium_Sulphate', phase = 's')

all_chemicals = tmo.Chemicals(['Water',
                               'Sulphuric_acid',
                               'Sodium_Hydroxide',
                               'Oleic_acid',
                               #'Tripalmitin',
                               'Methyl_oleate',
                                'Methanol',
                                Oil2,
                                CaO,
                                Calcium_Sulphate,
                                #'Methyl_palmitate',
                                '112-62-9',
                                Glycerol,
                                ])

all_chemicals.compile()
all_chemicals.show()
all_chemicals.set_synonym('Oleic_acid' ,'FFA')
all_chemicals.set_synonym('56-81-5' ,'Glycerol')
all_chemicals.set_synonym('122-32-7' ,'Oil2')
#all_chemicals.set_synonym('Methyl_palmitate' ,'fame1')
all_chemicals.set_synonym('112-62-9' ,'fame2')

##############################Settings#########################################3
bst.settings.set_thermo(all_chemicals)
bst.Stream.display_units.flow = 'kg/hr'
bst.Stream.display_units.composition = True
GWP = 'GWP 100yr'
bst.settings.define_impact_indicator(key=GWP, units='kg*CO2e')
Calcium_oxide = bst.Chemical('Calcium_oxide')
Calcium_oxide.at_state('s')
### TODO.xxx
Calcium_oxide.V.add_model(1e-6)
#############################UNITS#############################################
class PReactor(bst.BatchBioreactor):
    _N_ins = 1
    _N_outs = 1

    @property
    def effluent(self):
        return self.outs[0]
    @effluent.setter
    def effluent(self,effluent):
        self.outs[0]=effluent

    def __init__(self, ID='', ins=None, outs=(), thermo=None,
                  tau=None, N=None, V=None, T= 65 + 275, P=101325,
                  Nmin=2, Nmax=36):

        bst.BatchBioreactor.__init__(self, ID, ins, outs, thermo,
                                    tau = tau , N = N, V = V, T = T,
                                    P = P ,Nmin = Nmin , Nmax = Nmax)
        self.PR_conversion = 0.999

    def _setup(self):
        self.reactions = tmo.SeriesReaction([
            tmo.Rxn('Oleic_acid + Methanol -> Methyl_oleate + Water ',
                'Oleic_acid', X =  self.PR_conversion )
        ])
    # def _cost(self):
    #     density_in_kg_per_m3 =   913
    #     batch_vol = self.design_results['Reactor volume'] * self.V_wf
    #     R = (batch_vol/(0.002))**0.33 #here 2 is 2L
    #     N2 = 600*((1/R)**(2/3))
    #     D2 = 0.4*((batch_vol*4/3.14)**0.33)
    #     Np = 1.27
    #     power = Np * density_in_kg_per_m3*(N2**3)*(D2**5)
    #     self.power_utility(power)

    def _run(self):
        feed = self.ins[0]
        effluent = self.outs[0]
        effluent.copy_like(feed)
        self.reactions(effluent)
        effluent.T = self.T
        effluent.P = self.P

# Calcium oxide addition for Glycerol recovery
class Glycerol_recovery(bst.BatchBioreactor):
    _N_ins = 1
    _N_outs = 1

    @property
    def effluent(self):
        return self.outs[0]
    @effluent.setter
    def effluent(self,effluent):
        self.outs[0]=effluent

    def __init__(self, ID='', ins=None, outs=(), thermo=None,
                  tau=None, N=None, V=None, T= 60 + 275, P=101325,
                  Nmin=2, Nmax=36):

        bst.BatchBioreactor.__init__(self, ID, ins, outs, thermo,
                                    tau = tau , N = N, V = V, T = T,
                                    P = P ,Nmin = Nmin , Nmax = Nmax)

    def _setup(self):
            self.reactions = tmo.SeriesReaction([
                tmo.Rxn('Calcium_oxide + Sulphuric_acid -> Calcium_Sulphate + Water',
                          'Sulphuric_acid', X = 0.999)])
    # def _cost(self):
    #    density_in_kg_per_m3 = 873.9
    #    batch_vol = self.design_results['Reactor volume'] * self.V_wf
    #    R = (batch_vol/(0.002))**0.33 #here 2 is 2L
    #    N2 = 600*((1/R)**(2/3))
    #    D2 = 0.4*((batch_vol*4/3.14)**0.33)
    #    Np = 1.27
    #    power = Np * density_in_kg_per_m3*(N2**3)*(D2**5)
    #    self.power_utility(power)

    def _run(self):
        feed = self.ins[0]
        effluent = self.outs[0]
        effluent.copy_like(feed)
        self.reactions(effluent)
        effluent.T = self.T
        effluent.P = self.P

# Transesterification reactor
class TReactor(bst.BatchBioreactor):
    _N_ins = 1
    _N_outs = 1

    @property
    def effluent(self):
        return self.outs[0]
    @effluent.setter
    def effluent(self,effluent):
        self.outs[0]=effluent

    def __init__(self, ID='', ins=None, outs=(), thermo=None,
                  tau=None, N=None, V=None, T= 60 + 275, P=101325,
                  Nmin=2, Nmax=36):

        bst.BatchBioreactor.__init__(self, ID, ins, outs, thermo,
                                    tau = tau , N = N, V = V, T = T,
                                    P = P ,Nmin = Nmin , Nmax = Nmax)
        self.TR_conversion = 0.999
    def _setup(self):
            self.reactions = tmo.ParallelReaction([
                # tmo.Rxn('Tripalmitin + 3Methanol -> Methyl_palmitate + Glycerol',
                #         'Tripalmitin',
                #         X = 0.9),
                tmo.Rxn('Oil2 + 3Methanol -> Methyl_oleate + Glycerol',
                        'Oil2', X = self.TR_conversion)
            ])

    # def _cost(self):
    #      density_in_kg_per_m3 =   913
    #      batch_vol = self.design_results['Reactor volume'] * self.V_wf
    #      R = (batch_vol/(0.002))**0.33 #here 2 is 2L
    #      N2 = 600*((1/R)**(2/3))
    #      D2 = 0.4*((batch_vol*4/3.14)**0.33)
    #      Np = 1.27
    #      power = Np * density_in_kg_per_m3*(N2**3)*(D2**5)
    #      self.power_utility(power)

    def _run(self):
        feed = self.ins[0]
        effluent = self.outs[0]
        effluent.copy_like(feed)
        self.reactions(effluent)
        effluent.T = self.T
        effluent.P = self.P
############################ STREAMS ##########################################
WCO_range = [4947,5038,5128,
             5219,5310,5400,
             5491,5581,5672,
             5763,5853,5944,
             6035,6125,6216,
             6307,6397,6488,
             6579,6669]
GWP_obtained = [68.03354950443708, -18.60145643725397,]
Total_Collected_WCO = 6669
#Three values 4800,5525,6475
FFA = tmo.Chemical('Oleic_acid')
# Oil1 = tmo.Chemical('Tripalmitin')
Oil2 = all_chemicals['Oil2']
# FFA content is 6% of the total Oil
Collected_WCO = bst.Stream('Collected_WCO',
                            #Oil1 = 940/2 ,
                            Oil2 = 940,
                            FFA = 60,
                            units = 'kg/hr')
# http://dx.doi.org/10.17576/jkukm-2018-si1(2)-10

Collected_WCO.set_total_flow(Total_Collected_WCO,
                              units='kg/hr')
Fresh_water = bst.Stream('Fresh_water')
Washed_WCO = bst.Stream('Washed_WCO')
#TODO.xxx change this depending on the scale up ratios
Fresh_Methanol_1 = bst.Stream('Methanol1')
Fresh_Methanol_2 = bst.Stream('Methanol2')
Fresh_Sulacid = bst.Stream('Sulphuric_acid')
Fresh_Glycerol = bst.Stream('Glycerol')
Recycled_Methanol = bst.Stream('Recycled_Methanol')
Fresh_CaO = bst.Stream('Fresh_CaO')

# Calcium_Sulphate = bst.Chemical('Calcium_Sulphate')
Calcium_Sulphate.at_state('s')
Calcium_Sulphate.V.add_model(1e-6)
Reuse_Calcium_Sulphate = bst.Stream('Reuse_Calcium_Sulphate')

############################# SYSTEMS ##########################################
#Tank to store Collected WCO and Water
T101 = bst.StorageTank('T101_WCO',
                        ins = (Collected_WCO),
                        outs = 'WCO_to_pump')
# P101 = bst.Pump('P101_WCO',
#                   ins = T101-0,
#                   outs = 'WCO_for_waterwash')

T102 = bst.StorageTank('T102_Water',
                        ins = Fresh_water,
                        outs = 'Washing_water_to_pump')
P102 = bst.Pump('P102_Water',
                  ins = T102-0,
                  outs = 'water_for_wash')

# T103 = bst.MixTank('T103_Methanol',
#                     ins = (Fresh_Methanol),
#                     outs = 'Methanol_to_pump')
# #TODO.xxx Check
# P103 = bst.Pump('P103_Methanol',
#                   ins = T103-0,
#                   outs ='Methanol_to_reactor')

# M103_Methanol = bst.Mixer('M103_Methanol_mixing',
#                    ins = (Fresh_Methanol_1,Fresh_Methanol_2),
#                    outs = ('Methanol_to_splitter'))

# S1031 = bst.Splitter('S1031',
#                       ins = M103_Methanol-0,
#                       outs = ('Methanol_for_pretreatment',
#                             'Methanol_for_transesterification'),
#                       split = 0.5)

T104 = bst.StorageTank('T104_Sulacid',
                        ins = Fresh_Sulacid,
                        outs='Sulacid_to_pump')

P104 = bst.Pump('P104_Sulacid',
                  ins = T104-0,
                  outs='Sulacid_to_reactor')

T105 = bst.StorageTank('T105_Glycerol',
                        ins = Fresh_Glycerol,
                        outs = 'Glycerol_to_pump')

P105 = bst.Pump('P105_Glycerol',
                  P = 110300,
                  ins = T105-0,
                  outs = 'Glycerol_to_scrubber')

# T106 = bst.StorageTank('T106_CaO',
#                         ins = Fresh_CaO,
#                         vessel_type='Field erected',
#                         outs = 'CaO_to_pump')

# P106 = bst.Pump('P106_CaO',
#                   ins = Fresh_CaO,
#                   outs = 'CaO_for_deacidification')

#############################SYTEMS AND UNITS########################
#Mixtank to mix WCO and water
M101 = bst.MixTank('M101', ins = (T101-0,P102-0),
                    outs = 'Mixedfeed_for_waterwashing' )

@M101.add_specification(run=True)
def adjust_Washing_water_flow():
        Fresh_water.imass['Water'] = 10 * Collected_WCO.F_mass

#Splitter to seperate the water slurry from the effluent
#Considers 2% loss of oil
S101 = bst.units.Splitter('S101', ins=M101-0,
                    outs=('Water_slurry',
                          'WCO_pretreatment',
                          ),
                    split={'Water': 1,
                            'Oleic_acid': 0,
                            #'Oil1': 0.01,
                            'Oil2': 0.02,
                            })

#Mixer to mix all the streams before adding the stream to the reactor
M102 = bst.Mixer('M102',
                 ins = (S101-1,Fresh_Methanol_1,P104-0),
                 outs = 'mixedfeed_to_pretreatment')
Sul_acid_ratio = 0.0015
#######################################################################
@M102.add_specification(run=True)
def adjust_methanol():
    b = S101.ins[0].imass['Oleic_acid']
    Fresh_Methanol_1.imass['Methanol'] = 3*b
    Fresh_Sulacid.imass['Sulphuric_acid'] = Sul_acid_ratio * S101.outs[1].F_mass

#Reactor for the first esterification to reduce FFA content
#Assumes 90% conversion

R101 = PReactor('R101',
                  ins = M102.outs[0],
                  outs = 'feed_to_Glycerol_scrubber',
                  N = 4,
                  T =  48.5 + 273.15,
                  P = 400000,
                  tau = 2
                  )

#Glycerol scrubber to scrub Sulacid, Methanol and Water
L101_H = bst.units.HXutility('L101_H',
                              ins = R101-0,
                              outs ='feed_to_Glycerol_scrubber',
                              T = 65+273)

L101 = bst.MultiStageMixerSettlers('L101_Glycerol_scrubber',
                            ins = (L101_H-0,P105-0),
                            outs=('Methanol_extract',
                                  'raffinate_with_pretreated_WCO',
                                  ),
                            N_stages = 2,
                                  )
@L101.add_specification(run=True)
def adjust_glycerol():
    b = M102.outs[0].imass['Methanol']
    Fresh_Glycerol.imass['Glycerol'] = 0.547*b

#Methanol recovery
D101 = bst.units.ShortcutColumn('D101',
                                    ins = L101-0,
                                    outs = ('Recycled_methanol',
                                            'Glycerol_for_recovery'),
                                    LHK = ('Methanol','Water'),
                                    k = 2,
                                    P = 90000,
                                    #P = 101325./20,
                                    Lr = 0.99,
                                    Hr = 0.99,
                                    partial_condenser= False,
                                    )

@D101.add_specification
def D101_spec():
    oil_mol = D101.ins[0].imol['122-32-7']
    D101.ins[0].imol['122-32-7'] = 0
    D101._run()
    # D101.ins[0].imol['122-32-7'] = oil_mol
    # D101.outs[1].imol['122-32-7'] = oil_mol

M103 = bst.Mixer('M103',
                  ins = (D101-1,
                         Fresh_CaO),
                  outs = 'mixedfeed_for_deacidification'
                  )

@M103.add_specification(run = True)
def adjust_CaO():
    c = D101.outs[1].imass['Sulphuric_acid']
    print(c)
    # P106.outs[0].imass['Calcium_oxide'] = 5.4*c
    Fresh_CaO.imass['Calcium_oxide'] = 5.4*c

v_for_Calcium_sulphate_prod = 1.2* M103.outs[0].F_vol
#5.4 Kg of CaO/Kg of Sul acid

R102 = Glycerol_recovery('R102',
                          ins = M103-0,
                          outs='effluent_to_splitter',
                          T =  50 + 273.15,
                          tau = 3,
                          N = 3
                          )
# Splitter to remove Calcium sulphate
S103 = bst.units.Splitter('S103',
                          R102-0,
                          ['Reuse_Calcium_Sulphate',
                            'Glycerol_for_recovery'],
                      split={'Calcium_Sulphate': 1,
                              'Methanol': 0,
                              'Glycerol': 0,
                              'Sulphuric_acid': 0,
                              'Calcium_oxide': 0,
                              'Water': 0,}
                          )

T107 = bst.StorageTank('T107_Calcium_Sulphate',
                        ins = S103-0,
                        outs ='co_product_CaSO4')

#Reactor for transesterification

M104 = bst.Mixer('M104',
                 ins = (L101-1,D101-0,Fresh_Methanol_2),
                 outs = 'feed_to_transesterification')

R103 = TReactor('R103',
                 ins = M104.outs[0],
                 outs = 'mixedfeed_to_biodiesel_rectification',
                 T =  50 + 273.15,
                 tau = 4,
                 N = 3
                )

@M104.add_specification(run = True)
def adjust_methanol_for_transesterification():
    b =  M104.ins[0].imass['Oil2'] + M104.ins[1].imass['Oil2'] #M104.ins[0].imass['Tripalmitin'] + M104.ins[1].imass['Tripalmitin'] +
    a = M104.ins[0].imass['Methanol'] / b
    if a < 1.25:
        c = 1.25/a
        Fresh_Methanol_2.imass['Methanol'] = c*b

#Methanol recovery after transesterification
#TODO.xxx check how to recycle
D102 = bst.units.ShortcutColumn('D102',
                                ins = R103-0,
                                outs = ('Recycled_methanol',
                                        'Biodiesel_for_rectification'),
                                LHK = ('Methanol','Water'),
                                k = 4.6,
                                # P = 80000,
                                P = 101325/20,
                                Lr = 0.99999,
                                Hr = 0.99999,
                                partial_condenser= False,
                                    )

#Splitter to remove glycerol from the biodisel
S104 = bst.units.Splitter('S104',
                          D102-1,
                          ['FAME_for_recovery',
                            'Glycerol_for_recovery'],
                      split={'Water':0,
                              'Sulphuric_acid':0,
                              'Oleic_acid':1,
                              #'Tripalmitin':1,
                              'Methyl_oleate':1,
                              #'Methyl_palmitate':1,
                              'Methanol':0,
                              '122-32-7':1,
                              'Glycerol':1})

D103_H = bst.units.HXutility('Heat_exchanger_for_biodiesel_rectification',
                        ins = S104-0, T = 460
                        )

D103 = bst.units.ShortcutColumn('D103',
                               ins = D103_H-0,
                               outs = ('Glycerol_for_recovery',
                                       'BIODIESEL',
                                       ),
                               LHK = ('Glycerol',
                                        'Methyl_oleate'),
                               k = 2,
                               # P = 80000,
                               P=101325/20,
                               Lr = 0.95,
                               Hr = 0.9,
                               partial_condenser= False
                                  )
# Splitter_Final = bst.units.Splitter('S_final',
#                                     D103-1,
#                                     ['Biodiesel',
#                                      'Waste_oil'],
#                                     split={'Water':0,
#                                            'Sulphuric_acid':0,
#                                            'Oleic_acid':0,
#                                            #'Tripalmitin':1,
#                                            'Methyl_oleate':1,
#                                            #'Methyl_palmitate':1,
#                                            'Methanol':0,
#                                            '122-32-7':0,
#                                            'Glycerol':0})

##################################STORAGE TANKS ###############################
T108 = bst.MixTank('Waste_Glycerol',ins=(D103-0,S103-1))
T109 = bst.MixTank('Waste_Methanol',ins= (D102-0,S104-1))
T110 = bst.units.StorageTank('Biodiesel',ins=D103-1)
#################################System simulation ############################
WCO_sys = bst.main_flowsheet.create_system('WCO_sys')
WCO_sys.operating_hours = 345 * 24
WCO_sys.diagram(number=True)
WCO_sys.simulate()

#############################Characterisation_factors###########################
#Collected waste cooking oil impact
#https://v36.ecoquery.ecoinvent.org/Details/LCIA/d1f82144-4856-44f2-a56a-f01a40cc4d51/290c1f85-4cc4-4fa1-b0c8-2cb7f4276dce
Collected_WCO.set_CF(GWP, 0.26797)
#Sulphuric_acid impact
Fresh_Sulacid.set_CF(GWP,0.084321)
#Water
#https://v36.ecoquery.ecoinvent.org/Details/LCIA/979d50d1-eca2-4399-9805-e9fb36c9e463/290c1f85-4cc4-4fa1-b0c8-2cb7f4276dce
Fresh_water.set_CF(GWP, 0.00079796)
#Methanol_GWP
Fresh_Methanol_1.set_CF(GWP,0.65534)
##above is an assump
Fresh_Methanol_2.set_CF(GWP,0.65534)
#Electricity impact from biosteam
bst.PowerUtility.set_CF(GWP,0.38)
#STEAM: https://v36.ecoquery.ecoinvent.org/Details/LCIA/50cb253a-839c-494a-b98e-5ac2f3dc8e07/290c1f85-4cc4-4fa1-b0c8-2cb7f4276dce
bst.HeatUtility.set_CF('low_pressure_steam', GWP, 88.44, basis='MMBtu')
# NaOH GWP
#NaOH.set_CF(GWP,1.3106)
#CaO GWO
#below is an assum
Fresh_CaO.set_CF(GWP,0.1396)
#Waste_gypsum CF = 0.0091162
#Wastewater_CF =  11.432
#Glycerol GWP
Fresh_Glycerol.set_CF(GWP, 0.2588)

waste_methanol_value = 0.65534*WCO_sys.get_mass_flow(T109.outs[0])

waste_Glycerol_value = 0.2588*WCO_sys.get_mass_flow(T108.outs[0])

#https://v36.ecoquery.ecoinvent.org/Details/LCIA/f38b101d-77d0-4940-b775-5a3cd38f9f6d/290c1f85-4cc4-4fa1-b0c8-2cb7f4276dce
#https://v36.ecoquery.ecoinvent.org/Details/LCIA/ddf1244d-1051-47b3-96a3-5404d1036524/290c1f85-4cc4-4fa1-b0c8-2cb7f4276dce
waste_Water_value = 4.186* WCO_sys.get_mass_flow(S101.outs[0])
waste_CaSO4_value = 0.1396*WCO_sys.get_mass_flow(T107.outs[0])
total_out_mass = WCO_sys.get_mass_flow(D103.outs[1])
net_waste_GWP = waste_methanol_value + waste_Glycerol_value + waste_Water_value #+ waste_CaSO4_value
WCO_sys_net = WCO_sys.get_net_electricity_impact(GWP) + WCO_sys.get_total_feeds_impact(GWP)+WCO_sys.get_net_heat_utility_impact('low_pressure_steam', GWP)
GWP_total_displacement = WCO_sys_net - net_waste_GWP
GWP_total_transport = 127*4
a = GWP_total_displacement + GWP_total_transport
b = a/total_out_mass
print(b)

#b = 0.035435
# print(GWP_total_displacement/total_out_mass)
# print('GWP/mmBTU:',a/b)

#Biodiesel from WCO 37.3 (MJ/Kg)
#1 MJ equals 947.82 BTU
#Calorific value of 1 Kg WCO based biodiesel is 35448.36 BTU
#Biodiesel calorific value is  0.035435 mmBTU/Kg
#
#B100: 119,550 Btu/ga
#1 gal is 3.31Kg
#119550/3.31

#Residential #6 Oil; 166.7 lb/mmBtu
#166.7*0.45360.4536/mmBtu
# # # ####################################UNCERTAINITY ANALYSIS ########################################3

# model = bst.Model(WCO_sys)

# # @model.parameter(name='Total_input_WCO',
# #                   distribution=shape.Uniform(5000, 6000))
# # def set_value(X):
# #     Total_Collected_WCO = X
# @model.parameter(name='Pretreatment_conversion',
#                   distribution=shape.Uniform(0.9997, 0.9999))
# def set_conversion(X):
#     R101.PR_conversion = X
# @model.parameter(name='Transesterification_conversion',
#                   distribution=shape.Uniform(0.9997, 0.9999))
# def set_conversion(X):
#     R103.TR_conversion = X
#     # reactor is the object, reactant conversion is not defined

# @model.parameter(name='Sulphuric_acid_ratio',
#                   distribution=shape.Uniform(0.0015, 0.003))
# def set_SULACID_VALUE(X):
#     Sul_acid_ratio = X

# @model.metric(name = 'total_GWP_value')
# def LCA():
#     waste_Water_value = 4.186* WCO_sys.get_mass_flow(S101.outs[0])
#     waste_CaSO4_value = 0.1396*WCO_sys.get_mass_flow(T107.outs[0])
#     total_out_mass = WCO_sys.get_mass_flow(D103.outs[1])
#     net_waste_GWP = waste_methanol_value + waste_Glycerol_value + waste_Water_value #+ waste_CaSO4_value
#     WCO_sys_net = WCO_sys.get_net_electricity_impact(GWP) + WCO_sys.get_total_feeds_impact(GWP)+WCO_sys.get_net_heat_utility_impact('low_pressure_steam', GWP)
#     GWP_total_displacement = WCO_sys_net - net_waste_GWP
#     GWP_total_transport = 127*2
#     a = GWP_total_displacement + GWP_total_transport
#     b = a/total_out_mass
#     return(b)

# np.random.seed(1234) # For consistent results
# N_samples = 50
# rule = 'L' # For Latin-Hypercube sampling
# samples = model.sample(N_samples, rule)
# model.load_samples(samples)
# model.evaluate()
# model.table # All evaluations are stored as a pandas DataFrame
# print(model.table)

# df_rho, df_p = model.spearman_r()
# bst.plots.plot_spearman_1d(df_rho['Biorefinery', 'total_GWP_value'],
#                             index=[i.describe() for i in model.parameters],
#                             name= 'total_GWP_value')
# =============================================================================
#Residential fuel oil GWP 100: 0.44
# https://v36.ecoquery.ecoinvent.org/Details/LCIA/ed9b5526-d54a-4129-bbe0-52a36f02ae99/290c1f85-4cc4-4fa1-b0c8-2cb7f4276dce

#Calculating yield
#MW of AA = Actual yield/Theoritical yield(188 * 3.54)
#Density of WCO
#1 kilogram / cubic meter = 0.001 kilograms / liter