import os
import pickle
from contextlib import redirect_stderr, redirect_stdout
from logging import getLogger
from time import sleep
from typing import IO, Any, Dict

import numpy as np
import optuna
import pytest
import scipy.sparse as sps

from irspack import evaluator
from irspack.definitions import DenseScoreArray, InteractionMatrix, UserIndexArray
from irspack.optimizers.base_optimizer import BaseOptimizerWithEarlyStopping
from irspack.parameter_tuning import UniformSuggestion
from irspack.recommenders.base import BaseRecommender
from irspack.recommenders.base_earlystop import (
    BaseRecommenderWithEarlyStopping,
    TrainerBase,
)
from irspack.split import rowwise_train_test_split

X_small = sps.csr_matrix(
    (np.random.RandomState(42).rand(100, 32) > 0.5).astype(np.float64)
)
X_large = sps.csr_matrix(
    (np.random.RandomState(42).rand(50, 200) > 0.5).astype(np.float64)
)


class MockTrainer(TrainerBase):
    def __init__(self) -> None:
        self.epoch: int = 0

    def load_state(self, ifs: IO) -> None:
        self.epoch = pickle.load(ifs)["epoch"]

    def save_state(self, ofs: IO) -> None:
        pickle.dump(dict(epoch=self.epoch), ofs)

    def run_epoch(self) -> None:
        self.epoch += 1


class MockRecommender(BaseRecommenderWithEarlyStopping):
    trainer_class = MockTrainer

    def __init__(
        self,
        X: InteractionMatrix,
        target_epoch: int = 20,
        target_score: float = 0.0,
        **kwargs: Any,
    ):
        super().__init__(X, **kwargs)
        self.target_epoch = target_epoch
        self.target_score = target_score
        self.rns = np.random.RandomState(42)

    def _create_trainer(self) -> MockTrainer:
        sleep(0.01)
        return MockTrainer()

    def get_score(self, user_indices: UserIndexArray) -> DenseScoreArray:
        score = np.zeros(*self.X_train_all.shape)
        return score

    def _current_score(self) -> float:
        assert isinstance(self.trainer, MockTrainer)
        coeff: float
        if self.trainer.epoch > self.target_epoch:
            coeff = 1 - (self.trainer.epoch - self.target_epoch) / self.target_epoch
        else:
            coeff = self.trainer.epoch / self.target_epoch
        return self.target_score * coeff


class MockEvaluator(evaluator.Evaluator):
    def __init__(self, X: sps.csr_matrix) -> None:
        super().__init__(X, offset=0)
        self.target_metric = evaluator.TargetMetric.ndcg
        self.cutoff = 30

    def get_score(self, model: BaseRecommender) -> Dict[str, float]:
        assert isinstance(model, MockRecommender)
        return {self.target_metric.name: model._current_score()}


class MockOptimizer(BaseOptimizerWithEarlyStopping):
    recommender_class = MockRecommender
    default_tune_range = [UniformSuggestion("target_score", 0, 1)]


@pytest.mark.parametrize("X, target_epoch", [(X_small, 20)])
def test_optimizer_by_mock(X: InteractionMatrix, target_epoch: int) -> None:

    evaluator = MockEvaluator(X)
    optimizer = MockOptimizer(
        X,
        evaluator,
        fixed_params=dict(target_epoch=target_epoch),
        logger=getLogger("IGNORE"),
    )
    with redirect_stdout(open(os.devnull, "w")):
        with redirect_stderr(open(os.devnull, "w")):
            config, history = optimizer.optimize(n_trials=20, random_seed=42)
    assert len(config) == 2
    assert config["max_epoch"] == target_epoch
    best_index = np.nanargmax(-history.value.values)
    best_target_score_inferred = history.target_score.iloc[best_index]
    best_ndcg = history["ndcg@30"].iloc[best_index]

    assert best_target_score_inferred == pytest.approx(best_ndcg)
    assert config["target_score"] == pytest.approx(best_target_score_inferred)
    assert np.all(best_target_score_inferred >= history.target_score.values)
