"""Economy-level structuring of BLP problem results."""

import time
from typing import Any, Callable, Dict, Hashable, List, Optional, Sequence, TYPE_CHECKING, Tuple

import numpy as np
import scipy.linalg

from .results import Results
from .. import exceptions, options
from ..configurations.iteration import Iteration
from ..markets.results_market import ResultsMarket
from ..parameters import Parameters
from ..utilities.algebra import approximately_solve, multiply_matrix_and_tensor, precisely_compute_eigenvalues
from ..utilities.basics import (
    Array, Bounds, Error, TableFormatter, format_number, format_seconds, generate_items, output, output_progress
)
from ..utilities.statistics import compute_gmm_parameter_covariances, compute_gmm_weights

# only import objects that create import cycles when checking types
if TYPE_CHECKING:
    from .bootstrapped_results import BootstrappedResults  # noqa
    from .optimal_instrument_results import OptimalInstrumentResults  # noqa
    from ..economies.problem import Progress  # noqa

class ProblemResults(Results):
    r"""Results of a solved BLP problem.

    Many results are class attributes. Other post-estimation outputs be computed by calling class methods.

    .. note::

       All methods in this class support :func:`parallel` processing. If multiprocessing is used, market-by-market
       computation of each post-estimation output will be distributed among the processes.

    Attributes
    ----------
    problem : `Problem`
        :class:`Problem` that created these results.
    last_results : `ProblemResults`
        :class:`ProblemResults` from the last GMM step.
    step : `int`
        GMM step that created these results.
    optimization_time : `float`
        Number of seconds it took the optimization routine to finish.
    cumulative_optimization_time : `float`
        Sum of :attr:`ProblemResults.optimization_time` for this step and all prior steps.
    total_time : `float`
        Sum of :attr:`ProblemResults.optimization_time` and the number of seconds it took to set up the GMM step and
        compute results after optimization had finished.
    cumulative_total_time : `float`
        Sum of :attr:`ProblemResults.total_time` for this step and all prior steps.
    optimization_iterations : `int`
        Number of major iterations completed by the optimization routine.
    cumulative_optimization_iterations : `int`
        Sum of :attr:`ProblemResults.optimization_iterations` for this step and all prior steps.
    objective_evaluations : `int`
        Number of GMM objective evaluations.
    cumulative_objective_evaluations : `int`
        Sum of :attr:`ProblemResults.objective_evaluations` for this step and all prior steps.
    fp_converged : `ndarray`
        Flags for convergence of the iteration routine used to compute :math:`\delta(\hat{\theta})` in each market
        during each objective evaluation. Rows are in the same order as :attr:`Problem.unique_market_ids` and column
        indices correspond to objective evaluations.
    cumulative_fp_converged : `ndarray`
        Concatenation of :attr:`ProblemResults.fp_converged` for this step and all prior steps.
    fp_iterations : `ndarray`
        Number of major iterations completed by the iteration routine used to compute :math:`\delta(\hat{\theta})` in
        each market during each objective evaluation. Rows are in the same order as
        :attr:`Problem.unique_market_ids` and column indices correspond to objective evaluations.
    cumulative_fp_iterations : `ndarray`
        Concatenation of :attr:`ProblemResults.fp_iterations` for this step and all prior steps.
    contraction_evaluations : `ndarray`
        Number of times the contraction used to compute :math:`\delta(\hat{\theta})` was evaluated in each market during
        each objective evaluation. Rows are in the same order as :attr:`Problem.unique_market_ids` and column
        indices correspond to objective evaluations.
    cumulative_contraction_evaluations : `ndarray`
        Concatenation of :attr:`ProblemResults.contraction_evaluations` for this step and all prior steps.
    converged : `bool`
        Whether the optimization routine converged.
    cumulative_converged : `bool`
        Whether the optimization routine converged for this step and all prior steps.
    parameters : `ndarray`
        Stacked parameters in the following order: :math:`\hat{\theta}`, concentrated out elements of
        :math:`\hat{\beta}`, and concentrated out elements of :math:`\hat{\gamma}`.
    parameter_covariances : `ndarray`
        Estimated covariance matrix of the stacked parameters, from which standard errors are extracted.
    theta : `ndarray`
        Estimated unfixed parameters, :math:`\hat{\theta}` in the following order: :math:`\hat{\Sigma}`,
        :math:`\hat{\Pi}`, :math:`\hat{\rho}`, non-concentrated out elements from :math:`\hat{\beta}`, and
        non-concentrated out elements from :math:`\hat{\gamma}`.
    sigma : `ndarray`
        Estimated Cholesky decomposition of the covariance matrix that measures agents' random taste distribution,
        :math:`\hat{\Sigma}`.
    pi : `ndarray`
        Estimated parameters that measures how agent tastes vary with demographics, :math:`\hat{\Pi}`.
    rho : `ndarray`
        Estimated parameters that measure within nesting group correlations, :math:`\hat{\rho}`.
    beta : `ndarray`
        Estimated demand-side linear parameters, :math:`\hat{\beta}`.
    gamma : `ndarray`
        Estimated supply-side linear parameters, :math:`\hat{\gamma}`.
    sigma_se : `ndarray`
        Estimated standard errors for :math:`\hat{\Sigma}`.
    pi_se : `ndarray`
        Estimated standard errors for :math:`\hat{\Pi}`.
    rho_se : `ndarray`
        Estimated standard errors for :math:`\hat{\rho}`.
    beta_se : `ndarray`
        Estimated standard errors for :math:`\hat{\beta}`.
    gamma_se : `ndarray`
        Estimated standard errors for :math:`\hat{\gamma}`.
    sigma_bounds : `tuple`
        Bounds for :math:`\Sigma` that were used during optimization, which are of the form ``(lb, ub)``.
    pi_bounds : `tuple`
        Bounds for :math:`\Pi` that were used during optimization, which are of the form ``(lb, ub)``.
    rho_bounds : `tuple`
        Bounds for :math:`\rho` that were used during optimization, which are of the form ``(lb, ub)``.
    beta_bounds : `tuple`
        Bounds for :math:`\beta` that were used during optimization, which are of the form ``(lb, ub)``.
    gamma_bounds : `tuple`
        Bounds for :math:`\gamma` that were used during optimization, which are of the form ``(lb, ub)``.
    delta : `ndarray`
        Estimated mean utility, :math:`\delta(\hat{\theta})`.
    tilde_costs : `ndarray`
        Estimated transformed marginal costs, :math:`\tilde{c}(\hat{\theta})`. Transformed marginal costs are simply
        :math:`\tilde{c} = c`, marginal costs, under a linear cost specification, and are :math:`\tilde{c} = \log c`
        under a log-linear specification. If ``costs_bounds`` were specified in :meth:`Problem.solve`, :math:`c` may
        have been clipped.
    clipped_costs : `ndarray`
        Vector of booleans indicating whether the associated marginal costs were clipped. All elements will be ``False``
        if ``costs_bounds`` in :meth:`Problem.solve` was not specified.
    xi : `ndarray`
        Estimated unobserved demand-side product characteristics, :math:`\xi(\hat{\theta})`, or equivalently, the
        demand-side structural error term.
    omega : `ndarray`
        Estimated unobserved supply-side product characteristics, :math:`\omega(\hat{\theta})`, or equivalently, the
        supply-side structural error term.
    objective : `float`
        GMM objective value.
    xi_by_theta_jacobian : `ndarray`
        Estimated :math:`\partial\xi / \partial\theta = \partial\delta / \partial\theta`.
    omega_by_theta_jacobian : `ndarray`
        Estimated :math:`\partial\omega / \partial\theta = \partial\tilde{c} / \partial\theta`.
    gradient : `ndarray`
        Estimated gradient of the GMM objective with respect to :math:`\theta`, which is computed after the optimization
        routine finishes even if the routine was configured to not use analytic gradients.
    sigma_gradient : `ndarray`
        Estimated gradient of the GMM objective with respect to :math:`\Sigma` elements in :math:`\theta`.
    pi_gradient : `ndarray`
        Estimated gradient of the GMM objective with respect to :math:`\Pi` elements in :math:`\theta`.
    rho_gradient : `ndarray`
        Estimated gradient of the GMM objective with respect to :math:`\rho` elements in :math:`\theta`.
    beta_gradient : `ndarray`
        Estimated gradient of the GMM objective with respect to :math:`\beta` elements in :math:`\theta`.
    gamma_gradient : `ndarray`
        Estimated gradient of the GMM objective with respect to :math:`\gamma` elements in :math:`\theta`.
    gradient_norm : `ndarray`
        Infinity norm of :attr:`ProblemResults.gradient`.
    hessian : `ndarray`
        Estimated Hessian of the GMM objective with respect to :math:`\theta`. By default, this is computed with finite
        central differences after the optimization routine finishes.
    hessian_eigenvalues : `ndarray`
        Eigenvalues of :attr:`ProblemResults.hessian`.
    W : `ndarray`
        Weighting matrix, :math:`W`, used to compute these results.
    updated_W : `ndarray`
        Updated weighting matrix.

    Examples
    --------
        - :doc:`Tutorial </tutorial>`

    """

    last_results: Optional['ProblemResults']
    step: int
    optimization_time: float
    cumulative_optimization_time: float
    total_time: float
    cumulative_total_time: float
    optimization_iterations: int
    cumulative_optimization_iterations: int
    objective_evaluations: int
    cumulative_objective_evaluations: int
    fp_converged: Array
    cumulative_fp_converged: Array
    fp_iterations: Array
    cumulative_fp_iterations: Array
    contraction_evaluations: Array
    cumulative_contraction_evaluations: Array
    converged: bool
    cumulative_converged: bool
    parameters: Array
    parameter_covariances: Array
    theta: Array
    sigma: Array
    pi: Array
    rho: Array
    beta: Array
    gamma: Array
    sigma_se: Array
    pi_se: Array
    rho_se: Array
    beta_se: Array
    gamma_se: Array
    sigma_bounds: Bounds
    pi_bounds: Bounds
    rho_bounds: Bounds
    beta_bounds: Bounds
    gamma_bounds: Bounds
    delta: Array
    tilde_costs: Array
    clipped_costs: Array
    xi: Array
    omega: Array
    objective: Array
    xi_by_theta_jacobian: Array
    omega_by_theta_jacobian: Array
    gradient: Array
    gradient_norm: Array
    hessian: Array
    hessian_eigenvalues: Array
    sigma_gradient: Array
    pi_gradient: Array
    rho_gradient: Array
    beta_gradient: Array
    gamma_gradient: Array
    W: Array
    updated_W: Array
    _costs_type: str
    _se_type: str
    _errors: List[Error]
    _parameters: Parameters

    def __init__(
            self, progress: 'Progress', last_results: Optional['ProblemResults'], step_start_time: float,
            optimization_start_time: float, optimization_end_time: float, iterations: int, evaluations: int,
            converged_mappings: Sequence[Dict[Hashable, bool]], iteration_mappings: Sequence[Dict[Hashable, int]],
            evaluation_mappings: Sequence[Dict[Hashable, int]], converged: bool, costs_type: str, costs_bounds: Bounds,
            center_moments: bool, W_type: str, se_type: str) -> None:
        """Compute cumulative progress statistics, update weighting matrices, and estimate standard errors."""

        # initialize values from the progress structure
        super().__init__(progress.problem)
        self._errors = progress.errors
        self.problem = progress.problem
        self.W = progress.W
        self.theta = progress.theta
        self.delta = progress.delta
        self.tilde_costs = progress.tilde_costs
        self.xi_by_theta_jacobian = progress.xi_jacobian
        self.omega_by_theta_jacobian = progress.omega_jacobian
        self.xi = progress.xi
        self.omega = progress.omega
        self.beta = progress.beta
        self.gamma = progress.gamma
        self.objective = progress.objective
        self.gradient = progress.gradient
        self.gradient_norm = progress.gradient_norm
        self.hessian = progress.hessian

        # if the Hessian was computed, compute its eigenvalues and the ratio of the smallest to largest ones
        self.hessian_eigenvalues = np.full(progress.parameters.P, np.nan, options.dtype)
        if progress.parameters.P > 0 and np.isfinite(self.hessian).all():
            self.hessian_eigenvalues, successful = precisely_compute_eigenvalues(self.hessian)
            if not successful:
                self._errors.append(exceptions.HessianEigenvaluesError(self.hessian))

        # store information about cost bounds
        self._costs_bounds = costs_bounds
        self.clipped_costs = progress.clipped_costs

        # initialize counts, times, and convergence
        self.step = 1
        self.total_time = self.cumulative_total_time = time.time() - step_start_time
        self.optimization_time = self.cumulative_optimization_time = optimization_end_time - optimization_start_time
        self.optimization_iterations = self.cumulative_optimization_iterations = iterations
        self.objective_evaluations = self.cumulative_objective_evaluations = evaluations
        self.fp_converged = self.cumulative_fp_converged = np.array(
            [[m[t] if m else True for m in converged_mappings] for t in self.problem.unique_market_ids],
            dtype=np.int
        )
        self.fp_iterations = self.cumulative_fp_iterations = np.array(
            [[m[t] if m else 0 for m in iteration_mappings] for t in self.problem.unique_market_ids],
            dtype=np.int
        )
        self.contraction_evaluations = self.cumulative_contraction_evaluations = np.array(
            [[m[t] if m else 0 for m in evaluation_mappings] for t in self.problem.unique_market_ids],
            dtype=np.int
        )
        self.converged = self.cumulative_converged = converged

        # initialize last results and add to cumulative values
        self.last_results = last_results
        if last_results is not None:
            self.step += last_results.step
            self.cumulative_total_time += last_results.cumulative_total_time
            self.cumulative_optimization_time += last_results.cumulative_optimization_time
            self.cumulative_optimization_iterations += last_results.cumulative_optimization_iterations
            self.cumulative_objective_evaluations += last_results.cumulative_objective_evaluations
            self.cumulative_fp_converged = np.c_[
                last_results.cumulative_fp_converged, self.cumulative_fp_converged
            ]
            self.cumulative_fp_iterations = np.c_[
                last_results.cumulative_fp_iterations, self.cumulative_fp_iterations
            ]
            self.cumulative_contraction_evaluations = np.c_[
                last_results.cumulative_contraction_evaluations, self.cumulative_contraction_evaluations
            ]
            self.cumulative_converged = last_results.converged and converged

        # store estimated parameters and information about them (beta and gamma have already been stored above)
        self._parameters = progress.parameters
        self.sigma, self.pi, self.rho, _, _ = self._parameters.expand(self.theta)
        self.parameters = np.c_[np.r_[
            self.theta,
            self.beta[self._parameters.eliminated_beta_index],
            self.gamma[self._parameters.eliminated_gamma_index]
        ]]
        self.sigma_bounds = self._parameters.sigma_bounds
        self.pi_bounds = self._parameters.pi_bounds
        self.rho_bounds = self._parameters.rho_bounds
        self.beta_bounds = self._parameters.beta_bounds
        self.gamma_bounds = self._parameters.gamma_bounds

        # collect inputs to weighting matrix and standard error computation
        u_list = [self.xi]
        Z_list = [self.problem.products.ZD]
        jacobian_list = [np.c_[
            self.xi_by_theta_jacobian,
            -self.problem.products.X1[:, self._parameters.eliminated_beta_index.flat],
            np.zeros_like(self.problem.products.X3[:, self._parameters.eliminated_gamma_index.flat])
        ]]
        if self.problem.K3 > 0:
            u_list.append(self.omega)
            Z_list.append(self.problem.products.ZS)
            jacobian_list.append(np.c_[
                self.omega_by_theta_jacobian,
                np.zeros_like(self.problem.products.X1[:, self._parameters.eliminated_beta_index.flat]),
                -self.problem.products.X3[:, self._parameters.eliminated_gamma_index.flat]
            ])

        # update the weighting matrix
        with np.errstate(invalid='ignore'):
            self.updated_W, W_errors = compute_gmm_weights(
                u_list, Z_list, W_type, self.problem.products.clustering_ids, center_moments
            )
        self._errors.extend(W_errors)

        # compute parameter covariances (if this is the first step, an unadjusted weighting matrix needs to be used so
        #   that unadjusted covariances are scaled properly)
        update_W = se_type == 'unadjusted' and self.step == 1
        with np.errstate(all='ignore'):
            self.parameter_covariances, covariance_errors = compute_gmm_parameter_covariances(
                jacobian_list, u_list, Z_list, self.W, se_type, self.problem.products.clustering_ids, update_W
            )
        self._errors.extend(covariance_errors)

        # compute standard errors
        with np.errstate(invalid='ignore'):
            se = np.sqrt(np.c_[self.parameter_covariances.diagonal()] / self.problem.N)
        if np.isnan(se).any():
            self._errors.append(exceptions.InvalidParameterCovariancesError())

        # expand standard errors
        theta_se, eliminated_beta_se, eliminated_gamma_se = np.split(se, [
            self._parameters.P,
            self._parameters.P + self._parameters.eliminated_beta_index.sum()
        ])
        self.sigma_se, self.pi_se, self.rho_se, self.beta_se, self.gamma_se = (
            self._parameters.expand(theta_se, nullify=True)
        )
        self.beta_se[self._parameters.eliminated_beta_index] = eliminated_beta_se.flatten()
        self.gamma_se[self._parameters.eliminated_gamma_index] = eliminated_gamma_se.flatten()

        # expand gradients
        self.sigma_gradient, self.pi_gradient, self.rho_gradient, self.beta_gradient, self.gamma_gradient = (
            self._parameters.expand(self.gradient, nullify=True)
        )

        # store types that are used in other methods
        self._costs_type = costs_type
        self._se_type = se_type

    def __str__(self) -> str:
        """Format problem results (including parameters estimates) as a string."""

        # construct a standard error description
        if self._se_type == 'unadjusted':
            se_description = "Unadjusted SEs"
        elif self._se_type == 'robust':
            se_description = "Robust SEs"
        else:
            assert self._se_type == 'clustered'
            se_description = f'Robust SEs Adjusted for {np.unique(self.problem.products.clustering_ids).size} Clusters'

        # combine a summary table section and another with formatted estimates into one string
        return "\n\n".join([
            self._format_summary(),
            self._parameters.format_estimates(
                f"Estimates ({se_description} in Parentheses)", self.sigma, self.pi, self.rho, self.beta, self.gamma,
                self.sigma_se, self.pi_se, self.rho_se, self.beta_se, self.gamma_se
            )
        ])

    def _format_summary(self) -> str:
        """Format a summary table of problem results."""

        # at a minimum include time and the GMM step
        floats_index = 3
        header = [("", "Computation", "Time"), ("", "GMM", "Step")]
        values = [format_seconds(self.cumulative_total_time), self.step]

        # include any optimization information
        if self._parameters.P > 0:
            floats_index += 1
            header.append(("", "Optimization", "Iterations"))
            values.append(self.cumulative_optimization_iterations)

        # include objective evaluation information
        header.append(("", "Objective", "Evaluations"))
        values.append(self.cumulative_objective_evaluations)

        # include any fixed point information
        if np.any(self.cumulative_contraction_evaluations > 0):
            floats_index += 2
            header.extend([("", "Fixed Point", "Iterations"), ("", "Contraction", "Evaluations")])
            values.extend([self.cumulative_fp_iterations.sum(), self.cumulative_contraction_evaluations.sum()])

        # include any information about the final objective value
        header.append(("", "Objective", "Value"))
        values.append(format_number(float(self.objective)))
        if np.isfinite(self.gradient_norm):
            header.append(("", "Gradient", "Infinity Norm"))
            values.append(format_number(float(self.gradient_norm)))
        if np.isfinite(self.hessian_eigenvalues).any():
            if self.hessian_eigenvalues.size == 1:
                header.append(("", "Hessian", "Eigenvalue"))
                values.append(format_number(float(self.hessian_eigenvalues)))
            else:
                header.extend([
                    ("Smallest", "Hessian", "Eigenvalue"),
                    ("Largest", "Hessian", "Eigenvalue")
                ])
                values.extend([
                    format_number(float(np.min(self.hessian_eigenvalues))),
                    format_number(float(np.max(self.hessian_eigenvalues)))
                ])

        # include any information about clipped marginal costs
        if np.isfinite(self._costs_bounds).any():
            header.append(("Clipped", "Marginal", "Costs"))
            values.append(self.clipped_costs.sum())

        # format the table
        widths = [max(options.digits + 6 if i >= floats_index else 0, *map(len, k)) for i, k in enumerate(header)]
        formatter = TableFormatter(widths)
        lines = [
            "Problem Results Summary:",
            formatter.line()
        ]
        if any(k[0] for k in header):
            lines.append(formatter([k[0] for k in header]))
        lines.extend([
            formatter([k[1] for k in header]),
            formatter([k[2] for k in header], underline=True),
            formatter(values),
            formatter.line()
        ])
        return "\n".join(lines)

    def bootstrap(
            self, draws: int = 1000, seed: Optional[int] = None, iteration: Optional[Iteration] = None) -> (
            'BootstrappedResults'):
        r"""Use a parametric bootstrap to create an empirical distribution of results.

        The constructed :class:`BootstrappedResults` can be used just like :class:`ProblemResults` to compute various
        post-estimation outputs. The only difference is that :class:`BootstrappedResults` methods return arrays with an
        extra first dimension, along which bootstrapped results are stacked. These stacked results can be used to
        construct, for example, confidence intervals for post-estimation outputs.

        For each bootstrap draw, parameters are drawn from the estimated multivariate normal distribution of all
        parameters defined by :attr:`ProblemResults.parameters` and :attr:`ProblemResults.parameter_covariances`. Note
        that any bounds configured during the optimization routine will be used to bound parameter draws. These
        parameters are used to compute the implied mean utility, :math:`\delta`, and shares, :math:`s`. If a supply side
        was estimated, the implied marginal costs, :math:`c`, and prices, :math:`p`, are computed as well. Specifically,
        if a supply side was estimated, equilibrium prices and shares are computed by iterating over the
        :math:`\zeta`-markup equation from :ref:`references:Morrow and Skerlos (2011)`.

        .. note::

           By default, the bootstrapping procedure may use a lot of memory. This is because it stores in memory all
           bootstrapped results (for all ``draws``) at the same time. To reduce the memory footprint of the procedure,
           call this method in a loop with ``draws`` set to ``1``. In each iteration of the loop, compute the desired
           post-estimation output with the proper method of the returned :class:`BootstrappedResults` class and store
           these outputs.

        Parameters
        ----------
        draws : `int, optional`
            The number of draws that will be taken from the joint distribution of the parameters. The default is
            ``1000``.
        seed : `int, optional`
            Passed to :class:`numpy.random.RandomState` to seed the random number generator before any draws are taken.
            By default, a seed is not passed to the random number generator.
        iteration : `Iteration, optional`
            :class:`Iteration` configuration used to compute bootstrapped prices by iterating over the
            :math:`\zeta`-markup equation from :ref:`references:Morrow and Skerlos (2011)`. By default, if a supply side
            was estimated, this is ``Iteration('simple', {'tol': 1e-12})``. Analytic Jacobians are not supported for
            this contraction mapping and this configuration is not used if a supply side was not estimated.

        Returns
        -------
        `BootstrappedResults`
            Computed :class:`BootstrappedResults`.

        Examples
        --------
            - :doc:`Tutorial </tutorial>`

        """
        errors: List[Error] = []

        # keep track of long it takes to bootstrap results
        output("Bootstrapping results ...")
        start_time = time.time()

        # validate the number of draws
        if not isinstance(draws, int) or draws < 1:
            raise ValueError("draws must be a positive int.")

        # validate the iteration configuration
        if self.problem.K3 == 0:
            iteration = None
        elif iteration is None:
            iteration = Iteration('simple', {'tol': 1e-12})
        elif not isinstance(iteration, Iteration):
            raise TypeError("iteration must be None or an iteration instance.")
        elif iteration._compute_jacobian:
            raise ValueError("Analytic Jacobians are not supported for this contraction mapping.")

        # draw from the asymptotic distribution implied by the estimated parameters
        state = np.random.RandomState(seed)
        bootstrapped_parameters = np.atleast_3d(state.multivariate_normal(
            self.parameters.flatten(), self.parameter_covariances, draws
        ))

        # extract the parameters
        bootstrapped_sigma = np.zeros((draws, self.sigma.shape[0], self.sigma.shape[1]), options.dtype)
        bootstrapped_pi = np.zeros((draws, self.pi.shape[0], self.pi.shape[1]), options.dtype)
        bootstrapped_rho = np.zeros((draws, self.rho.shape[0], self.rho.shape[1]), options.dtype)
        bootstrapped_beta = np.zeros((draws, self.beta.shape[0], self.beta.shape[1]), options.dtype)
        bootstrapped_gamma = np.zeros((draws, self.gamma.shape[0], self.gamma.shape[1]), options.dtype)
        bootstrapped_theta, bootstrapped_eliminated_beta, bootstrapped_eliminated_gamma = np.split(
            bootstrapped_parameters,
            [self._parameters.P, self._parameters.P + self._parameters.eliminated_beta_index.sum()],
            axis=1
        )
        bootstrapped_beta[:, self._parameters.eliminated_beta_index.flat] = bootstrapped_eliminated_beta
        bootstrapped_gamma[:, self._parameters.eliminated_gamma_index.flat] = bootstrapped_eliminated_gamma
        for d in range(draws):
            bootstrapped_sigma[d], bootstrapped_pi[d], bootstrapped_rho[d], beta_d, gamma_d = self._parameters.expand(
                bootstrapped_theta[d]
            )
            bootstrapped_beta[d] = np.where(self._parameters.eliminated_beta_index, bootstrapped_beta[d], beta_d)
            bootstrapped_gamma[d] = np.where(self._parameters.eliminated_gamma_index, bootstrapped_gamma[d], gamma_d)
            bootstrapped_sigma[d] = np.clip(bootstrapped_sigma[d], *self.sigma_bounds)
            bootstrapped_pi[d] = np.clip(bootstrapped_pi[d], *self.pi_bounds)
            bootstrapped_rho[d] = np.clip(bootstrapped_rho[d], *self.rho_bounds)
            bootstrapped_beta[d] = np.clip(bootstrapped_beta[d], *self.beta_bounds)
            bootstrapped_gamma[d] = np.clip(bootstrapped_gamma[d], *self.gamma_bounds)

        # compute bootstrapped prices, shares, delta and marginal costs
        converged_mappings: List[Dict[Hashable, bool]] = []
        iteration_mappings: List[Dict[Hashable, int]] = []
        evaluation_mappings: List[Dict[Hashable, int]] = []
        bootstrapped_prices = np.zeros((draws, self.problem.N, 1), options.dtype)
        bootstrapped_shares = np.zeros((draws, self.problem.N, 1), options.dtype)
        bootstrapped_delta = np.zeros((draws, self.problem.N, 1), options.dtype)
        bootstrapped_costs = np.zeros((draws, self.problem.N, int(self.problem.K3 > 0)), options.dtype)
        for d in output_progress(range(draws), draws, start_time):
            prices_d, shares_d, delta_d, costs_d, converged_d, iterations_d, evaluations_d, errors_d = (
                self._compute_bootstrap(
                    iteration, bootstrapped_sigma[d], bootstrapped_pi[d], bootstrapped_rho[d], bootstrapped_beta[d],
                    bootstrapped_gamma[d]
                )
            )
            bootstrapped_prices[d] = prices_d
            bootstrapped_shares[d] = shares_d
            bootstrapped_delta[d] = delta_d
            bootstrapped_costs[d] = costs_d
            converged_mappings.append(converged_d)
            iteration_mappings.append(iterations_d)
            evaluation_mappings.append(evaluations_d)
            errors.extend(errors_d)

        # structure the results
        from .bootstrapped_results import BootstrappedResults  # noqa
        results = BootstrappedResults(
            self, bootstrapped_sigma, bootstrapped_pi, bootstrapped_rho, bootstrapped_beta, bootstrapped_gamma,
            bootstrapped_prices, bootstrapped_shares, bootstrapped_delta, bootstrapped_costs, start_time, time.time(),
            draws, converged_mappings, iteration_mappings, evaluation_mappings
        )
        output(f"Bootstrapped results after {format_seconds(results.computation_time)}.")
        output("")
        output(results)
        return results

    def _compute_bootstrap(
            self, iteration: Optional[Iteration], sigma: Array, pi: Array, rho: Array, beta: Array, gamma: Array) -> (
            Tuple[
                Array, Array, Array, Array, Dict[Hashable, bool], Dict[Hashable, int], Dict[Hashable, int], List[Error]
            ]):
        """Compute the equilibrium prices, shares, marginal costs, and delta associated with bootstrapped parameters
        market-by-market
        """
        errors: List[Error] = []

        # compute delta (which will change under equilibrium prices) and marginal costs (which won't change)
        delta = self.delta + self.problem._compute_true_X1() @ (beta - self.beta)
        costs = self.tilde_costs + self.problem._compute_true_X3() @ (gamma - self.gamma)
        if self._costs_type == 'log':
            costs = np.exp(costs)

        # prices will only change if there is an iteration configuration
        prices = self.problem.products.prices if iteration is None else None

        # define a factory for computing bootstrapped prices, shares, and delta in markets
        def market_factory(s: Hashable) -> Tuple[ResultsMarket, Array, Optional[Array], Optional[Iteration]]:
            """Build a market along with arguments used to compute equilibrium prices and shares along with delta."""
            market_s = ResultsMarket(self.problem, s, sigma, pi, rho, beta, delta)
            costs_s = costs[self.problem._product_market_indices[s]]
            prices_s = prices[self.problem._product_market_indices[s]] if prices is not None else None
            return market_s, costs_s, prices_s, iteration

        # compute bootstrapped prices, shares, and delta market-by-market
        converged_mapping: Dict[Hashable, bool] = {}
        iteration_mapping: Dict[Hashable, int] = {}
        evaluation_mapping: Dict[Hashable, int] = {}
        equilibrium_prices = np.zeros_like(self.problem.products.prices)
        equilibrium_shares = np.zeros_like(self.problem.products.shares)
        generator = generate_items(self.problem.unique_market_ids, market_factory, ResultsMarket.solve_equilibrium)
        for t, (prices_t, shares_t, delta_t, errors_t, converged_t, iterations_t, evaluations_t) in generator:
            equilibrium_prices[self.problem._product_market_indices[t]] = prices_t
            equilibrium_shares[self.problem._product_market_indices[t]] = shares_t
            delta[self.problem._product_market_indices[t]] = delta_t
            errors.extend(errors_t)
            converged_mapping[t] = converged_t
            iteration_mapping[t] = iterations_t
            evaluation_mapping[t] = evaluations_t

        # return all of the information associated with this bootstrap draw
        return (
            equilibrium_prices, equilibrium_shares, delta, costs, converged_mapping, iteration_mapping,
            evaluation_mapping, errors
        )

    def compute_optimal_instruments(
            self, method: str = 'normal', draws: int = 100, seed: Optional[int] = None,
            expected_prices: Optional[Any] = None, iteration: Optional[Iteration] = None) -> 'OptimalInstrumentResults':
        r"""Estimate optimal or efficient instruments, :math:`\mathscr{Z}_D` and :math:`\mathscr{Z}_S`.

        Optimal instruments have been shown, for example, by :ref:`references:Reynaert and Verboven (2014)`, to reduce
        bias, improve efficiency, and enhance stability of BLP estimates.

        In the spirit of :ref:`references:Chamberlain (1987)`, optimal instruments for :math:`\theta` are

        .. math::

           \begin{bmatrix}
               \mathscr{Z}_D \\
               \mathscr{Z}_S
           \end{bmatrix}_{jt}
           = \text{Var}(\xi, \omega)^{-1}\operatorname{\mathbb{E}}\left[
           \begin{matrix}
               \frac{\partial\xi_{jt}}{\partial\theta} \\
               \frac{\partial\omega_{jt}}{\partial\theta}
           \end{matrix}
           \mathrel{\Bigg|} Z \right],

        The expectation is taken by integrating over the joint density of :math:`\xi` and :math:`\omega`. For each error
        term realization, if not already estimated, equilibrium prices are computed via iteration over the
        :math:`\zeta`-markup equation from :ref:`references:Morrow and Skerlos (2011)`. Associated shares and
        :math:`\delta` are then computed before each Jacobian is evaluated.

        The expected Jacobians are estimated with the average over all computed Jacobian realizations. The normalizing
        matrix :math:`\text{Var}(\xi, \omega)^{-1}` is estimated with the sample covariance matrix of the error terms.

        Optimal instruments for linear parameters not included in :math:`\theta` are simple product characteristics, so
        they are not computed here but are rather included in the set of instruments by
        :meth:`OptimalInstrumentResults.to_problem`.

        Parameters
        ----------
        method : `str, optional`
            The method by which the integral over the joint density of :math:`\xi` and :math:`\omega` is computed. The
            following methods are supported:

                - ``'normal'`` (default) - Draw from the normal approximation to the joint distribution of the error
                  terms and take the average over the computed Jacobians (``draws`` determines the number of draws).

                - ``'empirical'`` - Draw with replacement from the empirical joint distribution of the error terms and
                  take the average over the computed Jacobians (``draws`` determines the number of draws).

                - ``'approximate'`` - Evaluate the Jacobians at the expected value of the error terms: zero (``draws``
                  will be ignored).

        draws : `int, optional`
            The number of draws that will be taken from the joint distribution of the error terms. This is ignored if
            ``method`` is ``'approximate'``. The default is ``100``.
        seed : `int, optional`
            Passed to :class:`numpy.random.RandomState` to seed the random number generator before any draws are taken.
            By default, a seed is not passed to the random number generator.
        expected_prices : `array-like, optional`
            Vector of expected prices conditional on all exogenous variables,
            :math:`\operatorname{\mathbb{E}}[p \mid Z]`. By default, if a supply side was estimated, ``iteration`` is
            used. If only a demand side was estimated, this is by default estimated with the fitted values from a
            reduced form regression of endogenous prices onto :math:`Z_D`: all exogenous variables including excluded
            instruments.
        iteration : `Iteration, optional`
            :class:`Iteration` configuration used to estimate expected prices by iterating over the :math:`\zeta`-markup
            equation from :ref:`references:Morrow and Skerlos (2011)`. By default, if a supply side was estimated, this
            is ``Iteration('simple', {'tol': 1e-12})``. Analytic Jacobians are not supported for this contraction
            mapping and this configuration is not used if ``expected_prices`` is specified.

        Returns
        -------
        `OptimalInstrumentResults`
           Computed :class:`OptimalInstrumentResults`.

        Examples
        --------
            - :doc:`Tutorial </tutorial>`

        """
        errors: List[Error] = []

        # keep track of long it takes to compute optimal instruments for theta
        output("Computing optimal instruments for theta ...")
        start_time = time.time()

        # validate the method and create a function that samples from the error distribution
        if method == 'approximate':
            sample = lambda: (np.zeros_like(self.xi), np.zeros_like(self.omega))
        else:
            state = np.random.RandomState(seed)
            if method == 'normal':
                if self.problem.K3 == 0:
                    variance = np.var(self.xi)
                    sample = lambda: (np.c_[state.normal(0, variance, self.problem.N)], self.omega)
                else:
                    covariances = np.cov(self.xi, self.omega, rowvar=False)
                    sample = lambda: np.hsplit(state.multivariate_normal([0, 0], covariances, self.problem.N), 2)
            elif method == 'empirical':
                if self.problem.K3 == 0:
                    sample = lambda: (self.xi[state.choice(self.problem.N, self.problem.N)], self.omega)
                else:
                    joint = np.c_[self.xi, self.omega]
                    sample = lambda: np.hsplit(joint[state.choice(self.problem.N, self.problem.N)], 2)
            else:
                raise ValueError("method must be 'approximate', 'normal', or 'empirical'.")

        # validate the number of draws (there will be only one for the approximate method)
        if method == 'approximate':
            draws = 1
        if not isinstance(draws, int) or draws < 1:
            raise ValueError("draws must be a positive int.")

        # validate expected prices or their integration configuration (or compute expected prices with a reduced form
        #   regression if unspecified and only a demand side)
        if expected_prices is not None:
            iteration = None
            expected_prices = np.c_[np.asarray(expected_prices, options.dtype)]
            if expected_prices.shape != (self.problem.N, 1):
                raise ValueError(f"expected_prices must be a {self.problem.N}-vector.")
        elif self.problem.K3 > 0:
            if iteration is None:
                iteration = Iteration('simple', {'tol': 1e-12})
            elif not isinstance(iteration, Iteration):
                raise TypeError("iteration must be None or an Iteration instance.")
        else:
            prices = self.problem.products.prices
            if self.problem._absorb_demand_ids is not None:
                prices, absorption_errors = self.problem._absorb_demand_ids(prices)
                errors.extend(absorption_errors)
            covariances = self.problem.products.ZD.T @ self.problem.products.ZD
            parameters, replacement = approximately_solve(covariances, self.problem.products.ZD.T @ prices)
            if replacement:
                errors.append(exceptions.FittedValuesInversionError(covariances, replacement))
            expected_prices = self.problem.products.ZD @ parameters + self.problem.products.prices - prices

        # validate expected prices or their iteration configuration
        if expected_prices is None:
            if self.problem.K3 == 0:
                raise TypeError("A supply side was not estimated, so expected_prices must be specified.")
            if iteration is None:
                iteration = Iteration('simple', {'tol': 1e-12})
            elif not isinstance(iteration, Iteration):
                raise TypeError("iteration must be None or an Iteration instance.")
            elif iteration._compute_jacobian:
                raise ValueError("Analytic Jacobians are not supported for this contraction mapping.")
        else:
            iteration = None
            expected_prices = np.c_[np.asarray(expected_prices, options.dtype)]
            if expected_prices.shape != (self.problem.N, 1):
                raise ValueError(f"expected_prices must be a {self.problem.N}-vector.")

        # average over Jacobian realizations
        converged_mappings: List[Dict[Hashable, bool]] = []
        iteration_mappings: List[Dict[Hashable, int]] = []
        evaluation_mappings: List[Dict[Hashable, int]] = []
        expected_xi_jacobian = np.zeros_like(self.xi_by_theta_jacobian)
        expected_omega_jacobian = np.zeros_like(self.omega_by_theta_jacobian)
        for _ in output_progress(range(draws), draws, start_time):
            xi_jacobian_i, omega_jacobian_i, converged_i, iterations_i, evaluations_i, errors_i = (
                self._compute_realizations(expected_prices, iteration, *sample())
            )
            expected_xi_jacobian += xi_jacobian_i / draws
            expected_omega_jacobian += omega_jacobian_i / draws
            converged_mappings.append(converged_i)
            iteration_mappings.append(iterations_i)
            evaluation_mappings.append(evaluations_i)
            errors.extend(errors_i)

        # output a warning about any errors
        if errors:
            output("")
            output(exceptions.MultipleErrors(errors))
            output("")

        # compute the optimal instruments
        if self.problem.K3 == 0:
            inverse_covariance_matrix = np.c_[1 / np.var(self.xi)]
            demand_instruments = inverse_covariance_matrix * expected_xi_jacobian
            supply_instruments = np.full((self.problem.N, 0), np.nan, options.dtype)
        else:
            inverse_covariance_matrix = np.c_[scipy.linalg.inv(np.cov(self.xi, self.omega, rowvar=False))]
            instruments = multiply_matrix_and_tensor(
                inverse_covariance_matrix,
                np.stack([expected_xi_jacobian, expected_omega_jacobian], axis=1)
            )
            demand_instruments, supply_instruments = np.split(instruments.reshape((self.problem.N, -1)), 2, axis=1)

        # structure the results
        from .optimal_instrument_results import OptimalInstrumentResults  # noqa
        results = OptimalInstrumentResults(
            self, demand_instruments, supply_instruments, inverse_covariance_matrix, expected_xi_jacobian,
            expected_omega_jacobian, expected_prices, start_time, time.time(), draws, converged_mappings,
            iteration_mappings, evaluation_mappings
        )
        output(f"Computed optimal instruments after {format_seconds(results.computation_time)}.")
        output("")
        output(results)
        return results

    def _compute_realizations(
            self, expected_prices: Optional[Array], iteration: Optional[Iteration], xi: Array, omega: Array) -> (
            Tuple[Array, Array, Dict[Hashable, bool], Dict[Hashable, int], Dict[Hashable, int], List[Error]]):
        """If they have not already been estimated, compute the equilibrium prices, shares, and delta associated with a
        realization of xi and omega market-by-market. Then, compute realizations of Jacobians of xi and omega with
        respect to theta.
        """
        errors: List[Error] = []

        # compute delta (which will change under equilibrium prices) and marginal costs (which won't change)
        delta = self.delta - self.xi + xi
        costs = tilde_costs = self.tilde_costs - self.omega + omega
        if self._costs_type == 'log':
            costs = np.exp(costs)

        # define a factory for computing realizations of prices, shares, and delta in markets
        def market_factory(s: Hashable) -> Tuple[ResultsMarket, Array, Optional[Array], Optional[Iteration]]:
            """Build a market along with arguments used to compute equilibrium prices and shares along with delta."""
            market_s = ResultsMarket(self.problem, s, self.sigma, self.pi, self.rho, self.beta, delta)
            costs_s = costs[self.problem._product_market_indices[s]]
            prices_s = expected_prices[self.problem._product_market_indices[s]] if expected_prices is not None else None
            return market_s, costs_s, prices_s, iteration

        # compute realizations of prices, shares, and delta market-by-market
        converged_mapping: Dict[Hashable, bool] = {}
        iteration_mapping: Dict[Hashable, int] = {}
        evaluation_mapping: Dict[Hashable, int] = {}
        equilibrium_prices = np.zeros_like(self.problem.products.prices)
        equilibrium_shares = np.zeros_like(self.problem.products.shares)
        generator = generate_items(self.problem.unique_market_ids, market_factory, ResultsMarket.solve_equilibrium)
        for t, (prices_t, shares_t, delta_t, errors_t, converged_t, iterations_t, evaluations_t) in generator:
            equilibrium_prices[self.problem._product_market_indices[t]] = prices_t
            equilibrium_shares[self.problem._product_market_indices[t]] = shares_t
            delta[self.problem._product_market_indices[t]] = delta_t
            errors.extend(errors_t)
            converged_mapping[t] = converged_t
            iteration_mapping[t] = iterations_t
            evaluation_mapping[t] = evaluations_t

        # compute the Jacobian of xi with respect to theta
        xi_jacobian, demand_errors = self._compute_demand_realization(equilibrium_prices, equilibrium_shares, delta)
        errors.extend(demand_errors)

        # compute the Jacobian of omega with respect to theta
        omega_jacobian = np.full((self.problem.N, self._parameters.P), np.nan, options.dtype)
        if self.problem.K3 > 0:
            omega_jacobian, supply_errors = self._compute_supply_realization(
                equilibrium_prices, equilibrium_shares, delta, tilde_costs, xi_jacobian
            )
            errors.extend(supply_errors)

        # return all of the information associated with this realization
        return xi_jacobian, omega_jacobian, converged_mapping, iteration_mapping, evaluation_mapping, errors

    def _compute_demand_realization(
            self, equilibrium_prices: Array, equilibrium_shares: Array, delta: Array) -> Tuple[Array, List[Error]]:
        """Compute a realization of the Jacobian of xi with respect to theta market-by-market. If necessary, revert
        problematic elements to their estimated values.
        """
        errors: List[Error] = []

        # check if the Jacobian does not need to be computed
        xi_jacobian = np.full((self.problem.N, self._parameters.P), np.nan, options.dtype)
        if self._parameters.P == 0:
            return xi_jacobian, errors

        # define a factory for computing the Jacobian of xi with respect to theta in markets
        def market_factory(s: Hashable) -> Tuple[ResultsMarket, Parameters]:
            """Build a market with the data realization along with arguments used to compute the Jacobian."""
            data_override_s = {
                'prices': equilibrium_prices[self.problem._product_market_indices[s]],
                'shares': equilibrium_shares[self.problem._product_market_indices[s]]
            }
            market_s = ResultsMarket(self.problem, s, self.sigma, self.pi, self.rho, self.beta, delta, data_override_s)
            return market_s, self._parameters

        # compute the Jacobian market-by-market
        generator = generate_items(
            self.problem.unique_market_ids, market_factory, ResultsMarket.compute_xi_by_theta_jacobian
        )
        for t, (xi_jacobian_t, errors_t) in generator:
            xi_jacobian[self.problem._product_market_indices[t]] = xi_jacobian_t
            errors.extend(errors_t)

        # replace invalid elements
        bad_jacobian_index = ~np.isfinite(xi_jacobian)
        if np.any(bad_jacobian_index):
            xi_jacobian[bad_jacobian_index] = self.xi_by_theta_jacobian[bad_jacobian_index]
            errors.append(exceptions.XiByThetaJacobianReversionError(bad_jacobian_index))
        return xi_jacobian, errors

    def _compute_supply_realization(
            self, equilibrium_prices: Array, equilibrium_shares: Array, delta: Array, tilde_costs: Array,
            xi_jacobian: Array) -> Tuple[Array, List[Error]]:
        """Compute a realization of the Jacobian of omega with respect to theta market-by-market. If necessary, revert
        problematic elements to their estimated values.
        """
        errors: List[Error] = []

        # define a factory for computing the Jacobian of omega with respect to theta in markets
        def market_factory(s: Hashable) -> Tuple[ResultsMarket, Array, Array, Parameters, str]:
            """Build a market with the data realization along with arguments used to compute the Jacobians."""
            data_override_s = {
                'prices': equilibrium_prices[self.problem._product_market_indices[s]],
                'shares': equilibrium_shares[self.problem._product_market_indices[s]]
            }
            market_s = ResultsMarket(self.problem, s, self.sigma, self.pi, self.rho, self.beta, delta, data_override_s)
            tilde_costs_s = tilde_costs[self.problem._product_market_indices[s]]
            xi_jacobian_s = xi_jacobian[self.problem._product_market_indices[s]]
            return market_s, tilde_costs_s, xi_jacobian_s, self._parameters, self._costs_type

        # compute the Jacobian market-by-market
        omega_jacobian = np.full((self.problem.N, self._parameters.P), np.nan, options.dtype)
        generator = generate_items(
            self.problem.unique_market_ids, market_factory, ResultsMarket.compute_omega_by_theta_jacobian
        )
        for t, (omega_jacobian_t, errors_t) in generator:
            omega_jacobian[self.problem._product_market_indices[t]] = omega_jacobian_t
            errors.extend(errors_t)

        # the Jacobian should be zero for any clipped marginal costs
        omega_jacobian[self.clipped_costs.flat] = 0

        # replace invalid elements
        bad_jacobian_index = ~np.isfinite(omega_jacobian)
        if np.any(bad_jacobian_index):
            omega_jacobian[bad_jacobian_index] = self.omega_by_theta_jacobian[bad_jacobian_index]
            errors.append(exceptions.OmegaByThetaJacobianReversionError(bad_jacobian_index))
        return omega_jacobian, errors

    def _coerce_matrices(self, matrices: Any) -> Array:
        """Coerce array-like stacked matrices into a stacked matrix and validate it."""
        matrices = np.c_[np.asarray(matrices, options.dtype)]
        if matrices.shape != (self.problem.N, self.problem._max_J):
            raise ValueError(f"matrices must be {self.problem.N} by {self.problem._max_J}.")
        return matrices

    def _coerce_optional_costs(self, costs: Optional[Any]) -> Array:
        """Coerce optional array-like costs into a column vector and validate it."""
        if costs is not None:
            costs = np.c_[np.asarray(costs, options.dtype)]
            if costs.shape != (self.problem.N, 1):
                raise ValueError(f"costs must be None or a {self.problem.N}-vector.")
        return costs

    def _coerce_optional_prices(self, prices: Optional[Any]) -> Array:
        """Coerce optional array-like prices into a column vector and validate it."""
        if prices is not None:
            prices = np.c_[np.asarray(prices, options.dtype)]
            if prices.shape != (self.problem.N, 1):
                raise ValueError(f"prices must be None or a {self.problem.N}-vector.")
        return prices

    def _coerce_optional_shares(self, shares: Optional[Any]) -> Array:
        """Coerce optional array-like shares into a column vector and validate it."""
        if shares is not None:
            shares = np.c_[np.asarray(shares, options.dtype)]
            if shares.shape != (self.problem.N, 1):
                raise ValueError(f"shares must be None or a {self.problem.N}-vector.")
        return shares

    def _combine_arrays(
            self, compute_market_results: Callable, fixed_args: Sequence = (), market_args: Sequence = ()) -> Array:
        """Compute an array for each market and stack them into a single matrix. An array for a single market is
        computed by passing fixed_args (identical for all markets) and market_args (matrices with as many rows as there
        are products that are restricted to the market) to compute_market_results, a ResultsMarket method that returns
        the output for the market and a set of any errors encountered during computation.
        """
        errors: List[Error] = []

        # keep track of how long it takes to compute the arrays
        start_time = time.time()

        # define a factory for computing arrays in markets
        def market_factory(s: Hashable) -> tuple:
            """Build a market along with arguments used to compute arrays."""
            indices_s = self.problem._product_market_indices[s]
            market_s = ResultsMarket(self.problem, s, self.sigma, self.pi, self.rho, self.beta, self.delta)
            args_s = [None if a is None else a[indices_s] for a in market_args]
            return (market_s, *fixed_args, *args_s)

        # construct a mapping from market IDs to market-specific arrays
        matrix_mapping: Dict[Hashable, Array] = {}
        generator = output_progress(
            generate_items(self.problem.unique_market_ids, market_factory, compute_market_results), self.problem.T,
            start_time
        )
        for t, (array_t, errors_t) in generator:
            matrix_mapping[t] = np.c_[array_t]
            errors.extend(errors_t)

        # output a warning about any errors
        if errors:
            output("")
            output(exceptions.MultipleErrors(errors))
            output("")

        # determine the number of rows and columns
        row_count = sum(matrix_mapping[t].shape[0] for t in self.problem.unique_market_ids)
        column_count = max(matrix_mapping[t].shape[1] for t in self.problem.unique_market_ids)

        # preserve the original product order or the sorted market order when stacking the arrays
        combined = np.full((row_count, column_count), np.nan, options.dtype)
        for t, matrix_t in matrix_mapping.items():
            if row_count == self.problem.N:
                combined[self.problem._product_market_indices[t], :matrix_t.shape[1]] = matrix_t
            else:
                combined[self.problem.unique_market_ids == t, :matrix_t.shape[1]] = matrix_t

        # output how long it took to compute the arrays
        end_time = time.time()
        output(f"Finished after {format_seconds(end_time - start_time)}.")
        output("")
        return combined


import os
import time
import re
from contextlib import contextmanager
from copy import copy
from pathlib import Path

import numpy as np
import math
import torch
import torch.nn as nn
from helper.torch_utils import is_parallel

@contextmanager
def torch_distributed_zero_first(local_rank: int):
    """
    Decorator to make all processes in distributed training wait for each local_master to do something.
    """
    if local_rank not in [-1, 0]:
        torch.distributed.barrier()
    yield
    if local_rank == 0:
        torch.distributed.barrier()

def xywh2xyxy(x):
    # Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
    y = torch.zeros_like(x) if isinstance(x, torch.Tensor) else np.zeros_like(x)
    y[:, 0] = x[:, 0] - x[:, 2] / 2  # top left x
    y[:, 1] = x[:, 1] - x[:, 3] / 2  # top left y
    y[:, 2] = x[:, 0] + x[:, 2] / 2  # bottom right x
    y[:, 3] = x[:, 1] + x[:, 3] / 2  # bottom right y
    return y

def xyxy2xywh(x):
    # Convert nx4 boxes from [x1, y1, x2, y2] to [x, y, w, h] where xy1=top-left, xy2=bottom-right
    y = torch.zeros_like(x) if isinstance(x, torch.Tensor) else np.zeros_like(x)
    y[:, 0] = (x[:, 0] + x[:, 2]) / 2  # x center
    y[:, 1] = (x[:, 1] + x[:, 3]) / 2  # y center
    y[:, 2] = x[:, 2] - x[:, 0]  # width
    y[:, 3] = x[:, 3] - x[:, 1]  # height
    return y

def smooth_BCE(eps=0.1):  # https://github.com/ultralytics/yolov3/issues/238#issuecomment-598028441
    # return positive, negative label smoothing BCE targets
    return 1.0 - 0.5 * eps, 0.5 * eps

def compute_ap(recall, precision):
    """ Compute the average precision, given the recall and precision curves.
    Source: https://github.com/rbgirshick/py-faster-rcnn.
    # Arguments
        recall:    The recall curve (list).
        precision: The precision curve (list).
    # Returns
        The average precision as computed in py-faster-rcnn.
    """

    # Append sentinel values to beginning and end
    mrec = np.concatenate(([0.], recall, [min(recall[-1] + 1E-3, 1.)]))
    mpre = np.concatenate(([0.], precision, [0.]))

    # Compute the precision envelope
    mpre = np.flip(np.maximum.accumulate(np.flip(mpre)))

    # Integrate area under curve
    method = 'interp'  # methods: 'continuous', 'interp'
    if method == 'interp':
        x = np.linspace(0, 1, 101)  # 101-point interp (COCO)
        ap = np.trapz(np.interp(x, mrec, mpre), x)  # integrate
    else:  # 'continuous'
        i = np.where(mrec[1:] != mrec[:-1])[0]  # points where x axis (recall) changes
        ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])  # area under curve

    return ap

class FocalLoss(nn.Module):
    # Wraps focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5)
    def __init__(self, loss_fcn, gamma=1.5, alpha=0.25):
        super(FocalLoss, self).__init__()
        self.loss_fcn = loss_fcn  # must be nn.BCEWithLogitsLoss()
        self.gamma = gamma
        self.alpha = alpha
        self.reduction = loss_fcn.reduction
        self.loss_fcn.reduction = 'none'  # required to apply FL to each element

    def forward(self, pred, true):
        loss = self.loss_fcn(pred, true)
        # p_t = torch.exp(-loss)
        # loss *= self.alpha * (1.000001 - p_t) ** self.gamma  # non-zero power for gradient stability

        # TF implementation https://github.com/tensorflow/addons/blob/v0.7.1/tensorflow_addons/losses/focal_loss.py
        pred_prob = torch.sigmoid(pred)  # prob from logits
        p_t = true * pred_prob + (1 - true) * (1 - pred_prob)
        alpha_factor = true * self.alpha + (1 - true) * (1 - self.alpha)
        modulating_factor = (1.0 - p_t) ** self.gamma
        loss *= alpha_factor * modulating_factor

        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        else:  # 'none'
            return loss

def build_targets(p, targets, model):
    # Build targets for compute_loss(), input targets(image,class,x,y,w,h)
    det = model.module.model[-1] if is_parallel(model) else model.model[-1]  # Detect() module
    na, nt = det.na, targets.shape[0]  # number of anchors, targets
    tcls, tbox, indices, anch = [], [], [], []
    gain = torch.ones(7, device=targets.device)  # normalized to gridspace gain
    ai = torch.arange(na, device=targets.device).float().view(na, 1).repeat(1, nt)  # same as .repeat_interleave(nt)
    targets = torch.cat((targets.repeat(na, 1, 1), ai[:, :, None]), 2)  # append anchor indices

    g = 0.5  # bias
    off = torch.tensor([[0, 0],
                        [1, 0], [0, 1], [-1, 0], [0, -1],  # j,k,l,m
                        # [1, 1], [1, -1], [-1, 1], [-1, -1],  # jk,jm,lk,lm
                        ], device=targets.device).float() * g  # offsets

    for i in range(det.nl):
        anchors = det.anchors[i]
        gain[2:6] = torch.tensor(p[i].shape)[[3, 2, 3, 2]]  # xyxy gain

        # Match targets to anchors
        t = targets * gain
        if nt:
            # Matches
            r = t[:, :, 4:6] / anchors[:, None]  # wh ratio
            j = torch.max(r, 1. / r).max(2)[0] < model.hyp['anchor_t']  # compare
            # j = wh_iou(anchors, t[:, 4:6]) > model.hyp['iou_t']  # iou(3,n)=wh_iou(anchors(3,2), gwh(n,2))
            t = t[j]  # filter

            # Offsets
            gxy = t[:, 2:4]  # grid xy
            gxi = gain[[2, 3]] - gxy  # inverse
            j, k = ((gxy % 1. < g) & (gxy > 1.)).T
            l, m = ((gxi % 1. < g) & (gxi > 1.)).T
            j = torch.stack((torch.ones_like(j), j, k, l, m))
            t = t.repeat((5, 1, 1))[j]
            offsets = (torch.zeros_like(gxy)[None] + off[:, None])[j]
        else:
            t = targets[0]
            offsets = 0

        # Define
        b, c = t[:, :2].long().T  # image, class
        gxy = t[:, 2:4]  # grid xy
        gwh = t[:, 4:6]  # grid wh
        gij = (gxy - offsets).long()
        gi, gj = gij.T  # grid xy indices

        # Append
        a = t[:, 6].long()  # anchor indices
        indices.append((b, a, gj, gi))  # image, anchor, grid indices
        tbox.append(torch.cat((gxy - gij, gwh), 1))  # box
        anch.append(anchors[a])  # anchors
        tcls.append(c)  # class

    return tcls, tbox, indices, anch

def bbox_iou(box1, box2, x1y1x2y2=True, GIoU=False, DIoU=False, CIoU=False, eps=1e-9):
    # Returns the IoU of box1 to box2. box1 is 4, box2 is nx4
    box2 = box2.T

    # Get the coordinates of bounding boxes
    if x1y1x2y2:  # x1, y1, x2, y2 = box1
        b1_x1, b1_y1, b1_x2, b1_y2 = box1[0], box1[1], box1[2], box1[3]
        b2_x1, b2_y1, b2_x2, b2_y2 = box2[0], box2[1], box2[2], box2[3]
    else:  # transform from xywh to xyxy
        b1_x1, b1_x2 = box1[0] - box1[2] / 2, box1[0] + box1[2] / 2
        b1_y1, b1_y2 = box1[1] - box1[3] / 2, box1[1] + box1[3] / 2
        b2_x1, b2_x2 = box2[0] - box2[2] / 2, box2[0] + box2[2] / 2
        b2_y1, b2_y2 = box2[1] - box2[3] / 2, box2[1] + box2[3] / 2

    # Intersection area
    inter = (torch.min(b1_x2, b2_x2) - torch.max(b1_x1, b2_x1)).clamp(0) * \
            (torch.min(b1_y2, b2_y2) - torch.max(b1_y1, b2_y1)).clamp(0)

    # Union Area
    w1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1 + eps
    w2, h2 = b2_x2 - b2_x1, b2_y2 - b2_y1 + eps
    union = w1 * h1 + w2 * h2 - inter + eps

    iou = inter / union
    if GIoU or DIoU or CIoU:
        cw = torch.max(b1_x2, b2_x2) - torch.min(b1_x1, b2_x1)  # convex (smallest enclosing box) width
        ch = torch.max(b1_y2, b2_y2) - torch.min(b1_y1, b2_y1)  # convex height
        if CIoU or DIoU:  # Distance or Complete IoU https://arxiv.org/abs/1911.08287v1
            c2 = cw ** 2 + ch ** 2 + eps  # convex diagonal squared
            rho2 = ((b2_x1 + b2_x2 - b1_x1 - b1_x2) ** 2 +
                    (b2_y1 + b2_y2 - b1_y1 - b1_y2) ** 2) / 4  # center distance squared
            if DIoU:
                return iou - rho2 / c2  # DIoU
            elif CIoU:  # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47
                v = (4 / math.pi ** 2) * torch.pow(torch.atan(w2 / h2) - torch.atan(w1 / h1), 2)
                with torch.no_grad():
                    alpha = v / ((1 + eps) - iou + v)
                return iou - (rho2 / c2 + v * alpha)  # CIoU
        else:  # GIoU https://arxiv.org/pdf/1902.09630.pdf
            c_area = cw * ch + eps  # convex area
            return iou - (c_area - union) / c_area  # GIoU
    else:
        return iou  # IoU

def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, fname='precision-recall_curve.png'):
    """ Compute the average precision, given the recall and precision curves.
    Source: https://github.com/rafaelpadilla/Object-Detection-Metrics.
    # Arguments
        tp:  True positives (nparray, nx1 or nx10).
        conf:  Objectness value from 0-1 (nparray).
        pred_cls:  Predicted object classes (nparray).
        target_cls:  True object classes (nparray).
        plot:  Plot precision-recall curve at mAP@0.5
        fname:  Plot filename
    # Returns
        The average precision as computed in py-faster-rcnn.
    """

    # Sort by objectness
    i = np.argsort(-conf)
    tp, conf, pred_cls = tp[i], conf[i], pred_cls[i]

    # Find unique classes
    unique_classes = np.unique(target_cls)

    # Create Precision-Recall curve and compute AP for each class
    px, py = np.linspace(0, 1, 1000), []  # for plotting
    pr_score = 0.1  # score to evaluate P and R https://github.com/ultralytics/yolov3/issues/898
    s = [unique_classes.shape[0], tp.shape[1]]  # number class, number iou thresholds (i.e. 10 for mAP0.5...0.95)
    ap, p, r = np.zeros(s), np.zeros(s), np.zeros(s)
    for ci, c in enumerate(unique_classes):
        i = pred_cls == c
        n_gt = (target_cls == c).sum()  # Number of ground truth objects
        n_p = i.sum()  # Number of predicted objects

        if n_p == 0 or n_gt == 0:
            continue
        else:
            # Accumulate FPs and TPs
            fpc = (1 - tp[i]).cumsum(0)
            tpc = tp[i].cumsum(0)

            # Recall
            recall = tpc / (n_gt + 1e-16)  # recall curve
            r[ci] = np.interp(-pr_score, -conf[i], recall[:, 0])  # r at pr_score, negative x, xp because xp decreases

            # Precision
            precision = tpc / (tpc + fpc)  # precision curve
            p[ci] = np.interp(-pr_score, -conf[i], precision[:, 0])  # p at pr_score

            # AP from recall-precision curve
            py.append(np.interp(px, recall[:, 0], precision[:, 0]))  # precision at mAP@0.5
            for j in range(tp.shape[1]):
                ap[ci, j] = compute_ap(recall[:, j], precision[:, j])

    # Compute F1 score (harmonic mean of precision and recall)
    f1 = 2 * p * r / (p + r + 1e-16)

    return p, r, ap, f1, unique_classes.astype('int32')

def box_iou(box1, box2):
    # https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py
    """
    Return intersection-over-union (Jaccard index) of boxes.
    Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
    Arguments:
        box1 (Tensor[N, 4])
        box2 (Tensor[M, 4])
    Returns:
        iou (Tensor[N, M]): the NxM matrix containing the pairwise
            IoU values for every element in boxes1 and boxes2
    """

    def box_area(box):
        # box = 4xn
        return (box[2] - box[0]) * (box[3] - box[1])

    area1 = box_area(box1.T)
    area2 = box_area(box2.T)

    # inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2)
    inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) - torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2)
    return inter / (area1[:, None] + area2 - inter)  # iou = inter / (area1 + area2 - inter)

def clip_coords(boxes, img_shape):
    # Clip bounding xyxy bounding boxes to image shape (height, width)
    boxes[:, 0].clamp_(0, img_shape[1])  # x1
    boxes[:, 1].clamp_(0, img_shape[0])  # y1
    boxes[:, 2].clamp_(0, img_shape[1])  # x2
    boxes[:, 3].clamp_(0, img_shape[0])  # y2

def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, merge=False, classes=None, agnostic=False):
    """Performs Non-Maximum Suppression (NMS) on inference results

    Returns:
         detections with shape: nx6 (x1, y1, x2, y2, conf, cls)
    """

    nc = prediction[0].shape[1] - 5  # number of classes
    xc = prediction[..., 4] > conf_thres  # candidates

    # Settings
    min_wh, max_wh = 2, 4096  # (pixels) minimum and maximum box width and height
    max_det = 300  # maximum number of detections per image
    time_limit = 10.0  # seconds to quit after
    redundant = True  # require redundant detections
    multi_label = nc > 1  # multiple labels per box (adds 0.5ms/img)

    t = time.time()
    output = [None] * prediction.shape[0]
    for xi, x in enumerate(prediction):  # image index, image inference
        # Apply constraints
        # x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0  # width-height
        x = x[xc[xi]]  # confidence

        # If none remain process next image
        if not x.shape[0]:
            continue

        # Compute conf
        x[:, 5:] *= x[:, 4:5]  # conf = obj_conf * cls_conf

        # Box (center x, center y, width, height) to (x1, y1, x2, y2)
        box = xywh2xyxy(x[:, :4])

        # Detections matrix nx6 (xyxy, conf, cls)
        if multi_label:
            i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T
            x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1)
        else:  # best class only
            conf, j = x[:, 5:].max(1, keepdim=True)
            x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres]

        # Filter by class
        if classes:
            x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]

        # Apply finite constraint
        # if not torch.isfinite(x).all():
        #     x = x[torch.isfinite(x).all(1)]

        # If none remain process next image
        n = x.shape[0]  # number of boxes
        if not n:
            continue

        # Sort by confidence
        # x = x[x[:, 4].argsort(descending=True)]

        # Batched NMS
        c = x[:, 5:6] * (0 if agnostic else max_wh)  # classes
        boxes, scores = x[:, :4] + c, x[:, 4]  # boxes (offset by class), scores
        i = torch.ops.torchvision.nms(boxes, scores, iou_thres)
        if i.shape[0] > max_det:  # limit detections
            i = i[:max_det]
        if merge and (1 < n < 3E3):  # Merge NMS (boxes merged using weighted mean)
            try:  # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)
                iou = box_iou(boxes[i], boxes) > iou_thres  # iou matrix
                weights = iou * scores[None]  # box weights
                x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True)  # merged boxes
                if redundant:
                    i = i[iou.sum(1) > 1]  # require redundancy
            except:  # possible CUDA error https://github.com/ultralytics/yolov3/issues/1139
                print(x, i, x.shape, i.shape)
                pass

        output[xi] = x[i]
        if (time.time() - t) > time_limit:
            break  # time limit exceeded

    return output

def compute_loss(p, targets, model):  # predictions, targets, model
    device = targets.device
    lcls, lbox, lobj = torch.zeros(1, device=device), torch.zeros(1, device=device), torch.zeros(1, device=device)
    tcls, tbox, indices, anchors = build_targets(p, targets, model)  # targets
    h = model.hyp  # hyperparameters

    # Define criteria
    BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.Tensor([h['cls_pw']])).to(device)
    BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.Tensor([h['obj_pw']])).to(device)

    # Class label smoothing https://arxiv.org/pdf/1902.04103.pdf eqn 3
    cp, cn = smooth_BCE(eps=0.0)

    # Focal loss
    g = h['fl_gamma']  # focal loss gamma
    if g > 0:
        BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g)

    # Losses
    nt = 0  # number of targets
    np = len(p)  # number of outputs
    balance = [4.0, 1.0, 0.4] if np == 3 else [4.0, 1.0, 0.4, 0.1]  # P3-5 or P3-6
    for i, pi in enumerate(p):  # layer index, layer predictions
        b, a, gj, gi = indices[i]  # image, anchor, gridy, gridx
        tobj = torch.zeros_like(pi[..., 0], device=device)  # target obj

        n = b.shape[0]  # number of targets
        if n:
            nt += n  # cumulative targets
            ps = pi[b, a, gj, gi]  # prediction subset corresponding to targets

            # Regression
            pxy = ps[:, :2].sigmoid() * 2. - 0.5
            pwh = (ps[:, 2:4].sigmoid() * 2) ** 2 * anchors[i]
            pbox = torch.cat((pxy, pwh), 1).to(device)  # predicted box
            iou = bbox_iou(pbox.T, tbox[i], x1y1x2y2=False, CIoU=True)  # iou(prediction, target)
            lbox += (1.0 - iou).mean()  # iou loss

            # Objectness
            tobj[b, a, gj, gi] = (1.0 - model.gr) + model.gr * iou.detach().clamp(0).type(tobj.dtype)  # iou ratio

            # Classification
            if model.nc > 1:  # cls loss (only if multiple classes)
                t = torch.full_like(ps[:, 5:], cn, device=device)  # targets
                t[range(n), tcls[i]] = cp
                lcls += BCEcls(ps[:, 5:], t)  # BCE

            # Append targets to text file
            # with open('targets.txt', 'a') as file:
            #     [file.write('%11.5g ' * 4 % tuple(x) + '\n') for x in torch.cat((txy[i], twh[i]), 1)]

        lobj += BCEobj(pi[..., 4], tobj) * balance[i]  # obj loss

    s = 3 / np  # output count scaling
    lbox *= h['box'] * s
    lobj *= h['obj'] * s * (1.4 if np == 4 else 1.)
    lcls *= h['cls'] * s
    bs = tobj.shape[0]  # batch size

    loss = lbox + lobj + lcls
    return loss * bs, torch.cat((lbox, lobj, lcls, loss)).detach()


# -*- coding: utf-8 -*-
"""
Created on Tue May  5 21:08:51 2015

@author: richard

TODO: implemement unit tests with nose
"""

import sys
import numpy as np
import pandas as pd
from flight import Flight
from decisions import Decisions
from observations import Observations
from random import choice as choose
from roboskeeter.math.math_toolbox import generate_random_unit_vector

class Simulator:
    """Our simulated mosquito.
    """

    def __init__(self, experiment, agent_kwargs):
        """ Load params
        """
        # dump kwarg dictionary into the agent object
        for key, value in agent_kwargs.iteritems():
            setattr(self, key, value)

        self.decisions = Decisions(self.decision_policy, self.stimulus_memory_n_timesteps)

        if experiment.is_simulation:
            self._simulated_agent_init(experiment)

    def _simulated_agent_init(self, experiment):
        # defaults
        self.mass = 2.88e-6  # avg. mass of our colony (kg) =2.88 mg,
        self.time_max = 15.
        self.dt = 0.01
        self.max_bins = int(np.ceil(self.time_max / self.dt))  # N bins

        # from gassian fit to experimental control data
        self.initial_velocity_mu = 0.18
        self.initial_velocity_stdev = 0.08

        # useful aliases
        self.experiment = experiment
        self.windtunnel = self.experiment.environment.windtunnel
        self.bounded = self.experiment.experiment_conditions['bounded']
        self.boundary = self.windtunnel.boundary
        self.heat = self.experiment.environment.heat

        # useful lists TODO: get rid of?
        self.kinematics_list = ['position', 'velocity', 'acceleration']  # curvature?
        self.forces_list = ['total_f', 'random_f', 'stim_f']
        self.other_list = ['tsi', 'times', 'decision', 'heat_signal', 'in_plume']

        # mk forces
        self.flight = Flight(self.random_f_strength,
                             self.stim_f_strength,
                             self.damping_coeff)

        # turn thresh, in units deg s-1.
        # From Sharri:
        # it is the stdev of the broader of two Gaussians that fit the distribution of angular velocity

        # # create repulsion landscape
        # self._repulsion_funcs = repulsion_landscape3D.landscape(boundary=self.boundary)

    def fly(self, n_trajectories=1):
        """ runs _generate_flight n_trajectories times
        """
        df_list = []
        traj_i = 0
        try:
            if self.verbose:
                print """Starting simulations with {} heat model and {} decision policy.
                If you run out of patience, press <CTL>-C to stop generating simulations and
                cut to the chase scene.""".format(
                self.heat.heat_model_name, self.decision_policy)

            while traj_i < n_trajectories:
                # print updates
                if self.verbose:
                    sys.stdout.write("\rTrajectory {}/{}".format(traj_i + 1, n_trajectories))
                    sys.stdout.flush()

                array_dict = self._generate_flight()

                # if len(array_dict['velocity_x']) < 5:  # hack to catch when optimizer makes trajectories explode
                #     print "catching explosion"
                #     break

                # add label column to enumerate the trajectories
                array_len = len(array_dict['tsi'])
                array_dict['trajectory_num'] = [traj_i] * array_len

                # mk df, add to list of dfs
                df = pd.DataFrame(array_dict)
                # df = df.set_index(['trajectory_num'])
                df_list.append(df)

                traj_i += 1

                if traj_i == n_trajectories:
                    if self.verbose:
                        sys.stdout.write("\rSimulations finished. Performing deep magic.")
                        sys.stdout.flush()

        except KeyboardInterrupt:
            print "\n Simulations interrupted at iteration {}. Moving along...".format(traj_i)
            pass

        observations = Observations()
        observations.kinematics = pd.concat(df_list)  # concatenate all the data frames at once for performance boost.

        return observations

    def _generate_flight(self):
        """Generate a single trajectory using our model.

        First put everything into np arrays stored inside of a dictionary
        """
        dt = self.dt
        m = self.mass
        vector_dict = self._initialize_vector_dict()

        # # dynamically create easy-to-read aliases for the contents of vector_dict
        # for key, value in vector_dict.iteritems():
        #     exec(key + " = vector_dict['" + key + "']")
        # unpack vector dict into nicer aliases
        in_plume = vector_dict['in_plume']
        heat_signal = vector_dict['heat_signal']
        position = vector_dict['position']
        velocity = vector_dict['velocity']
        acceleration = vector_dict['acceleration']
        random_f = vector_dict['random_f']
        stim_f = vector_dict['stim_f']
        total_f = vector_dict['total_f']
        decision = vector_dict['decision']

        position[0] = self._set_init_position()
        velocity[0] = self._set_init_velocity()

        for tsi in vector_dict['tsi']:
            in_plume[tsi] = self.heat.check_in_plume_bounds(position[tsi])  # returns False for non-Bool plume

            decision[tsi], heat_signal[tsi] = self.decisions.make_decision(in_plume[tsi], velocity[tsi][1])

            if heat_signal[tsi] == 'X':  # this is an awful hack telling us to look up the gradient
                heat_signal[tsi] = self.heat.get_nearest_gradient(position[tsi])

            stim_f[tsi], random_f[tsi], total_f[tsi] = self.flight.calc_forces(velocity[tsi], decision[tsi], heat_signal[tsi])

            # calculate current acceleration
            acceleration[tsi] = total_f[tsi] / m

            # check if time is out, end loop before we solve for future velo, position
            if tsi == self.max_bins-1: # -1 because of how range() works
                vector_dict = self._land(tsi, vector_dict)
                break

            ################################################
            # Calculate candidate velocity and positions
            ################################################
            candidate_velo = velocity[tsi] + acceleration[tsi] * dt

            # make sure velocity doesn't diverge to infinity if system is unstable
            # this stops the optimizer from crashing
            candidate_velo = self._velocity_ceiling(candidate_velo)

            candidate_pos = position[tsi] + candidate_velo * dt

            ################################################
            # test candidates
            ################################################
            if self.bounded:
                candidate_pos, candidate_velo = self._collide_with_wall(candidate_pos, candidate_velo)

            position[tsi + 1] = candidate_pos
            velocity[tsi + 1] = candidate_velo

        # once flight is finished, make dictionary ready to be loaded into DF
        vector_dict = self._fix_vector_dict(vector_dict)

        return vector_dict

    def _land(self, tsi, V):
        ''' trim excess timebins in arrays
        '''
        if tsi == 0:  # hack for if we need to chop a trajectory at the very start
            for k, array in V.iteritems():
                V[k] = array[:1]
        else:
            for k, array in V.iteritems():
                V[k] = array[:tsi - 1]
                V[k] = array[:tsi - 1]

        return V

    def _collide_with_wall(self, candidate_pos, candidate_velo):
        walls = self.windtunnel.walls
        xpos, ypos, zpos = candidate_pos
        xvelo, yvelo, zvelo = candidate_velo
        teleport_distance = 0.005  # this is arbitrary
        crash = False

        # print "test", candidate_velo

        # x dim
        if xpos < walls.downwind:  # too far behind
            crash = True
            xpos = walls.downwind + teleport_distance  # teleport back inside
            if self.collision_type == 'elastic':
                xvelo *= -1.
            elif self.collision_type == 'part_elastic':
                xvelo *= -self.restitution_coeff
            elif self.collision_type == 'crash':
                xvelo = 0.
            else:
                raise ValueError("unknown collision type {}".format(self.collision_type))
        if xpos > walls.upwind:  # reached far (upwind) wall (end)
            crash = True
            xpos = walls.upwind - teleport_distance  # teleport back inside
            if self.collision_type == 'elastic':
                xvelo *= -1.
            elif self.collision_type == 'part_elastic':
                xvelo *= -self.restitution_coeff
            elif self.collision_type == 'crash':
                xvelo = 0.

        # y dim
        if ypos < walls.left:  # too left
            crash = True
            ypos = walls.left + teleport_distance
            if self.collision_type == 'elastic':
                yvelo *= -1.
            elif self.collision_type == 'part_elastic':
                yvelo *= -self.restitution_coeff
            elif self.collision_type == "crash":
                yvelo = 0.

        if ypos > walls.right:  # too far right
            crash = True
            ypos = walls.right - teleport_distance
            if self.collision_type == 'elastic':
                yvelo *= -1.
            elif self.collision_type == 'part_elastic':
                yvelo *= -self.restitution_coeff
            elif self.collision_type == 'crash':
                yvelo = 0.

        # z dim
        if zpos > walls.ceiling:  # too far above
            crash = True
            zpos = walls.ceiling - teleport_distance
            if self.collision_type == 'elastic':
                zvelo *= -1.
            elif self.collision_type == 'part_elastic':
                zvelo *= -self.restitution_coeff
            elif self.collision_type == "crash":
                zvelo = 0.
        if zpos < walls.floor:  # too far below
            crash = True
            zpos = walls.floor + teleport_distance
            if self.collision_type == 'elastic':
                zvelo *= -1.
            elif self.collision_type == 'part_elastic':
                zvelo *= -self.restitution_coeff
            elif self.collision_type == 'crash':
                zvelo = 0.

        try:
            candidate_pos, candidate_velo = np.array([xpos, ypos, zpos]), np.array([xvelo, yvelo, zvelo])
        except:
            print " cand velo", [xvelo, yvelo, zvelo], "before", candidate_velo

        return candidate_pos, candidate_velo

    def _initialize_vector_dict(self):
        """
        initialize np arrays, store in dictionary
        """
        V = {}

        for name in self.kinematics_list + self.forces_list:
            V[name] = np.full((self.max_bins, 3), np.nan)

        V['tsi'] = np.arange(self.max_bins)
        V['times'] = np.linspace(0, self.time_max, self.max_bins)
        V['in_plume'] = np.zeros(self.max_bins, dtype=bool)
        V['heat_signal'] = np.array([None] * self.max_bins)
        V['decision'] = np.array([None] * self.max_bins)

        return V

    def _set_init_velocity(self):
        initial_velocity_norm = np.random.normal(self.initial_velocity_mu, self.initial_velocity_stdev, 1)

        unit_vector = generate_random_unit_vector()
        velocity_vec = initial_velocity_norm * unit_vector

        return velocity_vec

    def _set_init_position(self):
        ''' puts the agent in an initial position, usually within the bounds of the
        cage

        Options: [the cage] door, or anywhere in the plane at x=.1 meters

        set initial velocity from fitted distribution
        '''

        # generate random intial velocity condition using normal distribution fitted to experimental data
        if self.initial_position_selection == 'realistic':
            """these were calculated by taking selecting the initial positions of all observed trajectories in all
            conditions. Then, for each dimension, I calculated the distance of each initial position to the nearest wall
            in that dimension (i.e. for each z I calculated the distance to the floor and ceiling and selected
            the smallest distance. Then, Decisions aand"""
            downwind, upwind, left, right, floor, ceiling = self.boundary
            x_avg_dist_to_wall = 0.268
            y_avg_dist_to_wall = 0.044
            z_avg_dist_to_wall = 0.049
            x = choose([(downwind + x_avg_dist_to_wall), (upwind - x_avg_dist_to_wall)])
            y = choose([left + y_avg_dist_to_wall, (right - y_avg_dist_to_wall)])
            z = ceiling - z_avg_dist_to_wall
            initial_position = np.array([x,y,z])
        elif self.initial_position_selection == 'downwind_high':
            initial_position = np.array(
                [0.05, np.random.uniform(-0.127, 0.127), 0.2373])  # 0.2373 is mode of z pos distribution
        elif type(self.initial_position_selection) is list:
            initial_position = np.array(self.initial_position_selection)
        elif self.initial_position_selection == "door":  # start trajectories as they exit the front door
            initial_position = np.array([0.1909, np.random.uniform(-0.0381, 0.0381), np.random.uniform(0., 0.1016)])
            # FIXME cage is actually suspending above floor
        elif self.initial_position_selection == 'downwind_plane':
            initial_position = np.array([0.1, np.random.uniform(-0.127, 0.127), np.random.uniform(0., 0.254)])
        else:
            raise Exception('invalid agent position specified: {}'.format(self.initial_position_selection))

        return initial_position

    def _velocity_ceiling(self, candidate_velo):
        """check if we're seeing enormous velocities, which sometimes happens when running the optimization
         algoirithm. if so, cap the velocity instead of landing. this allows the optimizer to keep running.
        """
        for i, velo in enumerate(candidate_velo):
            if velo > 20:
                candidate_velo[i] = 20.
            elif velo < -20:
                candidate_velo[i] = -20.

        return candidate_velo

    def _fix_vector_dict(self, dct):
        # prepare dict for loading into pandas (dataframe only accepts 1D vectors)
        # split xyz dicts into separate x, y, z vectors for dataframe

        fixed_dct = {}
        for kinematic in self.kinematics_list + self.forces_list:
            fixed_dct[kinematic + '_x'], fixed_dct[kinematic + '_y'], fixed_dct[kinematic + '_z'] = np.split(
                dct[kinematic], 3, axis=1)

        # migrate rest of dict
        for v in self.other_list:
            fixed_dct[v] = dct[v]

        # fix pandas bug when trying to load (R,1) arrays when it expects (R,) arrays
        for key, dct in fixed_dct.iteritems():
            fixed_dct[key] = fixed_dct[key].reshape(len(dct))
            if fixed_dct[key].size == 0:
                fixed_dct[key] = np.array([0.])  # hack so that kde calculation doesn't freeze on empty arrays

        return fixed_dct


import pdb  # noqa: F401

import numpy as np
import cupy as cp
import scipy.signal as scsig

import os
from os import path
import time
from glob import glob
import deepdish as dd

from typing import Tuple, NamedTuple, TypeVar, List

import soundfile as sf

import multiprocessing as mp
import logging  # noqa: F401

N_CUDA_DEV = 4
NDArray = TypeVar('NDArray', np.ndarray, cp.ndarray)

def search_all_files(DIR_WAVFILE: str, ID: str) -> List[str]:
    result = []
    for folder, _, _ in os.walk(DIR_WAVFILE):
        files = glob(path.join(folder, ID))
        if not files:
            continue
        result += files

    return result

class SFTData(NamedTuple):
    """
    Constant Matrices/Vectors for Spherical Fourier Analysis
    """
    bEQspec: NDArray
    Yenc: NDArray
    Wnv: NDArray
    Wpv: NDArray
    Vv: NDArray

    def get_triags(self) -> Tuple:
        return (self.Wnv, self.Wpv, self.Vv)

class PreProcessor:
    def __init__(self, RIRs: np.ndarray, Ys: np.ndarray, sftdata: SFTData,
                 L_WIN_MS=20., RIRs_0: np.ndarray=None):
        # Bug Fix
        np.fft.restore_all()
        # From Parameters
        self.RIRs = RIRs
        self.N_LOC, self.N_MIC, self.L_RIR = RIRs.shape
        self.Ys = Ys
        self.RIRs_0 = RIRs_0
        self.sftdata = sftdata

        self.L_WIN_MS = L_WIN_MS

        # Determined during process
        self.DIR_IV = ''
        self.all_files = []

        # Common for all wave file
        self.Fs = 0
        self.N_wavfile = 0
        self.N_fft = 0
        self.N_freq = 0
        self.L_frame = 0
        self.L_hop = 0
        self.win = None

    def process(self, DIR_WAVFILE: str, ID: str, idx_start: int,
                DIR_IV: str, FORM: str, N_CORES=mp.cpu_count()//4):
        if not path.exists(DIR_IV):
            os.makedirs(DIR_IV)
        self.DIR_IV = DIR_IV

        if self.N_LOC < mp.cpu_count():
            N_CORES = self.N_LOC
        n_loc_per_core = int(np.ceil(self.N_LOC//N_CORES))

        max_n_pool = 1
        while max_n_pool*int(np.ceil(n_loc_per_core*N_CORES/N_CUDA_DEV)) < 30:
            max_n_pool += 1
        max_n_pool -= 1

        print(f'Start processing from the {idx_start}-th wave file')

        # Search all wave files
        # self.all_files = search_all_files(DIR_WAVFILE, ID)
        self.Fs = 16000
        data = np.random.normal(size=(16000*10))
        self.L_frame = int(self.Fs*self.L_WIN_MS//1000)
        self.N_fft = self.L_frame
        if self.N_fft % 2 == 0:
            self.N_freq = self.N_fft//2 + 1
        else:
            self.N_freq = self.N_fft//2
        self.L_hop = self.L_frame//2

        self.win = scsig.hamming(self.L_frame, sym=False)

        pdb.set_trace()
        # Main Process
        for i_proc in range(N_CORES):
            if (i_proc + 1) * n_loc_per_core <= self.N_LOC:
                range_loc = range(i_proc * n_loc_per_core,
                                  (i_proc+1) * n_loc_per_core)
            elif i_proc * n_loc_per_core < self.N_LOC:
                range_loc = range(i_proc * n_loc_per_core, self.N_LOC)
            else:
                break
            self.save_IV(i_proc % N_CUDA_DEV,
                         data,
                         range_loc,
                         FORM, self.N_wavfile+1)

    def save_IV(self, i_dev: int, data: NDArray, range_loc: iter,
                FORM: str, *args):
        """
        Save IV files.

        i_dev: GPU Device No.
        range_loc: RIR Index Range(S/M Location Index Range)
        FORM: format of filename
        args: format string arguments

        return: None
        """
        # CUDA Ready
        cp.cuda.Device(i_dev).use()
        data = cp.array(data)
        win = cp.array(self.win)
        Ys = cp.array(self.Ys)
        sftdata = SFTData(*[cp.array(item) for item in self.sftdata])

        N_frame_free = data.shape[0]//self.L_hop - 1
        N_frame_room = (data.shape[0]+self.L_RIR-1)//self.L_hop - 1

        for i_loc in range_loc:
            # RIR Filtering
            data_0 \
                = cp.array(scsig.fftconvolve(cp.asnumpy(data.reshape(1, -1)),
                                             self.RIRs_0[i_loc]))

            # using 0-th Order RIR
            iv_0 = cp.zeros((self.N_freq, N_frame_room, 16), dtype=complex)
            for i_frame in range(N_frame_room):
                interval = i_frame*self.L_hop + np.arange(self.L_frame)
                fft = cp.fft.fft(data_0[:, interval]*win, n=self.N_fft)
                anm = (sftdata.Yenc @ fft) * sftdata.bEQspec

                iv_0[:, i_frame, :] = anm[:, :self.N_freq].T

            # Free-field Intensity Vector Image
            iv_free = cp.zeros((self.N_freq, N_frame_free, 16), dtype=complex)
            for i_frame in range(N_frame_free):
                interval = i_frame*self.L_hop + np.arange(self.L_frame)
                fft = cp.fft.fft(data[interval]*win, n=self.N_fft)
                anm = cp.outer(Ys[i_loc].conj(), fft)

                iv_free[:, i_frame, :] = anm[:, :self.N_freq].T

            # Save
            dict_to_save = {'IV_free': cp.asnumpy(iv_free),
                            # 'IV_room': cp.asnumpy(iv_room),
                            'IV_0': cp.asnumpy(iv_0),
                            # 'data': cp.asnumpy(data),
                            # 'norm_factor_free': norm_factor_free,
                            # 'norm_factor_room': norm_factor_room,
                            }
            FNAME = FORM % (*args, i_loc)
            dd.io.save(path.join(self.DIR_IV, FNAME),
                       dict_to_save,
                       compression=None)

            print(FORM % (*args, i_loc))

    def __str__(self):
        return ('Wave Files Processed/Total: '
                f'{self.N_wavfile}/{len(self.all_files)}\n'
                f'Sample Rate: {self.Fs}\n'
                f'Number of source location: {self.N_LOC}\n'
                )

    def print_save_info(self):
        """
        Print __str__ and save metadata.
        """
        print(self)

        metadata = {'N_wavfile': self.N_wavfile,
                    'Fs': self.Fs,
                    # 'N_fft': self.N_fft,
                    'N_freq': self.N_freq,
                    'L_frame': self.L_frame,
                    'L_hop': self.L_hop,
                    'N_LOC': self.N_LOC,
                    'path_wavfiles': self.all_files,
                    }

        dd.io.save(path.join(self.DIR_IV, 'metadata.h5'), metadata)

    @staticmethod
    def seltriag(Ain: NDArray, nrord: int, shft: Tuple[int, int]) -> NDArray:
        xp = cp.get_array_module(Ain)
        N_freq = 1 if Ain.ndim == 1 else Ain.shape[1]
        N = int(np.ceil(np.sqrt(Ain.shape[0]))-1)
        idx = 0
        len_new = (N-nrord+1)**2

        Aout = xp.zeros((len_new, N_freq), dtype=Ain.dtype)
        for ii in range(N-nrord+1):
            for jj in range(-ii, ii+1):
                n = shft[0] + ii
                m = shft[1] + jj
                idx_from = m + n*(n+1)
                if -n <= m and m <= n and 0 <= n and n <= N \
                        and idx_from < Ain.shape[0]:
                    Aout[idx] = Ain[idx_from]
                idx += 1
        return Aout

    @classmethod
    def calc_intensity(cls, Asv: NDArray,
                       Wnv: NDArray, Wpv: NDArray, Vv: NDArray) -> NDArray:
        """
        Asv(anm) -> IV
        """
        xp = cp.get_array_module(Asv)

        aug1 = cls.seltriag(Asv, 1, (0, 0))
        aug2 = cls.seltriag(Wpv, 1, (1, -1))*cls.seltriag(Asv, 1, (1, -1)) \
            - cls.seltriag(Wnv, 1, (0, 0))*cls.seltriag(Asv, 1, (-1, -1))
        aug3 = cls.seltriag(Wpv, 1, (0, 0))*cls.seltriag(Asv, 1, (-1, 1)) \
            - cls.seltriag(Wnv, 1, (1, 1))*cls.seltriag(Asv, 1, (1, 1))
        aug4 = cls.seltriag(Vv, 1, (0, 0))*cls.seltriag(Asv, 1, (-1, 0)) \
            + cls.seltriag(Vv, 1, (1, 0))*cls.seltriag(Asv, 1, (1, 0))

        dx = (aug1.conj()*(aug2+aug3)/2).sum(axis=0)
        dy = (aug1.conj()*(aug2-aug3)/2j).sum(axis=0)
        dz = (aug1.conj()*aug4).sum(axis=0)

        return 0.5*xp.real(xp.stack((dx, dy, dz), axis=1))

if __name__ == '__main__':
    pass