import pandas as pd
from IPython.display import display
from sklearn.preprocessing import OneHotEncoder
from sklearn.preprocessing import LabelEncoder
import numpy as np

df = pd.read_csv('data/air2000_test.csv', header=0, index_col=0)
df.head()

# count the number of missing values per column
display(df.isnull().sum())

# drop columns with ALL NaN
df_drop_col = df.dropna(axis=1, thresh=1)
df_drop_col.head()

# drop rows with ANY NaN
df_drop_col_row = df_drop_col.dropna(axis=0, thresh=df_drop_col.shape[1])
df_drop_col_row.head()

X = df_drop_col_row.drop('ArrDelay', 1)
y = [int(arrDelay<=0) for arrDelay in df_drop_col_row['ArrDelay']]
#cols = df_drop_col_row.columns
#new_cols = ['ArrDelay'] + list(set(cols)-set(['ArrDelay']))
#X = df_drop_col_row[new_cols]
X.head()

print('Shape of input before one-hot: {}'.format(X.shape))

# Recognize non-numeric columns as categorical columns
cols = X.columns
num_cols = X._get_numeric_data().columns
catego_cols = list(set(cols) - set(num_cols))

# Add other categorical columns
catego_cols.extend(['Year', 'Month', 'DayofMonth', 'DayOfWeek', 'FlightNum'])#, 'Origin', 'Dest'])

print('Categorical Columns: {}'.format(catego_cols))

# encode label first
catego_le = LabelEncoder()

for i in catego_cols:
    X[i] = catego_le.fit_transform(X[i].values)
    classes_list = catego_le.classes_.tolist()
    
X.head()

# find the index of the categorical feature
catego_cols_idx = []
for str in catego_cols:
    catego_cols_idx.append(X.columns.tolist().index(str))

# give the column index you want to do one-hot encoding
ohe = OneHotEncoder(categorical_features = catego_cols_idx)

# fit one-hot encoder
onehot_data = ohe.fit_transform(X.values).toarray()
print('Shape of input after one-hot: {}'.format(onehot_data.shape))

data = pd.DataFrame(onehot_data, index=X.index)
data.head()

data['ArrDelay'] = y
data.head()

data.to_csv('logistic_input', header=False)