#!/usr/bin/env python
"""
Optuna Sampler 직렬화/역직렬화 테스트

주요 Optuna Sampler가 올바르게 직렬화되고 복원되는지 검증합니다.
Callable 타입 파라미터(gamma, weights 등)는 직렬화에서 제외됨을 확인합니다.
"""

import json
import os
import sys

import optuna
import pytest

# Add parent directory to path for imports
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../../runners"))

from runner_create_study import SAMPLER_WHITELIST
from runner_trial import EXIT_CODE_EXHAUSTED
from runner_trial import SAMPLER_WHITELIST as TRIAL_RUNNER_SAMPLER_WHITELIST

from aiauto.sampler import ExhaustiveCategoricalSampler, SearchSpaceExhausted
from aiauto.serializer import from_json, object_to_json


def test_random_sampler():
    """RandomSampler 테스트 - seed 파라미터"""
    sampler = optuna.samplers.RandomSampler(seed=42)
    json_str = object_to_json(sampler)

    parsed = json.loads(json_str)
    assert parsed["cls"] == "RandomSampler"
    # seed는 저장되지만 다른 형태로 저장될 수 있음
    if "seed" in parsed["kwargs"]:
        assert parsed["kwargs"]["seed"] == 42

    restored = from_json(json_str, SAMPLER_WHITELIST)
    assert type(restored).__name__ == "RandomSampler"


def test_tpe_sampler():
    """TPESampler 테스트 - 주요 파라미터"""
    sampler = optuna.samplers.TPESampler(
        n_startup_trials=5,
        n_ei_candidates=10,
        seed=42,
        multivariate=True,
        warn_independent_sampling=False,
    )
    json_str = object_to_json(sampler)

    parsed = json.loads(json_str)
    assert parsed["cls"] == "TPESampler"
    assert parsed["kwargs"]["n_startup_trials"] == 5
    assert parsed["kwargs"]["n_ei_candidates"] == 10
    assert parsed["kwargs"]["multivariate"] == True
    assert parsed["kwargs"]["warn_independent_sampling"] == False
    if "seed" in parsed["kwargs"]:
        assert parsed["kwargs"]["seed"] == 42

    # Callable 파라미터들(gamma, weights)은 제외되어야 함
    assert "gamma" not in parsed["kwargs"]
    assert "weights" not in parsed["kwargs"]

    restored = from_json(json_str, SAMPLER_WHITELIST)
    assert type(restored).__name__ == "TPESampler"
    assert restored._n_startup_trials == 5
    assert restored._n_ei_candidates == 10
    if hasattr(restored, "_seed"):
        assert restored._seed == 42


def test_tpe_sampler_default():
    """TPESampler 테스트 - 기본값"""
    sampler = optuna.samplers.TPESampler()
    json_str = object_to_json(sampler)

    parsed = json.loads(json_str)
    assert parsed["cls"] == "TPESampler"

    # 기본값들이 올바르게 직렬화되는지 확인
    assert "n_startup_trials" in parsed["kwargs"]
    assert "n_ei_candidates" in parsed["kwargs"]

    restored = from_json(json_str, SAMPLER_WHITELIST)
    assert type(restored).__name__ == "TPESampler"


def test_nsgaii_sampler():
    """NSGAIISampler 테스트 - 다목적 최적화용"""
    sampler = optuna.samplers.NSGAIISampler(
        population_size=10, mutation_prob=0.1, crossover_prob=0.8, swapping_prob=0.4, seed=42
    )
    json_str = object_to_json(sampler)

    parsed = json.loads(json_str)
    assert parsed["cls"] == "NSGAIISampler"
    assert parsed["kwargs"]["population_size"] == 10
    if "mutation_prob" in parsed["kwargs"]:
        assert parsed["kwargs"]["mutation_prob"] == 0.1
    if "crossover_prob" in parsed["kwargs"]:
        assert parsed["kwargs"]["crossover_prob"] == 0.8
    if "swapping_prob" in parsed["kwargs"]:
        assert parsed["kwargs"]["swapping_prob"] == 0.4
    if "seed" in parsed["kwargs"]:
        assert parsed["kwargs"]["seed"] == 42

    # Callable 파라미터들은 제외되거나 null이어야 함
    assert parsed["kwargs"].get("constraints_func") is None
    assert "elite_population_selection_strategy" not in parsed["kwargs"]

    restored = from_json(json_str, SAMPLER_WHITELIST)
    assert type(restored).__name__ == "NSGAIISampler"
    assert restored._population_size == 10
    if hasattr(restored, "_seed"):
        assert restored._seed == 42


def test_grid_sampler():
    """GridSampler 테스트 - search_space 파라미터"""
    search_space = {"x": [-10, -5, 0, 5, 10], "y": [1, 2, 3]}
    sampler = optuna.samplers.GridSampler(search_space=search_space)
    json_str = object_to_json(sampler)

    parsed = json.loads(json_str)
    assert parsed["cls"] == "GridSampler"
    assert parsed["kwargs"]["search_space"] == search_space

    restored = from_json(json_str, SAMPLER_WHITELIST)
    assert type(restored).__name__ == "GridSampler"
    assert restored._search_space == search_space


def test_bruteforce_sampler():
    """BruteForceSampler 테스트"""
    sampler = optuna.samplers.BruteForceSampler(seed=42)
    json_str = object_to_json(sampler)

    parsed = json.loads(json_str)
    assert parsed["cls"] == "BruteForceSampler"
    if "seed" in parsed["kwargs"]:
        assert parsed["kwargs"]["seed"] == 42

    restored = from_json(json_str, SAMPLER_WHITELIST)
    assert type(restored).__name__ == "BruteForceSampler"
    if hasattr(restored, "_seed"):
        assert restored._seed == 42


def test_cmaes_sampler():
    """CmaEsSampler 테스트"""
    sampler = optuna.samplers.CmaEsSampler(n_startup_trials=5, seed=42)
    json_str = object_to_json(sampler)

    parsed = json.loads(json_str)
    assert parsed["cls"] == "CmaEsSampler"
    assert parsed["kwargs"]["n_startup_trials"] == 5
    if "seed" in parsed["kwargs"]:
        assert parsed["kwargs"]["seed"] == 42
    # restart_strategy는 deprecated되어 저장되지 않음

    restored = from_json(json_str, SAMPLER_WHITELIST)
    assert type(restored).__name__ == "CmaEsSampler"
    assert restored._n_startup_trials == 5
    if hasattr(restored, "_seed"):
        assert restored._seed == 42


def test_qmc_sampler():
    """QMCSampler 테스트"""
    sampler = optuna.samplers.QMCSampler(qmc_type="sobol", scramble=True, seed=42)
    json_str = object_to_json(sampler)

    parsed = json.loads(json_str)
    assert parsed["cls"] == "QMCSampler"
    assert parsed["kwargs"]["qmc_type"] == "sobol"
    assert parsed["kwargs"]["scramble"] == True
    if "seed" in parsed["kwargs"]:
        assert parsed["kwargs"]["seed"] == 42

    restored = from_json(json_str, SAMPLER_WHITELIST)
    assert type(restored).__name__ == "QMCSampler"
    assert restored._qmc_type == "sobol"
    assert restored._scramble == True
    if hasattr(restored, "_seed"):
        assert restored._seed == 42


def test_runner_trial_exit_code_exhausted_is_3():
    """runner_trial.EXIT_CODE_EXHAUSTED는 반드시 3이어야 한다.

    operator는 exit code 3을 "search space 소진"으로 해석한다.
    이 값이 바뀌면 operator의 trial 상태 분류가 깨지므로 계약 보호 테스트다.
    """
    assert EXIT_CODE_EXHAUSTED == 3


def test_runner_trial_sampler_whitelist_includes_exhaustive_categorical():
    """runner_trial.SAMPLER_WHITELIST에 ExhaustiveCategoricalSampler가 포함되어야 한다.

    SAMPLER_JSON 환경변수로 ExhaustiveCategoricalSampler가 전달될 때
    runner_trial이 역직렬화할 수 있어야 한다.
    whitelist에서 제외되면 from_json()이 ValueError를 raise하여 trial이 실패한다.
    """
    assert "ExhaustiveCategoricalSampler" in TRIAL_RUNNER_SAMPLER_WHITELIST
    assert (
        TRIAL_RUNNER_SAMPLER_WHITELIST["ExhaustiveCategoricalSampler"]
        is ExhaustiveCategoricalSampler
    )


def test_exhaustive_categorical_sampler():
    """ExhaustiveCategoricalSampler 직렬화/역직렬화 테스트.

    constructor에 search_space가 없으므로 kwargs는 비어있거나 seed만 포함.
    직렬화 cls 이름과 역직렬화 type 검증.
    """
    sampler = ExhaustiveCategoricalSampler()
    json_str = object_to_json(sampler)

    parsed = json.loads(json_str)
    assert parsed["cls"] == "ExhaustiveCategoricalSampler"
    assert "search_space" not in parsed["kwargs"]

    restored = from_json(json_str, SAMPLER_WHITELIST)
    assert type(restored).__name__ == "ExhaustiveCategoricalSampler"


def test_exhaustive_categorical_sampler_rejects_non_categorical():
    """sample_independent는 CategoricalDistribution 외 타입을 거부해야 한다.

    ExhaustiveCategoricalSampler는 categorical-only sampler이므로
    FloatDistribution/IntDistribution 파라미터가 들어오면 ValueError를 raise한다.
    이 경로가 실제로 동작하는지 in-memory study를 통해 검증한다.
    """
    from optuna.distributions import FloatDistribution

    sampler = ExhaustiveCategoricalSampler()
    study = optuna.create_study(sampler=sampler, storage=None)
    trial = study.ask()

    # FrozenTrial을 얻기 위해 ask 결과를 storage에서 조회
    frozen = study._storage.get_trial(trial._trial_id)

    with pytest.raises(ValueError, match="CategoricalDistribution"):
        sampler.sample_independent(study, frozen, "y", FloatDistribution(0.0, 1.0))


def test_exhaustive_categorical_sampler_raises_search_space_exhausted():
    """모든 조합이 소진되면 SearchSpaceExhausted를 raise해야 한다.

    search_space {"x": ["a", "b"]}로 study를 생성하고
    두 trial을 COMPLETE 처리한 뒤, 세 번째 trial의 sample_relative 호출에서
    SearchSpaceExhausted가 발생해야 한다.

    study.tell()은 after_trial을 트리거하여 study.stop()을 호출하므로
    storage를 직접 조작하여 trial을 COMPLETE 상태로 만든다.
    """
    from optuna.distributions import CategoricalDistribution
    from optuna.trial import TrialState

    sampler = ExhaustiveCategoricalSampler()
    study = optuna.create_study(sampler=sampler, storage=None)

    dist = CategoricalDistribution(["a", "b"])

    # trial 0: ask -> storage에서 직접 COMPLETE 처리 (after_trial 호출 없이)
    trial0 = study.ask()
    study._storage.set_trial_param(trial0._trial_id, "x", dist.to_internal_repr("a"), dist)
    study._storage.set_trial_state_values(trial0._trial_id, TrialState.COMPLETE, [0.5])

    # trial 1: ask -> storage에서 직접 COMPLETE 처리
    trial1 = study.ask()
    study._storage.set_trial_param(trial1._trial_id, "x", dist.to_internal_repr("b"), dist)
    study._storage.set_trial_state_values(trial1._trial_id, TrialState.COMPLETE, [0.3])

    # trial 2: 모든 조합(a, b)이 소진된 상태에서 sample_relative 호출
    trial2 = study.ask()
    frozen2 = study._storage.get_trial(trial2._trial_id)

    search_space = {"x": CategoricalDistribution(["a", "b"])}
    with pytest.raises(SearchSpaceExhausted):
        sampler.sample_relative(study, frozen2, search_space)


def test_exhaustive_categorical_sampler_sample_independent_returns_valid_value():
    """sample_independent는 조합이 남아있을 때 valid한 값을 반환해야 한다.

    search_space {"x": ["a", "b"]}에서 아직 아무 trial도 없을 때
    sample_independent가 "a" 또는 "b" 중 하나를 반환하는지 검증한다.
    이 테스트는 SearchSpaceExhausted 경로가 아닌 정상 샘플링 경로를 커버한다.
    """
    from optuna.distributions import CategoricalDistribution

    sampler = ExhaustiveCategoricalSampler()
    study = optuna.create_study(sampler=sampler, storage=None)

    trial = study.ask()
    frozen = study._storage.get_trial(trial._trial_id)

    result = sampler.sample_independent(study, frozen, "x", CategoricalDistribution(["a", "b"]))
    assert result in ["a", "b"], f"Expected 'a' or 'b', got {result!r}"


def test_exhaustive_categorical_sampler_after_trial_stops_study():
    """after_trial은 모든 조합이 소진되면 study.stop()을 호출해야 한다.

    ExhaustiveCategoricalSampler.after_trial은 자체적으로 sorted tree를 구성하여
    count_unexpanded == 0이면 study.stop()을 호출한다.

    study.optimize()로 단일 조합({"x": ["a"]})을 실행하면 after_trial에서
    study.stop()이 호출되어 n_trials 관계없이 루프가 종료된다.
    결과적으로 trials가 정확히 1개만 생성됐다면 stop()이 호출된 것이다.

    주의: study.optimize()에 catch=(SearchSpaceExhausted,)를 지정하여
    sample_independent에서 raise된 예외를 Optuna가 swallow하도록 한다.
    """
    sampler = ExhaustiveCategoricalSampler()  # 단일 조합
    study = optuna.create_study(sampler=sampler, storage=None)

    # study.optimize()를 통해 after_trial 트리거
    # n_trials=100이지만 after_trial에서 study.stop()이 호출되어 1회만 실행된다
    def objective(trial):
        trial.suggest_categorical("x", ["a"])
        return 0.5

    study.optimize(objective, n_trials=100, catch=(SearchSpaceExhausted,))

    # after_trial이 study.stop()을 호출했다면 trials=1이어야 한다
    assert len(study.trials) == 1, (
        f"Expected 1 trial (study.stop() called by after_trial), got {len(study.trials)}"
    )


def test_exhaustive_categorical_sampler_after_trial_sets_total_combinations():
    """after_trial은 trial 완료 후 _total_combinations study attr을 설정해야 한다.

    ExhaustiveCategoricalSampler(seed=None)로 study를 생성하고
    search_space {"x": ["a", "b"], "y": [1, 2, 3]} (조합 6개)으로 objective 실행.
    첫 trial 완료 후 study.user_attrs["_total_combinations"] == "6" 검증.
    """
    sampler = ExhaustiveCategoricalSampler()
    study = optuna.create_study(sampler=sampler, storage=None)

    def objective(trial):
        trial.suggest_categorical("x", ["a", "b"])
        trial.suggest_categorical("y", [1, 2, 3])
        return 0.5

    study.optimize(objective, n_trials=1, catch=(SearchSpaceExhausted,))

    assert "_total_combinations" in study.user_attrs
    assert study.user_attrs["_total_combinations"] == "6"


def test_exhaustive_categorical_sampler_sorted_distributions():
    """분산 환경 시뮬레이션: 다른 dict 순서의 distributions를 가진 과거 trials가 있을 때
    crash 없이 정상 샘플링되는지 검증.

    gRPC storage에서 protobuf map은 순서가 보장되지 않아 trial마다
    distributions의 key 순서가 다를 수 있다. sample_relative()가 sorted key ordering
    으로 이를 처리하여 정상 동작하는지 확인.
    """
    from optuna.distributions import CategoricalDistribution
    from optuna.trial import TrialState

    sampler = ExhaustiveCategoricalSampler()
    study = optuna.create_study(sampler=sampler, storage=None)

    dist_x = CategoricalDistribution(["a", "b"])
    dist_y = CategoricalDistribution([1, 2])

    # trial 0: distributions 순서 {x, y}
    trial0 = study.ask()
    study._storage.set_trial_param(trial0._trial_id, "x", dist_x.to_internal_repr("a"), dist_x)
    study._storage.set_trial_param(trial0._trial_id, "y", dist_y.to_internal_repr(1), dist_y)
    study._storage.set_trial_state_values(trial0._trial_id, TrialState.COMPLETE, [0.5])

    # trial 1: distributions 순서 {y, x} (의도적으로 다른 순서)
    # storage의 set_trial_param은 호출 순서대로 dict에 추가되므로
    # y를 먼저 설정하여 순서를 뒤바꿈
    trial1 = study.ask()
    study._storage.set_trial_param(trial1._trial_id, "y", dist_y.to_internal_repr(2), dist_y)
    study._storage.set_trial_param(trial1._trial_id, "x", dist_x.to_internal_repr("a"), dist_x)
    study._storage.set_trial_state_values(trial1._trial_id, TrialState.COMPLETE, [0.3])

    # trial 2: 남은 조합에서 정상 샘플링되어야 함 (crash 없이)
    trial2 = study.ask()
    frozen2 = study._storage.get_trial(trial2._trial_id)

    # sample_relative: sorted key ordering으로 dict 순서 무관하게 동작
    search_space = {"x": dist_x, "y": dist_y}
    result = sampler.sample_relative(study, frozen2, search_space)
    assert result["x"] in ["a", "b"], f"Expected 'a' or 'b', got {result['x']!r}"
    assert result["y"] in [1, 2], f"Expected 1 or 2, got {result['y']!r}"
    # ("a", 1)과 ("a", 2)가 사용됨 -> 남은 조합은 ("b", 1) 또는 ("b", 2)
    assert result["x"] == "b", (
        f"Expected 'b' since all 'a' combinations are used, got {result['x']!r}"
    )


def test_exhaustive_categorical_sampler_after_trial_no_stop_when_combinations_remain():
    """after_trial은 조합이 남아있으면 study.stop()을 호출하지 않아야 한다.

    search_space {"x": ["a", "b"]}에서 첫 trial만 완료한 뒤 after_trial 호출.
    아직 "b" 조합이 남아있으므로 study.stop()이 호출되지 않아야 한다.
    study.optimize(n_trials=2)로 실행하여 2개 trial이 모두 생성됨을 확인.
    """
    sampler = ExhaustiveCategoricalSampler()
    study = optuna.create_study(sampler=sampler, storage=None)

    def objective(trial):
        trial.suggest_categorical("x", ["a", "b"])
        return 0.5

    study.optimize(objective, n_trials=2, catch=(SearchSpaceExhausted,))

    # 조합이 2개이므로 2개 trial이 모두 생성되어야 한다
    # 만약 after_trial이 첫 trial 후 잘못 stop()을 호출하면 1개만 생성됨
    assert len(study.trials) == 2, (
        f"Expected 2 trials (combinations remain after 1st trial), got {len(study.trials)}"
    )


def test_exhaustive_categorical_sampler_after_trial_empty_distributions():
    """after_trial에 distributions가 비어있는 trial을 전달하면
    _total_combinations study attr이 설정되지 않아야 한다.

    objective에서 suggest를 한 번도 호출하지 않으면 trial.distributions가 빈 dict.
    이 경우 after_trial의 `if trial.distributions:` 분기가 False -> skip.
    """
    from optuna.trial import TrialState, create_trial

    sampler = ExhaustiveCategoricalSampler()
    study = optuna.create_study(sampler=sampler, storage=None)

    # distributions가 비어있는 FrozenTrial을 직접 생성
    empty_trial = create_trial(
        state=TrialState.RUNNING,
        values=None,
        params={},
        distributions={},
    )

    # after_trial 직접 호출
    sampler.after_trial(study, empty_trial, TrialState.COMPLETE, [0.5])

    # _total_combinations가 설정되지 않았어야 한다
    assert "_total_combinations" not in study.user_attrs, (
        f"Expected no _total_combinations attr for empty distributions, "
        f"got {study.user_attrs.get('_total_combinations')}"
    )


def test_exhaustive_categorical_sampler_multi_param_exhaustion():
    """다중 파라미터 조합이 모두 소진되면 SearchSpaceExhausted를 raise해야 한다.

    search_space {"x": ["a", "b"], "y": [1, 2]} = 4 조합.
    4개 trial을 storage에 직접 COMPLETE 처리 후 5번째 sample_relative에서
    SearchSpaceExhausted가 발생해야 한다.
    """
    from optuna.distributions import CategoricalDistribution
    from optuna.trial import TrialState

    sampler = ExhaustiveCategoricalSampler()
    study = optuna.create_study(sampler=sampler, storage=None)

    dist_x = CategoricalDistribution(["a", "b"])
    dist_y = CategoricalDistribution([1, 2])

    # 4개 조합 모두 COMPLETE 처리
    combinations = [("a", 1), ("a", 2), ("b", 1), ("b", 2)]
    for x_val, y_val in combinations:
        t = study.ask()
        study._storage.set_trial_param(t._trial_id, "x", dist_x.to_internal_repr(x_val), dist_x)
        study._storage.set_trial_param(t._trial_id, "y", dist_y.to_internal_repr(y_val), dist_y)
        study._storage.set_trial_state_values(t._trial_id, TrialState.COMPLETE, [0.5])

    # 5번째 trial: 모든 조합 소진 -> SearchSpaceExhausted
    trial5 = study.ask()
    frozen5 = study._storage.get_trial(trial5._trial_id)

    search_space = {"x": dist_x, "y": dist_y}
    with pytest.raises(SearchSpaceExhausted):
        sampler.sample_relative(study, frozen5, search_space)


def test_exhaustive_categorical_sampler_seed_determinism():
    """동일 seed를 지정하면 동일한 샘플링 순서를 보장해야 한다.

    ExhaustiveCategoricalSampler(seed=42)로 2회 실행하여
    첫 trial의 sample_independent 결과가 동일한지 검증.
    """
    from optuna.distributions import CategoricalDistribution

    results = []
    for _ in range(2):
        sampler = ExhaustiveCategoricalSampler(seed=42)
        study = optuna.create_study(sampler=sampler, storage=None)
        trial = study.ask()
        frozen = study._storage.get_trial(trial._trial_id)
        result = sampler.sample_independent(
            study, frozen, "x", CategoricalDistribution(["a", "b", "c", "d"])
        )
        results.append(result)

    assert results[0] == results[1], (
        f"Expected deterministic sampling with seed=42, got {results[0]} and {results[1]}"
    )


if __name__ == "__main__":
    # 각 테스트 실행
    test_runner_trial_exit_code_exhausted_is_3()
    print("EXIT_CODE_EXHAUSTED contract test passed")

    test_runner_trial_sampler_whitelist_includes_exhaustive_categorical()
    print("SAMPLER_WHITELIST ExhaustiveCategoricalSampler inclusion test passed")

    test_random_sampler()
    print("RandomSampler test passed")

    test_tpe_sampler()
    print("TPESampler test passed")

    test_tpe_sampler_default()
    print("TPESampler (default) test passed")

    test_nsgaii_sampler()
    print("NSGAIISampler test passed")

    test_grid_sampler()
    print("GridSampler test passed")

    test_bruteforce_sampler()
    print("BruteForceSampler test passed")

    test_cmaes_sampler()
    print("CmaEsSampler test passed")

    test_qmc_sampler()
    print("QMCSampler test passed")

    test_exhaustive_categorical_sampler()
    print("ExhaustiveCategoricalSampler test passed")

    test_exhaustive_categorical_sampler_rejects_non_categorical()
    print("ExhaustiveCategoricalSampler non-categorical rejection test passed")

    test_exhaustive_categorical_sampler_raises_search_space_exhausted()
    print("ExhaustiveCategoricalSampler SearchSpaceExhausted test passed")

    test_exhaustive_categorical_sampler_sample_independent_returns_valid_value()
    print("ExhaustiveCategoricalSampler sample_independent normal path test passed")

    test_exhaustive_categorical_sampler_after_trial_stops_study()
    print("ExhaustiveCategoricalSampler after_trial study.stop() test passed")

    test_exhaustive_categorical_sampler_after_trial_sets_total_combinations()
    print("ExhaustiveCategoricalSampler after_trial _total_combinations test passed")

    test_exhaustive_categorical_sampler_sorted_distributions()
    print("ExhaustiveCategoricalSampler sorted distributions test passed")

    test_exhaustive_categorical_sampler_after_trial_no_stop_when_combinations_remain()
    print("ExhaustiveCategoricalSampler after_trial no stop when combinations remain test passed")

    test_exhaustive_categorical_sampler_after_trial_empty_distributions()
    print("ExhaustiveCategoricalSampler after_trial empty distributions test passed")

    test_exhaustive_categorical_sampler_multi_param_exhaustion()
    print("ExhaustiveCategoricalSampler multi-param exhaustion test passed")

    test_exhaustive_categorical_sampler_seed_determinism()
    print("ExhaustiveCategoricalSampler seed determinism test passed")

    print("\nAll Sampler tests passed!")
