#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
负荷预测模型评估脚本
功能：获取模型列表、评估模型性能、生成评估报告
"""

import requests
import argparse
import json
import time
import os
import subprocess
from datetime import datetime, timedelta
import numpy as np

# 正确的路径计算
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
SKILL_DIR = os.path.dirname(SCRIPT_DIR)
# WORKSPACE_DIR应该是skills的父目录
WORKSPACE_DIR = os.path.dirname(os.path.dirname(SKILL_DIR))

# API基础地址
API_BASE_URL = "http://47.96.16.84:5005"


def get_time_range():
    """
    获取数据的时间范围
    """
    url = f"{API_BASE_URL}/api/load-forecast/data/date-range"
    try:
        response = requests.get(url)
        response.raise_for_status()
        data = response.json()
        if data.get('code') == 200:
            return data['data']
        else:
            print(f"获取时间范围失败: {data.get('error')}")
            return None
    except Exception as e:
        print(f"获取时间范围异常: {e}")
        return None


def list_models():
    """
    获取模型列表
    """
    url = f"{API_BASE_URL}/api/load-forecast/model/list"
    try:
        response = requests.get(url)
        response.raise_for_status()
        data = response.json()
        if data.get('code') == 200:
            models = data['data']
            print("\n可用模型列表:")
            print("-" * 80)
            for model in models:
                print(f"模型名称: {model.get('name')}")
                print(f"模型路径: {model.get('path')}")
                print(f"模型类型: {model.get('model_type')}")
                print(f"修改时间: {model.get('mtime')}")
                print("-" * 80)
            return models
        else:
            print(f"获取模型列表失败: {data.get('error')}")
            return []
    except Exception as e:
        print(f"获取模型列表异常: {e}")
        return []


def load_model(model_dir):
    """
    加载模型
    """
    url = f"{API_BASE_URL}/api/load-forecast/model/load"
    data = {
        "model_dir": model_dir,
        "model_type": "cnn_rbilstm"
    }
    try:
        response = requests.post(url, json=data)
        response.raise_for_status()
        result = response.json()
        if result.get('code') == 200:
            print("模型加载成功")
            return True
        else:
            print(f"模型加载失败: {result.get('error')}")
            return False
    except Exception as e:
        print(f"模型加载异常: {e}")
        return False


def predict_multi(start_time, prediction_hours=24):
    """
    多步预测
    """
    url = f"{API_BASE_URL}/api/load-forecast/predict/multi"
    data = {
        "start_time": start_time,
        "prediction_hours": prediction_hours,
        "model_type": "cnn_rbilstm"
    }
    try:
        response = requests.post(url, json=data)
        response.raise_for_status()
        result = response.json()
        if result.get('code') == 200:
            return result['data']
        else:
            print(f"预测失败: {result.get('error')}")
            return None
    except Exception as e:
        print(f"预测异常: {e}")
        return None


def calculate_metrics(actual, predicted):
    """
    计算评估指标
    """
    actual = np.array(actual)
    predicted = np.array(predicted)
    
    # 计算MAE
    mae = np.mean(np.abs(actual - predicted))
    
    # 计算RMSE
    rmse = np.sqrt(np.mean((actual - predicted) ** 2))
    
    # 计算R²
    mean_actual = np.mean(actual)
    ss_res = np.sum((actual - predicted) ** 2)
    ss_tot = np.sum((actual - mean_actual) ** 2)
    r2 = 1 - (ss_res / ss_tot) if ss_tot != 0 else 0
    
    # 计算MAPE
    non_zero_actual = actual[actual != 0]
    non_zero_predicted = predicted[actual != 0]
    if len(non_zero_actual) > 0:
        mape = np.mean(np.abs((non_zero_actual - non_zero_predicted) / non_zero_actual)) * 100
    else:
        mape = 0
    
    return {
        'mae': mae,
        'rmse': rmse,
        'r2': r2,
        'mape': mape
    }


def generate_report(metrics, predictions, user_query="模型评估报告"):
    """
    生成评估报告
    """
    # 准备报告数据
    report_data = []
    for pred in predictions:
        if pred.get('actual') is not None:
            report_data.append({
                "time": pred.get('time'),
                "actual": pred.get('actual'),
                "predicted": pred.get('power'),
                "difference": pred.get('power') - pred.get('actual')
            })
    
    # 调用step1_select.py选择模板
    step1_script = f"{WORKSPACE_DIR}/report-generator/scripts/step1_select.py"
    if not os.path.exists(step1_script):
        # 尝试使用相对路径
        step1_script = os.path.join(os.path.dirname(SKILL_DIR), "report-generator", "scripts", "step1_select.py")
        if not os.path.exists(step1_script):
            raise Exception(f"Step1 script not found at {step1_script}")
    
    step1_cmd = ["python", step1_script, "--query", user_query, "--data", json.dumps(report_data, ensure_ascii=False)]
    step1_result = subprocess.run(step1_cmd, capture_output=True, text=True, timeout=60)
    
    if step1_result.returncode != 0:
        raise Exception(f"Step 1 failed: {step1_result.stderr}")
    
    lines = step1_result.stdout.strip().split('\n')
    step1_output = None
    for line in reversed(lines):
        try:
            step1_output = json.loads(line)
            break
        except:
            continue
    
    if step1_output is None or step1_output.get("status") != "success":
        raise Exception(f"Template selection failed")
    
    template_name = step1_output.get("template_name", "dataTable")
    
    # 调用step2_render.py生成报告
    step2_script = f"{WORKSPACE_DIR}/report-generator/scripts/step2_render.py"
    if not os.path.exists(step2_script):
        # 尝试使用相对路径
        step2_script = os.path.join(os.path.dirname(SKILL_DIR), "report-generator", "scripts", "step2_render.py")
        if not os.path.exists(step2_script):
            raise Exception(f"Step2 script not found at {step2_script}")
    
    step2_cmd = ["python", step2_script, "--query", user_query, "--template", template_name, "--data", json.dumps(report_data, ensure_ascii=False)]
    step2_result = subprocess.run(step2_cmd, capture_output=True, text=True, timeout=60)
    
    if step2_result.returncode != 0:
        raise Exception(f"Step 2 failed: {step2_result.stderr}")
    
    lines = step2_result.stdout.strip().split('\n')
    step2_output = None
    for line in reversed(lines):
        try:
            step2_output = json.loads(line)
            break
        except:
            continue
    
    if step2_output is None or step2_output.get("status") != "success":
        raise Exception(f"Report generation failed")
    
    report_address = step2_output.get('address', '')
    summary = step2_output.get('summary', '')
    
    print("\nGENERATING REPORT")
    print("=" * 60)
    print("PROGRESSION: ⏳ 正在分析数据，选择合适的可视化模板...")
    print(f"Selected template: {template_name}")
    print("PROGRESSION: ⏳ 正在将数据渲染为图表报告...")
    print("")
    print("✓ Report generated successfully!")
    
    print(f"Summary: {summary}")
    print("")
    print(f"@@@@@{report_address}@@@@@")
    print("")
    print("#####%查看详细数据%￥点击上方链接查看完整的数据分析报告￥ %继续深入分析%￥基于当前数据进行更深度的挖掘￥ %导出报告%￥将当前报告导出为PDF或图片格式￥#####")
    
    return report_address


def evaluate_model(model_id, mae_threshold=100.0, rmse_threshold=150.0, r2_threshold=0.70, generate_report=True):
    """
    评估模型
    """
    print("Step 1: Getting time range...")
    time_range = get_time_range()
    if not time_range:
        return False
    
    # 获取最后一天作为预测目标时间
    end_time = time_range['end_time']
    target_date = datetime.strptime(end_time, '%Y-%m-%d %H:%M:%S')
    target_time = target_date.strftime('%Y-%m-%d 00:00:00')
    print(f"Target time for prediction: {target_time}")
    
    print(f"Step 2: Loading model {model_id}...")
    # 这里需要根据model_id找到对应的模型路径
    models = list_models()
    model_path = None
    for model in models:
        if model_id in model['name']:
            model_path = model['path']
            break
    
    if not model_path:
        print(f"未找到模型: {model_id}")
        return False
    
    if not load_model(model_path):
        return False
    
    print(f"Step 3: Predicting for {target_time} with 96 steps...")
    prediction_result = predict_multi(target_time, 24)
    if not prediction_result:
        return False
    
    print("Step 4: Calculating metrics...")
    predictions = prediction_result['predictions']
    
    # 提取实际值和预测值
    actual_values = []
    predicted_values = []
    
    for pred in predictions:
        if pred.get('actual') is not None:
            actual_values.append(pred['actual'])
            predicted_values.append(pred['power'])
    
    if not actual_values:
        print("没有实际值用于评估")
        return False
    
    # 计算指标
    metrics = calculate_metrics(actual_values, predicted_values)
    print("Metrics:")
    print(f"  MAE:  {metrics['mae']:.2f} kW")
    print(f"  RMSE: {metrics['rmse']:.2f} kW")
    print(f"  R2:   {metrics['r2']:.2f}")
    print(f"  MAPE: {metrics['mape']:.2f} %")
    print("")
    
    # 评估预测效果
    if metrics['mae'] <= mae_threshold and metrics['rmse'] <= rmse_threshold and metrics['r2'] >= r2_threshold:
        print("  OK: Prediction performance is acceptable")
    else:
        print("  WARNING: Prediction performance is not acceptable")
        print("  Suggestion: Consider retraining the model")
    
    print("")
    print("=" * 60)
    print("EVALUATION SUMMARY")
    print("=" * 60)
    print("")
    
    if metrics['mae'] <= mae_threshold and metrics['rmse'] <= rmse_threshold and metrics['r2'] >= r2_threshold:
        print("Prediction performance is acceptable.")
    else:
        print("Prediction performance is not acceptable. Consider retraining.")
    
    print("")
    print("=" * 60)
    
    # 生成报告
    if generate_report:
        try:
            generate_report(metrics, predictions, f"模型{model_id}的评估报告")
        except Exception as e:
            print(f"报告生成失败: {e}")
    
    return True


def main():
    parser = argparse.ArgumentParser(description='负荷预测模型评估脚本')
    parser.add_argument('--action', type=str, required=True, choices=['list_models', 'evaluate_model'],
                        help='操作类型')
    parser.add_argument('--model_id', type=str, help='模型ID')
    parser.add_argument('--mae_threshold', type=float, default=100.0, help='MAE阈值')
    parser.add_argument('--rmse_threshold', type=float, default=150.0, help='RMSE阈值')
    parser.add_argument('--r2_threshold', type=float, default=0.70, help='R²阈值')
    parser.add_argument('--generate_report', type=bool, default=True, help='是否生成报告')
    
    args = parser.parse_args()
    
    if args.action == 'list_models':
        list_models()
    elif args.action == 'evaluate_model':
        if not args.model_id:
            print('错误：必须指定模型ID')
            return
        evaluate_model(args.model_id, args.mae_threshold, args.rmse_threshold, args.r2_threshold, args.generate_report)


if __name__ == '__main__':
    main()