from pygoose import *

import datetime

import lightgbm as lgb

from sklearn.model_selection import StratifiedKFold

"""
Explanation: Final Prediction: LightGBM
Train a GBM using K-fold CV and use the mean test prediction across the folds for the final submission.
Imports
This utility package imports numpy, pandas, matplotlib and a helper kg module into the root namespace.
End of explanation
"""


project = kg.Project.discover()

"""
Explanation: Config
Automatically discover the paths to various data folders and compose the project structure.
End of explanation
"""


NUM_FOLDS = 5

"""
Explanation: Number of CV folds.
End of explanation
"""


RANDOM_SEED = 2017

np.random.seed(RANDOM_SEED)

"""
Explanation: Make subsequent runs reproducible.
End of explanation
"""


feature_lists = [
    'simple_summaries',
    'jaccard_ngrams',
    'fuzzy',
    'tfidf',
    'lda',
    'nlp_tags',
    'wordnet_similarity',
    'phrase_embedding',
    'wmd',
    'wm_intersect',
    
    '3rdparty_abhishek',
    '3rdparty_dasolmar_whq',
    '3rdparty_mephistopheies',
    '3rdparty_image_similarity',
    
    'magic_pagerank',
    'magic_frequencies',
    'magic_cooccurrence_matrix',
    
    'oofp_nn_mlp_with_magic',
    'oofp_nn_cnn_with_magic',
    'oofp_nn_bi_lstm_with_magic',
    'oofp_nn_siamese_lstm_attention',
]

df_train, df_test, feature_list_ix = project.load_feature_lists(feature_lists)

X_train = df_train.values
X_test = df_test.values

y_train = kg.io.load(project.features_dir + 'y_train.pickle')

"""
Explanation: Read Data
Load all features we extracted earlier.
End of explanation
"""


print('X train:', X_train.shape)
print('X test: ', X_test.shape)
print('y train:', y_train.shape)

pd.DataFrame(feature_list_ix, columns=['feature_list', 'start_index', 'end_index'])

"""
Explanation: View feature summary.
End of explanation
"""


kfold = StratifiedKFold(
    n_splits=NUM_FOLDS,
    shuffle=True,
    random_state=RANDOM_SEED
)

y_test_pred = np.zeros((len(X_test), NUM_FOLDS))

"""
Explanation: Train models & compute test predictions from each fold
Calculate partitions.
End of explanation
"""


cv_scores = []

%%time

for fold_num, (ix_train, ix_val) in enumerate(kfold.split(X_train, y_train)):
    print(f'Fitting fold {fold_num + 1} of {kfold.n_splits}')
    
    X_fold_train = X_train[ix_train]
    X_fold_val = X_train[ix_val]

    y_fold_train = y_train[ix_train]
    y_fold_val = y_train[ix_val]
    
    lgb_params = {
        'objective': 'binary',
        'metric': 'binary_logloss',
        'boosting': 'gbdt',
        'device': 'cpu',
        'feature_fraction': 0.486,
        'num_leaves': 158,
        'lambda_l2': 50,
        'learning_rate': 0.01,
        'num_boost_round': 5000,
        'early_stopping_rounds': 10,
        'verbose': 1,
        'bagging_fraction_seed': RANDOM_SEED,
        'feature_fraction_seed': RANDOM_SEED,
    }
    
    lgb_data_train = lgb.Dataset(X_fold_train, y_fold_train)
    lgb_data_val = lgb.Dataset(X_fold_val, y_fold_val)    
    evals_result = {}
    
    model = lgb.train(
        lgb_params,
        lgb_data_train,
        valid_sets=[lgb_data_train, lgb_data_val],
        evals_result=evals_result,
        num_boost_round=lgb_params['num_boost_round'],
        early_stopping_rounds=lgb_params['early_stopping_rounds'],
        verbose_eval=False,
    )
    
    fold_train_scores = evals_result['training'][lgb_params['metric']]
    fold_val_scores = evals_result['valid_1'][lgb_params['metric']]
    
    print('Fold {}: {} rounds, training loss {:.6f}, validation loss {:.6f}'.format(
        fold_num + 1,
        len(fold_train_scores),
        fold_train_scores[-1],
        fold_val_scores[-1],
    ))
    print()
    
    cv_scores.append(fold_val_scores[-1])
    y_test_pred[:, fold_num] = model.predict(X_test).reshape(-1)

"""
Explanation: Fit all folds.
End of explanation
"""


pd.DataFrame({
    'column': list(df_train.columns),
    'importance': model.feature_importance(),
}).sort_values(by='importance')

final_cv_score = np.mean(cv_scores)

print('Final CV score:', final_cv_score)

"""
Explanation: Print CV score and feature importance.
End of explanation
"""


y_test = np.mean(y_test_pred, axis=1)

submission_id = datetime.datetime.now().strftime('%Y-%m-%d-%H%M')

df_submission = pd.DataFrame({
    'test_id': range(len(y_test)),
    'is_duplicate': y_test
})

"""
Explanation: Generate submission
End of explanation
"""


def recalibrate_prediction(pred, train_pos_ratio=0.3692, test_pos_ratio=0.165):
    a = test_pos_ratio / train_pos_ratio
    b = (1 - test_pos_ratio) / (1 - train_pos_ratio)
    return a * pred / (a * pred + b * (1 - pred))

df_submission['is_duplicate'] = df_submission['is_duplicate'].map(recalibrate_prediction)

df_submission = df_submission[['test_id', 'is_duplicate']]

"""
Explanation: Recalibrate predictions for a different target balance on test
Based on Mike Swarbrick Jones' blog.
$\alpha = \frac{p_{test}}{p_{train}}$
$\beta = \frac{1 - p_{test}}{1 - p_{train}}$
$\hat{y}{test}^{\prime} = \frac{\alpha \hat{y}{test}}{\alpha \hat{y}{test} + \beta(1 - \hat{y}{test})}$
Training set balance is 36.92%, test set balance is ~16.5%.
End of explanation
"""


pd.DataFrame(y_test).plot.hist()

print('Test duplicates with >0.9 confidence:', len(df_submission[df_submission.is_duplicate > 0.9]))
print('Test mean prediction:', np.mean(y_test))
print('Calibrated mean prediction:', df_submission['is_duplicate'].mean())

df_submission.to_csv(
    project.submissions_dir + f'{submission_id}-submission-draft-cv-{final_cv_score:.6f}.csv',
    header=True,
    float_format='%.8f',
    index=None,
)

"""
Explanation: Explore and save submission
End of explanation
"""


