import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn import metrics

class DecisionNode:
    pass

class LeafNode(DecisionNode):
    def __init__(self, value):
        self.value = value
    def nodetype(self):
        return "leaf"

class InternalNode(DecisionNode):
    def __init__(self, i, v, left, right, gini=0, bins=[]):
        self.i = i
        self.v = v
        self.left = left
        self.right = right
        self.gini = gini
        self.bins = bins
    def nodetype(self):
        return "internal"


def decide(data: np.array, classes: np.array, i : int, v : float):
    mask0 = data[:,i]>=v
    mask1 = data[:,i]<v
    data0 = data[mask0]
    classes0 = classes[mask0]
    data1 = data[mask1]
    classes1 = classes[mask1]
    return data0, classes0, data1, classes1

def gini(classes: np.array):
    N = classes.shape[0]
    T = (np.bincount(classes)/N)**2
    return 1-np.sum(T)

def gain(data: np.array, classes: np.array, i: int, v : float):
    d0,c0,d1,c1 = decide(data, classes, i, v)
    return gini(classes) - (c0.shape[0]/classes.shape[0])*gini(c0) - (c1.shape[0]/classes.shape[0])*gini(c1)

def maxgain(data: np.array, classes: np.array):
    best_idx = 0
    best_value = data[0][best_idx]
    currentg = gain(data, classes, best_idx, best_value)

    for i in range(data.shape[1]):
        candv = np.unique(data[:,i])
        for v in candv:
            if gain(data, classes, i, v) > currentg:
                best_idx = i
                best_value = v
                currentg = gain(data, classes, best_idx, best_value)
    return (best_idx, best_value)

def fit(data: np.array, classes: np.array):
    if gini(classes) == 0: # Il n'y a plus qu'une seule classes représentée dans le jeu de données
        return LeafNode(classes[0])
    else: # Sinon, on choisit la meilleure question et on construit l'arbre récursivement
        (i,v) = maxgain(data, classes)
        d0,c0,d1,c1 = decide(data, classes, i, v)
        left = fit(d0,c0)
        right = fit(d1,c1)
        
        return InternalNode(i, v, left, right, gini=gini(classes), bins=np.bincount(classes))

def predict_v(t, d):
    if t.nodetype() == "leaf":
        return t.value
    else:
        if d[t.i] >= t.v:
            return predict_v(t.left, d)
        else:
            return predict_v(t.right,d)

def predict(t, data):
    r = np.zeros(data.shape[0])
    for i in range(data.shape[0]):
        r[i] = predict_v(t, data[i])    
    return r

def printTree(t, tab=0):
    if t.nodetype() == "leaf":
        print(tab*"\t" + f"Return classe {t.value}")
    else:
        print(tab*"\t" + f"Current gini = {t.gini}")
        print(tab*"\t" + f"Current bins = {t.bins}")
        print(tab*"\t" + f"if (x_{t.i} < {t.v})")
        printTree(t.left, tab=tab+1)
        print(tab*"\t" + f"else")
        printTree(t.right, tab=tab+1)

def treeSize(t):
    if t.nodetype() == "leaf":
        return 1
    else:
        return 1+treeSize(t.left)+treeSize(t.right)
# Exemple jouet
X = np.array([[1, 2],
              [3, 4],
              [1, 4],
              [5, 2]
              ])
y = np.array([0,1,1,0])

t = fit(X,y)

# Iris

df = pd.read_csv("iris.csv")

feature_names = df.columns.to_numpy()[:4]
data = df.to_numpy()[:, :4]

d = {"Virginica": 0, "Versicolor": 1, "Setosa": 2}
cname = ["Virginica", "Versicolor", "Setosa"]
classes = df["variety"].replace(d).to_numpy(dtype=int)

data_train, data_test, class_train, class_test = train_test_split(data, classes, test_size=0.5, random_state=52)

print(gain(data_train, class_train, 3, 0.7))
print(gain(data_train, class_train, 2, 3.3))
t = fit(data_train, class_train)
printTree(t)
print(treeSize(t))

print(predict(t, data_test))
print(class_test)
accuracy = metrics.accuracy_score(class_test, predict(t, data_test))
print(f"Accurracy: {accuracy}")