"""
@Time ： 2022/9/20 16:17
@Auth ： 东
@File ：AlgorithmController.py
@IDE ：PyCharm
@Motto：ABC(Always Be Coding)
@Desc：算法接口

"""
import json
from functools import wraps
from threading import Thread

from flask import Blueprint, request

from app.schemas.TrainResult import Report, ProcessValueList
from app.utils.RedisMQTool import Task
from app.utils.StandardizedOutput import output_wrapped
from app.utils.redis_config import redis_client
from app.utils.websocket_tool import manager
import sys
from pathlib import Path

from pynvml import *
# FILE = Path(__file__).resolve()
# ROOT = FILE.parents[0]  # YOLOv5 root directory
# if str(ROOT) not in sys.path:
#     sys.path.append(str(ROOT))  # add ROOT to PATH
# sys.path.append("/mnt/sdc/algorithm/AICheck-MaskRCNN/app/maskrcnn_ppx")
# import ppx as pdx

bp = Blueprint('AlgorithmController', __name__)

ifKillDict = {}

def start_train_algorithm():
    """
    调用训练算法
    """

    def wrapTheFunction(func):
        @wraps(func)
        @bp.route('/start_train_algorithm', methods=['get'])
        def wrapped_function():
            param = request.args.get('param')
            id = request.args.get('id')
            t = Thread(target=func, args=(param, id))
            t.start()
            return output_wrapped(0, 'success', '成功')

        return wrapped_function

    return wrapTheFunction


def start_test_algorithm():
    """
    调用验证算法
    """

    def wrapTheFunction(func):
        @wraps(func)
        @bp.route('/start_test_algorithm', methods=['get'])
        def wrapped_function_test():
            param = request.args.get('param')
            id = request.args.get('id')
            t = Thread(target=func, args=(param, id))
            t.start()
            return output_wrapped(0, 'success', '成功')

        return wrapped_function_test

    return wrapTheFunction


def start_detect_algorithm():
    """
    调用检测算法
    """

    def wrapTheFunction(func):
        @wraps(func)
        @bp.route('/start_detect_algorithm', methods=['get'])
        def wrapped_function_detect():
            param = request.args.get('param')
            id = request.args.get('id')
            t = Thread(target=func, args=(param, id))
            t.start()
            return output_wrapped(0, 'success', '成功')

        return wrapped_function_detect

    return wrapTheFunction


def start_download_pt():
    """
    下载模型
    """

    def wrapTheFunction(func):
        @wraps(func)
        @bp.route('/start_download_pt', methods=['get'])
        def wrapped_function_start_download_pt():
            param = request.args.get('param')
            data = func(param)
            return output_wrapped(0, 'success', data)

        return wrapped_function_start_download_pt

    return wrapTheFunction


def algorithm_process_value():
    """
    获取中间值, redis订阅发布
    """

    def wrapTheFunction(func):
        @wraps(func)
        def wrapped_function(*args, **kwargs):
            data = func(*args, **kwargs)
            print(data)
            Task(redis_conn=redis_client.get_redis(), channel="ceshi").publish_task(
                data={'code': 0, 'msg': 'success', 'data': data})
            return output_wrapped(0, 'success', data)

        return wrapped_function

    return wrapTheFunction


def algorithm_process_value_websocket():
    """
    获取中间值, websocket发布
    """

    def wrapTheFunction(func):
        @wraps(func)
        def wrapped_function(*args, **kwargs):
            data = func(*args, **kwargs)
            id = data["id"]
            data_res = {'code': 0, "type": 'connected', 'msg': 'success', 'data': data}
            manager.send_message_proj_json(message=data_res, id=id)
            return data

        return wrapped_function

    return wrapTheFunction

def algorithm_kill_value_websocket():
    """
    获取kill值, websocket发布
    """

    def wrapTheFunction(func):
        @wraps(func)
        def wrapped_function(*args, **kwargs):
            data = func(*args, **kwargs)
            id = data["id"]
            data_res = {'code': 1, "type": 'kill', 'msg': 'success', 'data': data}
            manager.send_message_proj_json(message=data_res, id=id)
            return data

        return wrapped_function

    return wrapTheFunction


def algorithm_error_value_websocket():
    """
    获取error值, websocket发布
    """

    def wrapTheFunction(func):
        @wraps(func)
        def wrapped_function(*args, **kwargs):
            data = func(*args, **kwargs)
            id = data["id"]
            data_res = {'code': 2, "type": 'error', 'msg': 'fail', 'data': data}
            manager.send_message_proj_json(message=data_res, id=id)
            return data

        return wrapped_function

    return wrapTheFunction

def obtain_train_param():
    """
       获取训练参数
       """

    def wrapTheFunction(func):
        @wraps(func)
        @bp.route('/obtain_train_param', methods=['get'])
        def wrapped_function_train_param(*args, **kwargs):
            data = func(*args, **kwargs)
            return output_wrapped(0, 'success', data)

        return wrapped_function_train_param

    return wrapTheFunction

def obtain_test_param():
    """
       获取验证参数
       """

    def wrapTheFunction(func):
        @wraps(func)
        @bp.route('/obtain_test_param', methods=['get'])
        def wrapped_function_test_param(*args, **kwargs):
            data = func(*args, **kwargs)
            return output_wrapped(0, 'success', data)

        return wrapped_function_test_param

    return wrapTheFunction


def obtain_detect_param():
    """
       获取测试参数
       """

    def wrapTheFunction(func):
        @wraps(func)
        @bp.route('/obtain_detect_param', methods=['get'])
        def wrapped_function_inf_param(*args, **kwargs):
            data = func(*args, **kwargs)
            return output_wrapped(0, 'success', data)

        return wrapped_function_inf_param

    return wrapTheFunction


def obtain_download_pt_param():
    """
       获取下载模型参数
       """

    def wrapTheFunction(func):
        @wraps(func)
        @bp.route('/obtain_download_pt_param', methods=['get'])
        def wrapped_function_obtain_download_pt_param(*args, **kwargs):
            data = func(*args, **kwargs)
            return output_wrapped(0, 'success', data)

        return wrapped_function_obtain_download_pt_param

    return wrapTheFunction

@bp.route('/change_ifKillDIct', methods=['get'])
def change_ifKillDIct():
    """
    修改全局变量
    """
    id = request.args.get('id')
    type = request.args.get('type')
    global ifKillDict
    ifKillDict[id] = False
    return output_wrapped(0, 'success')

# @start_train_algorithm()
# def start(param: str):
#     """
#     例子
#     """
#     print(param)
#     process_value_list = ProcessValueList(name='1', value=[])
#     report = Report(rate_of_progess=0, process_value=[process_value_list], id='1')
#
#     @algorithm_process_value_websocket()
#     def process(v: int):
#         print(v)
#         report.rate_of_progess = ((v + 1) / 10) * 100
#         report.precision[0].value.append(v)
#         return report.dict()
#
#     for i in range(10):
#         process(i)
#     return report.dict()
from setparams import TrainParams
import os
from app.schemas.TrainResult import DetectProcessValueDice, DetectReport
from app import file_tool


def error_return(id: str, data):
    """
    算法出错，返回
    """
    data_res = {'code': 2, "type": 'error', 'msg': 'fail', 'data': data}
    manager.send_message_proj_json(message=data_res, id=id)

# 启动训练
@start_train_algorithm()
def train_R0DY(params_str, id):
    from app.yolov5.train_server import train_start
    params = TrainParams()
    params.read_from_str(params_str)
    print(params.get('device').default)
    data_list = file_tool.get_file(ori_path=params.get('DatasetDir').value, type_list=params.get('CLASS_NAMES').value)
    weights = params.get('resumeModPath').value  # 初始化模型绝对路径
    img_size = params.get('img_size').value
    savemodel = os.path.splitext(params.get('saveModDir').value)[0] + '_' + str(img_size) + '.pt'  # 模型命名加上图像参数
    epoches = params.get('epochnum').value
    batch_size = params.get('batch_size').value
    device = params.get('device').value
    try:
        train_start(weights, savemodel, epoches, img_size, batch_size, device, data_list, id)
        print("train down!")
    except Exception as e:
        print(repr(e))
        error_return(id=id,data=repr(e))


# 启动验证程序

@start_test_algorithm()
def validate_RODY(params_str, id):
    from app.yolov5.validate_server import validate_start
    params = TrainParams()
    params.read_from_str(params_str)
    weights = params.get('modPath').value  # 验证模型绝对路径
    (filename, extension) = os.path.splitext(weights)  # 文件名与后缀名分开
    img_size = int(filename.split('ROD')[1].split('_')[2])  # 获取图像参数
    # v_num = int(filename.split('ROD')[1].split('_')[1]) #获取版本号
    output = params.get('outputPath').value
    batch_size = params.get('batch_size').default
    device = params.get('device').value

    validate_start(weights, img_size, batch_size, device, output, id)


@start_detect_algorithm()
def detect_RODY(params_str, id):
    from app.yolov5.detect_server import detect_start
    params = TrainParams()
    params.read_from_str(params_str)
    weights = params.get('modPath').value  # 检测模型绝对路径
    input = params.get('inputPath').value
    outpath = params.get('outputPath').value
    # (filename, extension) = os.path.splitext(weights)  # 文件名与后缀名分开
    # img_size = int(filename.split('ROD')[1].split('_')[2]) #获取图像参数
    # v_num = int(filename.split('ROD')[1].split('_')[1]) #获取版本号
    # batch_size = params.get('batch_size').default
    device = params.get('device').value

    detect_start(input, weights, outpath, device, id)


@start_download_pt()
def Export_model_RODY(params_str):
    from app.yolov5.export import Start_Model_Export
    import zipfile
    params = TrainParams()
    params.read_from_str(params_str)
    exp_inputPath = params.get('exp_inputPath').value  # 模型路径
    print('输入模型：', exp_inputPath)
    exp_device = params.get('device').value
    imgsz = params.get('imgsz').value
    modellist = Start_Model_Export(exp_inputPath, exp_device, imgsz)
    exp_outputPath = exp_inputPath.replace('pt', 'zip')  # 压缩文件
    print('模型路径：',exp_outputPath)
    zipf = zipfile.ZipFile(exp_outputPath, 'w')
    for file in modellist:
        zipf.write(file, arcname=Path(file).name)  # 将torchscript和onnx模型压缩

    return exp_outputPath

@obtain_train_param()
def returnTrainParams():
    nvmlInit()
    gpuDeviceCount = nvmlDeviceGetCount()  # 获取Nvidia GPU块数
    _kernel = [f"cuda:{a}" for a in range(gpuDeviceCount)]
    params_list = [
        {"index": 0, "name": "epochnum", "value": 10, "description": '训练轮次', "default": 100, "type": "I", 'show': True},
        {"index": 1, "name": "batch_size", "value": 4, "description": '批次图像数量', "default": 1, "type": "I",
         'show': True},
        {"index": 2, "name": "img_size", "value": 640, "description": '训练图像大小', "default": 640, "type": "I",
         'show': True},
        {"index": 3, "name": "device", "value": 'CUDA', "description": '训练核心', "default": 'CUDA', "type": "E",
         "items": _kernel, 'show': True},  # _kernel
        {"index": 4, "name": "saveModDir", "value": "E:/alg_demo-master/alg_demo/app/yolov5/best.pt",
         "description": '保存模型路径',
         "default": "./app/maskrcnn/saved_model/test.pt", "type": "S", 'show': False},
        {"index": 5, "name": "resumeModPath", "value": '/yolov5s.pt',
         "description": '继续训练路径', "default": '', "type": "S",
         'show': False},
        {"index": 6, "name": "resumeMod", "value": '', "description": '继续训练模型', "default": '', "type": "E", "items": '',
         'show': True},
        {"index": 7, "name": "CLASS_NAMES", "value": ['hole', '456'], "description": '类别名称', "default": '', "type": "L",
         "items": '',
         'show': False},
        {"index": 8, "name": "DatasetDir", "value": "E:/aicheck/data_set/11442136178662604800/ori",
         "description": '数据集路径',
         "default": "./app/maskrcnn/datasets/test", "type": "S", 'show': False}  # ORI_PATH
    ]
    # {"index": 9, "name": "saveEpoch", "value": 2, "description": '保存模型轮次', "default": 2, "type": "I", 'show': True}]
    params_str = json.dumps(params_list)
    return params_str


@obtain_test_param()
def returnValidateParams():
    # nvmlInit()
    # gpuDeviceCount = nvmlDeviceGetCount()  # 获取Nvidia GPU块数
    # _kernel = [f"cuda:{a}" for a in range(gpuDeviceCount)]
    params_list = [
        {"index": 0, "name": "modPath", "value": "E:/alg_demo-master/alg_demo/app/yolov5/圆孔_123_RODY_1_640.pt",
         "description": '验证模型路径', "default": "./app/maskrcnn/saved_model/test.pt", "type": "S", 'show': False},
        {"index": 1, "name": "batch_size", "value": 1, "description": '批次图像数量', "default": 1, "type": "I",
         'show': False},
        {"index": 2, "name": "img_size", "value": 640, "description": '训练图像大小', "default": 640, "type": "I",
         'show': False},
        {"index": 3, "name": "outputPath", "value": 'E:/aicheck/data_set/11442136178662604800/val_results/',
         "description": '输出结果路径',
         "default": './app/maskrcnn/datasets/M006B_waibi/res', "type": "S", 'show': False},
        {"index": 4, "name": "device", "value": "0", "description": '训练核心', "default": "cuda", "type": "S",
         "items": '', 'show': False}  # _kernel
    ]
    # {"index": 9, "name": "saveEpoch", "value": 2, "description": '保存模型轮次', "default": 2, "type": "I", 'show': True}]
    params_str = json.dumps(params_list)
    return params_str


@obtain_detect_param()
def returnDetectParams():
    # nvmlInit()
    # gpuDeviceCount = nvmlDeviceGetCount()  # 获取Nvidia GPU块数
    # _kernel = [f"cuda:{a}" for a in range(gpuDeviceCount)]
    params_list = [
        {"index": 0, "name": "inputPath", "value": 'E:/aicheck/data_set/11442136178662604800/input/',
         "description": '输入图像路径', "default": './app/maskrcnn/datasets/M006B_waibi/JPEGImages', "type": "S",
         'show': False},
        {"index": 1, "name": "outputPath", "value": 'E:/aicheck/data_set/11442136178662604800/val_results/',
         "description": '输出结果路径',
         "default": './app/maskrcnn/datasets/M006B_waibi/res', "type": "S", 'show': False},
        {"index": 2, "name": "modPath", "value": "E:/alg_demo-master/alg_demo/app/yolov5/圆孔_123_RODY_1_640.pt",
         "description": '模型路径', "default": "./app/maskrcnn/saved_model/test.pt", "type": "S", 'show': False},
        {"index": 3, "name": "device", "value": "0", "description": '推理核', "default": "cpu", "type": "S",
         'show': False},
    ]
    # {"index": 9, "name": "saveEpoch", "value": 2, "description": '保存模型轮次', "default": 2, "type": "I", 'show': True}]
    params_str = json.dumps(params_list)
    return params_str


@obtain_download_pt_param()
def returnDownloadParams():
    params_list = [
        {"index": 0, "name": "exp_inputPath", "value": 'E:/alg_demo-master/alg_demo/app/yolov5/圆孔_123_RODY_1_640.pt',
         "description": '转化模型输入路径',
         "default": 'E:/alg_demo-master/alg_demo/app/yolov5/圆孔_123_RODY_1_640.pt/',
         "type": "S", 'show': False},
        {"index": 1, "name": "device", "value": 'gpu', "description": 'CPU或GPU', "default": 'gpu', "type": "S",
         'show': False},
        {"index": 2, "name": "imgsz", "value": 640, "description": '图像大小', "default": 640, "type": "I",
         'show': True}
    ]
    params_str = json.dumps(params_list)
    return params_str


if __name__ == '__main__':
    par = returnTrainParams()
    print(par)
    id='1'
    train_R0DY(par,id)