from __future__ import annotations

from collections.abc import Sequence
from typing import NamedTuple

import numpy as np
from tqdm import tqdm

import optuna
from optuna._hypervolume import compute_hypervolume
from optuna.study import Study
from optuna.study._study_direction import StudyDirection
from optuna.trial import TrialState
from optuna.visualization._plotly_imports import _imports


def this_pr(
    values_array: np.ndarray, ref_point: np.ndarray,
) -> np.ndarray:
    n_trials = values_array.shape[0]
    hypervolume_values = np.empty(n_trials, dtype=float)
    best_trials_values_normalized: np.ndarray | None = None
    hypervolume = 0.0
    for i in tqdm(range(n_trials)):
        values_normalized = values_array[i, np.newaxis, :]
        if best_trials_values_normalized is not None:
            if (best_trials_values_normalized <= values_normalized).all(axis=1).any(axis=0):
                hypervolume_values[i] = hypervolume
                continue

        hypervolume += np.prod(ref_point - values_normalized)
        if best_trials_values_normalized is None:
            best_trials_values_normalized = values_normalized
        else:
            limited_sols = np.maximum(best_trials_values_normalized, values_normalized)
            hypervolume -= compute_hypervolume(limited_sols, ref_point)
            is_kept = (best_trials_values_normalized < values_normalized).any(axis=1)
            best_trials_values_normalized = np.concatenate(
                [best_trials_values_normalized[is_kept, :], values_normalized], axis=0
            )
        hypervolume_values[i] = hypervolume

    return hypervolume_values


def master(
    values_array: np.ndarray, ref_point: np.ndarray,
) -> np.ndarray:
    n_trials = values_array.shape[0]
    hypervolume_values = np.empty(n_trials, dtype=float)
    best_trials_values_normalized: np.ndarray | None = None
    for i in tqdm(range(n_trials)):
        values_normalized = values_array[i, np.newaxis, :]
        if best_trials_values_normalized is not None:
            if (best_trials_values_normalized <= values_normalized).all(axis=1).any(axis=0):
                hypervolume_values[i] = hypervolume_values[i - 1]
                continue

        if best_trials_values_normalized is None:
            best_trials_values_normalized = values_normalized
        else:
            is_kept = (best_trials_values_normalized < values_normalized).any(axis=1)
            best_trials_values_normalized = np.concatenate(
                [best_trials_values_normalized[is_kept, :], values_normalized], axis=0
            )
        hypervolume_values[i] = compute_hypervolume(best_trials_values_normalized, ref_point)

    return hypervolume_values


if __name__ == "__main__":
    rng = np.random.RandomState(42)
    values_array = rng.normal(size=(5000, 4))
    ref_point = np.full(4, 10.0)
    out = this_pr(values_array.copy(), ref_point)
    ans = master(values_array.copy(), ref_point)
    assert np.allclose(out, ans)