import os
import sys

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from io import StringIO
from sklearn.datasets import fetch_openml
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from _treedraw import draw_tree  # noqa: E402

# The lab's field key for Palmer penguins: the same cleaning, split and chosen settings as the lab
# (max_depth=3, ccp_alpha=0.01), refit on the 256 development penguins and drawn as a dichotomous key.
peng = fetch_openml(data_id=42585, as_frame=True).frame
num = ['culmen_length_mm', 'culmen_depth_mm', 'flipper_length_mm', 'body_mass_g']
peng = peng.dropna(subset=num)
X = peng[num].copy()
X['is_male'] = peng['sex'].astype(object).map({'MALE': 1.0, 'FEMALE': 0.0})
X = X.join(pd.get_dummies(peng['island'], prefix='island', dtype=float))
y = peng['species'].astype(str)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, stratify=y, random_state=42)
key = DecisionTreeClassifier(max_depth=3, ccp_alpha=0.01, random_state=0).fit(X_train, y_train)

BLUE, ORANGE, GREEN = '#58a6ff', '#f0883e', '#3fb950'          # Adelie, Chinstrap, Gentoo
NICE = {'flipper_length_mm': 'flipper', 'culmen_length_mm': 'bill length', 'culmen_depth_mm': 'bill depth'}


def question(node, name, threshold):
    if name.startswith('island_'):
        return f'not from {name[7:]} island?'
    return f'{NICE.get(name, name)} ≤ {threshold:g} mm?'


plt.rcParams.update({'text.color': 'gray'})
fig, ax = plt.subplots(figsize=(11, 5.6))
xy = draw_tree(ax, key, list(X.columns), list(key.classes_), [BLUE, ORANGE, GREEN], xscale=1.75, yscale=1.45,
               fontsize=9, show_gini=False, question=question, gap=0.27)
xs = [x for x, _ in xy.values()]
ax.set_xlim(min(xs) - 1.0, max(xs) + 1.0)
ax.set_ylim(-3 * 1.45 - 0.9, 0.6)
ax.axis('off')
ax.text(np.mean([min(xs), max(xs)]), -3 * 1.45 - 0.75,
        'numbers in each box: Adelie · Chinstrap · Gentoo among the 256 development penguins  ·  '
        f'{key.get_n_leaves()} leaves, test accuracy {key.score(X_test, y_test):.3f}',
        ha='center', fontsize=9, style='italic')
plt.tight_layout()

buffer = StringIO()
plt.savefig(buffer, format="svg", transparent=True)
print(buffer.getvalue())

plt.close()