# -*- coding: utf-8 -*-

import sys
from collections import namedtuple

import numpy as np

from .matrix import BLOSUM62

__all__ = ["AlignmentResult", "aligner"]

IS_PY2 = False
if sys.version_info.major == 2:
    IS_PY2 = True

# Container for alignment result
AlignmentResult = namedtuple(
    'AlignmentResult',
    ['seq1', 'seq2', 'start1', 'start2',
     'end1', 'end2', 'n_gaps1', 'n_gaps2',
     'n_mismatches', 'score'])

def aligner(seqj, seqi, method='global', gap_open=-7, gap_extend=-7,
            gap_double=-7, matrix=BLOSUM62, max_hits=1):
    '''Calculates the alignment of two sequences.

    The supported 'methods' are:
        * 'global' for a global Needleman-Wunsh algorithm
        * 'local' for a local Smith-Waterman alignment
        * 'global_cfe' for a global alignment with cost-free ends
        * 'glocal' for an alignment which is 'global' only with respect to
          the shorter sequence (also known as a 'semi-global' alignment)

    Returns the aligned (sub)sequences as character arrays.

    Gotoh, O. (1982). J. Mol. Biol. 162, 705-708.
    Needleman, S. & Wunsch, C. (1970). J. Mol. Biol. 48(3), 443-53.
    Smith, T.F. & Waterman M.S. (1981). J. Mol. Biol. 147, 195-197.

    Arguments:

        - seqj (``sequence``) First aligned iterable object of symbols.
        - seqi (``sequence``) Second aligned iterable object of symbols.
        - method (``str``) Type of alignment: 'global', 'global_cfe', 'local',
          'glocal'.
        - gap_open (``float``) The gap-opening cost.
        - gap_extend (``float``) The cost of extending an open gap.
        - gap_double (``float``) The gap-opening cost if a gap is already open
          in the other sequence.
        - matrix (``dict``) A score matrix dictionary.
        - max_hits (``int``) The maximum number of results to return in
          case multiple alignments with the same score are found. If set to 1,
          a single ``AlignmentResult`` object is returned. If set to values
          larger than 1, a list containing ``AlignmentResult`` objects are
          returned. If set to `None`, all alignments with the maximum score
          are returned.
    '''
    assert max_hits is None or max_hits > 0
    NONE, LEFT, UP, DIAG = range(4)  # NONE is 0
    GAP_CHAR = ord('-') if not IS_PY2 else '-'
    max_j = len(seqj)
    max_i = len(seqi)

    if max_j > max_i:
        flip = 1
        seqi, seqj = seqj, seqi
        max_i, max_j = max_j, max_i
    else:
        flip = 0

    seqi = seqi.encode() if not isinstance(seqi, bytes) else seqi
    seqj = seqj.encode() if not isinstance(seqj, bytes) else seqj

    F = np.zeros((max_i + 1, max_j + 1), dtype=np.float32)
    I = np.ndarray((max_i + 1, max_j + 1), dtype=np.float32)
    I.fill(-np.inf)
    J = np.ndarray((max_i + 1, max_j + 1), dtype=np.float32)
    J.fill(-np.inf)
    pointer = np.zeros((max_i + 1, max_j + 1), dtype=np.uint)  # NONE

    if method == 'global':
        pointer[0, 1:] = LEFT
        pointer[1:, 0] = UP
        F[0, 1:] = gap_open + gap_extend * \
            np.arange(0, max_j, dtype=np.float32)
        F[1:, 0] = gap_open + gap_extend * \
            np.arange(0, max_i, dtype=np.float32)
    elif method == 'global_cfe':
        pointer[0, 1:] = LEFT
        pointer[1:, 0] = UP
    elif method == 'glocal':
        pointer[0, 1:] = LEFT
        F[0, 1:] = gap_open + gap_extend * \
            np.arange(0, max_j, dtype=np.float32)

    for i in range(1, max_i + 1):
        ci = seqi[i - 1:i]
        for j in range(1, max_j + 1):
            cj = seqj[j - 1:j]
            # I
            I[i, j] = max(
                         F[i, j - 1] + gap_open,
                         I[i, j - 1] + gap_extend,
                         J[i, j - 1] + gap_double)
            # J
            J[i, j] = max(
                         F[i - 1, j] + gap_open,
                         J[i - 1, j] + gap_extend,
                         I[i - 1, j] + gap_double)
            # F
            diag_score = F[i - 1, j - 1] + matrix[cj][ci]
            left_score = I[i, j]
            up_score = J[i, j]
            max_score = max(diag_score, up_score, left_score)

            F[i, j] = max(0, max_score) if method == 'local' else max_score

            if method == 'local':
                if F[i, j] == 0:
                    pass  # point[i,j] = NONE
                elif max_score == diag_score:
                    pointer[i, j] = DIAG
                elif max_score == up_score:
                    pointer[i, j] = UP
                elif max_score == left_score:
                    pointer[i, j] = LEFT
            elif method == 'glocal':
                # In a semi-global alignment we want to consume as much as
                # possible of the longer sequence.
                if max_score == up_score:
                    pointer[i, j] = UP
                elif max_score == diag_score:
                    pointer[i, j] = DIAG
                elif max_score == left_score:
                    pointer[i, j] = LEFT
            else:
                # global
                if max_score == up_score:
                    pointer[i, j] = UP
                elif max_score == left_score:
                    pointer[i, j] = LEFT
                else:
                    pointer[i, j] = DIAG

    # container for traceback coordinates
    ij_pairs = []
    if method == 'local':
        # max anywhere
        maxv_indices = np.argwhere(F == F.max())[:max_hits]
        for index in maxv_indices:
            ij_pairs.append(index)
    elif method == 'glocal':
        # max in last col
        max_score = F[:, -1].max()
        maxi_indices = np.argwhere(F[:, -1] == F[:, -1].max())\
            .flatten()[:max_hits]
        for i in maxi_indices:
            ij_pairs.append((i, max_j))
    elif method == 'global_cfe':
        # from i,j to max(max(last row), max(last col)) for free
        row_max = F[-1].max()
        col_max = F[:, -1].max()
        # expecting max to exist on either last column or last row
        if row_max > col_max:
            col_idces = np.argwhere(F[-1] == row_max).flatten()[:max_hits]
            for cid in col_idces:
                ij_pairs.append((i, cid))
        elif row_max < col_max:
            row_idces = np.argwhere(F[:, -1] == col_max).flatten()[:max_hits]
            for rid in row_idces:
                ij_pairs.append((rid, j))
        # special case: max is on last row, last col
        elif row_max == col_max == F[i, j]:
            # check if max score also exist on other cells in last row
            # or last col. we expect only one of the case.
            col_idces = np.argwhere(F[-1] == row_max).flatten()
            row_idces = np.argwhere(F[:, -1] == col_max).flatten()
            ncol_idces = len(col_idces)
            nrow_idces = len(row_idces)

            # tiebreaker between row/col is whichever has more max scores
            if ncol_idces > nrow_idces:
                for cid in col_idces[:max_hits]:
                    ij_pairs.append((i, cid))
            elif ncol_idces < nrow_idces:
                for rid in row_idces[:max_hits]:
                    ij_pairs.append((rid, j))
            elif ncol_idces == nrow_idces == 1:
                ij_pairs.append((i, j))
            else:
                raise RuntimeError('Unexpected multiple maximum global_cfe'
                                   ' scores.')
        else:
            raise RuntimeError('Unexpected global_cfe scenario.')
    else:
        # method must be global at this point
        ij_pairs.append((i, j))

    results = []
    for i, j in ij_pairs:
        align_j = []
        align_i = []
        score = F[i, j]
        p = pointer[i, j]
        # mimic Python's coord system
        if method.startswith("global"):
            end_i, end_j = max_i, max_j
        else:
            end_i, end_j = i, j
        n_gaps_i, n_gaps_j, n_mmatch = 0, 0, 0

        # special case for global_cfe ~ one cell may contain multiple pointer
        # directions
        if method == 'global_cfe':
            if i < max_i:
                align_i.extend([c for c in seqi[i:][::-1]])
                align_j.extend([GAP_CHAR] * (max_i - i))
                n_gaps_j += 1
            elif j < max_j:
                align_i.extend([GAP_CHAR] * (max_j - j))
                align_j.extend([c for c in seqj[j:][::-1]])
                n_gaps_i += 1

        while p != NONE:
            if p == DIAG:
                i -= 1
                j -= 1
                ichar = seqi[i]
                jchar = seqj[j]
                if ichar != jchar:
                    n_mmatch += 1
                align_j.append(jchar)
                align_i.append(ichar)
            elif p == LEFT:
                j -= 1
                align_j.append(seqj[j])
                if not align_i or align_i[-1] != GAP_CHAR:
                    n_gaps_i += 1
                align_i.append(GAP_CHAR)
            elif p == UP:
                i -= 1
                align_i.append(seqi[i])
                if not align_j or align_j[-1] != GAP_CHAR:
                    n_gaps_j += 1
                align_j.append(GAP_CHAR)
            else:
                raise Exception('wtf!')
            p = pointer[i, j]

        align_i = bytes(align_i[::-1]) \
            if not IS_PY2 else ''.join(align_i[::-1])
        align_j = bytes(align_j[::-1]) \
            if not IS_PY2 else ''.join(align_j[::-1])

        aln = (AlignmentResult(align_i, align_j, i, j, end_i, end_j,
                               n_gaps_i, n_gaps_j, n_mmatch, score)
               if flip else
               AlignmentResult(align_j, align_i, j, i, end_j, end_i,
                               n_gaps_j, n_gaps_i, n_mmatch, score))

        results.append(aln)

    return results


#!/usr/bin/env python
# -*- coding: utf-8 -*-
r"""Provide the manipulability task.

The manipulability task implements a task that tries to maximize the manipulability measure given in [1]:

.. math:: w(q) = \sqrt{ \det( J(q) W J(q)^\top ) }

where :math:`W` is a constant weight matrix, :math:`q` are the joint positions, and :math:`J(q)` is the jacobian.
The gradient of :math:`w` is then computed and projected using the gradient projection method [2].

The quadratic cost being minimized is:

.. math:: ||\dot{q} - \dot{q}_0||^2

where :math:`\dot{q}` are the joint velocities being optimized,
:math:`\dot{q}_0 = k_0 \left( \frac{\partial w(q)}{\partial q} \right)^\top` where :math:`k_0 > 0` and
:math:`w(q)` is an objective function of the joint variables, where in this case, the manipulability measure is
given by :math:`w(q) = \sqrt{\det( J(q) J^\top(q) )}`. By maximizing this measure, we move away from singularities.

The implementation of this class is inspired by [1] (which is licensed under the LGPLv2).

References:
    - [1] "OpenSoT: A whole-body control library for the compliant humanoid robot COMAN", Rocchi et al., 2015
"""

# TODO: finish to implement this

import numpy as np

from pyrobolearn.priorities.tasks import JointVelocityTask

__author__ = "Brian Delhaisse"
__copyright__ = "Copyright 2019, PyRoboLearn"
__credits__ = ["Enrico Mingo Hoffman (C++)", "Brian Delhaisse (Python + doc)"]
__license__ = "GNU GPLv3"
__version__ = "1.0.0"
__maintainer__ = "Brian Delhaisse"
__email__ = "briandelhaisse@gmail.com"
__status__ = "Development"

class ManipulabilityTask(JointVelocityTask):
    r"""Manipulability Task

    The manipulability task implements a task that tries to maximize the manipulability measure given in [1]:

    .. math:: w(q) = \sqrt{ \det( J(q) W J(q)^\top ) }

    where :math:`W` is a constant weight matrix, :math:`q` are the joint positions, and :math:`J(q)` is the jacobian.
    The gradient of :math:`w` is then computed and projected using the gradient projection method [2].

    The quadratic cost being minimized is:

    .. math:: ||\dot{q} - \dot{q}_0||^2

    where :math:`\dot{q}` are the joint velocities being optimized,
    :math:`\dot{q}_0 = k_0 \left( \frac{\partial w(q)}{\partial q} \right)^\top` where :math:`k_0 > 0` and
    :math:`w(q)` is an objective function of the joint variables, where in this case, the manipulability measure is
    given by :math:`w(q) = \sqrt{\det( J(q) J^\top(q) )}`. By maximizing this measure, we move away from singularities.

    References:
        - [1] "Robotics: Modelling, Planning, and Control", Siciliano et al., 2010
        - [2] "OpenSoT: A whole-body control library for the compliant humanoid robot COMAN", Rocchi et al., 2015
    """

    def __init__(self, model, weight=1., constraints=[]):
        """
        Initialize the task.

        Args:
            model (ModelInterface): model interface
            weight (float, np.array[float[N,N]]): weight scalar or matrix associated to the task.
            constraints (list[Constraint]): list of constraints associated with the task.
        """
        super(ManipulabilityTask, self).__init__(model=model, weight=weight, constraints=constraints)

        raise NotImplementedError("This class has not been implemented yet.")


# -*- coding: utf-8 -*-

import numpy as np

class AEP_load():
    """ Calculate AEP (and fatigue load) of a given wind farm.

    Parameters
    ----------
    site_conditions: class object (WindResource class)
        Storing the terrain flow and wind resource data, and providing a
        function to get the site conditions of given location(s) for certain
        inflow wind direction.

    wind_farm: class object (WindFarm class)
        Defining the design of the wind farm, i.e., the layout and turbine type
        and hub height information for all turbines.

    wake_model: class object (WakeModel class)
        Specifying the wake models to use in the calculation. The default ones
        are N.O. Jensen for wake deficit and G.C. Larsen for turbulence.

    ws_binned: array:float, np.linspace(1, 30, num=30)
        Discretized wind speed bins at a reference height above the ground for
        far field inflow. ) [m/s]

    wd_binned: array:float, np.linspace(0, 330, num=12)
        Discretized wind direction bins at a reference height above the ground
        for far field inflow [deg]. This should be equally spaced and contains
        at least two wind directions, as we derive the sector width from it.

    height_ref: float, 85
        Reference height above the ground for defining the far field inflow
        condition [m].

    k_star: float, 0.075
        Wake decay parameter used in the N.O. Jensen wake model for the
        wake_model [-].

    availability: float, 1.0
        Availability factor for the wind farm [-].

    num_evals: integer, 0
        Number of evluations of .cal_AEP_load() [-].

    z0: float, 0.001
        Used to transfer far field inflow wind speeds between different heights
        above the ground [m].

    Returns
    -------
    wind2load: wind2load class
        Specifying the load model for calculating mean equivalent fatigue loads
        for each turbine at different channels. Currently it is the surrogate
        load model implemented. Note that when wind2load is not provided as
        input, this class will only calculate AEP.

    Methods
    -------
    reset_num_evals()
        Reset the number of AEP (and load) evalutions to zero.

    cal_AEP_load(cal_load=True)
        Calculate gross AEP, net AEP (and mean load) of this wind farm.
    """

    def __init__(self,
                 site_conditions,  # include terrain and wind resource
                 wind_farm,  # defines the design of wind farm
                 wake_model,  # wake model to use
                 wind2load=None,  # optional input, calc only AEP when None
                 ws_binned=np.linspace(1, 30, num=30),  # [m/s]
                 wd_binned=np.linspace(0, 330, num=12),  # [deg]
                 height_ref=85,  # [m]
                 k_star=0.075,  # wake decay parameter
                 availability=1.0,  # availability factor for the wind farm
                 num_evals=0,    # number of evluations of .cal_AEP_load()
                 z0=0.001  # used to transfer wind speed between diff height
                 ):
        """ ws_binned, wd_binned and height_ref defines a set of discretized
        ideal far field inflow condition at the a height above the ground
        (height_ref), which controls and specifies the bin sizes of wind speed
        and wind direction in the AEP calculation.
        """
        # inputed attributes
        self.site_conditions = site_conditions
        self.wind_farm = wind_farm
        self.wake_model = wake_model
        self.wind2load = wind2load
        self.ws_binned = ws_binned
        self.wd_binned = wd_binned
        self.height_ref = height_ref
        self.k_star = k_star
        self.availability = availability
        self.num_evals = num_evals
        self.z0 = z0

        # calculated attributes (storing necessary data during calculation)
        self.num_ws_bins = len(ws_binned)
        self.num_wd_bins = len(wd_binned)
        self.wf_design = wind_farm.get_summary().values  # [x, y, z, D, H, P]
        self.num_turbines = len(self.wf_design[:, 0])
        if self.wind2load is not None:
            self.num_channels = self.wind2load.num_channel

        # Extend the wind dir bins to include an extra dir for integration
        self.num_wd_bins = self.num_wd_bins + 1
        self.wd_binned = np.hstack(
                (self.wd_binned,
                 self.wd_binned[-1] + self.wd_binned[1] - self.wd_binned[0]))
        # note we assume wd_binned is equally spaced.

        # index i, k, l denotes the ith turbine, the kth wind speed and the lth
        # wind direction, which is defined in self.ws_binned and self.wd_binned
        shape_ikl = [self.num_turbines, self.num_ws_bins, self.num_wd_bins]
        self.local_ws_ideal_ikl = np.zeros(shape_ikl)
        self.local_wd_ideal_ikl = np.zeros(shape_ikl)
        self.local_power_ideal_ikl = np.zeros(shape_ikl)
        self.local_ws_real_ikl = np.zeros(shape_ikl)
        self.local_wd_real_ikl = np.zeros(shape_ikl)
        self.local_power_real_ikl = np.zeros(shape_ikl)
        self.local_pdf_ikl = np.zeros(shape_ikl)
        self.local_TI_real_ikl = np.zeros(shape_ikl)
        self.local_Ct_ikl = np.zeros(shape_ikl)

        if self.wind2load is not None:
            self.load_iklm = np.zeros([self.num_turbines,
                                       self.num_ws_bins,
                                       self.num_wd_bins,
                                       self.num_channels])

    def reset_num_evals(self):
        """ Reset the number of evaluation to 0.
        """
        self.num_evals = 0

    def cal_AEP_load(self, cal_load=True):
        """ Calculate gross AEP, net AEP (and mean load) of this wind farm.

        If wind2load is provided and load calculation is turned on (cal_load=
        True), returns (AEP_gross, AEP_net and mean_loads).

        If wind2load is not provided or load calculation is turned off
        (cal_load=False), returns (AEP_gross, AEP_net)

        This is the vectorized version.

        Parameters
        ----------
        cal_load: boolean (default: True)
            A flag to turn on/off the load calculation.

        Returns
        -------
        AEP_gross: array:float
            Gross AEP values of each turbine in the wind farm.

        AEP_net: array:float
            Net AEP values of each turbine in the wind farm.

        mean_loads: array:float
            Mean equivalent fatigue loads for each turbine at different
            channels. Note this one is only returned when the wind2load
            instance is provided and the load calculation is turned on.
        """
        num_hrs_a_year = 8760

        #######################################################################
        # Step 1. Get and store sector (wind direction) wise site conditions
        speed_up_il = np.zeros([self.num_turbines, self.num_wd_bins])
        turning_il = np.zeros([self.num_turbines, self.num_wd_bins])
        inclination_il = np.zeros([self.num_turbines, self.num_wd_bins])
        Weibull_A_il = np.zeros([self.num_turbines, self.num_wd_bins])
        Weibull_k_il = np.zeros([self.num_turbines, self.num_wd_bins])
        frequency_il = np.zeros([self.num_turbines, self.num_wd_bins])
        turbulence_il = np.zeros([self.num_turbines, self.num_wd_bins])
        wind_shear_il = np.zeros([self.num_turbines, self.num_wd_bins])
        rho_il = np.zeros([self.num_turbines, self.num_wd_bins])

        for l_wd in range(self.num_wd_bins):
            wd = self.wd_binned[l_wd]
            # For this general inflow direcion, get wind resource and terrain
            # effect information for all turbine locations.

            conditions = self.site_conditions.get_site_conditions(
                self.wf_design[:, (0, 1, 4)], wd)

            speed_up_il[:, l_wd] = conditions['spd_up'].values
            turning_il[:, l_wd] = conditions['deviation'].values
            inclination_il[:, l_wd] = conditions['inflow_angle'].values
            Weibull_A_il[:, l_wd] = conditions['A'].values
            Weibull_k_il[:, l_wd] = conditions['k'].values
            frequency_il[:, l_wd] = conditions['freq_per_degree'].values
            turbulence_il[:, l_wd] = conditions['tke_amb'].values
            wind_shear_il[:, l_wd] = conditions['alpha'].values
            rho_il[:, l_wd] = conditions['rho'].values

        #######################################################################
        # Step 2. Calculate ideal local flow field (ws, wd) and related pdf
        for l_wd in range(self.num_wd_bins):
            wd = self.wd_binned[l_wd]

            for k_ws in range(self.num_ws_bins):
                ws = self.ws_binned[k_ws]

                for i_wt in range(self.num_turbines):
                    H_hub = self.wf_design[i_wt, 4]

                    # calculate local wind speed and wind direction without
                    # wake effects by considering terrain effect
                    local_ws_ideal = (ws * np.log(H_hub / self.z0) /
                                      np.log(self.height_ref / self.z0) *
                                      speed_up_il[i_wt, l_wd])

                    local_wd_ideal = (wd + turning_il[i_wt, l_wd])

                    self.local_ws_ideal_ikl[i_wt, k_ws, l_wd] = local_ws_ideal
                    self.local_wd_ideal_ikl[i_wt, k_ws, l_wd] = local_wd_ideal

                    # calculating related pdf for all wind speeds
                    self.local_pdf_ikl[i_wt, k_ws, l_wd] = \
                        self.cal_pdf_Weibull(self.local_ws_ideal_ikl[i_wt,
                                                                     k_ws,
                                                                     l_wd],
                                             Weibull_A_il[i_wt, l_wd],
                                             Weibull_k_il[i_wt, l_wd]) \
                        * frequency_il[i_wt, l_wd]

        #######################################################################
        # Step 3. Calculate ideal local Ct and power
        for i_wt in range(self.num_turbines):
            self.local_power_ideal_ikl[i_wt, :, :] = (
                    self.wind_farm.get_power(i_wt,
                                             self.local_ws_ideal_ikl[i_wt, :,
                                                                     :]))

            self.local_Ct_ikl[i_wt, :, :] = (
                    self.wind_farm.get_ct(i_wt,
                                          self.local_ws_ideal_ikl[i_wt, :,
                                                                  :]))

        #######################################################################
        # Step 4. Calculate real local wind speed and turbulence intensity

        # assuming same wake decay coefficients for all turbines
        k_star_list = [self.k_star] * self.num_turbines

        (self.local_ws_real_ikl, self.local_TI_real_ikl) = self.wake_model.cal_wake(
                    self.wf_design[:, 0],   # [x_i]
                    self.wf_design[:, 1],   # [y_i]
                    self.wf_design[:, 4],   # [H_i]]
                    self.wf_design[:, 3],   # [D_i]
                    self.local_ws_ideal_ikl,
                    self.local_wd_ideal_ikl,
                    self.local_Ct_ikl,
                    turbulence_il,
                    k_star_list)

        #######################################################################
        # Step 5. Calculate real power of each turbine
        for i_wt in range(self.num_turbines):
            self.local_power_real_ikl[i_wt, :, :] = (
                    self.wind_farm.get_power(i_wt,
                                             self.local_ws_real_ikl[i_wt, :,
                                                                    :]))

        #######################################################################
        # Step 6. Calculate loads

        # if the wind2load is available and load calculation is turned on
        if (self.wind2load is not None) and cal_load:
            # for the single value 0 dimension vectorized calculation
            # for l_wd in range(self.num_wd_bins):
            #     for k_ws in range(self.num_ws_bins):
            #         for i_wt in range(self.num_turbines):
            #
            #             self.load_iklm[i_wt, k_ws, l_wd, :] = (
            #                 self.wind2load.load_calculation(
            #                     self.local_ws_real_ikl[i_wt, k_ws, l_wd],
            #                     self.local_TI_real_ikl[i_wt, k_ws, l_wd],
            #                     wind_shear_il[i_wt, l_wd],
            #                     inclination_il[i_wt, l_wd],
            #                     rho_il[i_wt, l_wd],
            #                     self.wf_design[i_wt, 4],      # H
            #                     self.wf_design[i_wt, 3],      # D
            #                     self.wf_design[i_wt, 5]))     # P_rated
            # 2d dimensional data passing of vectorized calculation
            for i_wt in range(self.num_turbines):
                self.load_iklm[i_wt, :, :, :] = (
                    self.wind2load.load_calculation_2d(
                        self.local_ws_real_ikl[i_wt, :, :],
                        self.local_TI_real_ikl[i_wt, :, :],
                        np.tile(wind_shear_il[i_wt, :],
                                (self.num_ws_bins, 1)),
                        np.tile(inclination_il[i_wt, :],
                                (self.num_ws_bins, 1)),
                        np.tile(rho_il[i_wt, :], (self.num_ws_bins, 1)),
                        self.wf_design[i_wt, 4],  # H
                        self.wf_design[i_wt, 3],  # D
                        self.wf_design[i_wt, 5]))  # P_rated

        #######################################################################
        # Step 7. Calculate mean power and AEP values using numerical integ.

        delta_ws = (self.local_ws_ideal_ikl[:, 1:, 1:] -
                    self.local_ws_ideal_ikl[:, :-1, 1:])

        delta_wd = (self.local_wd_ideal_ikl[:, 1:, 1:] -
                    self.local_wd_ideal_ikl[:, 1:, :-1])

        pdf_array = (self.local_pdf_ikl[:, 1:, 1:] +
                     self.local_pdf_ikl[:, :-1, 1:]) / 2

        power_ideal_array = (self.local_power_ideal_ikl[:, 1:, 1:] +
                             self.local_power_ideal_ikl[:, :-1, 1:]) / 2

        power_real_array = (self.local_power_real_ikl[:, 1:, 1:] +
                            self.local_power_real_ikl[:, :-1, 1:]) / 2

        mean_power_ideal = np.sum(
            delta_ws * delta_wd * pdf_array * power_ideal_array, (1, 2))

        mean_power_real = np.sum(
            delta_ws * delta_wd * pdf_array * power_real_array, (1, 2))

        AEP_gross = num_hrs_a_year * mean_power_ideal * self.availability
        AEP_net = num_hrs_a_year * mean_power_real * self.availability

        #######################################################################
        # Step 8. Calculate loads

        # if the wind2load is available and load calculation is turned on
        if (self.wind2load is not None) and cal_load:
            load_array = (self.load_iklm[:, 1:, 1:, :]
                          + self.load_iklm[:, :-1, 1:, :])/2
            slope_array = self.wind2load.pce_slopes
            mean_loads = np.zeros([self.num_turbines, self.num_channels])

            for m_channel in range(self.num_channels):
                mean_loads[:, m_channel] = (np.sum(
                    delta_ws * delta_wd * pdf_array *
                    (load_array[:, :, :, m_channel] **
                     slope_array[m_channel]), (1, 2))
                     / self.wind2load.frequence) ** \
                    (1 / slope_array[m_channel])

        #################################################################
        # updating number of evluations
        self.num_evals = self.num_evals + 1

        if (self.wind2load is not None) and cal_load:
            return (AEP_gross, AEP_net, mean_loads)
        else:
            return (AEP_gross, AEP_net)

    def cal_AEP_load_naive(self, cal_load=True):
        """ Calculate gross AEP, net AEP and mean load of this wind farm

        This is the naive version with all in for loops. This naive version is
        only kept here for understanding the process and possible testing
        usages.
        """
        num_hrs_a_year = 8760

        #######################################################################
        # Step 1. Calculate flow field, pdf , power and load

        for l_wd in range(self.num_wd_bins):
            wd = self.wd_binned[l_wd]
            # For this general inflow direcion, get wind resource and terrain
            # effect information for all turbine locations. All thess lists are
            # one dimensional array with self.num_turbs elements.

            conditions = self.site_conditions.get_site_conditions(
                self.wf_design[:, (0, 1, 4)], wd)

            speed_up_list = conditions['spd_up'].values
            turning_list = conditions['deviation'].values
            inclination_list = conditions['inflow_angle'].values
            Weibull_A_list = conditions['A'].values
            Weibull_k_list = conditions['k'].values
            frequency_list = conditions['freq_per_degree'].values
            turbulence_list = conditions['tke_amb'].values
            wind_shear_list = conditions['alpha'].values
            rho_list = conditions['rho'].values

            for k_ws in range(self.num_ws_bins):
                ws = self.ws_binned[k_ws]
                Ct_list = np.zeros(self.num_turbines)

                ###########################################################
                # ideal case: without wake effects
                for i_wt in range(self.num_turbines):
                    H_hub = self.wf_design[i_wt, 4]

                    # calculate local wind speed and wind direction without
                    # wake effects by considering terrain effect
                    local_ws_ideal = (ws * np.log(H_hub / self.z0) /
                                      np.log(self.height_ref / self.z0) *
                                      speed_up_list[i_wt])
                    local_wd_ideal = (wd + turning_list[i_wt])

                    self.local_ws_ideal_ikl[i_wt, k_ws, l_wd] = local_ws_ideal
                    self.local_wd_ideal_ikl[i_wt, k_ws, l_wd] = local_wd_ideal

                    # calculating relating pdf and power
                    self.local_pdf_ikl[i_wt, k_ws, l_wd] = (
                        (self.cal_pdf_Weibull(
                            local_ws_ideal,
                            Weibull_A_list[i_wt],
                            Weibull_k_list[i_wt])) *
                        frequency_list[i_wt])

                    self.local_power_ideal_ikl[i_wt, k_ws, l_wd] = (
                        self.wind_farm.get_power(i_wt, local_ws_ideal))

                    # get Ct
                    Ct_list[i_wt] = self.wind_farm.get_ct(i_wt,
                                                          local_ws_ideal)

                ###########################################################
                # Real case: with wake effects
                # caculate wake influced flow field
                k_star_list = [self.k_star] * self.num_turbines

                (ws_eff, TI_eff) = self.wake_model.cal_wake(
                    self.wf_design[:, 0],
                    self.wf_design[:, 1],
                    self.wf_design[:, 4],
                    self.wf_design[:, 3],
                    self.local_ws_ideal_ikl[:, k_ws, l_wd],
                    self.local_wd_ideal_ikl[:, k_ws, l_wd],
                    Ct_list,
                    turbulence_list,
                    k_star_list)

                self.local_ws_real_ikl[:, k_ws, l_wd] = ws_eff

                if (self.wind2load is not None) and cal_load:

                    for i_wt in range(self.num_turbines):
                        self.local_power_real_ikl[i_wt, k_ws, l_wd] = (
                            self.wind_farm.get_power(i_wt,
                                                     ws_eff[i_wt]))

                        # Calculate laod for each WT
                        wind_condition = np.array([ws_eff[i_wt],
                                                   TI_eff[i_wt],
                                                   wind_shear_list[i_wt],
                                                   inclination_list[i_wt],
                                                   rho_list[i_wt]])
                        turbine_parameter = np.array([self.wf_design[i_wt, 4],
                                                      self.wf_design[i_wt, 3],
                                                      self.wf_design[i_wt, 5]])
                        print(cal_load)
                        self.load_iklm[i_wt, k_ws, l_wd, :] = (
                            self.wind2load.load_calculation(
                                wind_condition,
                                turbine_parameter))

        #######################################################################
        # Step 4. Calculate mean power and AEP values using numerical integ.
        delta_ws = (self.local_ws_ideal_ikl[:, 1:, 1:] -
                    self.local_ws_ideal_ikl[:, :-1, 1:])

        delta_wd = (self.local_wd_ideal_ikl[:, 1:, 1:] -
                    self.local_wd_ideal_ikl[:, 1:, :-1])

        pdf_array = (self.local_pdf_ikl[:, 1:, 1:] +
                     self.local_pdf_ikl[:, :-1, 1:]) / 2

        power_ideal_array = (self.local_power_ideal_ikl[:, 1:, 1:] +
                             self.local_power_ideal_ikl[:, :-1, 1:]) / 2

        power_real_array = (self.local_power_real_ikl[:, 1:, 1:] +
                            self.local_power_real_ikl[:, :-1, 1:]) / 2

        mean_power_ideal = np.sum(
            delta_ws * delta_wd * pdf_array * power_ideal_array, (1, 2))

        mean_power_real = np.sum(
            delta_ws * delta_wd * pdf_array * power_real_array, (1, 2))

        AEP_gross = num_hrs_a_year * mean_power_ideal
        AEP_net = num_hrs_a_year * mean_power_real

        #######################
        # mean loads
        if (self.wind2load is not None) and cal_load:
            load_array = (self.load_iklm[:, 1:, 1:, :]
                          + self.load_iklm[:, :-1, 1:, :])/2

            mean_loads = np.zeros([self.num_turbines, self.num_channels])

            for m_channel in range(self.num_channels):
                mean_loads[:, m_channel] = np.sum(
                    delta_ws * delta_wd * pdf_array *
                    load_array[:, :, :, m_channel], (1, 2))

        if (self.wind2load is not None) and cal_load:
            return (AEP_gross, AEP_net, mean_loads)
        else:
            return (AEP_gross, AEP_net)

    def cal_pdf_Weibull(self, ws, Weibull_A, Weibull_k):
        """ calculate pdf of a given wind speed based on Weibull distribution.

        Parameters
        ----------
        ws: array:float
            Wind speed [m/s]

        Weibull_A: array:float
            Scale parameter of Weibull distribution [m/s]
        Weibull_k: array:float
            Shape parameter of Weibull distribution [-].

        Returns
        -------
        pdf: array:float
            Probability density function calculated using Weibull distribution.

        """
        pdf = ((Weibull_k / Weibull_A) * (ws / Weibull_A) ** (Weibull_k - 1) *
               np.exp(-(ws / Weibull_A) ** Weibull_k))

        return pdf


"""Implementations of interpolators, which representgeometric paths.

"""
import logging
import warnings
import numpy as np
from scipy.interpolate import UnivariateSpline, CubicSpline, PPoly

logger = logging.getLogger(__name__)

try:
    import openravepy as orpy
except ImportError as err:
    logger.warning("Unable to import openravepy. Exception: %s", err.args[0])
except SyntaxError as err:
    logger.warning("Unable to import openravepy. Exception: %s", err.args[0])

def normalize(gridpoints):
    # type: (np.ndarray) -> np.ndarray
    """Normalize the path discretization.

    Parameters
    ----------
    gridpoints: Path position array.

    Returns
    -------
    out: Normalized path position array.
    """
    return np.array(gridpoints) / gridpoints[-1]

def _find_left_index(gridpoints, s):
    # type: (np.ndarray, float) -> int
    """Find the least lowest entry that is larger or equal.

    Parameters
    ----------
    gridpoints:
        Array of path positions.
    s:
        A path position.

    Returns
    -------
    out:
        The desired index.

    """
    for i in range(1, len(gridpoints)):
        if gridpoints[i - 1] <= s < gridpoints[i]:
            return i - 1
    return len(gridpoints) - 2

class Interpolator(object):
    """Abstract class for interpolators."""

    def __init__(self):
        pass

    def get_dof(self):
        # type: () -> int
        """Return the degree-of-freedom of the path.

        Returns
        -------
        out:
            Degree-of-freedom of the path.
        """
        raise NotImplementedError

    @property
    def duration(self):
        """Return the duration of the path."""
        raise NotImplementedError

    @property
    def dof(self):
        """Return the degrees-of-freedom of the path."""
        raise NotImplementedError

    def get_path_interval(self):
        # type: () -> np.ndarray
        """Return the starting and ending path positions.

        Returns
        -------
        out:
            The starting and ending path positions.

        """
        return np.array([self.s_start, self.s_end])

    def eval(self, ss_sam):
        # type: (any[np.ndarray, float]) -> np.ndarray
        """Evaluate joint positions at specified path positions.

        Parameters
        ----------
        ss_sam :
            Shape (m,) or float. The path positions to sample at.

        Returns
        -------
        out :
            Shape (m, dof) if input is an array: evaluated values at positions.
            Shape (dof,) if input is a float.
        """
        raise NotImplementedError

    def evald(self, ss_sam):
        # type: (any[np.ndarray, float]) -> np.ndarray
        """Evaluate first derivative at specified path positions.

        Parameters
        ----------
        ss_sam :
            Shape (m,) or float. The path positions to sample at.

        Returns
        -------
        out :
            Shape (m, dof) if input is an array: evaluated values at positions.
            Shape (dof,) if input is a float.
        """
        raise NotImplementedError

    def evaldd(self, ss_sam):
        # type: (Union[np.ndarray, float]) -> np.ndarray
        """Evaluate second derivative at specified path positions.

        Parameters
        ----------
        ss_sam :
            Shape (m,) or float. The path positions to sample at.

        Returns
        -------
        out :
            Shape (m, dof) if input is an array: evaluated values at positions.
            Shape (dof,) if input is a float.
        """
        raise NotImplementedError

    def compute_rave_trajectory(self, robot):
        """Return the corresponding Openrave Trajectory."""
        raise NotImplementedError

    def compute_ros_trajectory(self):
        """Return the corresponding ROS trajectory."""
        raise NotImplementedError

class RaveTrajectoryWrapper(Interpolator):
    """An interpolator that wraps OpenRAVE's :class:`GenericTrajectory`.

    Only trajectories using quadratic interpolation or cubic
    interpolation are supported.  The trajectory is represented as a
    piecewise polynomial. The polynomial could be quadratic or cubic
    depending the interpolation method used by the input trajectory
    object.

    """

    def __init__(self, traj, robot):
        # type: (orpy.RaveTrajectory, orpy.Robot) -> None
        """Initialize the Trajectory Wrapper.

        Parameters
        ----------
        traj:
            An OpenRAVE joint trajectory.
        robot:
            An OpenRAVE robot.
        """
        super(RaveTrajectoryWrapper, self).__init__()
        self.traj = traj  #: init
        self.spec = traj.GetConfigurationSpecification()
        self._dof = robot.GetActiveDOF()

        self._interpolation = self.spec.GetGroupFromName('joint').interpolation
        if self._interpolation not in ['quadratic', 'cubic']:
            raise ValueError(
                "This class only handles trajectories with quadratic or cubic interpolation"
            )
        self._duration = traj.GetDuration()
        all_waypoints = traj.GetWaypoints(0, traj.GetNumWaypoints()).reshape(
            traj.GetNumWaypoints(), -1)
        valid_wp_indices = [0]
        self.ss_waypoints = [0.0]
        for i in range(1, traj.GetNumWaypoints()):
            dt = self.spec.ExtractDeltaTime(all_waypoints[i])
            if dt > 1e-5:  # If delta is too small, skip it.
                valid_wp_indices.append(i)
                self.ss_waypoints.append(self.ss_waypoints[-1] + dt)

        self.n_waypoints = len(valid_wp_indices)
        self.ss_waypoints = np.array(self.ss_waypoints)
        self.s_start = self.ss_waypoints[0]
        self.s_end = self.ss_waypoints[-1]

        self.waypoints = np.array([
            self.spec.ExtractJointValues(all_waypoints[i], robot,
                                         robot.GetActiveDOFIndices())
            for i in valid_wp_indices
        ])
        self.waypoints_d = np.array([
            self.spec.ExtractJointValues(all_waypoints[i], robot,
                                         robot.GetActiveDOFIndices(), 1)
            for i in valid_wp_indices
        ])

        # Degenerate case: there is only one waypoint.
        if self.n_waypoints == 1:
            pp_coeffs = np.zeros((1, 1, self.dof))
            for idof in range(self.dof):
                pp_coeffs[0, 0, idof] = self.waypoints[0, idof]
            # A constant function
            self.ppoly = PPoly(pp_coeffs, [0, 1])

        elif self._interpolation == "quadratic":
            self.waypoints_dd = []
            for i in range(self.n_waypoints - 1):
                qdd = ((self.waypoints_d[i + 1] - self.waypoints_d[i]) /
                       (self.ss_waypoints[i + 1] - self.ss_waypoints[i]))
                self.waypoints_dd.append(qdd)
            self.waypoints_dd = np.array(self.waypoints_dd)

            # Fill the coefficient matrix for scipy.PPoly class
            pp_coeffs = np.zeros((3, self.n_waypoints - 1, self.dof))
            for idof in range(self.dof):
                for iseg in range(self.n_waypoints - 1):
                    pp_coeffs[:, iseg, idof] = [
                        self.waypoints_dd[iseg, idof] / 2,
                        self.waypoints_d[iseg, idof],
                        self.waypoints[iseg, idof]
                    ]
            self.ppoly = PPoly(pp_coeffs, self.ss_waypoints)

        elif self._interpolation == "cubic":
            self.waypoints_dd = np.array([
                self.spec.ExtractJointValues(all_waypoints[i], robot,
                                             robot.GetActiveDOFIndices(), 2)
                for i in valid_wp_indices
            ])
            self.waypoints_ddd = []
            for i in range(self.n_waypoints - 1):
                qddd = ((self.waypoints_dd[i + 1] - self.waypoints_dd[i]) /
                        (self.ss_waypoints[i + 1] - self.ss_waypoints[i]))
                self.waypoints_ddd.append(qddd)
            self.waypoints_ddd = np.array(self.waypoints_ddd)

            # Fill the coefficient matrix for scipy.PPoly class
            pp_coeffs = np.zeros((4, self.n_waypoints - 1, self.dof))
            for idof in range(self.dof):
                for iseg in range(self.n_waypoints - 1):
                    pp_coeffs[:, iseg, idof] = [
                        self.waypoints_ddd[iseg, idof] / 6,
                        self.waypoints_dd[iseg, idof] / 2,
                        self.waypoints_d[iseg, idof],
                        self.waypoints[iseg, idof]
                    ]
            self.ppoly = PPoly(pp_coeffs, self.ss_waypoints)

        self.ppoly_d = self.ppoly.derivative()
        self.ppoly_dd = self.ppoly.derivative(2)

    def get_duration(self):
        warnings.warn(
            "`get_duration` method is deprecated, use `duration` property instead",
            PendingDeprecationWarning)
        return self.duration

    def get_dof(self):  # type: () -> int
        warnings.warn("This method is deprecated, use the property instead",
                      PendingDeprecationWarning)
        return self.dof

    @property
    def duration(self):
        return self._duration

    @property
    def dof(self):
        return self._dof

    def eval(self, ss_sam):
        return self.ppoly(ss_sam)

    def evald(self, ss_sam):
        return self.ppoly_d(ss_sam)

    def evaldd(self, ss_sam):
        return self.ppoly_dd(ss_sam)

class SplineInterpolator(Interpolator):
    """Interpolate the given waypoints by cubic spline.

    This interpolator is implemented as a simple wrapper over scipy's
    CubicSpline class.

    Parameters
    ----------
    ss_waypoints: array
        Shaped (N+1,). Path positions of the waypoints.
    waypoints: array
        Shaped (N+1, dof). Waypoints.
    bc_type: str, optional
        Boundary condition. Can be 'not-a-knot', 'clamped', 'natural' or 'periodic'.
        See scipy.CubicSpline documentation for more details.

    Attributes
    ----------
    dof : int
        Output dimension of the function
    cspl : :class:`scipy.interpolate.CubicSpline`
        The path.
    cspld : :class:`scipy.interpolate.CubicSpline`
        The path 1st derivative.
    cspldd : :class:`scipy.interpolate.CubicSpline`
        The path 2nd derivative.

    """

    def __init__(self, ss_waypoints, waypoints, bc_type='clamped'):
        super(SplineInterpolator, self).__init__()
        assert ss_waypoints[0] == 0, "First index must equals zero."
        self.ss_waypoints = np.array(ss_waypoints)
        self.waypoints = np.array(waypoints)
        self.bc_type = bc_type

        assert self.ss_waypoints.shape[0] == self.waypoints.shape[0]
        self.s_start = self.ss_waypoints[0]
        self.s_end = self.ss_waypoints[-1]

        if len(ss_waypoints) == 1:

            def _1dof_cspl(s):
                try:
                    ret = np.zeros((len(s), self.dof))
                    ret[:, :] = self.waypoints[0]
                except TypeError:
                    ret = self.waypoints[0]
                return ret

            def _1dof_cspld(s):
                try:
                    ret = np.zeros((len(s), self.dof))
                except TypeError:
                    ret = np.zeros(self.dof)
                return ret

            self.cspl = _1dof_cspl
            self.cspld = _1dof_cspld
            self.cspldd = _1dof_cspld
        else:
            self.cspl = CubicSpline(ss_waypoints, waypoints, bc_type=bc_type)
            self.cspld = self.cspl.derivative()
            self.cspldd = self.cspld.derivative()

    def get_waypoints(self):
        """Return the appropriate scaled waypoints."""
        return self.ss_waypoints, self.waypoints

    def get_duration(self):
        warnings.warn(
            "get_duration is deprecated, use duration (property) instead",
            PendingDeprecationWarning)
        return self.duration

    @property
    def duration(self):
        return self.ss_waypoints[-1] - self.ss_waypoints[0]

    @property
    def dof(self):
        if np.isscalar(self.waypoints[0]):
            return 1
        return self.waypoints[0].shape[0]

    def get_dof(self):  # type: () -> int
        warnings.warn("get_dof is deprecated, use dof (property) instead",
                      PendingDeprecationWarning)
        return self.dof

    def eval(self, ss_sam):
        return self.cspl(ss_sam)

    def evald(self, ss_sam):
        return self.cspld(ss_sam)

    def evaldd(self, ss_sam):
        return self.cspldd(ss_sam)

    def compute_rave_trajectory(self, robot):
        """Compute an OpenRAVE trajectory equivalent to this trajectory.

        Parameters
        ----------
        robot:
            Openrave robot.

        Returns
        -------
        trajectory:
            Equivalent openrave trajectory.
        """

        traj = orpy.RaveCreateTrajectory(robot.GetEnv(), "")
        spec = robot.GetActiveConfigurationSpecification('cubic')
        spec.AddDerivativeGroups(1, False)
        spec.AddDerivativeGroups(2, True)

        traj.Init(spec)
        deltas = [0]
        for i in range(len(self.ss_waypoints) - 1):
            deltas.append(self.ss_waypoints[i + 1] - self.ss_waypoints[i])
        if len(self.ss_waypoints) == 1:
            q = self.eval(0)
            qd = self.evald(0)
            qdd = self.evaldd(0)
            traj.Insert(traj.GetNumWaypoints(),
                        list(q) + list(qd) + list(qdd) + [0])
        else:
            qs = self.eval(self.ss_waypoints)
            qds = self.evald(self.ss_waypoints)
            qdds = self.evaldd(self.ss_waypoints)
            for (q, qd, qdd, dt) in zip(qs, qds, qdds, deltas):
                traj.Insert(traj.GetNumWaypoints(),
                            q.tolist() + qd.tolist() + qdd.tolist() + [dt])
        return traj

class UnivariateSplineInterpolator(Interpolator):
    """ Smooth given wayspoints by a cubic spline.

    This is a simple wrapper over `scipy.UnivariateSplineInterpolator`
    class.

    Parameters
    ----------
    ss_waypoints: ndarray
        Path positions of the waypoints.
    waypoints: ndarray
        The waypoints.
    """

    def __init__(self, ss_waypoints, waypoints):
        super(UnivariateSplineInterpolator, self).__init__()
        assert ss_waypoints[0] == 0, "First index must equals zero."
        self.ss_waypoints = np.array(ss_waypoints)
        self.waypoints = np.array(waypoints)
        if np.isscalar(waypoints[0]):
            self.dof = 1
        else:
            self.dof = waypoints[0].shape[0]
        self.duration = ss_waypoints[-1]
        assert self.ss_waypoints.shape[0] == self.waypoints.shape[0]
        self.s_start = self.ss_waypoints[0]
        self.s_end = self.ss_waypoints[-1]

        self.uspl = []
        for i in range(self.dof):
            self.uspl.append(
                UnivariateSpline(self.ss_waypoints, self.waypoints[:, i]))
        self.uspld = [spl.derivative() for spl in self.uspl]
        self.uspldd = [spl.derivative() for spl in self.uspld]

    def get_duration(self):
        return self.duration

    def eval(self, ss_sam):
        data = []
        for spl in self.uspl:
            data.append(spl(ss_sam))
        return np.array(data).T

    def evald(self, ss_sam):
        data = []
        for spl in self.uspld:
            data.append(spl(ss_sam))
        return np.array(data).T

    def evaldd(self, ss_sam):
        data = []
        for spl in self.uspldd:
            data.append(spl(ss_sam))
        return np.array(data).T

class PolynomialPath(Interpolator):
    """ A class representing polynominal paths.

    If coeff is a 1d array, the polynomial's equation is given by

    .. math::

    coeff[0] + coeff[1] s + coeff[2] s^2 + ...

    If coeff is a 2d array, the i-th joint position is the polynomial

    .. math::

    coeff[i, 0] + coeff[i, 1] s + coeff[i, 2] s^2 + ...
    """

    def __init__(self, coeff, s_start=0.0, s_end=1.0):
        # type: (np.ndarray, float, float) -> None
        """Initialize the polynomial path.

        Parameters
        ----------
        coeff
            Coefficients of the polynomials.
        s_start
            Starting path position.
        s_end
            Ending path position.
        """
        super(PolynomialPath, self).__init__()
        self.coeff = np.array(coeff)
        self.s_end = s_end
        self.s_start = s_start
        if np.isscalar(self.coeff[0]):
            self.poly = [np.polynomial.Polynomial(self.coeff)]
            self.coeff = self.coeff.reshape(1, -1)
        else:
            self.poly = [
                np.polynomial.Polynomial(self.coeff[i])
                for i in range(self.dof)
            ]

        self.polyd = [poly.deriv() for poly in self.poly]
        self.polydd = [poly.deriv() for poly in self.polyd]

    @property
    def dof(self):
        return self.coeff.shape[0]

    @property
    def duration(self):
        return self.s_end - self.s_start

    def get_duration(self):
        warnings.warn("get_duration is deprecated, use duration",
                      PendingDeprecationWarning)
        return self.duration

    def get_dof(self):
        warnings.warn("get_dof is deprecated, use dof",
                      PendingDeprecationWarning)
        return self.dof

    def eval(self, ss_sam):
        res = [poly(np.array(ss_sam)) for poly in self.poly]
        if self.dof == 1:
            return np.array(res).flatten()
        return np.array(res).T

    def evald(self, ss_sam):
        res = [poly(np.array(ss_sam)) for poly in self.polyd]
        if self.dof == 1:
            return np.array(res).flatten()
        return np.array(res).T

    def evaldd(self, ss_sam):
        res = [poly(np.array(ss_sam)) for poly in self.polydd]
        if self.dof == 1:
            return np.array(res).flatten()
        return np.array(res).T


# Copyright 2021 Huawei Technologies Co., Ltd
#
# 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.
# ============================================================================
"""Vision Transformer implementation."""

from importlib import import_module
from easydict import EasyDict as edict
import numpy as np

import mindspore
from mindspore.common.initializer import initializer
from mindspore.common.parameter import Parameter
from mindspore.nn import Cell, Dense, Dropout, SequentialCell
from mindspore.ops import operations as P
import mindspore.common.dtype as mstype
from mindspore import Tensor

MIN_NUM_PATCHES = 4

class VitConfig:
    """
    VitConfig
    """
    def __init__(self, configs):
        self.configs = configs

        # network init
        self.network_norm = mindspore.nn.LayerNorm((configs.normalized_shape,))
        self.network_init = mindspore.common.initializer.Normal(sigma=1.0)
        self.network_dropout_rate = 0.1
        self.network_pool = 'cls'
        self.network = ViT

        # stem
        self.stem_init = mindspore.common.initializer.XavierUniform()
        self.stem = VitStem

        # body
        self.body_norm = mindspore.nn.LayerNorm
        self.body_drop_path_rate = 0.1
        self.body = Transformer

        # body attention
        self.attention_init = mindspore.common.initializer.XavierUniform()
        self.attention_activation = mindspore.nn.Softmax()
        self.attention_dropout_rate = 0.1
        self.attention = Attention

        # body feedforward
        self.feedforward_init = mindspore.common.initializer.XavierUniform()
        self.feedforward_activation = mindspore.nn.GELU()
        self.feedforward_dropout_rate = 0.1
        self.feedforward = FeedForward

        # head
        self.head = origin_head
        self.head_init = mindspore.common.initializer.XavierUniform()
        self.head_dropout_rate = 0.1
        self.head_norm = mindspore.nn.LayerNorm((configs.normalized_shape,))
        self.head_activation = mindspore.nn.GELU()

class DropPath(Cell):
    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
    """

    def __init__(self, drop_prob=None, seed=0):
        super(DropPath, self).__init__()
        self.keep_prob = 1 - drop_prob
        seed = min(seed, 0) # always be 0
        self.rand = P.UniformReal(seed=seed) # seed must be 0, if set to other value, it's not rand for multiple call
        self.shape = P.Shape()
        self.floor = P.Floor()

    def construct(self, x):
        if self.training:
            x_shape = self.shape(x) # B N C
            random_tensor = self.rand((x_shape[0], 1, 1))
            random_tensor = random_tensor + self.keep_prob
            random_tensor = self.floor(random_tensor)
            x = x / self.keep_prob
            x = x * random_tensor
        return x

class BatchDense(Cell):
    """BatchDense module."""

    def __init__(self, in_features, out_features, initialization, has_bias=True):
        super().__init__()
        self.out_features = out_features
        self.dense = Dense(in_features, out_features, has_bias=has_bias)
        self.dense.weight.set_data(initializer(initialization, [out_features, in_features]))
        self.reshape = P.Reshape()

    def construct(self, x):
        bs, seq_len, d_model = x.shape
        out = self.reshape(x, (bs * seq_len, d_model))
        out = self.dense(out)
        out = self.reshape(out, (bs, seq_len, self.out_features))
        return out

class ResidualCell(Cell):
    """Cell which implements x + f(x) function."""
    def __init__(self, cell):
        super().__init__()
        self.cell = cell

    def construct(self, x, **kwargs):
        return self.cell(x, **kwargs) + x

def pretrain_head(vit_config):
    """Head for ViT pretraining."""
    d_model = vit_config.configs.d_model
    mlp_dim = vit_config.configs.mlp_dim
    num_classes = vit_config.configs.num_classes

    dropout_rate = vit_config.head_dropout_rate
    initialization = vit_config.head_init
    normalization = vit_config.head_norm
    activation = vit_config.head_activation

    dense1 = Dense(d_model, mlp_dim)
    dense1.weight.set_data(initializer(initialization, [mlp_dim, d_model]))
    dense2 = Dense(mlp_dim, num_classes)
    dense2.weight.set_data(initializer(initialization, [num_classes, mlp_dim]))

    return SequentialCell([
        normalization,
        dense1,
        activation,
        Dropout(keep_prob=(1. - dropout_rate)),
        dense2])

def origin_head(vit_config):
    """Head for ViT pretraining."""
    d_model = vit_config.configs.d_model
    num_classes = vit_config.configs.num_classes
    initialization = vit_config.head_init
    dense = Dense(d_model, num_classes)
    dense.weight.set_data(initializer(initialization, [num_classes, d_model]))
    return SequentialCell([dense])

class VitStem(Cell):
    """Stem layer for ViT."""

    def __init__(self, vit_config):
        super().__init__()
        d_model = vit_config.configs.d_model
        patch_size = vit_config.configs.patch_size
        image_size = vit_config.configs.image_size
        initialization = vit_config.stem_init
        channels = 3

        assert image_size % patch_size == 0, 'Image dimensions must be divisible by the patch size.'
        num_patches = (image_size // patch_size) ** 2
        assert num_patches > MIN_NUM_PATCHES, f'your number of patches {num_patches} is too small'
        patch_dim = channels * patch_size ** 2

        self.patch_size = patch_size
        self.reshape = P.Reshape()
        self.transpose = P.Transpose()
        self.patch_to_embedding = BatchDense(patch_dim, d_model, initialization, has_bias=True)

    def construct(self, img):
        p = self.patch_size
        bs, channels, h, w = img.shape
        x = self.reshape(img, (bs, channels, h // p, p, w // p, p))
        x = self.transpose(x, (0, 2, 4, 1, 3, 5))
        x = self.reshape(x, (bs, (h//p)*(w//p), channels*p*p))
        x = self.patch_to_embedding(x)
        return x

class ViT(Cell):
    """Vision Transformer implementation."""

    def __init__(self, vit_config):
        super().__init__()

        d_model = vit_config.configs.d_model
        patch_size = vit_config.configs.patch_size
        image_size = vit_config.configs.image_size

        initialization = vit_config.network_init
        pool = vit_config.network_pool
        dropout_rate = vit_config.network_dropout_rate
        norm = vit_config.network_norm

        stem = vit_config.stem(vit_config)
        body = vit_config.body(vit_config)
        head = vit_config.head(vit_config)

        assert pool in {'cls', 'mean'}, 'pool type must be either cls or mean'
        num_patches = (image_size // patch_size) ** 2

        if pool == "cls":
            self.cls_token = Parameter(initializer(initialization, (1, 1, d_model)),
                                       name='cls', requires_grad=True)
            self.pos_embedding = Parameter(initializer(initialization, (1, num_patches + 1, d_model)),
                                           name='pos_embedding', requires_grad=True)
            self.tile = P.Tile()
            self.cat_1 = P.Concat(axis=1)
        else:
            self.pos_embedding = Parameter(initializer(initialization, (1, num_patches, d_model)),
                                           name='pos_embedding', requires_grad=True)
            self.mean = P.ReduceMean(keep_dims=False)
        self.pool = pool

        self.cast = P.Cast()
        self.dropout = Dropout(keep_prob=(1. - dropout_rate))
        self.stem = stem
        self.body = body
        self.head = head
        self.norm = norm

    def construct(self, img):
        x = self.stem(img)
        bs, seq_len, _ = x.shape

        if self.pool == "cls":
            cls_tokens = self.tile(self.cls_token, (bs, 1, 1))
            x = self.cat_1((cls_tokens, x)) # now x has shape = (bs, seq_len+1, d)
            x += self.pos_embedding[:, :(seq_len + 1)]
        else:
            x += self.pos_embedding[:, :seq_len]

        y = self.cast(x, mstype.float32)
        y = self.dropout(y)
        x = self.cast(y, x.dtype)

        x = self.body(x)

        if self.norm is not None:
            x = self.norm(x)

        if self.pool == "cls":
            x = x[:, 0]
        else:
            x = self.mean(x, (-2,))

        return self.head(x)

class Attention(Cell):
    """Attention layer implementation."""

    def __init__(self, vit_config):
        super().__init__()
        d_model = vit_config.configs.d_model
        dim_head = vit_config.configs.dim_head
        heads = vit_config.configs.heads

        initialization = vit_config.attention_init
        activation = vit_config.attention_activation
        dropout_rate = vit_config.attention_dropout_rate

        inner_dim = heads * dim_head
        self.dim_head = dim_head
        self.heads = heads
        self.scale = Tensor([dim_head ** -0.5])

        self.to_q = Dense(d_model, inner_dim, has_bias=True)
        self.to_q.weight.set_data(initializer(initialization, [inner_dim, d_model]))
        self.to_k = Dense(d_model, inner_dim, has_bias=True)
        self.to_k.weight.set_data(initializer(initialization, [inner_dim, d_model]))
        self.to_v = Dense(d_model, inner_dim, has_bias=True)
        self.to_v.weight.set_data(initializer(initialization, [inner_dim, d_model]))

        self.to_out = Dense(inner_dim, d_model, has_bias=True)
        self.to_out.weight.set_data(initializer(initialization, [inner_dim, d_model]))
        self.dropout = Dropout(1 - dropout_rate)

        self.activation = activation

        #auxiliary functions
        self.reshape = P.Reshape()
        self.transpose = P.Transpose()
        self.cast = P.Cast()
        self.mul = P.Mul()
        self.q_matmul_k = P.BatchMatMul(transpose_b=True)
        self.attn_matmul_v = P.BatchMatMul()
        self.softmax_nz = True

    def construct(self, x):
        '''x size - BxNxd_model'''
        bs, seq_len, d_model, h, d = x.shape[0], x.shape[1], x.shape[2], self.heads, self.dim_head

        x_2d = self.reshape(x, (-1, d_model))
        q, k, v = self.to_q(x_2d), self.to_k(x_2d), self.to_v(x_2d)

        if self.softmax_nz:
            q = self.reshape(q, (bs, seq_len, h, d))
            q = self.transpose(q, (0, 2, 1, 3))
            q = self.cast(q, mstype.float32)
            q = self.mul(q, self.scale)

            k = self.reshape(k, (bs, seq_len, h, d))
            k = self.transpose(k, (0, 2, 1, 3))
            v = self.reshape(v, (bs, seq_len, h, d))
            v = self.transpose(v, (0, 2, 1, 3))

            q = self.cast(q, k.dtype)
            attn_scores = self.q_matmul_k(q, k) #bs x h x seq_len x seq_len
            attn_scores = self.cast(attn_scores, x.dtype)
            attn_scores = self.activation(attn_scores)
        else:
            q = self.reshape(q, (bs, seq_len, h, d))
            q = self.transpose(q, (0, 2, 1, 3))
            k = self.reshape(k, (bs, seq_len, h, d))
            k = self.transpose(k, (0, 2, 1, 3))
            v = self.reshape(v, (bs, seq_len, h, d))
            v = self.transpose(v, (0, 2, 1, 3))

            attn_scores = self.q_matmul_k(q, k) #bs x h x seq_len x seq_len
            attn_scores = self.cast(attn_scores, mstype.float32)
            attn_scores = self.mul(attn_scores, self.scale)
            attn_scores = self.cast(attn_scores, x.dtype)
            attn_scores = self.activation(attn_scores)

        out = self.attn_matmul_v(attn_scores, v) #bs x h x seq_len x dim_head
        out = self.transpose(out, (0, 2, 1, 3))
        out = self.reshape(out, (bs*seq_len, h*d))
        out = self.to_out(out)
        out = self.reshape(out, (bs, seq_len, d_model))
        #out = self.dropout(out)
        y = self.cast(out, mstype.float32)
        y = self.dropout(y)
        out = self.cast(y, out.dtype)
        #out = self.reshape(out, (bs, seq_len, d_model))
        return out

class FeedForward(Cell):
    """FeedForward layer implementation."""

    def __init__(self, vit_config):
        super().__init__()

        d_model = vit_config.configs.d_model
        hidden_dim = vit_config.configs.mlp_dim

        initialization = vit_config.feedforward_init
        activation = vit_config.feedforward_activation
        dropout_rate = vit_config.feedforward_dropout_rate

        self.ff1 = BatchDense(d_model, hidden_dim, initialization)
        self.activation = activation
        self.dropout = Dropout(keep_prob=1.-dropout_rate)
        self.ff2 = BatchDense(hidden_dim, d_model, initialization)
        self.cast = P.Cast()

    def construct(self, x):
        y = self.ff1(x)
        y = self.cast(y, mstype.float32)
        y = self.activation(y)
        y = self.dropout(y)
        y = self.cast(y, x.dtype)
        y = self.ff2(y)
        y = self.cast(y, mstype.float32)
        y = self.dropout(y)
        y = self.cast(y, x.dtype)
        return y

class Transformer(Cell):
    """Transformer implementation."""

    def __init__(self, vit_config):
        super().__init__()

        depth = vit_config.configs.depth
        drop_path_rate = vit_config.body_drop_path_rate

        dpr = [x.item() for x in np.linspace(0, drop_path_rate, depth)]
        att_seeds = [np.random.randint(1024) for _ in range(depth)]
        mlp_seeds = [np.random.randint(1024) for _ in range(depth)]

        layers = []
        for i in range(depth):
            normalization = vit_config.body_norm((vit_config.configs.normalized_shape,))
            normalization2 = vit_config.body_norm((vit_config.configs.normalized_shape,))
            attention = vit_config.attention(vit_config)
            feedforward = vit_config.feedforward(vit_config)

            if drop_path_rate > 0:
                layers.append(
                    SequentialCell([
                        ResidualCell(SequentialCell([normalization,
                                                     attention,
                                                     DropPath(dpr[i], att_seeds[i])])),
                        ResidualCell(SequentialCell([normalization2,
                                                     feedforward,
                                                     DropPath(dpr[i], mlp_seeds[i])]))
                    ])
                )
            else:
                layers.append(
                    SequentialCell([
                        ResidualCell(SequentialCell([normalization,
                                                     attention])),
                        ResidualCell(SequentialCell([normalization2,
                                                     feedforward]))
                    ])
                )

        self.layers = SequentialCell(layers)

    def construct(self, x):
        return self.layers(x)

def load_function(func_name):
    """Load function using its name."""
    modules = func_name.split(".")
    if len(modules) > 1:
        module_path = ".".join(modules[:-1])
        name = modules[-1]
        module = import_module(module_path)
        return getattr(module, name)
    return func_name

vit_cfg = edict({
    'd_model': 768,
    'depth': 12,
    'heads': 12,
    'mlp_dim': 3072,
    'dim_head': 64,
    'patch_size': 32,
    'normalized_shape': 768,
    'image_size': 224,
    'num_classes': 1001,
})

def vit_base_patch16(args):
    """vit_base_patch16"""
    vit_cfg.d_model = 768
    vit_cfg.depth = 12
    vit_cfg.heads = 12
    vit_cfg.mlp_dim = 3072
    vit_cfg.dim_head = vit_cfg.d_model // vit_cfg.heads
    vit_cfg.patch_size = 16
    vit_cfg.normalized_shape = vit_cfg.d_model
    vit_cfg.image_size = args.train_image_size
    vit_cfg.num_classes = args.class_num

    if args.vit_config_path != '':
        print("get vit_config_path")
        vit_config = load_function(args.vit_config_path)(vit_cfg)
    else:
        print("get default_vit_cfg")
        vit_config = VitConfig(vit_cfg)

    model = vit_config.network(vit_config)
    return model

def vit_base_patch32(args):
    """vit_base_patch32"""
    vit_cfg.d_model = 768
    vit_cfg.depth = 12
    vit_cfg.heads = 12
    vit_cfg.mlp_dim = 3072
    vit_cfg.dim_head = vit_cfg.d_model // vit_cfg.heads
    vit_cfg.patch_size = 32
    vit_cfg.normalized_shape = vit_cfg.d_model
    vit_cfg.image_size = args.train_image_size
    vit_cfg.num_classes = args.class_num

    if args.vit_config_path != '':
        print("get vit_config_path")
        vit_config = load_function(args.vit_config_path)(vit_cfg)
    else:
        print("get default_vit_cfg")
        vit_config = VitConfig(vit_cfg)

    model = vit_config.network(vit_config)

    return model

def get_network(backbone_name, args):
    """get_network"""
    if backbone_name == 'vit_base_patch32':
        backbone = vit_base_patch32(args=args)
    elif backbone_name == 'vit_base_patch16':
        backbone = vit_base_patch16(args=args)
    else:
        raise NotImplementedError
    return backbone


from functools import partial
from typing import Callable

import jax
import jax.numpy as jnp
import jax.random
from jax import ops
from jax.tree_util import tree_map

from numpyro import handlers
from numpyro.contrib.funsor import enum, config_enumerate
from numpyro.distributions import Distribution
from numpyro.distributions.transforms import IdentityTransform
from numpyro.infer import NUTS, MCMC, VI
from numpyro.infer.guide import ReinitGuide
from numpyro.infer.kernels import SteinKernel
from numpyro.infer.util import transform_fn, get_parameter_transform, _guess_max_plate_nesting
from numpyro.util import ravel_pytree

# TODO
# Fix MCMC updates to work reasonably with optimizer

# Lots of code based on SVI interface and commonalities should be refactored
class Stein(VI):

    def __init__(self, model, guide: ReinitGuide, optim, loss, kernel_fn: SteinKernel, num_particles: int = 10,
                 loss_temperature: float = 1.0, repulsion_temperature: float = 1.0,
                 classic_guide_params_fn: Callable[[str], bool] = lambda name: False,
                 enum=True, sp_mcmc_crit='infl',
                 sp_mode='local', num_mcmc_particles: int = 0, num_mcmc_warmup: int = 100, num_mcmc_updates: int = 10,
                 sampler_fn=NUTS, sampler_kwargs=None, mcmc_kwargs=None, **static_kwargs):
        """
        Stein Variational Gradient Descent for Non-parametric Inference.
        :param model: Python callable with Pyro primitives for the model.
        :param guide: Python callable with Pyro primitives for the guide
            (recognition network).
        :param optim: an instance of :class:`~numpyro.optim._NumpyroOptim`.
        :param loss: ELBO loss, i.e. negative Evidence Lower Bound, to minimize.
        :param kernel_fn: Function that produces a logarithm of the statistical kernel to use with Stein inference
        :param num_particles: number of particles for Stein inference.
            (More particles capture more of the posterior distribution)
        :param loss_temperature: scaling of loss factor
        :param repulsion_temperature: scaling of repulsive forces (Non-linear Stein)
        :param enum: whether to apply automatic marginalization of discrete variables
        :param classic_guide_param_fn: predicate on names of parameters in guide which should be optimized classically without Stein (E.g., parameters for large normal networks or other transformation)
        :param sp_mcmc_crit: Stein Point MCMC update selection criterion, either 'infl' for most influential or 'rand' for random (EXPERIMENTAL)
        :param sp_mode: Stein Point MCMC mode for calculating Kernelized Stein Discrepancy. Either 'local' for only the updated MCMC particles or 'global' for all particles. (EXPERIMENTAL)
        :param num_mcmc_particles: Number of particles that should be updated with Stein Point MCMC (should be a subset of number of Stein particles) (EXPERIMENTAL)
        :param num_mcmc_warmup: Number of warmup steps for the MCMC sampler (EXPERIMENTAL)
        :param num_mcmc_updates: Number of MCMC update steps at each iteration (EXPERIMENTAL)
        :param sampler_fn: The MCMC sampling kernel used for the Stein Point MCMC updates (EXPERIMENTAL)
        :param sampler_kwargs: Keyword arguments provided to the MCMC sampling kernel (EXPERIMENTAL)
        :param mcmc_kwargs: Keyword arguments provided to the MCMC interface (EXPERIMENTAL)
        :param static_kwargs: Static keyword arguments for the model / guide, i.e. arguments
            that remain constant during fitting.
        """
        super().__init__(model, guide, optim, loss, name='Stein', **static_kwargs)
        assert sp_mcmc_crit == 'infl' or sp_mcmc_crit == 'rand'
        assert sp_mode == 'local' or sp_mode == 'global'
        assert 0 <= num_mcmc_particles <= num_particles

        self._inference_model = model
        self.model = model
        self.guide = guide
        self.optim = optim
        self.loss = loss
        self.kernel_fn = kernel_fn
        self.static_kwargs = static_kwargs
        self.num_particles = num_particles
        self.loss_temperature = loss_temperature
        self.repulsion_temperature = repulsion_temperature
        self.enum = enum
        self.classic_guide_params_fn = classic_guide_params_fn
        self.sp_mcmc_crit = sp_mcmc_crit
        self.sp_mode = sp_mode
        self.num_mcmc_particles = num_mcmc_particles
        self.num_mcmc_warmup = num_mcmc_warmup
        self.num_mcmc_updates = num_mcmc_updates
        self.sampler_fn = sampler_fn
        self.sampler_kwargs = sampler_kwargs or dict()
        self.mcmc_kwargs = mcmc_kwargs or dict()
        self.mcmc: MCMC = None
        self.guide_param_names = None
        self.constrain_fn = None
        self.uconstrain_fn = None
        self.particle_transform_fn = None

    def _apply_kernel(self, kernel, x, y, v):
        if self.kernel_fn.mode == 'norm' or self.kernel_fn.mode == 'vector':
            return kernel(x, y) * v
        else:
            return kernel(x, y) @ v

    def _kernel_grad(self, kernel, x, y):
        if self.kernel_fn.mode == 'norm':
            return jax.grad(lambda x: kernel(x, y))(x)
        elif self.kernel_fn.mode == 'vector':
            return jax.vmap(lambda i: jax.grad(lambda x: kernel(x, y)[i])(x)[i])(jnp.arange(x.shape[0]))
        else:
            return jax.vmap(lambda l: jnp.sum(jax.vmap(lambda m: jax.grad(lambda x: kernel(x, y)[l, m])(x)[m])
                                              (jnp.arange(x.shape[0]))))(jnp.arange(x.shape[0]))

    def _param_size(self, param):
        if isinstance(param, tuple) or isinstance(param, list):
            return sum(map(self._param_size, param))
        return param.size

    def _calc_particle_info(self, uparams, num_particles):
        uparam_keys = list(uparams.keys())
        uparam_keys.sort()
        start_index = 0
        res = {}
        for k in uparam_keys:
            end_index = start_index + self._param_size(uparams[k]) // num_particles
            res[k] = (start_index, end_index)
            start_index = end_index
        return res

    def _svgd_loss_and_grads(self, rng_key, unconstr_params, *args, **kwargs):
        # 0. Separate model and guide parameters, since only guide parameters are updated using Stein
        classic_uparams = {p: v for p, v in unconstr_params.items() if
                           p not in self.guide_param_names or self.classic_guide_params_fn(p)}
        stein_uparams = {p: v for p, v in unconstr_params.items() if p not in classic_uparams}
        # 1. Collect each guide parameter into monolithic particles that capture correlations
        # between parameter values across each individual particle
        stein_particles, unravel_pytree = ravel_pytree(stein_uparams, batch_dims=1)
        unravel_pytree_batched = jax.vmap(unravel_pytree)
        particle_info = self._calc_particle_info(stein_uparams, stein_particles.shape[0])

        # 2. Calculate loss and gradients for each parameter
        def scaled_loss(rng_key, classic_params, stein_params):
            params = {**classic_params, **stein_params}
            loss_val = self.loss.loss(rng_key, params, handlers.scale(self._inference_model, self.loss_temperature),
                                      self.guide, *args, **kwargs)
            return - loss_val

        def kernel_particle_loss_fn(ps):
            return scaled_loss(rng_key, self.constrain_fn(classic_uparams),
                               self.constrain_fn(unravel_pytree(ps)))

        def particle_transform_fn(particle):
            params = unravel_pytree(particle)
            tparams = self.particle_transform_fn(params)
            tparticle, _ = ravel_pytree(tparams)
            return tparticle

        tstein_particles = jax.vmap(particle_transform_fn)(stein_particles)

        loss, particle_ljp_grads = jax.vmap(jax.value_and_grad(kernel_particle_loss_fn))(tstein_particles)
        classic_param_grads = jax.vmap(lambda ps: jax.grad(lambda cps:
                                                           scaled_loss(rng_key, self.constrain_fn(cps),
                                                                       self.constrain_fn(unravel_pytree(ps))))(
            classic_uparams))(stein_particles)
        classic_param_grads = tree_map(partial(jnp.mean, axis=0), classic_param_grads)

        # 3. Calculate kernel on monolithic particle
        kernel = self.kernel_fn.compute(stein_particles, particle_info, kernel_particle_loss_fn)

        # 4. Calculate the attractive force and repulsive force on the monolithic particles
        attractive_force = jax.vmap(lambda y: jnp.sum(
            jax.vmap(lambda x, x_ljp_grad: self._apply_kernel(kernel, x, y, x_ljp_grad))(tstein_particles,
                                                                                particle_ljp_grads), axis=0))(
            tstein_particles)
        repulsive_force = jax.vmap(lambda y: jnp.sum(
            jax.vmap(lambda x: self.repulsion_temperature * self._kernel_grad(kernel, x, y))(tstein_particles),
            axis=0))(
            tstein_particles)

        def single_particle_grad(particle, att_force, rep_force):
            reparam_jac = jax.jacfwd(particle_transform_fn)(particle)
            return (att_force + rep_force) @ reparam_jac

        particle_grads = jax.vmap(single_particle_grad)(stein_particles, attractive_force,
                                                        repulsive_force) / self.num_particles

        # 5. Decompose the monolithic particle forces back to concrete parameter values
        stein_param_grads = unravel_pytree_batched(particle_grads)

        # 6. Return loss and gradients (based on parameter forces)
        res_grads = tree_map(lambda x: -x, {**classic_param_grads, **stein_param_grads})
        return -jnp.mean(loss), res_grads

    def _score_sp_mcmc(self, rng_key, subset_idxs, stein_uparams, sp_mcmc_subset_uparams, classic_uparams,
                       *args, **kwargs):
        if self.sp_mode == 'local':
            _, ksd = self._svgd_loss_and_grads(rng_key, {**sp_mcmc_subset_uparams, **classic_uparams}, *args, **kwargs)
        else:
            stein_uparams = {p: ops.index_update(v, subset_idxs, sp_mcmc_subset_uparams[p]) for p, v in
                             stein_uparams.items()}
            _, ksd = self._svgd_loss_and_grads(rng_key, {**stein_uparams, **classic_uparams}, *args, **kwargs)
        ksd_res = jnp.sum(jnp.concatenate([jnp.ravel(v) for v in ksd.values()]))
        return ksd_res

    def _sp_mcmc(self, rng_key, unconstr_params, *args, **kwargs):
        # 0. Separate classical and stein parameters
        classic_uparams = {p: v for p, v in unconstr_params.items() if
                           p not in self.guide_param_names or self.classic_guide_params_fn(p)}
        stein_uparams = {p: v for p, v in unconstr_params.items() if p not in classic_uparams}

        # 1. Run warmup on a subset of particles to tune the MCMC state
        warmup_key, mcmc_key = jax.random.split(rng_key)
        sampler = self.sampler_fn(
            potential_fn=lambda params: self.loss.loss(warmup_key, {**params, **self.constrain_fn(classic_uparams)},
                                                       self._inference_model, self.guide, *args, **kwargs))
        mcmc = MCMC(sampler, self.num_mcmc_warmup, self.num_mcmc_updates, num_chains=self.num_mcmc_particles,
                    progress_bar=False, chain_method='vectorized',
                    **self.mcmc_kwargs)
        stein_params = self.constrain_fn(stein_uparams)
        stein_subset_params = {p: v[0:self.num_mcmc_particles] for p, v in stein_params.items()}
        mcmc.warmup(warmup_key, *args, init_params=stein_subset_params, **kwargs)

        # 2. Choose MCMC particles
        mcmc_key, choice_key = jax.random.split(mcmc_key)
        if self.num_mcmc_particles == self.num_particles:
            idxs = jnp.arange(self.num_particles)
        else:
            if self.sp_mcmc_crit == 'rand':
                idxs = jax.random.shuffle(choice_key, jnp.arange(self.num_particles))[:self.num_mcmc_particles]
            elif self.sp_mcmc_crit == 'infl':
                _, grads = self._svgd_loss_and_grads(choice_key, unconstr_params, *args, **kwargs)
                ksd = jnp.linalg.norm(
                    jnp.concatenate([jnp.reshape(grads[p], (self.num_particles, -1)) for p in stein_uparams.keys()],
                                    axis=-1),
                    ord=2, axis=-1)
                idxs = jnp.argsort(ksd)[:self.num_mcmc_particles]
            else:
                assert False, "Unsupported SP MCMC criterion: {}".format(self.sp_mcmc_crit)

        # 3. Run MCMC on chosen particles
        stein_params = self.constrain_fn(stein_uparams)
        stein_subset_params = {p: v[idxs] for p, v in stein_params.items()}
        mcmc.run(mcmc_key, *args, init_params=stein_subset_params, **kwargs)
        samples_subset_stein_params = mcmc.get_samples(group_by_chain=True)
        sss_uparams = self.uconstrain_fn(samples_subset_stein_params)

        # 4. Select best MCMC iteration to update particles
        scores = jax.vmap(
            lambda i: self._score_sp_mcmc(mcmc_key, idxs, stein_uparams, {p: v[:, i] for p, v in sss_uparams.items()},
                                          classic_uparams, *args, **kwargs))(jnp.arange(self.num_mcmc_particles))
        mcmc_idx = jnp.argmax(scores)
        stein_uparams = {p: ops.index_update(v, idxs, sss_uparams[p][:, mcmc_idx]) for p, v in stein_uparams.items()}
        return {**stein_uparams, **classic_uparams}

    def init(self, rng_key, *args, **kwargs):
        """
        :param jax.random.PRNGKey rng_key: random number generator seed.
        :param args: arguments to the model / guide (these can possibly vary during
            the course of fitting).
        :param kwargs: keyword arguments to the model / guide (these can possibly vary
            during the course of fitting).
        :return: initial :data:`CurrentState`
        """
        rng_key, model_seed, guide_seed = jax.random.split(rng_key, 3)
        model_init = handlers.seed(self.model, model_seed)
        guide_init = handlers.seed(self.guide, guide_seed)
        guide_trace = handlers.trace(guide_init).get_trace(*args, **kwargs, **self.static_kwargs)
        model_trace = handlers.trace(model_init).get_trace(*args, **kwargs, **self.static_kwargs)
        rng_key, particle_seed = jax.random.split(rng_key)
        particle_seeds = jax.random.split(particle_seed, num=self.num_particles)
        self.guide.find_params(particle_seeds, *args, **kwargs,
                               **self.static_kwargs)  # Get parameter values for each particle
        guide_init_params = self.guide.init_params()
        params = {}
        transforms = {}
        inv_transforms = {}
        particle_transforms = {}
        guide_param_names = set()
        should_enum = False
        for site in model_trace.values():
            if isinstance(site['fn'], Distribution) and site['fn'].is_discrete:
                if site['fn'].has_enumerate_support and self.enum:
                    should_enum = True
                else:
                    raise Exception("Cannot enumerate model with discrete variables without enumerate support")
        # NB: params in model_trace will be overwritten by params in guide_trace
        for site in list(model_trace.values()) + list(guide_trace.values()):
            if site['type'] == 'param':
                transform = get_parameter_transform(site)
                inv_transforms[site['name']] = transform
                transforms[site['name']] = transform.inv
                particle_transforms[site['name']] = site.get('particle_transform', IdentityTransform())
                if site['name'] in guide_init_params:
                    pval, _ = guide_init_params[site['name']]
                    if self.classic_guide_params_fn(site['name']):
                        pval = tree_map(lambda x: x[0], pval)
                else:
                    pval = site['value']
                params[site['name']] = transform.inv(pval)
                if site['name'] in guide_trace:
                    guide_param_names.add(site['name'])

        if should_enum:
            mpn = _guess_max_plate_nesting(model_trace)
            self._inference_model = enum(config_enumerate(self.model), - mpn - 1)
        self.guide_param_names = guide_param_names
        self.constrain_fn = partial(transform_fn, inv_transforms)
        self.uconstrain_fn = partial(transform_fn, transforms)
        self.particle_transform_fn = partial(transform_fn, particle_transforms)
        return VI.CurrentState(self.optim.init(params), rng_key)

    def get_params(self, state):
        """
        Gets values at `param` sites of the `model` and `guide`.
        :param svi_state: current state of the optimizer.
        """
        params = self.constrain_fn(self.optim.get_params(state.optim_state))
        return params

    def update(self, state, *args, **kwargs):
        """
        Take a single step of Stein (possibly on a batch / minibatch of data),
        using the optimizer.
        :param state: current state of Stein.
        :param args: arguments to the model / guide (these can possibly vary during
            the course of fitting).
        :param kwargs: keyword arguments to the model / guide (these can possibly vary
            during the course of fitting).
        :return: tuple of `(state, loss)`.
        """
        rng_key, rng_key_mcmc, rng_key_step = jax.random.split(state.rng_key, num=3)
        params = self.optim.get_params(state.optim_state)
        # Run Stein Point MCMC
        if self.num_mcmc_particles > 0:
            new_params = self._sp_mcmc(rng_key_mcmc, params, *args, **kwargs, **self.static_kwargs)
            grads = {p: new_params[p] - params[p] for p in params}
            optim_state = self.optim.update(grads, state.optim_state)
            params = self.optim.get_params(state.optim_state)
        else:
            optim_state = state.optim_state
        loss_val, grads = self._svgd_loss_and_grads(rng_key_step, params,
                                                    *args, **kwargs, **self.static_kwargs)
        optim_state = self.optim.update(grads, optim_state)
        return VI.CurrentState(optim_state, rng_key), loss_val

    def evaluate(self, state, *args, **kwargs):
        """
        Take a single step of Stein (possibly on a batch / minibatch of data).
        :param state: current state of Stein.
        :param args: arguments to the model / guide (these can possibly vary during
            the course of fitting).
        :param kwargs: keyword arguments to the model / guide.
        :return: evaluate loss given the current parameter values (held within `state.optim_state`).
        """
        # we split to have the same seed as `update_fn` given a state
        _, rng_key_eval = jax.random.split(state.rng_key)
        params = self.optim.get_params(state.optim_state)
        loss_val, _ = self._svgd_loss_and_grads(rng_key_eval, params,
                                                *args, **kwargs, **self.static_kwargs)
        return loss_val

    def predict(self, state, *args, num_samples=1, **kwargs):
        _, rng_key_predict = jax.random.split(state.rng_key)
        params = self.get_params(state)
        classic_params = {p: v for p, v in params.items() if
                          p not in self.guide_param_names or self.classic_guide_params_fn(p)}
        stein_params = {p: v for p, v in params.items() if p not in classic_params}
        if num_samples == 1:
            return jax.vmap(lambda sp: self._predict_model(rng_key_predict, {**sp, **classic_params}, *args, **kwargs)
                            )(stein_params)
        else:
            return jax.vmap(lambda rk: jax.vmap(lambda sp: self._predict_model(rk, {**sp, **classic_params},
                                                                               *args, **kwargs)
                                                )(stein_params))(jax.random.split(rng_key_predict, num_samples))