#!/usr/bin/env python3
# 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.

# pyre-strict

from typing import Dict, List, Optional, Set, Tuple

from ax.core.observation import (
    Observation,
    ObservationData,
    ObservationFeatures,
    separate_observations,
)
from ax.core.optimization_config import OptimizationConfig
from ax.core.parameter import ChoiceParameter, FixedParameter
from ax.core.search_space import SearchSpace
from ax.core.types import TParamValueList
from ax.exceptions.core import UserInputError
from ax.modelbridge.base import GenResults, ModelBridge
from ax.modelbridge.modelbridge_utils import (
    array_to_observation_data,
    get_fixed_features,
)
from ax.modelbridge.torch import (
    extract_objective_weights,
    extract_outcome_constraints,
    validate_optimization_config,
)
from ax.models.discrete_base import DiscreteModel
from ax.models.types import TConfig


FIT_MODEL_ERROR = "Model must be fit before {action}."


# pyre-fixme[13]: Attribute `model` is never initialized.
# pyre-fixme[13]: Attribute `outcomes` is never initialized.
# pyre-fixme[13]: Attribute `parameters` is never initialized.
# pyre-fixme[13]: Attribute `search_space` is never initialized.
class DiscreteModelBridge(ModelBridge):
    """A model bridge for using models based on discrete parameters.

    Requires that all parameters have been transformed to ChoiceParameters.
    """

    model: DiscreteModel
    outcomes: List[str]
    parameters: List[str]
    search_space: Optional[SearchSpace]

    def _fit(
        self,
        model: DiscreteModel,
        search_space: SearchSpace,
        observations: List[Observation],
    ) -> None:
        self.model = model
        # Convert observations to arrays
        self.parameters = list(search_space.parameters.keys())
        all_metric_names: Set[str] = set()
        observation_features, observation_data = separate_observations(observations)
        for od in observation_data:
            all_metric_names.update(od.metric_names)
        self.outcomes = list(all_metric_names)
        # Convert observations to arrays
        Xs_array, Ys_array, Yvars_array = self._convert_observations(
            observation_data=observation_data,
            observation_features=observation_features,
            outcomes=self.outcomes,
            parameters=self.parameters,
        )
        # Extract parameter values
        parameter_values = _get_parameter_values(search_space, self.parameters)
        self.model.fit(
            Xs=Xs_array,
            Ys=Ys_array,
            Yvars=Yvars_array,
            parameter_values=parameter_values,
            outcome_names=self.outcomes,
        )

    def _predict(
        self, observation_features: List[ObservationFeatures]
    ) -> List[ObservationData]:
        # Convert observations to array
        X = [
            [of.parameters[param] for param in self.parameters]
            for of in observation_features
        ]
        f, cov = self.model.predict(X=X)
        # Convert arrays to observations
        return array_to_observation_data(f=f, cov=cov, outcomes=self.outcomes)

    def _validate_gen_inputs(
        self,
        n: int,
        search_space: Optional[SearchSpace] = None,
        optimization_config: Optional[OptimizationConfig] = None,
        pending_observations: Optional[Dict[str, List[ObservationFeatures]]] = None,
        fixed_features: Optional[ObservationFeatures] = None,
        model_gen_options: Optional[TConfig] = None,
    ) -> None:
        """Validate inputs to `ModelBridge.gen`.

        Currently, this is only used to ensure that `n` is a positive integer or -1.
        """
        if n < 1 and n != -1:
            raise UserInputError(
                f"Attempted to generate n={n} points. Number of points to generate "
                "must be either a positive integer or -1."
            )

    def _gen(
        self,
        n: int,
        search_space: SearchSpace,
        pending_observations: Dict[str, List[ObservationFeatures]],
        fixed_features: Optional[ObservationFeatures],
        model_gen_options: Optional[TConfig] = None,
        optimization_config: Optional[OptimizationConfig] = None,
    ) -> GenResults:
        """Generate new candidates according to search_space and
        optimization_config.

        The outcome constraints should be transformed to no longer be relative.
        """
        # Validation
        if not self.parameters:
            raise ValueError(FIT_MODEL_ERROR.format(action="_gen"))
        # Extract parameter values
        parameter_values = _get_parameter_values(search_space, self.parameters)
        # Extract objective and outcome constraints
        if len(self.outcomes) == 0 or optimization_config is None:
            objective_weights = None
            outcome_constraints = None
        else:
            validate_optimization_config(optimization_config, self.outcomes)
            objective_weights = extract_objective_weights(
                objective=optimization_config.objective, outcomes=self.outcomes
            )
            outcome_constraints = extract_outcome_constraints(
                outcome_constraints=optimization_config.outcome_constraints,
                outcomes=self.outcomes,
            )

        # Get fixed features
        fixed_features_dict = get_fixed_features(
            fixed_features=fixed_features, param_names=self.parameters
        )

        # Pending observations
        if len(pending_observations) == 0:
            pending_array: Optional[List[List[TParamValueList]]] = None
        else:
            pending_array = [[] for _ in self.outcomes]
            for metric_name, po_list in pending_observations.items():
                pending_array[self.outcomes.index(metric_name)] = [
                    [po.parameters[p] for p in self.parameters] for po in po_list
                ]

        # Generate the candidates
        X, w, gen_metadata = self.model.gen(
            n=n,
            parameter_values=parameter_values,
            objective_weights=objective_weights,
            outcome_constraints=outcome_constraints,
            fixed_features=fixed_features_dict,  # pyre-ignore
            pending_observations=pending_array,
            model_gen_options=model_gen_options,
        )
        observation_features = []
        for x in X:
            observation_features.append(
                ObservationFeatures(
                    parameters={p: x[i] for i, p in enumerate(self.parameters)}
                )
            )
        # TODO[drfreund, bletham]: implement best_point identification and
        # return best_point instead of None
        return GenResults(
            observation_features=observation_features,
            weights=w,
            gen_metadata=gen_metadata,
        )

    def _cross_validate(
        self,
        search_space: SearchSpace,
        cv_training_data: List[Observation],
        cv_test_points: List[ObservationFeatures],
        use_posterior_predictive: bool = False,
    ) -> List[ObservationData]:
        """Make predictions at cv_test_points using only the data in obs_feats
        and obs_data.
        """
        observation_features, observation_data = separate_observations(cv_training_data)
        Xs_train, Ys_train, Yvars_train = self._convert_observations(
            observation_data=observation_data,
            observation_features=observation_features,
            outcomes=self.outcomes,
            parameters=self.parameters,
        )
        X_test = [
            [obsf.parameters[param] for param in self.parameters]
            for obsf in cv_test_points
        ]
        # Use the model to do the cross validation
        f_test, cov_test = self.model.cross_validate(
            Xs_train=Xs_train,
            Ys_train=Ys_train,
            Yvars_train=Yvars_train,
            X_test=X_test,
            use_posterior_predictive=use_posterior_predictive,
        )
        # Convert array back to ObservationData
        return array_to_observation_data(f=f_test, cov=cov_test, outcomes=self.outcomes)

    @classmethod
    def _convert_observations(
        cls,
        observation_data: List[ObservationData],
        observation_features: List[ObservationFeatures],
        outcomes: List[str],
        parameters: List[str],
    ) -> Tuple[List[List[TParamValueList]], List[List[float]], List[List[float]]]:
        Xs: List[List[TParamValueList]] = [[] for _ in outcomes]
        Ys: List[List[float]] = [[] for _ in outcomes]
        Yvars: List[List[float]] = [[] for _ in outcomes]
        for i, obsf in enumerate(observation_features):
            try:
                x = [obsf.parameters[param] for param in parameters]
            except (KeyError, TypeError):
                # Out of design point
                raise ValueError("Out of design points cannot be converted.")
            for j, m in enumerate(observation_data[i].metric_names):
                k = outcomes.index(m)
                Xs[k].append(x)
                Ys[k].append(observation_data[i].means[j])
                Yvars[k].append(observation_data[i].covariance[j, j])
        return Xs, Ys, Yvars


def _get_parameter_values(
    search_space: SearchSpace, param_names: List[str]
) -> List[TParamValueList]:
    """Extract parameter values from a search space of discrete parameters."""
    parameter_values: List[TParamValueList] = []
    for p_name in param_names:
        p = search_space.parameters[p_name]
        # Validation
        if isinstance(p, ChoiceParameter):
            # Set values
            parameter_values.append(p.values)
        elif isinstance(p, FixedParameter):
            parameter_values.append([p.value])
        else:
            raise ValueError(f"{p} not ChoiceParameter or FixedParameter")
    return parameter_values
