import pandas as pd
from pymongo import MongoClient
import pickle
from flask import Flask, request
import time
import logging
import traceback
from common.database_dml import get_data_from_mongo, insert_data_into_mongo

app = Flask('res_prediction——service')


def str_to_list(arg):
    if arg == '':
        return []
    else:
        return arg.split(',')


@app.route('/res_prediction', methods=['POST'])
def model_prediction_lightgbm():
    # 获取程序开始时间
    start_time = time.time()
    result = {}
    success = 0
    print("Program starts execution!")
    try:
        args = request.values.to_dict()
        print('args', args)
        logger.info(args)
        col_reserve = str_to_list(args['col_reserve'])
        power_df = get_data_from_mongo(args)
        if 'is_limit' in power_df.columns:
            power_df = power_df[power_df['is_limit'] == False]
        power_df['model'] = args['model']
        power_df['predict'] = power_df[args['col_pre']]
        features_reserve = col_reserve + ['model', 'predict']
        insert_data_into_mongo(power_df[set(features_reserve)], args)
        success = 1
    except Exception as e:
        my_exception = traceback.format_exc()
        my_exception.replace("\n", "\t")
        result['msg'] = my_exception
    end_time = time.time()

    result['success'] = success
    result['args'] = args
    result['start_time'] = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(start_time))
    result['end_time'] = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(end_time))
    print("Program execution ends!")
    return result


if __name__ == "__main__":
    print("Program starts execution!")
    logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    logger = logging.getLogger("res_prediction log")
    from waitress import serve

    serve(app, host="0.0.0.0", port=10105)
    print("server start!")