#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @FileName  :tf_lstm_train.py
# @Time      :2025/2/13 10:52
# @Author    :David
# @Company: shenyang JY
import json, os
import numpy as np
import traceback
import logging
from app.common.logs import params
from app.common.data_handler import DataHandler, write_number_to_file
import time
from app.common.tf_fmi import FMIHandler
from app.common.dbmg import MongoUtils
from app.common.logs import logger
from copy import deepcopy
np.random.seed(42)  # NumPy随机种子
# tf.set_random_seed(42)  # TensorFlow随机种子

mgUtils = MongoUtils(logger)

def model_training(train_data, input_file, cap):
    # 获取程序开始时间
    start_time = time.time()
    success = 0
    logger.info("Program starts execution!")
    farm_id = input_file.split('/')[-2]
    output_file = input_file.replace('IN', 'OUT')
    status_file = 'STATUS.TXT'
    # 创建线程独立的实例
    local_params = deepcopy(params)
    dh = DataHandler(logger, local_params)
    fmi = FMIHandler(logger, local_params)
    try:
        # ------------ 获取数据，预处理训练数据 ------------
        dh.opt.cap = cap
        train_x, valid_x, train_y, valid_y, scaled_train_bytes, scaled_target_bytes, scaled_cap = dh.train_data_handler(train_data)
        fmi.opt.cap = round(scaled_cap, 2)
        fmi.opt.Model['input_size'] = train_x.shape[2]
        # ------------ 训练模型，保存模型 ------------
        # 1. 如果是加强训练模式，先加载预训练模型特征参数，再预处理训练数据
        # 2. 如果是普通模式，先预处理训练数据，再根据训练数据特征加载模型
        model = fmi.train_init() if fmi.opt.Model['add_train'] else fmi.get_keras_model(fmi.opt)
        if fmi.opt.Model['add_train']:
            if model:
                feas = json.loads(fmi.model_params).get('features', dh.opt.features)
                if set(feas).issubset(set(dh.opt.features)):
                    dh.opt.features = list(feas)
                    train_x, train_y, valid_x, valid_y, scaled_train_bytes, scaled_target_bytes, scaled_cap = dh.train_data_handler(train_data)
                else:
                    model = fmi.get_keras_model(fmi.opt)
                    logger.info("训练数据特征，不满足，加强训练模型特征")
            else:
                model = fmi.get_keras_model(fmi.opt)

        ts_model = fmi.training(model, [train_x, valid_x, train_y, valid_y])
        success = 1
        # 更新算法状态：1. 启动成功
        write_number_to_file(os.path.join(output_file, status_file), 1, 1, 'rewrite')
        # ------------ 组装模型数据 ------------
        local_params['Model']['features'] = ','.join(dh.opt.features)
        local_params.update({
            'params': json.dumps(local_params),
            'descr': f'南网竞赛-{farm_id}',
            'gen_time': time.strftime('%Y-%m-%d %H:%M:%S', time.localtime()),
            'model_table': local_params['model_table'] + farm_id,
            'scaler_table': local_params['scaler_table'] + farm_id
        })
        mgUtils.insert_trained_model_into_mongo(ts_model, local_params)
        mgUtils.insert_scaler_model_into_mongo(scaled_train_bytes, scaled_target_bytes, local_params)
        # 更新算法状态：正常结束
        write_number_to_file(os.path.join(output_file, status_file), 2, 2)
    except Exception as e:
        # 如果算法状态没启动，不更新
        if success:
            write_number_to_file(os.path.join(output_file, status_file), 2, 3)
        my_exception = traceback.format_exc()
        my_exception.replace("\n", "\t")
    end_time = time.time()
    logger.info("fmi训练任务：用了 %s 秒 " % (end_time-start_time))


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

    serve(app, host="0.0.0.0", port=10103, threads=4)
    print("server start!")
    # args_dict = {"mongodb_database": 'realtimeDq', 'scaler_table': 'j00600_scaler', 'model_name': 'lstm1',
    # 'model_table': 'j00600_model', 'mongodb_read_table': 'j00600', 'col_time': 'dateTime',
    # 'features': 'speed10,direction10,speed30,direction30,speed50,direction50,speed70,direction70,speed90,direction90,speed110,direction110,speed150,direction150,speed170,direction170'}
    # args_dict['features'] = args_dict['features'].split(',')
    # args.update(args_dict)
    # dh = DataHandler(logger, args)
    # ts = TSHandler(logger, args)
    # opt = argparse.Namespace(**args)
    # opt.Model['input_size'] = len(opt.features)
    # train_data = get_data_from_mongo(args_dict)
    # train_x, train_y, valid_x, valid_y, scaled_train_bytes, scaled_target_bytes = dh.train_data_handler(train_data)
    # ts_model = ts.training([train_x, train_y, valid_x, valid_y])
    #
    # args_dict['gen_time'] = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time()))
    # args_dict['params'] = args
    # args_dict['descr'] = '测试'
    # insert_trained_model_into_mongo(ts_model, args_dict)
    # insert_scaler_model_into_mongo(scaled_train_bytes, scaled_target_bytes, args_dict)