import builtins
import re
import sys
import webbrowser
from pathlib import Path
from unittest.mock import Mock

import pytest
from sklearn.dummy import DummyClassifier
from sklearn.feature_selection import SelectKBest

import skrub
from skrub import datasets
from skrub._data_ops import _inspection


@pytest.mark.skipif(not _inspection._has_graphviz(), reason="report requires graphviz")
def test_output_dir(tmp_path):
    e = skrub.X()
    assert e.skb.full_report(open=False)["report_path"].is_relative_to(
        datasets.get_data_dir()
    )
    out = tmp_path / "report"
    assert (
        e.skb.full_report(open=False, output_dir=out)["report_path"]
        == out / "index.html"
    )
    with pytest.raises(FileExistsError, match=".*Set 'overwrite=True'"):
        e.skb.full_report(open=False, output_dir=out)

    assert (
        e.skb.full_report(open=False, output_dir=out, overwrite=True)["report_path"]
        == out / "index.html"
    )


@pytest.mark.skipif(not _inspection._has_graphviz(), reason="report requires graphviz")
def test_full_report():
    # smoke test for the full report
    # TODO we should have a private function that returns the JSON data so we
    #      can check the content before rendering with jinja
    # however that requires first settling on the content of the report etc.
    e = -(
        (skrub.var("a", 12345) + 1).skb.set_name("b").skb.set_description("this is b")
        / skrub.var("c", 1)
    )
    report = e.skb.full_report(open=False)
    assert report["error"] is None
    assert report["result"] == -12346.0
    assert "-12346.0" in (report["report_path"].parent / "node_4.html").read_text(
        "utf-8"
    )
    report = e.skb.full_report({"a": 12345, "c": 0}, open=False)
    assert isinstance(report["error"], (ZeroDivisionError, RuntimeError))
    assert report["result"] is None
    out = report["report_path"].parent
    text = (out / "node_1.html").read_text("utf-8")
    assert "12346" in text and "this is b" in text
    assert "ZeroDivisionError" in (out / "node_3.html").read_text("utf-8")
    assert "This step did not run" in (out / "node_4.html").read_text("utf-8")


@pytest.mark.skipif(not _inspection._has_graphviz(), reason="report requires graphviz")
def test_full_report_title():
    # TODO we should have a private function that returns the JSON data so we
    #      can check the content before rendering with jinja
    # however that requires first settling on the content of the report etc.
    data_op = skrub.var("a", 1)
    title = "small data ops"
    report = data_op.skb.full_report(open=False, title=title)
    assert title in report["report_path"].read_text("utf-8")


@pytest.mark.skipif(not _inspection._has_graphviz(), reason="report requires graphviz")
def test_preview_subsample():
    X = datasets.fetch_employee_salaries().X
    preview = skrub.X(X).skb.subsample(n=3)._repr_html_()
    assert "subsample" in preview


@pytest.mark.skipif(not _inspection._has_graphviz(), reason="report requires graphviz")
def test_full_report_failed_apply():
    # Somewhat contrived example for the corner case where an Apply does not
    # have an easily identifiable estimator.
    orders = skrub.datasets.toy_orders()
    e = (
        skrub.X()
        .skb.apply(SelectKBest())  # error! missing y
        .skb.apply(
            # never reached so there is no fitted estimator
            skrub.choose_from([DummyClassifier(), DummyClassifier()], name="classif"),
            y=skrub.y(),
        )
    )
    report = e.skb.full_report({"X": orders.X, "y": orders.y}, open=False)
    assert report["error"] is not None


@pytest.mark.skipif(not _inspection._has_graphviz(), reason="report requires graphviz")
def test_report_fit_mode():
    # non-regression: in fit mode the individual node pages used to show the
    # dataop itself as the output instead of the result of fit_transform for
    # intermediate nodes.
    learner = (
        skrub.var("a")
        .skb.apply_func(lambda x: x[::-1])
        .skb.apply_func(lambda x: x.upper())
        .skb.make_learner()
    )
    report = learner.report(environment={"a": "hello"}, mode="fit", open=False)
    report_dir = Path(report["report_path"]).parent
    node_1_report = (report_dir / "node_1.html").read_text("utf-8")
    # the page does display the actual output rather than the repr of the node
    # ('hello'[::-1] = 'olleh')
    assert "olleh" in node_1_report


@pytest.mark.skipif(not _inspection._has_graphviz(), reason="report requires graphviz")
def test_full_report_open(monkeypatch):
    mock = Mock()
    monkeypatch.setattr(webbrowser, "open", mock)
    skrub.as_data_op(0).skb.full_report()
    mock.assert_called_once()


@pytest.mark.skipif(not _inspection._has_graphviz(), reason="report requires graphviz")
def test_draw_graph():
    g = skrub.as_data_op(0).skb.draw_graph()
    assert repr(g) == "<GraphDrawing: use .open() to display>"
    assert b"<svg" in g.svg
    assert "<svg" in g._repr_html_()
    assert g.png.startswith(b"\x89PNG")
    assert g._repr_png_().startswith(b"\x89PNG")


@pytest.mark.skipif(not _inspection._has_graphviz(), reason="report requires graphviz")
def test_svg_anchor_google_colab(monkeypatch):
    """non-regression test for #1589"""
    monkeypatch.setitem(sys.modules, "google.colab", None)
    svg = skrub.as_data_op(0).skb.set_description("SOME TEXT").skb.draw_graph().svg
    assert re.search(rb'<a target="_blank" xlink:title=".*SOME TEXT', svg)


def test_no_pydot(monkeypatch):
    monkeypatch.delitem(sys.modules, "pydot", raising=False)
    builtin_import = builtins.__import__

    def _import(name, *args, **kwargs):
        if name == "pydot":
            raise ImportError(name)
        return builtin_import(name, *args, **kwargs)

    monkeypatch.setattr(builtins, "__import__", _import)
    with pytest.raises(ImportError, match="Please install pydot"):
        skrub.as_data_op(0).skb.draw_graph()


def test_no_graphviz(monkeypatch):
    pydot = pytest.importorskip("pydot")
    monkeypatch.setattr(pydot.Dot, "create_svg", Mock(side_effect=Exception()))
    with pytest.raises(ImportError, match="Please install pydot and graphviz"):
        skrub.as_data_op(0).skb.draw_graph()


@pytest.mark.skipif(not _inspection._has_graphviz(), reason="report requires graphviz")
def test_draw_graph_open(monkeypatch):
    mock = Mock()
    monkeypatch.setattr(_inspection, "open_in_browser", mock)
    skrub.as_data_op(0).skb.draw_graph().open()
    mock.assert_called_once()


def test_describe_param_grid():
    """
    >>> from sklearn.linear_model import LogisticRegression
    >>> from sklearn.ensemble import RandomForestClassifier
    >>> from sklearn.decomposition import PCA
    >>> from sklearn.feature_selection import SelectKBest
    >>> from sklearn.impute import SimpleImputer
    >>> from sklearn.preprocessing import StandardScaler, RobustScaler

    >>> import skrub

    >>> X = skrub.X()
    >>> y = skrub.y()

    >>> imputed = X.skb.apply(skrub.optional(SimpleImputer(), name="impute"))
    >>> dim_reduction = skrub.choose_from(
    ...     {
    ...         "PCA": PCA(),
    ...         "SelectKBest": SelectKBest(),
    ...     },
    ...     name="dim_reduction",
    ... )
    >>> selected = imputed.skb.apply(dim_reduction)
    >>> use_scaling = skrub.choose_bool(name="scaling")
    >>> scaling_kind = skrub.choose_from(["robust", "standard"], name="scaling_kind")
    >>> scaler = scaling_kind.match(
    ...     {
    ...         "robust": RobustScaler(
    ...             **skrub.choose_bool(name="robust_scaler__with_centering")
    ...         ),
    ...         "standard": StandardScaler(),
    ...     }
    ... )
    >>> scaled = selected.skb.apply(use_scaling.if_else(scaler, None))
    >>> classifier = skrub.choose_from(
    ...     {
    ...         "logreg": LogisticRegression(
    ...             **skrub.choose_float(0.001, 100, log=True, name="C")
    ...         ),
    ...         "rf": RandomForestClassifier(
    ...             n_estimators=skrub.choose_int(20, 400, name="N 🌴")
    ...         ),
    ...     },
    ...     name="classifier",
    ... )
    >>> pred = scaled.skb.apply(classifier, y=y)

    those need to be split into separate subgrids:

    - scaling or not because the nested choice scaling kind is only used if
      scaling is true
    - scaling kind because the robust scaler has a nested choice with_centering
    - the classifier because the random forest and logistic regression have
      nested choices (their hyperparameters)

    so we end up with 2 (rf or logreg) x 3 (no scaling, robust, standard) = 6
    subgrids.

    >>> print(pred.skb.describe_param_grid())
    - impute: [SimpleImputer(), None]
      dim_reduction: ['PCA', 'SelectKBest']
      scaling: True
      scaling_kind: 'robust'
      robust_scaler__with_centering: [True, False]
      classifier: 'logreg'
      C: choose_float(0.001, 100, log=True, name='C')
    - impute: [SimpleImputer(), None]
      dim_reduction: ['PCA', 'SelectKBest']
      scaling: True
      scaling_kind: 'robust'
      robust_scaler__with_centering: [True, False]
      classifier: 'rf'
      N 🌴: choose_int(20, 400, name='N 🌴')
    - impute: [SimpleImputer(), None]
      dim_reduction: ['PCA', 'SelectKBest']
      scaling: True
      scaling_kind: 'standard'
      classifier: 'logreg'
      C: choose_float(0.001, 100, log=True, name='C')
    - impute: [SimpleImputer(), None]
      dim_reduction: ['PCA', 'SelectKBest']
      scaling: True
      scaling_kind: 'standard'
      classifier: 'rf'
      N 🌴: choose_int(20, 400, name='N 🌴')
    - impute: [SimpleImputer(), None]
      dim_reduction: ['PCA', 'SelectKBest']
      scaling: False
      classifier: 'logreg'
      C: choose_float(0.001, 100, log=True, name='C')
    - impute: [SimpleImputer(), None]
      dim_reduction: ['PCA', 'SelectKBest']
      scaling: False
      classifier: 'rf'
      N 🌴: choose_int(20, 400, name='N 🌴')
    """


def test_describe_params():
    c1 = skrub.choose_float(0.0, 1.0)
    c2 = skrub.choose_from((5.5, c1), name="c2")
    c3 = skrub.choose_bool()
    c4 = skrub.choose_from({"2": c2, "1": c1, "3": c3})
    c5 = skrub.choose_int(10, 20, default=11, name="c5")
    c6 = skrub.choose_from(["a", "b"])
    c7 = skrub.choose_float(100.0, 200.0, default=110.5)
    e = c6.match({"a": [c4, c2.as_data_op() + c7], "b": c5}).as_data_op()
    print(e.skb.describe_defaults())
    expected = {
        "choose_from(['a', 'b'])": "a",
        "choose_from({'2': …, '1': …, '3': …})": "2",
        "c2": 5.5,
        "choose_float(100.0, 200.0, default=110.5)": 110.5,
    }

    assert e.skb.describe_defaults() == expected
    assert e.skb.make_learner().describe_params() == expected
    assert skrub.X().skb.describe_defaults() == {}
