# ==================== Energy 任务模块 ====================
"""
Energy 任务相关的所有函数和界面组件
支持多用户并发：使用 gr.State 管理每个用户会话的状态
使用统一进度管理模块存储数据
"""
import json
import os
from typing import List, Tuple, Optional, Dict, Any
import gradio as gr

# 导入统一进度管理模块
import progress_manager

# 导入 Energy 环境
import sys
current_dir = os.path.dirname(os.path.abspath(__file__))
energyenv_path = os.path.join(current_dir, "EnergyEnv")
if os.path.exists(energyenv_path):
    sys.path.insert(0, energyenv_path)
from EnergyEnv_v5 import DynamicEnergyGrid

# ------------------- 常量 -------------------
ENERGY_MAX_STEPS = 120

# ------------------- 示例文本 -------------------
ENERGY_EXAMPLE_TEXT = """
## 📖 Energy Environment Usage Example

### Scenario Description
You need to manage an energy grid, balancing generation, demand, and budget while meeting stability and carbon emission targets, completing at least 120 days of tasks. If demand default or budget default occurs for three consecutive days, the task will fail directly.

### Task Objectives
- **Completion Days**: Complete at least 120 days
- **Stability Target**: Final average stability must ≥ target value (shown in state)
- **Carbon Emission Target**: Final carbon emission ratio must ≤ target value (shown in state)
- **Default Limit**: 3 consecutive days of demand default or budget default will cause task failure

### Available Operations
- **🔥 Thermal**: Input thermal power generation (≥0)
- **💨 Wind**: Input wind power generation (≥0)
- **☀️ Solar**: Input solar power generation (≥0)
- **🔋 Battery**: Input battery operation
  - Negative value = Charge (e.g., -20)
  - Positive value = Discharge (e.g., 20)
  - 0 = Do not use battery
  - Battery has maximum capacity limit of 80

### Actual Generation Calculation
- Actual generation = Input generation × Efficiency coefficient
- After actual generation, store to battery, no loss at this stage
- For example, input thermal 10, wind 20, solar 30, battery storage 10. Thermal efficiency 0.9, wind efficiency 1.1, solar efficiency 1
- Then actual generation = 10×0.9 + 20×1.1 + 30×1 = 61
- Applied to grid (subtract battery storage): 61 - 10 = 51
- **Note**: Thermal efficiency fluctuates randomly around 1; wind and solar efficiencies repeat in periodic functions with small fluctuations

### Stability Requirements
- Daily generation configuration changes cannot be too large, otherwise it will cause grid instability
- Stability calculation considers: magnitude of generation configuration changes (ramping), budget default, demand default
- If budget default or demand default occurs, stability will be greatly reduced
- **Important**: Insufficient stability will not directly terminate the task, but will be used to judge task success after final completion. So you need to adjust strategy in time to improve stability

### Carbon Emission Requirements
- Carbon emission ratio = Historical cumulative thermal actual generation / Historical cumulative total actual generation
- When task is completed, carbon emission ratio must ≤ target value
- Need to control thermal power proportion of all generation throughout the task
- **Important**: Excessive carbon emissions will not directly terminate the task, but will be used to judge task success after final completion. So you need to adjust strategy in time to reduce carbon emissions

### Default Explanation
- **Demand Default**: Actual supply < Demand
- **Budget Default**: Actual cost > Budget
- Insufficient stability or excessive carbon emissions do not count as defaults
- Three consecutive days of defaults will directly terminate and fail the task
- **Important**: Only demand default and budget default will increase consecutive default days. Insufficient stability and excessive carbon emissions do not count as defaults but affect final results

### Initial Configuration
- Day 1 will show initial generation configuration, which is the system's initial state
- Your Day 1 operation should refer to this initial configuration to avoid excessive changes affecting stability

## Example
### Scenario Description
- Thermal, wind, solar unit prices are 2, 4, 6 per unit respectively, battery operation cost 0.1 per unit
- Carbon emission ratio target ≤ 0.81 (i.e., thermal proportion ≤ 0.19)
- Stability target ≥ 0.5
- This example demonstrates 6 days, actual task requires completing 120 days

### Example Logic (Only shown in examples. In actual tasks, these rules are hidden and need to be inferred by users)
- Thermal efficiency sequence: [1.0, 1.0, 1.0, 0.9, 1.1, 1.0] (randomly fluctuates around 1)
- Wind efficiency sequence: [1.1, 1.0, 1.1, 1.0, 1.1, 1.0] (cycle every 2 days)
- Solar efficiency sequence: [0.9, 1.0, 1.1, 0.9, 1.0, 1.1] (cycle every 3 days)

### Important Tips
- In actual tasks, efficiency coefficients are hidden and need to be inferred from historical data
- Need to balance cost, stability, carbon emissions, and demand satisfaction
- Insufficient stability and excessive carbon emissions will not directly terminate the task but will affect final task completion conditions
- Only demand default and budget default will increase consecutive default days. 3 consecutive days of defaults will cause task failure
- When defaults occur, need to adjust strategy in time to avoid consecutive defaults
- In actual problems, you cannot see the specific calculation process of stability coefficient, you can only see a result. Please adjust strategy based on this result
"""


# ------------------- 状态管理 -------------------

def create_energy_state() -> Dict[str, Any]:
    """创建初始的 Energy 任务状态（每个用户会话独立）"""
    return {
        'env': None,                    # DynamicEnergyGrid 实例
        'test_data': [],                # 测试数据
        'current_env_idx': 0,           # 当前环境索引
        'history_records': [],          # 操作历史记录
        'last_step_violations': {       # 上一步的违约信息
            "demand_violation": False,
            "budget_violation": False,
            "stability_violation": False,
            "violation_days_cont": 0
        }
    }


# ------------------- 工具函数 -------------------

def get_energy_steps_info(state: Dict[str, Any]) -> str:
    """获取 Energy 任务的步数信息（包含天数，天数从1开始显示）"""
    env = state.get('env')
    history_records = state.get('history_records', [])
    executed_steps = len(history_records)
    # 显示当前步数（已执行步数 + 1，但不超过最大值）
    current_step = min(executed_steps + 1, ENERGY_MAX_STEPS)
    if env is not None:
        # env.t 是当前所在的天数（从0开始），显示时加1
        current_day = env.t + 1
        return f"{current_step} / {ENERGY_MAX_STEPS} (Day {current_day})"
    else:
        return f"{current_step} / {ENERGY_MAX_STEPS} (Day 1)"


def calculate_estimated_cost(state: Dict[str, Any], thermal: float, wind: float, solar: float, battery: float) -> str:
    """计算预计支出"""
    env = state.get('env')
    test_data = state.get('test_data', [])
    current_env_idx = state.get('current_env_idx', 0)
    
    # 处理 None 值
    thermal = float(thermal) if thermal is not None else 0.0
    wind = float(wind) if wind is not None else 0.0
    solar = float(solar) if solar is not None else 0.0
    battery = float(battery) if battery is not None else 0.0
    
    # 确保非负
    thermal = max(0.0, thermal)
    wind = max(0.0, wind)
    solar = max(0.0, solar)
    
    # 获取价格信息
    prices = None
    battery_op_cost = 0.1
    
    if env is not None:
        prices = env.prices
        battery_op_cost = env.battery_op_cost
    elif test_data and current_env_idx < len(test_data):
        config = test_data[current_env_idx]
        prices = config.get("prices", {})
        battery_op_cost = 0.1
    
    if prices is None:
        prices = {"thermal": 3.0, "wind": 5.0, "solar": 6.0}
    
    # 计算各项成本
    thermal_cost = thermal * prices.get("thermal", 3.0)
    wind_cost = wind * prices.get("wind", 5.0)
    solar_cost = solar * prices.get("solar", 6.0)
    battery_cost = abs(battery) * battery_op_cost
    
    total_cost = thermal_cost + wind_cost + solar_cost + battery_cost
    total_generation = thermal + wind + solar + battery
    
    # 获取今日预算和需求
    budget_today = None
    demand_today = None
    if env is not None:
        t = min(env.t, env.horizon - 1)
        if t < len(env.budget_series):
            budget_today = env.budget_series[t]
        if t < len(env.demand_series):
            demand_today = env.demand_series[t]
    
    # 获取电池当前电量
    battery_cur = None
    if env is not None:
        battery_cur = env.battery_cur
    elif test_data and current_env_idx < len(test_data):
        battery_cur = 0.0
    
    lines = []
    if battery_cur is not None:
        lines.append(f"🔋 Battery current charge: {battery_cur:.2f} MW")
    
    lines.append("\n⚡ Current Total Generation:")
    lines.append(f"  Thermal: {thermal:.2f} MW")
    lines.append(f"  Wind: {wind:.2f} MW")
    lines.append(f"  Solar: {solar:.2f} MW")
    lines.append(f"  Battery: {battery:.2f} MW {'(charging)' if battery < 0 else '(discharging)' if battery > 0 else ''}")
    lines.append(f"  Total: {total_generation:.2f} MW")
    
    if demand_today is not None:
        lines.append(f"  Today's demand: {demand_today:.2f} MW")
        if total_generation < demand_today:
            lines.append(f"  ⚠️ Insufficient supply: {demand_today - total_generation:.2f} MW")
        elif total_generation > demand_today:
            lines.append(f"  ✅ Sufficient supply: excess {total_generation - demand_today:.2f} MW")
        else:
            lines.append(f"  ✅ Supply-demand balance")
    
    lines.append("\n💰 Estimated Cost:")
    lines.append(f"  Thermal: {thermal_cost:.2f} (Unit price: {prices.get('thermal', 3.0):.2f} × {thermal:.2f})")
    lines.append(f"  Wind: {wind_cost:.2f} (Unit price: {prices.get('wind', 5.0):.2f} × {wind:.2f})")
    lines.append(f"  Solar: {solar_cost:.2f} (Unit price: {prices.get('solar', 6.0):.2f} × {solar:.2f})")
    lines.append(f"  Battery: {battery_cost:.2f} (Unit price: {battery_op_cost:.2f} × {abs(battery):.2f})")
    lines.append(f"  Total: {total_cost:.2f}")
    
    if budget_today is not None:
        lines.append(f"\n📊 Today's budget: {budget_today:.2f}")
        if total_cost > budget_today:
            lines.append(f"⚠️ Over budget: {total_cost - budget_today:.2f}")
        else:
            lines.append(f"✅ Budget remaining: {budget_today - total_cost:.2f}")
    
    return "\n".join(lines)


def format_energy_history_record(step_num: int, day: int, action: Dict[str, Any], obs: Dict[str, Any], feedback_msg: str, reward: float, error: str = None, demand: float = None) -> str:
    """格式化单步历史记录
    Args:
        step_num: 步骤编号
        day: 当前天数（从0开始）
        action: 输入的动作字典
        obs: 执行动作后的观察（包含实际发电量）
        feedback_msg: 反馈消息
        reward: 奖励
        error: 错误信息（如果有）
        demand: 当天的需求（可选）
    """
    lines = []
    lines.append(f"Step {step_num} (Day {day + 1}):")
    
    # Input generation
    lines.append("Input generation:")
    lines.append(f"  Thermal: {action.get('thermal', 0):.2f}")
    lines.append(f"  Wind: {action.get('wind', 0):.2f}")
    lines.append(f"  Solar: {action.get('solar', 0):.2f}")
    lines.append(f"  Battery: {action.get('battery', 0):.2f}")
    
    # Actual generation (from obs if available)
    actual_prev = obs.get('actual_prev', {})
    if actual_prev:
        lines.append("Actual generation:")
        lines.append(f"  Thermal: {actual_prev.get('thermal', 0):.2f}")
        lines.append(f"  Wind: {actual_prev.get('wind', 0):.2f}")
        lines.append(f"  Solar: {actual_prev.get('solar', 0):.2f}")
        battery_flow = actual_prev.get('battery', 0)
        if battery_flow < 0:
            lines.append(f"  Battery: {abs(battery_flow):.2f} (charging)")
        elif battery_flow > 0:
            lines.append(f"  Battery: {battery_flow:.2f} (discharging)")
        else:
            lines.append(f"  Battery: 0.00")
        supply = actual_prev.get('supply', 0)
        lines.append(f"  Total supply: {supply:.2f}")
        # Display total demand
        if demand is not None:
            lines.append(f"  Total demand: {demand:.2f}")
            if supply < demand:
                lines.append(f"  ⚠️ Insufficient supply: {demand - supply:.2f}")
            elif supply > demand:
                lines.append(f"  ✅ Sufficient supply: excess {supply - demand:.2f}")
            else:
                lines.append(f"  ✅ Supply-demand balance")
    
    # Feedback
    if error:
        lines.append(f"Feedback: ❌ {error}")
    else:
        lines.append(f"Feedback: {feedback_msg}, Reward={reward:.2f}")
    
    return "\n".join(lines)


def format_energy_state(state: Dict[str, Any], obs: Dict[str, Any], last_violations: Optional[Dict[str, Any]] = None, 
                       thermal_input: Optional[float] = None, wind_input: Optional[float] = None, 
                       solar_input: Optional[float] = None) -> str:
    """格式化 Energy 环境状态显示
    Args:
        state: 状态字典
        obs: 观察字典
        last_violations: 上一步的违约情况
        thermal_input: 火电输入值（可选，用于实时显示碳排放比例）
        wind_input: 风电输入值（可选，用于实时显示碳排放比例）
        solar_input: 太阳能输入值（可选，用于实时显示碳排放比例）
    """
    env = state.get('env')
    if last_violations is None:
        last_violations = state.get('last_step_violations', {})
    
    lines = []
    current_day = obs.get('day', 0)
    
    # 显示电池当前电量（始终显示，让用户知道电池状态）
    battery_cur = None
    if env is not None:
        battery_cur = env.battery_cur
    elif obs.get('battery_cur') is not None:
        battery_cur = obs.get('battery_cur')
    
    if battery_cur is not None:
        battery_capacity = 80.0  # 电池最大容量
        if env is not None and hasattr(env, 'capacity'):
            battery_capacity = env.capacity.get('battery', 80.0)
        lines.append(f"🔋 Battery current charge: {battery_cur:.2f} / {battery_capacity:.2f} MW")
    
    # 第一天显示初始发电量配置
    if current_day == 0:
        if env is not None and hasattr(env, 'initial_rated_cfg'):
            initial_rated = env.initial_rated_cfg
            lines.append("\nInitial generation configuration (use this to ensure stability, your first step should not differ too much from this configuration):")
            lines.append(f"  Thermal: {initial_rated.get('thermal', 0):.2f}")
            lines.append(f"  Wind: {initial_rated.get('wind', 0):.2f}")
            lines.append(f"  Solar: {initial_rated.get('solar', 0):.2f}")
            lines.append(f"  Battery: 0.00")
    elif current_day > 0:
        rated_prev = obs.get('rated_prev', {})
        if rated_prev:
            lines.append("Previous moment input generation:")
            lines.append(f"  Thermal: {rated_prev.get('thermal', 0):.2f}")
            lines.append(f"  Wind: {rated_prev.get('wind', 0):.2f}")
            lines.append(f"  Solar: {rated_prev.get('solar', 0):.2f}")
            lines.append(f"  Battery: {rated_prev.get('battery', 0):.2f}")
    
    # Get previous moment actual generation
    if current_day > 0:
        actual_prev = obs.get('actual_prev', {})
        if actual_prev:
            lines.append("\nPrevious moment actual generation:")
            lines.append(f"  Thermal: {actual_prev.get('thermal', 0):.2f}")
            lines.append(f"  Wind: {actual_prev.get('wind', 0):.2f}")
            lines.append(f"  Solar: {actual_prev.get('solar', 0):.2f}")
            battery_flow = actual_prev.get('battery', 0)
            if battery_flow < 0:
                lines.append(f"  Battery: {abs(battery_flow):.2f} (charging)")
            elif battery_flow > 0:
                lines.append(f"  Battery: {battery_flow:.2f} (discharging)")
            else:
                lines.append(f"  Battery: 0.00")
            lines.append(f"  Total supply: {actual_prev.get('supply', 0):.2f}")
    
    # 显示上一天的违约情况
    if obs.get('day', 0) > 0:
        # lines.append("\n" + "="*30)
        demand_vio = last_violations.get('demand_violation', False)
        budget_vio = last_violations.get('budget_violation', False)
        
        # lines.append("📊 上一天违约情况:")
        violation_days = obs.get('violation_days_cont', 0)
        has_violation = demand_vio or budget_vio
        
        if has_violation:
            lines.append("  ❌ Previous day had violations")
            violation_reasons = []
            if demand_vio:
                violation_reasons.append("Demand not met")
            if budget_vio:
                violation_reasons.append("Budget exceeded")
            lines.append(f"  Violation reasons: {', '.join(violation_reasons)}")
        else:
            lines.append("  ✅ Previous day had no violations")
        
        if violation_days > 0:
            lines.append(f"  Consecutive violation days: {violation_days} days")
            if violation_days >= 3:
                lines.append("  ⚠️ Warning: Consecutive violations reached 3 days, task failed!")
        else:
            lines.append("  Consecutive violation days: 0 days")
        # lines.append("="*30)
    
    # 显示稳定性及目标
    stability_value = obs.get('stability', 0)
    target_stability = None
    if env is not None and hasattr(env, 'target_stability'):
        target_stability = env.target_stability
    
    if target_stability is not None:
        lines.append(f"\nStability: {stability_value:.3f} (Target: ≥{target_stability:.3f})")
    else:
        lines.append(f"\nStability: {stability_value:.3f}")
    
    # 计算碳排放比例
    target_carbon = None
    if env is not None and hasattr(env, 'target_carbon'):
        target_carbon = env.target_carbon
    
    carbon_value = obs.get('carbon', 0)  # 累计碳排放比例
    
    # 计算今天的实时碳排放比例
    today_carbon_ratio = None
    # 优先使用输入值计算实时碳排放比例（如果提供了输入值）
    if thermal_input is not None and wind_input is not None and solar_input is not None:
        thermal_val = float(thermal_input) if thermal_input is not None else 0.0
        wind_val = float(wind_input) if wind_input is not None else 0.0
        solar_val = float(solar_input) if solar_input is not None else 0.0
        total_generation = thermal_val + wind_val + solar_val
        if total_generation > 0:
            today_carbon_ratio = thermal_val / total_generation
    elif env is not None:
        # 如果没有输入值，使用实际发电量
        thermal_today = getattr(env, 'thermal_actual', 0)
        wind_today = getattr(env, 'wind_actual', 0)
        solar_today = getattr(env, 'solar_actual', 0)
        total_generation_today = thermal_today + wind_today + solar_today
        if total_generation_today > 0:
            today_carbon_ratio = thermal_today / total_generation_today
    
    if target_carbon is not None:
        if today_carbon_ratio is not None:
            if thermal_input is not None:
                lines.append(f"Carbon emission ratio: {carbon_value:.3f} (cumulative, target: ≤{target_carbon:.3f})")
                lines.append(f"Today's carbon emission ratio: {today_carbon_ratio:.3f}")
            else:
                lines.append(f"Carbon emission ratio: {carbon_value:.3f} (cumulative, target: ≤{target_carbon:.3f})")
                lines.append(f"Today's carbon emission ratio: {today_carbon_ratio:.3f}")
        else:
            lines.append(f"Carbon emission ratio: {carbon_value:.3f} (cumulative, target: ≤{target_carbon:.3f})")
    else:
        if today_carbon_ratio is not None:
            if thermal_input is not None:
                lines.append(f"Carbon emission ratio: {carbon_value:.3f} (cumulative)")
                lines.append(f"Today's carbon emission ratio: {today_carbon_ratio:.3f}")
            else:
                lines.append(f"Carbon emission ratio: {carbon_value:.3f} (cumulative)")
                lines.append(f"Today's carbon emission ratio: {today_carbon_ratio:.3f}")
        else:
            lines.append(f"Carbon emission ratio: {carbon_value:.3f} (cumulative)")
    
    return "\n".join(lines)


def load_energy_test_data(state: Dict[str, Any], current_dir: str) -> Tuple[Dict[str, Any], str]:
    """加载 Energy 测试数据"""
    test_file = os.path.join(
        current_dir, "test_data/energy/test_energy_lite_251207.json")
    if not os.path.exists(test_file):
        test_file = "test_data/energy/test_energy_lite_251207.json"
    
    try:
        with open(test_file, 'r', encoding='utf-8') as f:
            state['test_data'] = json.load(f)
        return state, f"✅ Successfully loaded {len(state['test_data'])} test environments"
    except FileNotFoundError:
        return state, f"❌ File not found: {test_file}"
    except Exception as e:
        return state, f"❌ Load failed: {str(e)}"


def energy_save_progress_internal(state: Dict[str, Any], current_user_id: str, save_dir: str) -> str:
    """保存 Energy 环境进度（使用统一进度管理模块）"""
    # Auto-generate user ID if not provided
    if not current_user_id:
        import uuid
        current_user_id = f"user_{uuid.uuid4().hex[:8]}"
    
    env = state.get('env')
    if env is None:
        return "⚠️ No progress to save"
    
    try:
        current_env_idx = state.get('current_env_idx', 0)
        history_records = state.get('history_records', [])
        test_data = state.get('test_data', [])
        last_step_violations = state.get('last_step_violations', {})
        
        # 保存环境状态变量
        prev_rated = getattr(env, 'prev_rated', {})
        if not isinstance(prev_rated, dict):
            prev_rated = {}
        prev_rated_dict = {
            "thermal": float(prev_rated.get("thermal", 0)),
            "wind": float(prev_rated.get("wind", 0)),
            "solar": float(prev_rated.get("solar", 0)),
            "battery": float(prev_rated.get("battery", 0)),
        }
        
        env_state = {
            "thermal_actual": float(getattr(env, 'thermal_actual', 0)),
            "wind_actual": float(getattr(env, 'wind_actual', 0)),
            "solar_actual": float(getattr(env, 'solar_actual', 0)),
            "battery_actual": float(getattr(env, 'battery_actual', 0)),
            "prev_rated": prev_rated_dict,
            "stability_avg": float(getattr(env, 'stability_avg', 1.0)),
            "share_thermal": float(getattr(env, 'share_thermal', 0.0)),
            "supply_total": float(getattr(env, 'supply_total', 0)),
            "cum_carbon": float(getattr(env, 'cum_carbon', 0)),  # 保存累计火电发电量，用于正确计算碳排放比例
            # 不再保存 stability_sta 列表以提升性能（stability_avg 已足够）
            # 加载时会根据 stability_avg 和步数重建一个近似列表
        }
        
        # 计算 success：需要同时满足所有成功条件
        # 根据 EnergyEnv_v5.py 第248行的逻辑：
        # success = done AND stability_avg > target_stability AND share_thermal < target_carbon AND violation_days_cont < 3
        done = env.done
        stability_avg = float(getattr(env, 'stability_avg', 1.0))
        share_thermal = float(getattr(env, 'share_thermal', 0.0))
        violation_days_cont = getattr(env, 'violation_days_cont', 0)
        target_stability = getattr(env, 'target_stability', 0.0)
        target_carbon = getattr(env, 'target_carbon', 1.0)
        
        success = (
            done and
            stability_avg > target_stability and
            share_thermal < target_carbon and
            violation_days_cont < 3
        )
        
        env_progress = {
            "user_id": current_user_id,
            "env_idx": current_env_idx,
            "env_idx_display": current_env_idx + 1,
            # 不再保存 config，因为可以从 test_data[env_idx] 获取
            "day": env.t,
            "battery_cur": float(env.battery_cur),
            "history": history_records,
            "num_steps": len(history_records),
            "done": done,
            "success": success,
            "violation_days_cont": violation_days_cont,
            "last_violations": last_step_violations,
            "env_state": env_state,
        }
        
        result = progress_manager.save_task_environment_progress(
            current_user_id, save_dir, "energy", current_env_idx, env_progress
        )
        
        return f"✅ Progress saved (Environment {current_env_idx + 1}, Steps {len(history_records)})"
    except Exception as e:
        return f"❌ Save failed: {str(e)}"


def energy_load_environment(state: Dict[str, Any], env_idx_display: int, current_user_id: str, save_dir: str) -> Tuple[Dict[str, Any], str, str, str, str, str, str]:
    """加载 Energy 环境（使用统一进度管理模块）
    Returns: (state, info, state_display, logic, history_display, progress, steps_info)
    """
    # Auto-generate user ID if not provided
    if not current_user_id:
        import uuid
        current_user_id = f"user_{uuid.uuid4().hex[:8]}"
    
    test_data = state.get('test_data', [])
    if not test_data:
        return state, "❌ Please load test data first", "", "", "", "Click 'View Uncompleted Problems' button to view progress", "0 / 120 (Day 1)"
    
    env_idx = env_idx_display - 1
    if env_idx < 0 or env_idx >= len(test_data):
        return state, f"❌ Environment index out of range (1-{len(test_data)})", "", "", "", "Click 'View Unfinished Problems' button to view progress", "0 / 120 (Day 1)"
    
    # 使用统一进度管理模块检查是否有保存的进度
    saved_progress_data = progress_manager.get_task_environment_progress(
        current_user_id, save_dir, "energy", env_idx
    )
    
    # 如果有保存的进度，加载它
    if saved_progress_data:
        state['current_env_idx'] = env_idx
        state['history_records'] = saved_progress_data.get("history", [])
        
        # 从 test_data 获取 config（不再从保存的数据中获取，以节省存储空间）
        # 为了向后兼容，如果保存的数据中有 config，优先使用（旧数据可能没有 test_data）
        config = saved_progress_data.get("config")
        if not config and env_idx < len(test_data):
            config = test_data[env_idx]
        
        if config:
            state['env'] = DynamicEnergyGrid(config)
            state['env'].t = saved_progress_data.get("day", 0)
            state['env'].battery_cur = saved_progress_data.get("battery_cur", 0.0)
            state['env'].done = saved_progress_data.get("done", False)
            if "violation_days_cont" in saved_progress_data:
                state['env'].violation_days_cont = saved_progress_data.get("violation_days_cont", 0)
            
            # 恢复环境状态变量
            if "env_state" in saved_progress_data:
                env_state = saved_progress_data.get("env_state", {})
                state['env'].thermal_actual = env_state.get("thermal_actual", 0)
                state['env'].wind_actual = env_state.get("wind_actual", 0)
                state['env'].solar_actual = env_state.get("solar_actual", 0)
                state['env'].battery_actual = env_state.get("battery_actual", 0)
                prev_rated_loaded = env_state.get("prev_rated", {})
                if isinstance(prev_rated_loaded, dict):
                    state['env'].prev_rated = {
                        "thermal": float(prev_rated_loaded.get("thermal", 0)),
                        "wind": float(prev_rated_loaded.get("wind", 0)),
                        "solar": float(prev_rated_loaded.get("solar", 0)),
                        "battery": float(prev_rated_loaded.get("battery", 0)),
                    }
                else:
                    state['env'].prev_rated = {"thermal": 0.0, "wind": 0.0, "solar": 0.0, "battery": 0.0}
                state['env'].stability_avg = env_state.get("stability_avg", 1.0)
                state['env'].share_thermal = env_state.get("share_thermal", 0.0)
                state['env'].supply_total = env_state.get("supply_total", 0)
                # 恢复累计火电发电量，用于正确计算碳排放比例
                # 如果旧数据中没有 cum_carbon，从 share_thermal 和 supply_total 反推
                if "cum_carbon" in env_state:
                    state['env'].cum_carbon = env_state.get("cum_carbon", 0)
                else:
                    # 兼容旧数据：从 share_thermal 和 supply_total 反推
                    share_thermal = env_state.get("share_thermal", 0.0)
                    supply_total = env_state.get("supply_total", 0)
                    state['env'].cum_carbon = share_thermal * supply_total if supply_total > 0 else 0
                # 恢复 stability_sta 列表（如果旧数据中有，就使用；否则重建）
                # 新版本不再保存 stability_sta 以提升性能，但为了兼容旧数据，先尝试加载
                if "stability_sta" in env_state:
                    # 旧数据中有 stability_sta，直接使用
                    state['env'].stability_sta = env_state.get("stability_sta", [])
                else:
                    # 新数据中没有 stability_sta，根据步数和平均值重建
                    # 这样可以在保持性能的同时，确保环境状态的一致性
                    num_steps = len(state['history_records'])
                    if num_steps > 0:
                        stability_avg = env_state.get("stability_avg", 1.0)
                        state['env'].stability_sta = [stability_avg] * num_steps
                    else:
                        state['env'].stability_sta = []
        
        # 恢复上一步的违约信息
        if "last_violations" in saved_progress_data:
            state['last_step_violations'] = saved_progress_data.get("last_violations", {
                "demand_violation": False, "budget_violation": False, "stability_violation": False, "violation_days_cont": 0
            })
        else:
            state['last_step_violations'] = {"demand_violation": False, "budget_violation": False, "stability_violation": False, "violation_days_cont": 0}
        
        if state['env'] is not None:
            obs = state['env']._get_obs()
            state_display = format_energy_state(state, obs)
        else:
            state_display = "Environment loading failed"
        history_display = "\n\n".join(state['history_records']) if state['history_records'] else "No history records"  # Add blank lines between steps
        
        info = f"✅ Environment {env_idx_display}/{len(test_data)} loaded\n"
        info += f"Steps: {len(state['history_records'])}"
        
        steps_info = get_energy_steps_info(state)
        
        return state, info, state_display, "", history_display, "Click 'View Unfinished Problems' button to view progress", steps_info
    
    # 没有保存的进度，初始化新环境
    state['current_env_idx'] = env_idx
    config = test_data[env_idx]
    state['env'] = DynamicEnergyGrid(config)
    state['history_records'] = []
    state['last_step_violations'] = {"demand_violation": False, "budget_violation": False, "stability_violation": False, "violation_days_cont": 0}
    energy_save_progress_internal(state, current_user_id, save_dir)
    
    obs = state['env']._get_obs()
    state_display = format_energy_state(state, obs)
    history_display = "Environment initialized (new environment)\n"
    
    info = f"✅ Environment {env_idx_display}/{len(test_data)} initialized (new environment)\n"
    
    steps_info = get_energy_steps_info(state)
    
    return state, info, state_display, "", history_display, "Click 'View Unfinished Problems' button to view progress", steps_info


def energy_step_environment_from_inputs(state: Dict[str, Any], thermal: float, wind: float, solar: float, battery: float, current_user_id: str, save_dir: str) -> Tuple[Dict[str, Any], str, str, str, bool, str]:
    """从输入框执行 Energy 环境一步动作"""
    thermal = float(thermal) if thermal is not None else 0.0
    wind = float(wind) if wind is not None else 0.0
    solar = float(solar) if solar is not None else 0.0
    battery = float(battery) if battery is not None else 0.0
    
    action = {"thermal": thermal, "wind": wind, "solar": solar, "battery": battery}
    action_str = json.dumps(action, ensure_ascii=False)
    return energy_step_environment(state, action_str, current_user_id, save_dir)


def energy_step_environment(state: Dict[str, Any], action_str: str, current_user_id: str, save_dir: str) -> Tuple[Dict[str, Any], str, str, str, bool, str]:
    """执行 Energy 环境一步动作
    Returns: (state, feedback, state_display, history_display, done, steps_info)
    """
    env = state.get('env')
    history_records = state.get('history_records', [])
    
    current_state_display = ""
    if env is not None:
        obs = env._get_obs()
        current_state_display = format_energy_state(state, obs)
    
    if env is None:
        return state, "❌ Please initialize environment first", current_state_display if current_state_display else "Please initialize environment first", "", False, "0 / 120 (Day 1)"
    
    # 检查 episode 是否已完成
    if env.done:
        history_display = "\n\n".join(history_records) if history_records else ""  # 每步之间加空行
        steps_info = get_energy_steps_info(state)
        current_steps = len(history_records)
        if current_steps < ENERGY_MAX_STEPS:
            feedback_info = "❌ Task failed (completed)!\n"
            feedback_info += f"Task ended at {current_steps} steps, did not reach required {ENERGY_MAX_STEPS} steps.\n"
        else:
            feedback_info = "🎉 Task completed!\n"
            feedback_info += f"Successfully completed {current_steps} steps.\n"
        feedback_info += "Task ended, cannot continue executing new steps.\n"
        return state, feedback_info, current_state_display, history_display, True, steps_info
    
    # Auto-generate user ID if not provided
    if not current_user_id:
        import uuid
        current_user_id = f"user_{uuid.uuid4().hex[:8]}"
    
    # 解析动作
    try:
        action = json.loads(action_str.strip())
    except json.JSONDecodeError:
        step_num = len(history_records) + 1
        obs = env._get_obs()
        current_day = obs.get('day', 0)
        # 获取当天的需求
        demand_today = None
        if current_day < len(env.demand_series):
            demand_today = env.demand_series[current_day]
        history_record = format_energy_history_record(
            step_num, current_day, {"thermal": 0, "wind": 0, "solar": 0, "battery": 0}, 
            obs, "", 0, "JSON format error", demand=demand_today
        )
        history_records.append(history_record)
        state['history_records'] = history_records
        history_display = "\n\n".join(history_records)  # 每步之间加空行
        energy_save_progress_internal(state, current_user_id, save_dir)
        feedback_info = f"Action: {action_str}\nFeedback: ❌ JSON format error\n"
        steps_info = get_energy_steps_info(state)
        return state, feedback_info, current_state_display, history_display, False, steps_info
    
    # 检查是否达到步骤上限
    if len(history_records) >= ENERGY_MAX_STEPS:
        history_display = "\n\n".join(history_records) if history_records else ""  # 每步之间加空行
        energy_save_progress_internal(state, current_user_id, save_dir)
        feedback_info = f"⚠️ Reached step limit ({ENERGY_MAX_STEPS} steps)\n"
        feedback_info += "Task ended (failed to complete within the specified number of steps)\n"
        steps_info = get_energy_steps_info(state)
        return state, feedback_info, current_state_display, history_display, True, steps_info
    
    # 执行动作
    try:
        # 在执行 step 前获取当前天数和当天的需求（执行后 env.t 会增加）
        current_day_before_step = env.t
        demand_before_step = None
        if current_day_before_step < len(env.demand_series):
            demand_before_step = env.demand_series[current_day_before_step]
        
        obs, reward, done, info = env.step(action)
        
        current_violations = {
            "demand_violation": info.get('demand_violation', False) if isinstance(info, dict) else False,
            "budget_violation": info.get('budget_violation', False) if isinstance(info, dict) else False,
            "stability_violation": info.get('stability_violation', False) if isinstance(info, dict) else False,
            "violation_days_cont": info.get('violation_days_cont', 0) if isinstance(info, dict) else 0
        }
        state['last_step_violations'] = current_violations
        
        state_display = format_energy_state(state, obs, last_violations=current_violations)
        
        # 获取实际发电量（从 env 对象中获取）
        actual_generation = {
            "thermal": getattr(env, 'thermal_actual', 0),
            "wind": getattr(env, 'wind_actual', 0),
            "solar": getattr(env, 'solar_actual', 0),
            "battery": getattr(env, 'battery_actual', 0),
            "supply": getattr(env, 'thermal_actual', 0) + getattr(env, 'wind_actual', 0) + 
                     getattr(env, 'solar_actual', 0) + getattr(env, 'battery_actual', 0)
        }
        
        # 更新 obs 以便历史记录可以显示实际发电量
        obs['actual_prev'] = actual_generation
        
        feedback_msg = info.get('last_message', '') if isinstance(info, dict) else str(info)
        step_num = len(history_records) + 1
        history_record = format_energy_history_record(
            step_num, current_day_before_step, action, obs, feedback_msg, reward, demand=demand_before_step
        )
        history_records.append(history_record)
        state['history_records'] = history_records
        history_display = "\n\n".join(history_records)  # 每步之间加空行
        
        energy_save_progress_internal(state, current_user_id, save_dir)
        
        feedback_info = f"Action: {action_str}\nFeedback: {feedback_msg}\nReward: {reward:.2f}\n"
        if done:
            current_steps = len(history_records)
            if current_steps < ENERGY_MAX_STEPS:
                feedback_info += "❌ Task failed!\n"
                feedback_info += f"Task ended at {current_steps} steps, did not reach required {ENERGY_MAX_STEPS} steps.\n"
            else:
                feedback_info += "🎉 Task completed!\n"
                feedback_info += f"Successfully completed {current_steps} steps.\n"
        
        steps_info = get_energy_steps_info(state)
        
        return state, feedback_info, state_display, history_display, done, steps_info
    except Exception as e:
        step_num = len(history_records) + 1
        obs = env._get_obs()
        current_day = obs.get('day', 0)
        # 获取当天的需求
        demand_today = None
        if current_day < len(env.demand_series):
            demand_today = env.demand_series[current_day]
        try:
            action_dict = json.loads(action_str.strip())
        except:
            action_dict = {"thermal": 0, "wind": 0, "solar": 0, "battery": 0}
        history_record = format_energy_history_record(
            step_num, current_day, action_dict, obs, "", 0, str(e), demand=demand_today
        )
        history_records.append(history_record)
        state['history_records'] = history_records
        history_display = "\n\n".join(history_records)  # 每步之间加空行
        energy_save_progress_internal(state, current_user_id, save_dir)
        feedback_info = f"Action: {action_str}\nFeedback: ❌ {str(e)}\n"
        steps_info = get_energy_steps_info(state)
        return state, feedback_info, current_state_display, history_display, False, steps_info


def energy_reset_environment(state: Dict[str, Any], current_user_id: str, save_dir: str) -> Tuple[Dict[str, Any], str, str, str, str, str]:
    """重置 Energy 环境
    Returns: (state, info, state_display, history_display, progress, steps_info)
    """
    env = state.get('env')
    
    if env is None:
        return state, "❌ Please initialize environment first", "", "", "Click 'View Uncompleted Problems' button to view progress", "0 / 120 (Day 1)"
    
    env.reset()
    if hasattr(env, 'violation_days_cont'):
        env.violation_days_cont = 0
    state['history_records'] = []
    state['last_step_violations'] = {"demand_violation": False, "budget_violation": False, "stability_violation": False, "violation_days_cont": 0}
    energy_save_progress_internal(state, current_user_id, save_dir)
    
    obs = env._get_obs()
    state_display = format_energy_state(state, obs)
    history_display = "Environment reset\n"
    
    steps_info = get_energy_steps_info(state)
    
    return state, "✅ Environment reset", state_display, history_display, "Click 'View Unfinished Problems' button to view progress", steps_info


def get_energy_current_env_idx(state: Dict[str, Any]) -> int:
    """获取当前 Energy 环境索引"""
    return state.get('current_env_idx', 0)


def get_energy_test_data(state: Dict[str, Any]) -> List[dict]:
    """获取 Energy 测试数据"""
    return state.get('test_data', [])


def get_energy_history_records(state: Dict[str, Any]) -> List[str]:
    """获取 Energy 历史记录"""
    return state.get('history_records', [])


def get_energy_env(state: Dict[str, Any]) -> Optional[DynamicEnergyGrid]:
    """获取当前的 Energy 环境对象"""
    return state.get('env')


def get_energy_progress_summary(state: Dict[str, Any], user_id: str, save_dir: str) -> str:
    """获取 Energy 任务用户进度摘要（使用统一进度管理模块）"""
    # Auto-generate user ID if not provided
    if not user_id or not user_id.strip():
        import uuid
        user_id = f"user_{uuid.uuid4().hex[:8]}"
    
    user_id = user_id.strip()
    test_data = state.get('test_data', [])
    
    # 使用统一进度管理模块加载进度
    task_data = progress_manager.load_task_progress(user_id, save_dir, "energy")
    environments = task_data.get("environments", {})
    
    completed_envs = set()
    for env_key, progress_data in environments.items():
        env_idx = progress_data.get("env_idx", -1)
        done = progress_data.get("done", False)
        success = progress_data.get("success", False)
        num_steps = progress_data.get("num_steps", 0)
        
        is_completed = False
        if success or done:
            is_completed = True
        elif num_steps >= ENERGY_MAX_STEPS:
            is_completed = True
        
        if is_completed:
            completed_envs.add(env_idx)
    
    total_envs = len(test_data) if test_data else 0
    if total_envs == 0:
        return "⚠️ Please load test data first"
    
    all_env_indices = set(range(total_envs))
    incomplete_envs = sorted(all_env_indices - completed_envs)
    
    summary_lines = []
    summary_lines.append(f"📊 Energy Task - Progress Summary for User {user_id}")
    summary_lines.append(f"Total environments: {total_envs}")
    summary_lines.append(f"Completed: {len(completed_envs)}/{total_envs}")
    summary_lines.append(f"Incomplete: {len(incomplete_envs)}/{total_envs}")
    
    if incomplete_envs:
        summary_lines.append("\n❌ Incomplete environments:")
        for i in range(0, len(incomplete_envs), 5):
            env_display_list = [str(env_idx + 1) for env_idx in incomplete_envs[i:i+5]]
            summary_lines.append("  " + ", ".join(env_display_list))
    else:
        summary_lines.append("\n🎉 Congratulations! All environments are completed!")
    
    return "\n".join(summary_lines)


def create_energy_interface(current_dir: str, save_dir: str, user_id_input: gr.Textbox) -> Tuple:
    """创建 Energy 任务界面组件
    Returns: (energy_interface, energy_env_idx_input, energy_init_btn, energy_reset_btn,
              energy_env_info, energy_state_display, energy_steps_info_text,
              energy_thermal_input, energy_wind_input, energy_solar_input, energy_battery_input,
              energy_cost_display, energy_step_btn, energy_feedback_display, energy_history_display)
    
    注意：环境控制组件（energy_env_idx_input, energy_init_btn, energy_reset_btn, energy_env_info）
    需要在主界面中手动添加到进度摘要下方，不包含在 energy_interface 中。
    为了保持函数签名一致，这里返回 None 作为占位符，主界面会忽略这些返回值。
    """
    # 创建主界面 Row（不包含环境控制）
    with gr.Row(visible=False) as energy_interface:
        with gr.Column(scale=1):
            energy_steps_info_text = gr.Textbox(
                label="Steps Info (Day)",
                value="0 / 120 (Day 1)",
                interactive=False,
                visible=True,
                lines=2
            )
            gr.Markdown("### 📜 Action History")
            energy_history_display = gr.Textbox(
                label="Action History",
                interactive=False,
                lines=10
            )
        
        with gr.Column(scale=2):
            gr.Markdown("### ⚡ Current State")
            with gr.Row():
                energy_state_display = gr.Textbox(
                    label="Energy State",
                    interactive=False,
                    lines=10,
                    value="Please load environment first"
                )
                energy_cost_display = gr.Textbox(
                    label="Total Generation & Estimated Cost",
                    interactive=False,
                    lines=10,
                    value="Please input generation to view total generation and estimated cost"
                )
            
            gr.Markdown("### 🎯 Energy Operations")
            with gr.Row():
                energy_thermal_input = gr.Number(
                    label="🔥 Thermal",
                    value=0.0,
                    minimum=0.0,
                    precision=2,
                    info="Thermal power generation (≥0)"
                )
                energy_wind_input = gr.Number(
                    label="💨 Wind",
                    value=0.0,
                    minimum=0.0,
                    precision=2,
                    info="Wind power generation (≥0)"
                )
                energy_solar_input = gr.Number(
                    label="☀️ Solar",
                    value=0.0,
                    minimum=0.0,
                    precision=2,
                    info="Solar power generation (≥0)"
                )
                energy_battery_input = gr.Number(
                label="🔋 Battery",
                value=0.0,
                precision=2,
                info="Battery operation: negative=charge, positive=discharge"
            )
            
            energy_step_btn = gr.Button("Execute Operation", variant="primary")
            
            # Environment feedback box removed, but keep variable for interface compatibility
            energy_feedback_display = gr.Textbox(
                label="Feedback Info",
                interactive=False,
                lines=5,
                visible=False
            )
    
    # 返回占位符（主界面会使用自己创建的环境控制组件）
    return (energy_interface, None, None, None,
            None, energy_state_display, energy_steps_info_text,
            energy_thermal_input, energy_wind_input, energy_solar_input, energy_battery_input,
            energy_cost_display, energy_step_btn, energy_feedback_display, energy_history_display)