from __future__ import annotations

from typing import TYPE_CHECKING, Any, Callable

if TYPE_CHECKING:
    import numpy as np
    from numpy.random import Generator

from .turbo_config import TurboENNConfig


class TurboENNImpl:
    def __init__(self, config: TurboENNConfig) -> None:
        self._config = config
        self._enn: Any | None = None
        self._fitted_params: Any | None = None
        self._fitted_n_obs: int = 0

    @property
    def always_clears_on_restart(self) -> bool:
        return True

    def get_x_center(
        self,
        x_obs_list: list,
        y_obs_list: list,
        rng: Generator,
        tr_state: Any = None,
    ) -> np.ndarray | None:
        import numpy as np

        from .impl_helpers import get_x_center_fallback
        from .turbo_utils import argmax_random_tie

        if len(y_obs_list) == 0:
            return None
        if self._enn is None or self._fitted_params is None:
            return get_x_center_fallback(
                self._config, x_obs_list, y_obs_list, rng, tr_state
            )
        if self._fitted_n_obs != len(x_obs_list):
            raise RuntimeError(
                f"ENN fitted on {self._fitted_n_obs} obs but get_x_center called with {len(x_obs_list)}"
            )

        y_array = np.asarray(y_obs_list, dtype=float)
        x_array = np.asarray(x_obs_list, dtype=float)
        k = self._config.k if self._config.k is not None else 10

        # For morbo: top-k per metric → union → scalarize mu
        if self._config.tr_type == "morbo" and tr_state is not None:
            if y_array.ndim == 1:
                y_array = y_array.reshape(-1, tr_state.num_metrics)
            num_metrics = y_array.shape[1]

            # Find top-k indices for each metric and take union
            union_indices = set()
            for m in range(num_metrics):
                num_top = min(k, len(y_array))
                top_m = np.argpartition(-y_array[:, m], num_top - 1)[:num_top]
                union_indices.update(top_m.tolist())
            union_indices = np.array(sorted(union_indices), dtype=int)

            x_union = x_array[union_indices]
            posterior = self._enn.posterior(x_union, params=self._fitted_params)
            mu = posterior.mu  # (len(union), num_metrics)

            scalarized = tr_state.scalarize(mu, clip=False)
            best_idx_in_union = argmax_random_tie(scalarized, rng=rng)
            return x_union[best_idx_in_union]
        else:
            # Single-objective: original logic
            num_top = min(k, len(y_array))
            top_indices = np.argpartition(-y_array, num_top - 1)[:num_top]

            x_top = x_array[top_indices]
            posterior = self._enn.posterior(x_top, params=self._fitted_params)
            mu = posterior.mu[:, 0]

            best_idx_in_top = argmax_random_tie(mu, rng=rng)
            return x_top[best_idx_in_top]

    def needs_tr_list(self) -> bool:
        return True

    def try_early_ask(
        self,
        num_arms: int,
        x_obs_list: list,
        draw_initial_fn: Callable[[int], np.ndarray],
        get_init_lhd_points_fn: Callable[[int], np.ndarray],
    ) -> np.ndarray | None:
        return None

    def prepare_ask(
        self,
        x_obs_list: list,
        y_obs_list: list,
        yvar_obs_list: list,
        num_dim: int,
        gp_num_steps: int,
        rng: Any | None = None,
    ) -> tuple[Any, float | None, float | None, np.ndarray | None]:
        from .proposal import mk_enn

        k = self._config.k if self._config.k is not None else 10
        self._enn, self._fitted_params = mk_enn(
            x_obs_list,
            y_obs_list,
            k,
            yvar_obs_list,
            num_fit_samples=self._config.num_fit_samples,
            num_fit_candidates=self._config.num_fit_candidates,
            scale_x=self._config.scale_x,
            rng=rng,
            params_warm_start=self._fitted_params,
        )
        self._fitted_n_obs = len(x_obs_list)
        return None, None, None, None

    def select_candidates(
        self,
        x_cand: np.ndarray,
        num_arms: int,
        num_dim: int,
        rng: Generator,
        fallback_fn: Callable[[np.ndarray, int], np.ndarray],
        from_unit_fn: Callable[[np.ndarray], np.ndarray],
        tr_state: Any = None,  # noqa: ARG002
    ) -> np.ndarray:
        import numpy as np

        from enn.enn.enn_params import ENNParams

        acq_type = self._config.acq_type
        k = self._config.k

        if self._enn is None:
            return fallback_fn(x_cand, num_arms)

        if self._fitted_params is not None:
            params = self._fitted_params
        else:
            k_val = k if k is not None else 10
            params = ENNParams(k=k_val, epi_var_scale=1.0, ale_homoscedastic_scale=0.0)

        posterior = self._enn.posterior(x_cand, params=params)
        mu = posterior.mu[:, 0]
        se = posterior.se[:, 0]

        if acq_type == "pareto":
            from enn.enn.enn_util import arms_from_pareto_fronts

            x_arms = arms_from_pareto_fronts(x_cand, mu, se, num_arms, rng)
        elif acq_type == "ucb":
            scores = mu + se
            shuffled_indices = rng.permutation(len(scores))
            shuffled_scores = scores[shuffled_indices]
            top_k_in_shuffled = np.argpartition(-shuffled_scores, num_arms - 1)[
                :num_arms
            ]
            idx = shuffled_indices[top_k_in_shuffled]
            x_arms = x_cand[idx]
        elif acq_type == "thompson":
            # One function sample per arm; each arm = argmax over all candidates
            base_seed = rng.integers(0, 2**31)
            function_seeds = np.arange(base_seed, base_seed + num_arms, dtype=np.int64)
            f_samples = self._enn.posterior_function_draw(
                x_cand, params, function_seeds=function_seeds
            )
            # f_samples: (num_arms, num_candidates, num_metrics)
            arm_indices = np.argmax(f_samples[:, :, 0], axis=1)
            x_arms = x_cand[arm_indices]
        else:
            raise ValueError(f"Unknown acq_type: {acq_type}")

        return from_unit_fn(x_arms)

    def estimate_y(self, x_unit: np.ndarray, y_observed: np.ndarray) -> np.ndarray:
        if self._enn is None or self._fitted_params is None:
            return y_observed
        posterior = self._enn.posterior(x_unit, params=self._fitted_params)
        return posterior.mu

    def get_mu_sigma(self, x_unit: np.ndarray) -> tuple[np.ndarray, np.ndarray] | None:
        if self._enn is None:
            return None
        k = self._config.k if self._config.k is not None else 10
        from enn.enn.enn_params import ENNParams

        params = (
            self._fitted_params
            if self._fitted_params is not None
            else ENNParams(k=k, epi_var_scale=1.0, ale_homoscedastic_scale=0.0)
        )
        posterior = self._enn.posterior(x_unit, params=params, observation_noise=False)
        # For multi-metric (morbo), return full mu/sigma; for single-metric, return 1D
        if posterior.mu.shape[1] > 1:
            return posterior.mu, posterior.se
        return posterior.mu[:, 0], posterior.se[:, 0]
