from sklearn.tree import DecisionTreeRegressor as sklDT
from sklearn.metrics import mean_poisson_deviance
import numpy as np
import pandas as pd
from scipy.stats import beta
import matplotlib
import matplotlib.pyplot as plt
import seaborn as sns

matplotlib.use("Agg")
sns.set()


def beta_dataset():
    np.random.seed(33)
    n = 1000

    X = np.random.random((n, 4)).astype(np.float32)
    a, b = 2.31, 0.627
    y = beta.rvs(a, b, size=n).astype(np.float32) * 10
    return X, y


rs = np.random.RandomState(92)
depths = range(1, 8)
bootstrap = None
max_features = 1.0
n_estimators = 1
min_impurity_decrease = 0  # 1e-5
algo = {
    "skl_dt_poisson": sklDT(
        random_state=rs,
        min_impurity_decrease=min_impurity_decrease,
        criterion="poisson",
    ),
    "skl_dt_mse": sklDT(
        random_state=rs, min_impurity_decrease=min_impurity_decrease, criterion="mse",
    ),
}

datasets = {
    "poisson": beta_dataset(),
}
for data_name, (X, y) in datasets.items():
    X = X.astype(np.float32)
    y = y.astype(np.float32)
    df = pd.DataFrame(columns=["algorithm", "accuracy"])
    for d in depths:
        for name, alg in algo.items():
            name, alg, d
            alg.set_params(max_depth=d)
            alg.fit(X, y)

            pred = alg.predict(X)
            accuracy = mean_poisson_deviance(y, (pred))
            df = df.append(
                {"algorithm": name, "accuracy": accuracy, "depth": d},
                ignore_index=True,
            )

    print(df)
    sns.lineplot(data=df, x="depth", y="accuracy", hue="algorithm")
    plt.ylabel("train poisson")
    plt.tight_layout()
    plt.savefig("poisson_convergence.png")
    plt.clf()