import os
import pandas as pd
import numpy as np
import random
import data_prep_aws 
import predictors
import logging
import time
import pickle
import datetime
from datetime import date
from imblearn.over_sampling import SMOTE
from sklearn.preprocessing import  OrdinalEncoder

import warnings
warnings.filterwarnings("ignore")



import snowflake.connector
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives.asymmetric import dsa
from cryptography.hazmat.primitives import serialization

from password import PRIVATE_KEY_PASSPHRASE

with open("/Users/impr001/Keys_snowflake/rsa_key.p8", "rb") as key: 
    p_key= serialization.load_pem_private_key(
        key.read(),
        password=PRIVATE_KEY_PASSPHRASE.encode(),
        backend=default_backend()
    )

pkb = p_key.private_bytes(
    encoding=serialization.Encoding.DER,
    format=serialization.PrivateFormat.PKCS8,
    encryption_algorithm=serialization.NoEncryption())

def timing_val(func):
    def wrapper(*arg, **kw):
        t1 = time.time()
        res = func(*arg, **kw)
        t2 = time.time()
        return (t2 - t1), res, func.__name__
    return wrapper

@timing_val
def load_full_data():
    ctx = snowflake.connector.connect(
        user='eimpara', 
        account='orchard',
        private_key=pkb,
        role= 'PROD_DATALYTICS_ROLE',
        warehouse = "DEV_ANALYTICS_ORCHARD"
        )
    print("And here...")

    try:
        cs = ctx.cursor()
        print("Here")
        sql = """
        select * from dev_engineering.eimpara.maze_country_table
        --limit 1000;
        """
        cs.execute(sql)
        original = cs.fetch_pandas_all()
        return original
        print(original.shape)

    finally:
        cs.close()
        ctx.close()

    

if __name__ == '__main__':

    timing, df, _ = load_full_data()
    print("{} rows loaded in {} seconds".format(df.shape[0], timing))
    (X_train, y_train, X_test, y_test) = data_prep_aws.split_train_test_data(df)
    log_model = predictors.LogisticRegressionModel(X_train, y_train)
    log_model.train()
    log_performances = predictors.estimate_predictor(log_model, X_test, y_test)
    logging.info(f"Perfromance metric for Logistic Model {log_performances}")
    matrix_log = predictors.confusion_matrix_calculation(log_model, X_test, y_test)
    logging.info(f"Perfromance metric for Logistic Model {matrix_log}")
    
    with open('/Users/impr001/Documents/Jupyter_Notebooks/Maze/log_model_XXX.pkl', 'wb') as file:
        pickle.dump(log_model, file)