#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
模型比较脚本
调用统一API的模型比较接口
"""
import requests
import json
from datetime import datetime
import argparse
from typing import Optional, Dict, Any, List

def compare_models(
    model_ids: List[str],
    metrics: Optional[List[str]] = None
) -> Optional[Dict[str, Any]]:
    """
    比较多个模型的性能指标
    
    参数:
    model_ids: 模型ID列表
    metrics: 要比较的指标列表 (r2, rmse, mae, mape)，默认: 全部
    
    返回:
    接口返回的JSON响应
    """
    # API 端点 - 预留位置，等待后续接口开发
    url = "http://47.96.16.84:5001/api/unified/models/compare"
    
    # 构建请求体
    request_body = {
        "model_ids": model_ids
    }
    
    if metrics:
        request_body["metrics"] = metrics
    
    headers = {
        "Content-Type": "application/json"
    }
    
    try:
        print(f"正在比较模型性能...")
        print(f"API 端点: {url}")
        print(f"模型ID: {', '.join(model_ids)}")
        if metrics:
            print(f"比较指标: {', '.join(metrics)}")
        print()
        
        response = requests.post(url, json=request_body, headers=headers, timeout=30)
        response.raise_for_status()
        result = response.json()
        
        if result.get("success"):
            print("✓ 模型比较成功！")
            print()
            
            comparison_data = result.get("data", {}).get("comparison", [])
            
            if comparison_data:
                # 获取所有可用的指标
                all_metrics = set()
                for model in comparison_data:
                    all_metrics.update([k for k in model.keys() if k not in ['model_id', 'station_name', 'training_date']])
                
                # 构建表头
                headers_list = ['模型ID', '站点', '训练日期']
                if metrics:
                    headers_list.extend(metrics)
                else:
                    headers_list.extend(sorted(all_metrics))
                
                # 计算列宽
                col_widths = [max(len(h), 30) for h in headers_list]
                
                # 打印表格
                print("=" * (sum(col_widths) + len(col_widths) * 3 + 1))
                header_row = " | ".join([h.center(col_widths[i]) for i, h in enumerate(headers_list)])
                print(header_row)
                print("=" * (sum(col_widths) + len(col_widths) * 3 + 1))
                
                for model in comparison_data:
                    row_data = [
                        model.get('model_id', 'N/A')[:30],
                        model.get('station_name', 'N/A'),
                        model.get('training_date', 'N/A')
                    ]
                    
                    if metrics:
                        for metric in metrics:
                            value = model.get(metric, 'N/A')
                            if isinstance(value, float):
                                row_data.append(f"{value:.4f}")
                            else:
                                row_data.append(str(value))
                    else:
                        for metric in sorted(all_metrics):
                            value = model.get(metric, 'N/A')
                            if isinstance(value, float):
                                row_data.append(f"{value:.4f}")
                            else:
                                row_data.append(str(value))
                    
                    row_str = " | ".join([str(v).center(col_widths[i]) for i, v in enumerate(row_data)])
                    print(row_str)
                
                print("=" * (sum(col_widths) + len(col_widths) * 3 + 1))
                print()
                
                # 性能排名
                if 'r2' in [m.lower() for m in all_metrics]:
                    print("R² 排名 (从高到低):")
                    sorted_by_r2 = sorted(comparison_data, key=lambda x: x.get('r2', 0), reverse=True)
                    for i, model in enumerate(sorted_by_r2, 1):
                        print(f"  {i}. {model.get('model_id', 'N/A')}: R² = {model.get('r2', 'N/A'):.6f}")
                    print()
                
                # 输出完整的 JSON 格式
                print("=" * 80)
                print("完整信息 (JSON 格式)")
                print("=" * 80)
                print(json.dumps(result, indent=2, ensure_ascii=False))
            else:
                print("未找到匹配的模型")
            
            return result
        else:
            print(f"✗ 模型比较失败: {result.get('message', '未知错误')}")
            print(f"返回内容: {json.dumps(result, indent=2, ensure_ascii=False)}")
            return None
            
    except requests.exceptions.Timeout:
        print(f"✗ 请求超时（30秒）")
        return None
    except requests.exceptions.ConnectionError:
        print(f"✗ 无法连接到 API 服务器: {url}")
        print(f"请确保服务器地址正确且服务正在运行")
        return None
    except requests.exceptions.HTTPError as e:
        print(f"✗ HTTP 错误: {e}")
        print(f"状态码: {response.status_code}")
        print(f"响应内容: {response.text}")
        return None
    except Exception as e:
        print(f"✗ 模型比较失败: {e}")
        import traceback
        traceback.print_exc()
        return None

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='模型比较脚本')
    parser.add_argument('--model_ids', type=str, required=True, 
                        help='逗号分隔的模型ID列表 (例如: model_v1,model_v2,model_v3)')
    parser.add_argument('--metrics', type=str, default=None,
                        help='逗号分隔的指标列表 (r2, rmse, mae, mape)，默认: 全部')
    
    args = parser.parse_args()
    
    model_ids = [m.strip() for m in args.model_ids.split(',')]
    metrics = None
    if args.metrics:
        metrics = [m.strip() for m in args.metrics.split(',')]
    
    result = compare_models(
        model_ids=model_ids,
        metrics=metrics
    )
    
    if result is None:
        exit(1)