import argparse
import requests
import json
import time
import sys
from datetime import datetime

API_BASE_URL = "http://47.96.16.84:5006"

def get_time_range():
    """获取数据的可用时间范围"""
    url = f"{API_BASE_URL}/api/data/timerange"
    try:
        response = requests.get(url, timeout=30)
        response.raise_for_status()
        return response.json()
    except Exception as e:
        print(f"Error getting time range: {e}")
        return None

def start_training(station_id, base_model_id=None, learning_rate=0.001, batch_size=64, 
                   max_epochs=100, hidden_dim=64, dropout=0.1, patience=15, train_ratio=0.7):
    """
    启动模型训练任务
    """
    url = f"{API_BASE_URL}/api/training/start"
    
    # 获取时间范围
    time_range_data = get_time_range()
    if not time_range_data or not time_range_data.get("success"):
        print("Failed to get time range.")
        return None
    
    start_time = time_range_data.get("start_time", "2024-01-01 00:00:00")
    end_time = time_range_data.get("end_time")
    
    if not end_time:
        print("Invalid time range data.")
        return None
    
    payload = {
        "stations": [station_id],
        "start_time": start_time,
        "end_time": end_time,
        "learning_rate": learning_rate,
        "batch_size": batch_size,
        "max_epochs": max_epochs,
        "hidden_dim": hidden_dim,
        "dropout": dropout,
        "patience": patience,
        "train_ratio": train_ratio
    }
    
    # 如果提供了基础模型ID，添加到payload中
    if base_model_id:
        payload["base_model_id"] = base_model_id
    
    try:
        response = requests.post(url, json=payload, timeout=30)
        response.raise_for_status()
        return response.json()
    except Exception as e:
        print(f"Error starting training: {e}")
        return None

def get_training_status():
    """获取当前训练任务的状态"""
    url = f"{API_BASE_URL}/api/training/status"
    try:
        response = requests.get(url, timeout=30)
        response.raise_for_status()
        return response.json()
    except Exception as e:
        print(f"Error getting training status: {e}")
        return None

def monitor_training(check_interval=10, max_wait_time=3600):
    """
    监控训练进度
    """
    start_time = time.time()
    
    while True:
        elapsed = time.time() - start_time
        if elapsed > max_wait_time:
            print(f"Training monitoring timeout after {max_wait_time} seconds")
            return False
            
        status = get_training_status()
        if not status:
            print("Failed to get training status")
            return False
            
        is_training = status.get("is_training", False)
        current_epoch = status.get("current_epoch", 0)
        max_epochs = status.get("max_epochs", 100)
        train_loss = status.get("train_loss", 0)
        val_loss = status.get("val_loss", 0)
        progress = status.get("progress", 0)
        
        print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] "
              f"Epoch {current_epoch}/{max_epochs} | "
              f"Progress: {progress:.1f}% | "
              f"Train Loss: {train_loss:.6f} | "
              f"Val Loss: {val_loss:.6f}")
        
        if not is_training:
            print("Training completed!")
            return True
            
        time.sleep(check_interval)

def format_training_params(station_id, base_model_id=None, learning_rate=0.001, batch_size=64, 
                          max_epochs=100, hidden_dim=64, dropout=0.1, patience=15, train_ratio=0.7, 
                          start_time=None, end_time=None):
    """
    格式化训练参数为JSON格式，用于用户二次确认
    包含所有默认值
    """
    params = {
        "stations": [station_id],
        "start_time": start_time,
        "end_time": end_time,
        "learning_rate": learning_rate,
        "batch_size": batch_size,
        "max_epochs": max_epochs,
        "hidden_dim": hidden_dim,
        "dropout": dropout,
        "patience": patience,
        "train_ratio": train_ratio
    }
    
    # 如果提供了基础模型ID，添加到参数中
    if base_model_id:
        params["base_model_id"] = base_model_id
    
    return params

def main():
    parser = argparse.ArgumentParser(description="Train PV prediction model for a specific station")
    parser.add_argument("--action", type=str, required=True, 
                       choices=["train_station", "check_status", "train_and_monitor", "preview_params"],
                       help="Action to perform")
    parser.add_argument("--station_id", type=str, help="Station ID to train")
    parser.add_argument("--base_model_id", type=str, help="Base model ID for transfer learning (optional)")
    parser.add_argument("--learning_rate", type=float, default=0.001, help="Learning rate (default: 0.001)")
    parser.add_argument("--batch_size", type=int, default=64, help="Batch size (default: 64)")
    parser.add_argument("--max_epochs", type=int, default=100, help="Maximum epochs (default: 100)")
    parser.add_argument("--hidden_dim", type=int, default=64, help="Hidden dimension (default: 64)")
    parser.add_argument("--dropout", type=float, default=0.1, help="Dropout rate (default: 0.1)")
    parser.add_argument("--patience", type=int, default=15, help="Early stopping patience (default: 15)")
    parser.add_argument("--train_ratio", type=float, default=0.7, help="Training data ratio (default: 0.7)")
    parser.add_argument("--check_interval", type=int, default=10, help="Status check interval in seconds (default: 10)")
    parser.add_argument("--max_wait_time", type=int, default=3600, help="Maximum wait time in seconds (default: 3600)")
    
    args = parser.parse_args()
    
    if args.action == "preview_params":
        """预览训练参数，供用户二次确认"""
        if not args.station_id:
            print("Error: --station_id is required for preview_params action")
            sys.exit(1)
            
        print(f"Step 1: Getting available time range...")
        time_range_data = get_time_range()
        if not time_range_data or not time_range_data.get("success"):
            print("Failed to get time range.")
            sys.exit(1)
            
        start_time = time_range_data.get("start_time")
        end_time = time_range_data.get("end_time")
        
        if not start_time or not end_time:
            print("Invalid time range data.")
            sys.exit(1)
        
        # 格式化参数（使用实际值，包含默认值）
        params = format_training_params(
            station_id=args.station_id,
            base_model_id=args.base_model_id,
            learning_rate=args.learning_rate,
            batch_size=args.batch_size,
            max_epochs=args.max_epochs,
            hidden_dim=args.hidden_dim,
            dropout=args.dropout,
            patience=args.patience,
            train_ratio=args.train_ratio,
            start_time=start_time,
            end_time=end_time
        )
        
        # 输出参数预览信息
        print("\n" + "="*70)
        print("TRAINING PARAMETERS PREVIEW")
        print("="*70)
        print("\nTraining parameters for station:", args.station_id)
        print(json.dumps(params, indent=2, ensure_ascii=False))
        print("\n" + "="*70)
        print("\nPlease confirm the parameters above.")
        print("\nConfirmation format (copy this to confirm):")
        print(f"&**&{json.dumps(params, ensure_ascii=False)}&**&")
        print("="*70)
        
    elif args.action == "train_station":
        if not args.station_id:
            print("Error: --station_id is required for train_station action")
            sys.exit(1)
            
        print(f"Step 1: Getting available time range...")
        time_range_data = get_time_range()
        if not time_range_data or not time_range_data.get("success"):
            print("Failed to get time range.")
            sys.exit(1)
            
        start_time = time_range_data.get("start_time")
        end_time = time_range_data.get("end_time")
        
        if not start_time or not end_time:
            print("Invalid time range data.")
            sys.exit(1)
            
        print(f"Time range: {start_time} to {end_time}")
        
        print(f"\nStep 2: Starting training for station {args.station_id}...")
        train_result = start_training(
            args.station_id,
            base_model_id=args.base_model_id,
            learning_rate=args.learning_rate,
            batch_size=args.batch_size,
            max_epochs=args.max_epochs,
            hidden_dim=args.hidden_dim,
            dropout=args.dropout,
            patience=args.patience,
            train_ratio=args.train_ratio
        )
        
        if not train_result or not train_result.get("success"):
            print("Failed to start training.")
            print(json.dumps(train_result, indent=2, ensure_ascii=False))
            sys.exit(1)
            
        print(f"Training started successfully!")
        print(json.dumps(train_result, indent=2, ensure_ascii=False))
        
    elif args.action == "check_status":
        print("Checking training status...")
        status = get_training_status()
        if status:
            print(json.dumps(status, indent=2, ensure_ascii=False))
        else:
            print("Failed to get training status.")
            sys.exit(1)
            
    elif args.action == "train_and_monitor":
        if not args.station_id:
            print("Error: --station_id is required for train_and_monitor action")
            sys.exit(1)
            
        print(f"Step 1: Getting available time range...")
        time_range_data = get_time_range()
        if not time_range_data or not time_range_data.get("success"):
            print("Failed to get time range.")
            sys.exit(1)
            
        start_time = time_range_data.get("start_time")
        end_time = time_range_data.get("end_time")
        
        print(f"Time range: {start_time} to {end_time}")
        
        print(f"\nStep 2: Starting training for station {args.station_id}...")
        train_result = start_training(
            args.station_id,
            base_model_id=args.base_model_id,
            learning_rate=args.learning_rate,
            batch_size=args.batch_size,
            max_epochs=args.max_epochs,
            hidden_dim=args.hidden_dim,
            dropout=args.dropout,
            patience=args.patience,
            train_ratio=args.train_ratio
        )
        
        if not train_result or not train_result.get("success"):
            print("Failed to start training.")
            print(json.dumps(train_result, indent=2, ensure_ascii=False))
            sys.exit(1)
            
        print("Training started successfully!")
        
        print(f"\nStep 3: Monitoring training progress...")
        success = monitor_training(args.check_interval, args.max_wait_time)
        
        if success:
            print("\nStep 4: Getting final training status...")
            final_status = get_training_status()
            if final_status:
                print(json.dumps(final_status, indent=2, ensure_ascii=False))
        else:
            print("Training monitoring failed or timed out.")
            sys.exit(1)

if __name__ == "__main__":
    main()