# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

from typing import Dict, List, Optional, Union

import numpy as np
import torch
from torch.distributions import Normal

from theseus.core.objective import Objective
from theseus.optimizer import OptimizerInfo
from theseus.optimizer.variable_ordering import VariableOrdering
from theseus.third_party.lml import LML

from .nonlinear_optimizer import (
    BackwardMode,
    EndIterCallbackType,
    NonlinearOptimizer,
    NonlinearOptimizerInfo,
    NonlinearOptimizerStatus,
)


class DCEM(NonlinearOptimizer):
    """
    DCEM optimizer for nonlinear optimization using sampling based techniques.
    The optimizer can be really sensitive to hypermeter tuning. Here are few tuning
    hints:
    1. If have to lower the max_iterations, then increase the n_sample.
    2. The higher the n_sample, the slowly with variance of samples will decrease.
    3. The higher the n_sample, more the chances of optimum being in the elite set.
    4. The higher the n_elite, the slower is convergence, but more accurate it might
       be, but would need more iterations. n_elite= 5 is good enough for most cases.
    """

    def __init__(
        self,
        objective: Objective,
        vectorize: bool = False,
        max_iterations: int = 50,
        n_sample: int = 100,
        n_elite: int = 5,
        temp: float = 1.0,
        init_sigma: Union[float, torch.Tensor] = 1.0,
        lb: float = None,
        ub: float = None,
        lml_verbose: bool = False,
        lml_eps: float = 1e-3,
        normalize: bool = True,
        abs_err_tolerance: float = 1e-6,
        rel_err_tolerance: float = 1e-4,
        **kwargs,
    ) -> None:
        super().__init__(
            objective,
            vectorize=vectorize,
            abs_err_tolerance=abs_err_tolerance,
            rel_err_tolerance=rel_err_tolerance,
            max_iterations=max_iterations,
            **kwargs,
        )

        self.objective = objective
        self.ordering = VariableOrdering(objective)
        self.n_samples = n_sample
        self.n_elite = n_elite
        self.lb = lb
        self.ub = ub
        self.temp = temp
        self.normalize = normalize
        self._tot_dof = sum([x.dof() for x in self.ordering])
        self.lml_eps = lml_eps
        self.lml_verbose = lml_verbose
        self.init_sigma = init_sigma

    def _mu_vec_to_dict(self, mu: torch.Tensor) -> Dict[str, torch.Tensor]:
        idx = 0
        mu_dic = {}
        for var in self.ordering:
            mu_dic[var.name] = mu[:, slice(idx, idx + var.dof())]
            idx += var.dof()
        return mu_dic

    def reset_sigma(self, init_sigma: Union[float, torch.Tensor]) -> None:
        self.sigma = (
            torch.ones(
                (self.objective.batch_size, self._tot_dof), device=self.objective.device
            )
            * init_sigma
        )

    def _CEM_step(self):
        """
        Performs one iteration of CEM.
        Updates the self.sigma and return the new mu.
        """
        device = self.objective.device
        n_batch = self.ordering[0].shape[0]

        mu = torch.cat([var.tensor for var in self.ordering], dim=-1)

        X = Normal(mu, self.sigma).rsample((self.n_samples,))

        X_samples: List[Dict[str, torch.Tensor]] = []
        for sample in X:
            X_samples.append(self._mu_vec_to_dict(sample))

        fX = torch.stack(
            [self.objective.error_metric(X_samples[i]) for i in range(self.n_samples)],
            dim=1,
        )

        assert fX.shape == (n_batch, self.n_samples)

        if self.temp is not None and self.temp < np.inf:
            if self.normalize:
                fX_mu = fX.mean(dim=1).unsqueeze(1)
                fX_sigma = fX.std(dim=1).unsqueeze(1)
                _fX = (fX - fX_mu) / (fX_sigma + 1e-6)
            else:
                _fX = fX

            if self.n_elite == 1:
                # indexes = LML(N=n_elite, verbose=lml_verbose, eps=lml_eps)(-_fX*temp)
                indexes = torch.softmax(-_fX * self.temp, dim=1)
            else:
                indexes = LML(
                    N=self.n_elite, verbose=self.lml_verbose, eps=self.lml_eps
                )(-_fX * self.temp)
            indexes = indexes.unsqueeze(2)
            eps = 0

        else:
            indexes_vals = fX.argsort(dim=1)[:, : self.n_elite]
            # Scatter 1.0 to the indexes using indexes_vals
            indexes = torch.zeros(n_batch, self.n_samples, device=device).scatter_(
                1, indexes_vals, 1.0
            )
            indexes = indexes.unsqueeze(2)
            eps = 1e-10
        # indexes.shape should be (n_batch, n_sample, 1)

        X = X.transpose(0, 1)

        assert indexes.shape[:2] == X.shape[:2]

        X_I = indexes * X

        mu = torch.sum(X_I, dim=1) / self.n_elite
        self.sigma = (
            (indexes * (X - mu.unsqueeze(1)) ** 2).sum(dim=1) / self.n_elite
        ).sqrt() + eps  # adding eps to avoid sigma=0, which is happening when temp=None

        assert self.sigma.shape == (n_batch, self._tot_dof)

        return self._mu_vec_to_dict(mu)

    def _optimize_loop(
        self,
        num_iter: int,
        info: NonlinearOptimizerInfo,
        verbose: bool,
        end_iter_callback: Optional[EndIterCallbackType] = None,
        **kwargs,
    ) -> int:
        converged_indices = torch.zeros_like(info.last_err).bool()
        iters_done = 0
        for it_ in range(num_iter):
            iters_done += 1
            try:
                mu = self._CEM_step()
            except RuntimeError as error:
                raise RuntimeError(f"There is an error in update {error}.")

            self.objective.update(mu)

            # check for convergence
            with torch.no_grad():
                err = self.objective.error_metric()
                self._update_info(info, it_, err, converged_indices)
                if verbose:
                    print(
                        f"Nonlinear optimizer. Iteration: {it_+1}. "
                        f"Error: {err.mean().item()} "
                    )
                converged_indices = self._check_convergence(err, info.last_err)
                info.status[
                    np.array(converged_indices.cpu().numpy())
                ] = NonlinearOptimizerStatus.CONVERGED

                if converged_indices.all():
                    break  # nothing else will happen at this point
                info.last_err = err

                if end_iter_callback is not None:
                    end_iter_callback(self, info, mu, it_)

        info.status[
            info.status == NonlinearOptimizerStatus.START
        ] = NonlinearOptimizerStatus.MAX_ITERATIONS

        return iters_done

    def _optimize_impl(
        self,
        track_best_solution: bool = False,
        track_err_history: bool = False,
        track_state_history: bool = False,
        verbose: bool = False,
        backward_mode: Union[str, BackwardMode] = BackwardMode.UNROLL,
        end_iter_callback: Optional[EndIterCallbackType] = None,
        **kwargs,
    ) -> OptimizerInfo:
        backward_mode = BackwardMode.resolve(backward_mode)
        init_sigma = kwargs.get("init_sigma", self.init_sigma)
        self.reset_sigma(init_sigma)

        with torch.no_grad():
            info = self._init_info(
                track_best_solution, track_err_history, track_state_history
            )

        if verbose:
            print(
                f"DCEM optimizer. Iteration: 0. "
                f"Error: {info.last_err.mean().item()}"
            )

        if backward_mode in [BackwardMode.UNROLL, BackwardMode.DLM]:
            self._optimize_loop(
                num_iter=self.params.max_iterations,
                info=info,
                verbose=verbose,
                end_iter_callback=end_iter_callback,
                **kwargs,
            )
            # If didn't coverge, remove misleading converged_iter value
            info.converged_iter[
                info.status == NonlinearOptimizerStatus.MAX_ITERATIONS
            ] = -1
            return info

        else:
            raise NotImplementedError(
                "DCEM currently only supports 'unroll' backward mode."
            )
