from __future__ import annotations

import numpy
from numpy.testing import assert_almost_equal

from pyearth import Earth
from pyearth._basis import (
    Basis,
    ConstantBasisFunction,
    HingeBasisFunction,
    LinearBasisFunction,
)
from pyearth._types import BOOL
from pyearth.export import export_python_function

numpy.random.seed(0)

basis = Basis(10)
constant = ConstantBasisFunction()
basis.append(constant)
bf1 = HingeBasisFunction(constant, 0.1, 10, 1, False, "x1")
bf2 = HingeBasisFunction(constant, 0.1, 10, 1, True, "x1")
bf3 = LinearBasisFunction(bf1, 2, "x2")
basis.append(bf1)
basis.append(bf2)
basis.append(bf3)
X = numpy.random.normal(size=(100, 10))
missing = numpy.zeros_like(X, dtype=BOOL)
B = numpy.empty(shape=(100, 4), dtype=numpy.float64)
basis.transform(X, missing, B)
beta = numpy.random.normal(size=4)
y = numpy.empty(shape=100, dtype=numpy.float64)
y[:] = numpy.dot(B, beta) + numpy.random.normal(size=100)
beta2 = numpy.random.normal(size=4)
y2 = numpy.empty(shape=100, dtype=numpy.float64)
y2[:] = numpy.dot(B, beta2) + numpy.random.normal(size=100)
Y = numpy.concatenate([y[:, None], y2[:, None]], axis=1)
default_params = {"penalty": 1}


def test_export_python_function():
    for smooth in (True, False):
        model = Earth(penalty=1, smooth=smooth, max_degree=2).fit(X, y)
        export_model = export_python_function(model)
        for exp_pred, model_pred in zip(model.predict(X), export_model(X), strict=False):
            assert_almost_equal(exp_pred, model_pred)


# def test_export_python_string():
#     for smooth in (True, False):
#         model = Earth(penalty=1, smooth=smooth, max_degree=2).fit(X, y)
#         export_model = export_python_string(model, "my_test_model")
#         six.exec_(export_model, globals())
#         for exp_pred, model_pred in zip(model.predict(X), my_test_model(X)):  # noqa: F821
#             assert_almost_equal(exp_pred, model_pred)


# @if_pandas
# @if_sympy
# def test_export_sympy():
#     import pandas as pd
#     from sympy.printing.lambdarepr import NumPyPrinter
#     from sympy.utilities.lambdify import lambdify

#     class PyEarthNumpyPrinter(NumPyPrinter):
#         def _print_Max(self, expr):
#             return "maximum(" + ",".join(self._print(i) for i in expr.args) + ")"

#         def _print_NaNProtect(self, expr):
#             return (
#                 "where(isnan("
#                 + ",".join(self._print(a) for a in expr.args)
#                 + "), 0, "
#                 + ",".join(self._print(a) for a in expr.args)
#                 + ")"
#             )

#         def _print_Missing(self, expr):
#             return (
#                 "isnan("
#                 + ",".join(self._print(a) for a in expr.args)
#                 + ").astype(float)"
#             )

#     for smooth, n_cols, allow_missing in product((True, False), (1, 2), (True, False)):
#         X_df = pd.DataFrame(X.copy(), columns=["x_%d" % i for i in range(X.shape[1])])
#         y_df = pd.DataFrame(Y[:, :n_cols])
#         if allow_missing:
#             # Randomly remove some values so that the fitted model contains MissingnessBasisFunctions
#             X_df["x_1"][
#                 numpy.random.binomial(n=1, p=0.1, size=X_df.shape[0]).astype(bool)
#             ] = numpy.nan

#         model = Earth(allow_missing=allow_missing, smooth=smooth, max_degree=2).fit(
#             X_df, y_df
#         )
#         expressions = export_sympy(model) if n_cols > 1 else [export_sympy(model)]
#         module_dict = {
#             "select": numpy.select,
#             "less_equal": numpy.less_equal,
#             "isnan": numpy.isnan,
#             "greater_equal": numpy.greater_equal,
#             "logical_and": numpy.logical_and,
#             "less": numpy.less,
#             "logical_not": numpy.logical_not,
#             "greater": numpy.greater,
#             "maximum": numpy.maximum,
#             "Missing": lambda x: numpy.isnan(x).astype(float),
#             "NaNProtect": lambda x: numpy.where(numpy.isnan(x), 0, x),
#             "nan": numpy.nan,
#             "float": float,
#             "where": numpy.where,
#         }

#         for i, expression in enumerate(expressions):
#             # The lambdified functions for smoothed basis functions only work with modules='numpy' and
#             # for regular basis functions with modules={'Max':numpy.maximum}.  This is a confusing situation
#             func = lambdify(
#                 X_df.columns,
#                 expression,
#                 printer=PyEarthNumpyPrinter,
#                 modules=module_dict,
#             )
#             y_pred_sympy = func(*[X_df.loc[:, var] for var in X_df.columns])

#             y_pred = model.predict(X_df)[:, i] if n_cols > 1 else model.predict(X_df)
#             assert_array_almost_equal(y_pred, y_pred_sympy)
