import numpy as np
import pandas as pd
from joblib import load


class Main:
    def __init__(self):
        self.preprocess = load('model/preprocess.joblib')
        self.model = load('model/model.joblib')

    def predict(self, json_df):
        print('json_df', json_df)
        df = pd.read_json(json_df, orient='records')
        print(df.info())
        print('df', df)
        preprocessed_df = self.preprocess.transform(df.copy())
        prediction_tuples = self.model.predict_proba(preprocessed_df)

        predictions = []
        confidences = []
        for prediction_tuple in prediction_tuples:
            prediction = np.argmax(prediction_tuple)
            confidence = prediction_tuple[prediction]
            predictions.append('Yes' if prediction == 1 else 'No')
            confidences.append(confidence)
        customer_ids = df['customerID'].values

        out_df = pd.DataFrame(np.dstack([customer_ids, predictions, confidences])[0], columns=['customerID', 'Churn', 'Confidence'])
        return out_df.to_json(orient='records')


if __name__ == '__main__':
    main = Main()
    df = pd.read_csv('data/WA_Fn-UseC_-Telco-Customer-Churn.csv')
    X = df[[col for col in df.columns if col != 'Churn']]
    X = X.loc[:4, :]
    json_df = X.to_json(orient='records')
    print(main.predict(json_df))
