# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
from contextlib import nullcontext
from importlib import import_module
from importlib.metadata import PackageNotFoundError, version
import inspect
from pathlib import Path
import pkgutil

import numpy as np
from numpy.testing import assert_almost_equal, assert_array_equal
from packaging.version import parse
import pandas as pd
import pytest
from sklearn.base import BaseEstimator, TransformerMixin

import sksurv
from sksurv.metrics import concordance_index_censored


def assert_cindex_almost_equal(event_indicator, event_time, estimate, expected):
    result = concordance_index_censored(event_indicator, event_time, estimate)
    assert_array_equal(result[1:], expected[1:])
    concordant, discordant, tied_risk = result[1:4]
    cc = (concordant + 0.5 * tied_risk) / (concordant + discordant + tied_risk)
    assert_almost_equal(result[0], cc)
    assert_almost_equal(result[0], expected[0])


def assert_survival_function_properties(surv_fns):
    if not np.isfinite(surv_fns).all():
        raise AssertionError("survival function contains values that are not finite")
    if np.any(surv_fns < 0.0):
        raise AssertionError("survival function contains negative values")
    if np.any(surv_fns > 1.0):
        raise AssertionError("survival function contains values larger 1")

    d = np.apply_along_axis(np.diff, 1, surv_fns)
    if np.any(d > 0):
        raise AssertionError("survival functions are not monotonically decreasing")

    # survival function at first time point
    num_closer_to_zero = np.sum(1.0 - surv_fns[:, 0] >= surv_fns[:, 0])
    if num_closer_to_zero / surv_fns.shape[0] > 0.5:
        raise AssertionError(f"most ({num_closer_to_zero}) probabilities at first time point are closer to 0 than 1")

    # survival function at last time point
    num_closer_to_one = np.sum(1.0 - surv_fns[:, -1] < surv_fns[:, -1])
    if num_closer_to_one / surv_fns.shape[0] > 0.5:
        raise AssertionError(f"most ({num_closer_to_one}) probabilities at last time point are closer to 1 than 0")


def assert_chf_properties(chf):
    if not np.isfinite(chf).all():
        raise AssertionError("chf contains values that are not finite")
    if np.any(chf < 0.0):
        raise AssertionError("chf contains negative values")

    d = np.apply_along_axis(np.diff, 1, chf)
    if np.any(d < 0):
        raise AssertionError("chf are not monotonically increasing")

    # chf at first time point
    num_closer_to_one = np.sum(1.0 - chf[:, 0] < chf[:, 0])
    if num_closer_to_one / chf.shape[0] > 0.5:
        raise AssertionError(f"most ({num_closer_to_one}) hazard rates at first time point are closer to 1 than 0")


def _is_survival_estimator(x):
    return (
        inspect.isclass(x)
        and issubclass(x, BaseEstimator)
        and not issubclass(x, TransformerMixin)
        and x.__module__.startswith("sksurv.")
        and not x.__name__.startswith("_")
        and x.__module__.split(".", 2)[1] not in {"metrics", "nonparametric"}
    )


def all_survival_estimators():
    root = str(Path(sksurv.__file__).parent)
    all_classes = []
    for _importer, modname, _ispkg in pkgutil.walk_packages(path=[root], prefix="sksurv."):
        # meta-estimators require base estimators
        if modname.startswith("sksurv.meta"):
            continue
        module = import_module(modname)
        for _name, cls in inspect.getmembers(module, _is_survival_estimator):
            if inspect.isabstract(cls):
                continue
            all_classes.append(cls)
    return set(all_classes)


class FixtureParameterFactory:
    def get_cases(self):
        cases = []
        for name, func in inspect.getmembers(self):
            if name.startswith("data_"):
                values = func()
                cases.append(pytest.param(*values, id=name))
        return cases

    def get_cases_func(self):
        cases = []
        for name, func in inspect.getmembers(self):
            if name.startswith("data_"):
                cases.append(pytest.param(func, id=name))
        return cases


def check_module_minimum_version(module, min_version_str, max_version_str=None):
    """
    Check whether a module of a specified minimum version is available.

    Parameters
    ----------
    module : str
        Name of the module.
    min_version_str : str
        Minimum version of the module.
    max_version_str : str, optional
        Maximum version of the module (excluding).

    Returns
    -------
    available : bool
        True if the module is available and its version is >= `version_str`.
    """
    try:
        module_version = parse(version(module))
        required_min_version = parse(min_version_str)
        if max_version_str is None:
            return module_version >= required_min_version
        required_max_version = parse(max_version_str)
        return required_min_version <= module_version < required_max_version
    except PackageNotFoundError:  # pragma: no cover
        return False


def get_pandas_infer_string_context():
    if check_module_minimum_version("pandas", "2.3.0", "3.0.0"):
        return (
            pytest.param(pd.option_context("future.infer_string", False), id="infer_string=False"),
            pytest.param(pd.option_context("future.infer_string", True), id="infer_string=True"),
        )
    return (
        pytest.param(nullcontext(), id="pandas default options"),
        pytest.param(nullcontext(), marks=pytest.mark.skip("no pandas 2.3.0")),
    )
