import datetime
import os
from unittest.mock import Mock, patch

import pytest

try:
    import optuna
except ImportError:
    optuna = None

if optuna is None:
    pytest.skip("optuna not installed", allow_module_level=True)

from confit import Config

from edsnlp.tune import (
    compute_importances,
    compute_n_trials,
    compute_remaining_n_trials_possible,
    compute_time_per_trial,
    is_plotly_install,
    load_config,
    optimize,
    process_results,
    tune,
)


def build_trial(number, value, params, datetime_start, datetime_complete):
    trial = Mock(spec=optuna.trial.FrozenTrial)
    trial.number = number
    trial.value = value
    trial.values = [value]
    trial.params = params
    trial.distributions = {
        "param1": optuna.distributions.FloatDistribution(
            high=0.3,
            log=False,
            low=0.0,
            step=0.05,
        ),
        "param2": optuna.distributions.FloatDistribution(
            high=0.3,
            log=False,
            low=0.0,
            step=0.05,
        ),
    }
    trial.datetime_start = datetime_start
    trial.datetime_complete = datetime_complete
    trial.state = optuna.trial.TrialState.COMPLETE
    trial.system_attrs = {}
    trial.user_attrs = {}
    return trial


@pytest.fixture
def study():
    study = Mock(spec=optuna.study.Study)
    study.study_name = "mock_study"
    study._is_multi_objective.return_value = False

    trials = []
    trial_0 = build_trial(
        number=0,
        value=0.9,
        params={"param1": 0.15, "param2": 0.3},
        datetime_start=datetime.datetime(2025, 1, 1, 12, 0, 0),
        datetime_complete=datetime.datetime(2025, 1, 1, 12, 5, 0),
    )
    trials.append(trial_0)

    trial_1 = build_trial(
        number=1,
        value=0.75,
        params={"param1": 0.05, "param2": 0.2},
        datetime_start=datetime.datetime(2025, 1, 1, 12, 5, 0),
        datetime_complete=datetime.datetime(2025, 1, 1, 12, 10, 0),
    )
    trials.append(trial_1)

    trial_2 = build_trial(
        number=2,
        value=0.99,
        params={"param1": 0.3, "param2": 0.25},
        datetime_start=datetime.datetime(2025, 1, 1, 12, 10, 0),
        datetime_complete=datetime.datetime(2025, 1, 1, 12, 15, 0),
    )
    trials.append(trial_2)

    study.trials = trials
    study.get_trials.return_value = trials
    study.best_trial = trials[2]
    return study


@pytest.mark.parametrize("ema", [True, False])
def test_compute_time_per_trial_with_ema(study, ema):
    result = compute_time_per_trial(study, ema=ema, alpha=0.1)
    assert result == pytest.approx(300.00)


@pytest.mark.parametrize(
    "gpu_hours, time_per_trial, expected_n_trials, raises_exception",
    [
        (1, 120, 30, False),
        (0.5, 3600, None, True),
    ],
)
def test_compute_n_trials(
    gpu_hours, time_per_trial, expected_n_trials, raises_exception
):
    if raises_exception:
        with pytest.raises(ValueError):
            compute_n_trials(gpu_hours, time_per_trial)
    else:
        result = compute_n_trials(gpu_hours, time_per_trial)
        assert result == expected_n_trials


def test_compute_importances(study):
    importance = compute_importances(study)
    assert importance == {"param2": 0.5239814153755754, "param1": 0.4760185846244246}


@pytest.mark.parametrize("viz", [True, False])
@pytest.mark.parametrize(
    "config_path", ["tests/tuning/config.yml", "tests/tuning/config.cfg"]
)
def test_process_results(study, tmpdir, viz, config_path):
    output_dir = tmpdir.mkdir("output")
    config = {
        "train": {
            "param1": None,
        },
        ".lr": {
            "param2": 0.01,
        },
    }
    hyperparameters = {
        "train.param1": {
            "type": "int",
            "alias": "param1",
            "low": 2,
            "high": 8,
            "step": 2,
        },
    }
    best_params, importances = process_results(
        study, output_dir, viz, config, config_path, hyperparameters
    )

    assert isinstance(best_params, dict)
    assert isinstance(importances, dict)

    results_file = os.path.join(output_dir, "results_summary.txt")
    assert os.path.exists(results_file)

    with open(results_file, "r") as f:
        content = f.read()
        assert "Study Summary" in content
        assert "Best trial" in content
        assert "Value" in content
        assert "Params" in content
        assert "Importances" in content

    if config_path.endswith("yml") or config_path.endswith("yaml"):
        config_file = os.path.join(output_dir, "config.yml")
    else:
        config_file = os.path.join(output_dir, "config.cfg")
    assert os.path.exists(config_file), f"Expected file {config_file} not found"

    with open(config_file, "r", encoding="utf-8") as f:
        content = f.read()
    assert (
        "# My usefull comment" in content
    ), f"Expected comment not found in {config_file}"

    if viz:
        optimization_history_file = os.path.join(
            output_dir, "optimization_history.html"
        )
        assert os.path.exists(
            optimization_history_file
        ), f"Expected file {optimization_history_file} not found"

        parallel_coord_file = os.path.join(output_dir, "parallel_coordinate.html")
        assert os.path.exists(
            parallel_coord_file
        ), f"Expected file {parallel_coord_file} not found"

        contour_file = os.path.join(output_dir, "contour.html")
        assert os.path.exists(contour_file), f"Expected file {contour_file} not found"

        edf_file = os.path.join(output_dir, "edf.html")
        assert os.path.exists(edf_file), f"Expected file {edf_file} not found"

        timeline_file = os.path.join(output_dir, "timeline.html")
        assert os.path.exists(timeline_file), f"Expected file {timeline_file} not found"


def test_compute_remaining_n_trials_possible(study):
    gpu_hours = 0.5
    remaining_trials = compute_remaining_n_trials_possible(study, gpu_hours)
    assert remaining_trials == 3


@patch("edsnlp.tune.objective_with_param")
@patch("optuna.study.Study.optimize")
@pytest.mark.parametrize("has_study", [True, False])
def test_optimize(mock_objective_with_param, mock_optimize_study, has_study, study):
    mock_objective_with_param.return_value = 0.9
    metric = ("ner", "micro", "f")
    checkpoint_dir = "./checkpoint"

    if has_study:

        def pass_fn(obj, n_trials, callbacks):
            pass

        study.optimize = pass_fn
        study = optimize(
            "config_path",
            tuned_parameters={},
            n_trials=1,
            metric=metric,
            checkpoint_dir=checkpoint_dir,
            study=study,
        )
        assert isinstance(study, Mock)
        assert len(study.trials) == 3

    else:
        study = optimize(
            "config_path",
            tuned_parameters={},
            n_trials=1,
            metric=metric,
            checkpoint_dir=checkpoint_dir,
            study=None,
        )
        assert isinstance(study, optuna.study.Study)
        assert len(study.trials) == 0


@patch("edsnlp.tune.optimize")
@patch("edsnlp.tune.process_results")
@patch("edsnlp.tune.load_config")
@patch("edsnlp.tune.compute_n_trials")
@patch("edsnlp.tune.update_config")
@pytest.mark.parametrize("n_trials", [10, None])
@pytest.mark.parametrize("two_phase_tuning", [False, True])
def test_tune(
    mock_update_config,
    mock_compute_n_trials,
    mock_load_config,
    mock_process_results,
    mock_optimize,
    study,
    n_trials,
    two_phase_tuning,
):
    mock_load_config.return_value = {"train": {}, "scorer": {}, "val_data": {}}
    mock_update_config.return_value = None, {"train": {}, "scorer": {}, "val_data": {}}
    mock_optimize.return_value = study
    mock_process_results.return_value = ({}, {})
    mock_compute_n_trials.return_value = 10
    config_meta = {"config_path": ["fake_path"]}
    hyperparameters = {
        "param1": {"type": "float", "low": 0.0, "high": 1.0},
        "param2": {"type": "float", "low": 0.0, "high": 1.0},
    }
    output_dir = "output_dir"
    checkpoint_dir = "checkpoint_dir"
    gpu_hours = 0.25
    seed = 42

    tune(
        config_meta=config_meta,
        hyperparameters=hyperparameters,
        output_dir=output_dir,
        checkpoint_dir=checkpoint_dir,
        gpu_hours=gpu_hours,
        n_trials=n_trials,
        two_phase_tuning=two_phase_tuning,
        seed=seed,
    )

    mock_load_config.assert_called_once()

    if two_phase_tuning:
        if n_trials is None:
            assert mock_compute_n_trials.call_count == 2  # 1 at begining + 1 at end
            assert mock_optimize.call_count == 3  # 1 at begining + 2 for tuning
        else:
            mock_compute_n_trials.assert_not_called()
            assert mock_optimize.call_count == 2  # 2 for tuning

        assert mock_process_results.call_count == 2  # one for each phase

    else:
        if n_trials is None:
            assert mock_compute_n_trials.call_count == 2  # 1 at begining + 1 at end
            assert (
                mock_optimize.call_count == 3
            )  # 1 at begining + 1 for tuning + 1 at end
        else:
            mock_compute_n_trials.assert_not_called()
            assert mock_optimize.call_count == 1  # 1 for tuning

        mock_process_results.assert_called_once()


@patch("importlib.util.find_spec")
def test_plotly(mock_importlib_util_find_spec):
    mock_importlib_util_find_spec.return_value = None
    assert not is_plotly_install()


def test_load_config(tmpdir):
    cfg = """\
    "a":
        "aa": 1
    "b": 2
    "c": "test"
    """
    config_dir = tmpdir.mkdir("configs")
    config_path = os.path.join(config_dir, "config.yml")
    Config.from_yaml_str(cfg).to_disk(config_path)
    config = load_config(config_path)
    assert isinstance(config, Config)
    with pytest.raises(FileNotFoundError):
        load_config("wrong_path")