import matplotlib
import matplotlib.pyplot as plt
import numpy as np
from pandas import read_csv, merge
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.neural_network import MLPClassifier
from sklearn.metrics import classification_report, confusion_matrix

# Read in the focus/dose data with the first row (i.e., row 0) as the header
exp_data = read_csv('exposure_data.csv', header=0)

# Check for missing values
print("Total number of rows in DataFrame: " + str(exp_data.shape[0]))
print("Number of rows with missing values by column:")
print(exp_data.isnull().sum())

# drop rows with missing values, printing the row counts before and after
exp_data.dropna(inplace=True)
# Check for missing values
print("Total number of rows in DataFrame: " + str(exp_data.shape[0]))
print("Number of rows with missing values by column:")
print(exp_data.isnull().sum())

# Read in the inspection data
insp_data = read_csv('inspection_data.csv', header=0)

# Check for missing values
print("Total number of rows in dataframe: " + str(insp_data.shape[0]))
print("Number of rows with missing values by column:")
print(insp_data.isnull().sum())

# drop rows with missing values, printing the row counts before and after
insp_data.dropna(inplace=True)
# Check for missing values
print("Total number of rows in dataframe: " + str(insp_data.shape[0]))
print("Number of rows with missing values by column:")
print(insp_data.isnull().sum())

# merge data into a single DataFrame
merged_data = merge(exp_data, insp_data, how='inner', on=['Lot', 'Wafer', 'FieldX', 'FieldY'], sort=False)

# Print a few rows to see what the results look like
print(merged_data[0:5])

data_array = merged_data.values

# Print a few rows for comparison with the DataFrame view in the previous step
print(data_array[0:4,:])

X = data_array[:,9:11].astype(float)
Y = data_array[:,12].astype(int)

mask = Y.astype(bool)
a = 20*np.ones(mask.shape)
a_good = np.ma.masked_where(mask, a)
a_bad = np.ma.masked_where(np.logical_not(mask), a)
fig = plt.figure(figsize=(6,6), dpi=100)
axes = fig.add_axes([0.1, 0.1, 0.9, 0.9])
axes.scatter(X[:,0], X[:,1], s=a_good, color='green', label='No Defect')
axes.scatter(X[:,0], X[:,1], s=a_bad, color='red', label='Defect')
axes.set_xlabel('Focus [nm]')
axes.set_ylabel('Dose [mJ]')
axes.set_title('Raw Data');
axes.legend();

# Split the data into training and test sets
X_train, X_test, y_train, y_test = train_test_split(X,Y,test_size=0.2, random_state = 2018)

# Create a 'StandardScaler' instance and adjust the normalization transforms to the training data
scaler = StandardScaler()
scaler.fit(X_train)

# Now apply the transformations to the the training and test sets
X_train = scaler.transform(X_train)
X_test = scaler.transform(X_test)

# Define a multi-layer perceptron classifier with three hidden layers, each with 40 neurons
mlp = MLPClassifier(hidden_layer_sizes=(40,40,40), random_state = 2018)

# Fit the classifier to the training data
mlp.fit(X_train,y_train);

predictions = mlp.predict(X_test)

print(confusion_matrix(y_test,predictions))

print(classification_report(y_test,predictions))