#!/usr/bin/env python3
import argparse
import subprocess
import sys
import os

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))

def get_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):
    """
    获取训练参数，返回字符串格式
    """
    params = []
    params.append(f"stations={station_id}")
    params.append(f"learning_rate={learning_rate}")
    params.append(f"batch_size={batch_size}")
    params.append(f"max_epochs={max_epochs}")
    params.append(f"hidden_dim={hidden_dim}")
    params.append(f"dropout={dropout}")
    params.append(f"patience={patience}")
    params.append(f"train_ratio={train_ratio}")
    
    if base_model_id:
        params.append(f"base_model_id={base_model_id}")
    
    return params

def get_time_range():
    """
    从API获取时间范围
    """
    import requests
    API_BASE_URL = "http://47.96.16.84:5006"
    url = f"{API_BASE_URL}/api/data/timerange"
    try:
        response = requests.get(url, timeout=30)
        response.raise_for_status()
        data = response.json()
        if data.get("success"):
            return data.get("start_time"), data.get("end_time")
    except Exception as e:
        pass
    return None, None

def main():
    parser = argparse.ArgumentParser(description="Handle user training requests with parameter confirmation")
    parser.add_argument("--station_id", type=str, required=True, 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)")
    
    args = parser.parse_args()
    
    # 获取时间范围
    start_time, end_time = get_time_range()
    
    # 获取训练参数
    params = get_training_params(
        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 start_time and end_time:
        params.append(f"start_time={start_time}")
        params.append(f"end_time={end_time}")
    
    # 输出格式：&**&key1=value1|key2=value2|...|keyN=valueN&**&
    param_string = "|".join(params)
    print(f"&**&{param_string}&**&")

if __name__ == "__main__":
    main()
