# Apply STG on XOR dataset
from stg import STG
import numpy as np
import scipy.stats # for creating a simple dataset 
import matplotlib.pyplot as plt 
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score

def generate_XOR_labels(X):
    y = np.exp(X[:,0]*X[:,1])

    prob_1 = np.expand_dims(1 / (1+y) ,1)
    prob_0 = np.expand_dims(y / (1+y) ,1)

    # y = np.concatenate((prob_0,prob_1), axis = 1)
    y = np.where(prob_0 > prob_1, 0, 1)
    return y

X_data = np.random.randn(100, 10)
y_i = generate_XOR_labels(X_data)
y_data = y_i.flatten()

print(X_data.shape)
print(y_i)
print(y_data.shape)

X_train, X_test, y_train, y_test = train_test_split(X_data, y_data, train_size=0.3)
X_train, X_valid, y_train, y_valid = train_test_split(X_train, y_train, train_size=0.8)

model = STG(task_type='classification',input_dim=X_train.shape[1], output_dim=2, hidden_dims=[60, 20], activation='tanh',
    optimizer='SGD', learning_rate=0.1, batch_size=X_train.shape[0], feature_selection=True, sigma=0.5, lam=0.5, random_state=1, device="cpu")

model.fit(X_train, y_train, nr_epochs=6000, valid_X=X_valid, valid_y=y_valid, print_interval=1000)

print(model.get_gates(mode='prob'))

print(model.get_gates(mode='raw'))

y_pred=model.predict(X_test)

print(accuracy_score(y_pred,y_test)) #
# Apply STG on XOR dataset
from stg import STG
import numpy as np
import scipy.stats # for creating a simple dataset 
import matplotlib.pyplot as plt 
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score

def generate_XOR_labels(X):
    y = np.exp(X[:,0]*X[:,1])

    prob_1 = np.expand_dims(1 / (1+y) ,1)
    prob_0 = np.expand_dims(y / (1+y) ,1)

    # y = np.concatenate((prob_0,prob_1), axis = 1)
    y = np.where(prob_0 > prob_1, 0, 1)
    return y

X_data = np.random.randn(100, 10)
y_i = generate_XOR_labels(X_data)
y_data = y_i.flatten()

print(X_data.shape)
print(y_i)
print(y_data.shape)

X_train, X_test, y_train, y_test = train_test_split(X_data, y_data, train_size=0.3)
X_train, X_valid, y_train, y_valid = train_test_split(X_train, y_train, train_size=0.8)

model = STG(task_type='classification',input_dim=X_train.shape[1], output_dim=2, hidden_dims=[60, 20], activation='tanh',
    optimizer='SGD', learning_rate=0.1, batch_size=X_train.shape[0], feature_selection=True, sigma=0.5, lam=0.5, random_state=1, device="cpu")

model.fit(X_train, y_train, nr_epochs=6000, valid_X=X_valid, valid_y=y_valid, print_interval=1000)

print(model.get_gates(mode='prob'))

print(model.get_gates(mode='raw'))

y_pred=model.predict(X_test)

print(accuracy_score(y_pred,y_test)) #
# Apply STG on XOR dataset # for creating a simple dataset # y = np.concatenate((prob_0,prob_1), axis = 1) #