import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from scipy.stats import spearmanr
import os
import time

# ==========================================
# 1. 配置与初始化
# ==========================================
def configure_environment():
    """
    配置绘图环境，解决中文乱码和负号显示问题，
    并设置符合学术规范的绘图风格。
    """
    # 设置字体为 Times New Roman (英文) 或 SimHei (中文兼容)
    # 为了美赛(MCM)标准，主要使用英文标签，但配置好中文支持以防万一
    plt.rcParams['font.family'] = 'sans-serif'
    plt.rcParams['font.sans-serif'] = ['SimHei', 'Arial'] # 优先使用黑体显示中文
    plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示为方块的问题
    
    # 学术风格配置
    plt.style.use('seaborn-v0_8-whitegrid')
    plt.rcParams.update({
        'font.size': 12,
        'axes.labelsize': 14,
        'axes.titlesize': 16,
        'xtick.labelsize': 12,
        'ytick.labelsize': 12,
        'legend.fontsize': 12,
        'figure.figsize': (10, 6),
        'figure.dpi': 150
    })

    # 确保输出目录存在
    if not os.path.exists('output_q2'):
        os.makedirs('output_q2')

# ==========================================
# 2. 核心物理模型类
# ==========================================
class StochasticBatteryModel:
    def __init__(self):
        # 电池物理参数
        self.Q_design = 4000 / 1000.0  # 4000 mAh -> 4.0 Ah
        self.R_int_base = 0.15         # 内阻 (Ohm)
        
        # Shepherd 模型参数 (OCV 曲线)
        self.E0 = 3.4
        self.K = 0.02
        self.A = 0.5
        self.B = 15.0

    def get_capacity_correction(self, temp_c):
        """
        根据温度修正有效容量 (Arrhenius 效应)
        Temp_c: 环境温度 (摄氏度)
        """
        temp_k = temp_c + 273.15
        ref_k = 298.15 # 25°C
        # 温度越低，容量越小；温度越高，容量略微增加但有限制
        factor = np.exp(0.8 * (1 - ref_k / temp_k))
        return np.clip(factor, 0.4, 1.1) # 限制修正因子范围

    def get_ocv(self, soc):
        """计算开路电压 (Open Circuit Voltage)"""
        soc = np.clip(soc, 0.001, 1.0)
        # Shepherd Model + 线性项
        term1 = self.E0
        term2 = -self.K / soc
        term3 = -self.R_int_base * 0.1 # 简化极化项
        term4 = self.A * np.exp(-self.B * (1 - soc))
        term5 = 0.7 * soc # 线性部分
        v_ocv = 3.0 + term5 + term4 + term2 
        return np.clip(v_ocv, 2.8, 4.35)

    def calculate_power_drain(self, base_power, signal_dbm, noise_scale=0.1):
        """
        计算瞬时功率，包含信号强度的非线性影响和随机噪声
        P_total = P_base + P_signal + Noise
        """
        # 1. 信号功耗模型：信号越弱(越负)，功耗呈指数上升
        # 假设 -60dBm 为基准，每下降 10dBm 功耗显著增加
        # 归一化信号强度：将 -120dBm 到 -50dBm 映射到 0-1 之间
        sig_norm = np.clip((-signal_dbm - 50) / 70, 0, 1)
        p_net = 2.0 * (sig_norm ** 3) # 三次幂关系，模拟弱信号时的急剧恶化
        
        # 2. 随机噪声 (模拟 CPU 动态调频)
        noise = np.random.normal(0, noise_scale)
        
        return max(0.1, base_power + p_net + noise)

# ==========================================
# 3. 蒙特卡洛模拟引擎
# ==========================================
def run_monte_carlo_simulation(n_simulations=2000):
    """
    执行 N 次蒙特卡洛模拟
    """
    print(f"Starting Monte Carlo Simulation with {n_simulations} samples...")
    
    model = StochasticBatteryModel()
    results = []
    
    # 定义时间步长 (秒)
    dt = 60.0  # 1分钟一步，平衡精度与速度
    max_time = 48 * 3600 # 最大模拟 48 小时
    
    start_time = time.time()

    for i in range(n_simulations):
        # --- A. 随机生成初始条件 (输入分布) ---
        
        # 1. 初始电量 SOC0: 均匀分布 U(0.2, 1.0)
        # 模拟用户在不同电量下开始使用手机
        soc_0 = np.random.uniform(0.2, 1.0)
        
        # 2. 环境温度 Temp: 正态分布 N(25, 8)，截断在 [-10, 45]
        temp_c = np.random.normal(25, 8)
        temp_c = np.clip(temp_c, -10, 45)
        
        # 3. 平均信号强度 Signal: 偏态分布
        # 大部分时候信号较好 (-80dBm), 偶尔很差 (-110dBm)
        signal_dbm = np.random.triangular(-120, -85, -50)
        
        # 4. 用户行为模式 (基准功率)
        # 混合高斯分布：待机为主(0.5W)，偶尔重度使用(4W)
        user_type = np.random.choice(['Light', 'Medium', 'Heavy'], p=[0.4, 0.4, 0.2])
        if user_type == 'Light':
            base_power_avg = 0.5 # 待机、阅读
        elif user_type == 'Medium':
            base_power_avg = 1.5 # 视频、社交
        else:
            base_power_avg = 4.0 # 游戏
            
        # --- B. 执行单次时间步进模拟 (Euler Integration) ---
        soc = soc_0
        t = 0
        q_eff = model.Q_design * model.get_capacity_correction(temp_c)
        
        while soc > 0.02 and t < max_time: # 截止电压设为 2%
            # 计算当前功率 (加入随机性)
            p_inst = model.calculate_power_drain(base_power_avg, signal_dbm, noise_scale=0.2)
            
            # 计算端电压
            v_term = model.get_ocv(soc) - (p_inst / 3.7) * model.R_int_base
            v_term = max(v_term, 3.0) # 防止电压过低
            
            # 计算电流 I = P / V
            current = p_inst / v_term
            
            # 更新 SOC: dS = -I * dt / Q
            d_soc = -current * dt / (q_eff * 3600)
            soc += d_soc
            t += dt
            
        # --- C. 记录结果 ---
        tte_hours = t / 3600.0
        results.append({
            'Initial_SOC': soc_0,
            'Temperature_C': temp_c,
            'Signal_dBm': signal_dbm,
            'User_Profile': user_type,
            'Base_Power': base_power_avg,
            'TTE_Hours': tte_hours
        })
        
        if (i+1) % 500 == 0:
            print(f"  Processed {i+1}/{n_simulations} simulations...")

    print(f"Simulation completed in {time.time() - start_time:.2f} seconds.")
    return pd.DataFrame(results)

# ==========================================
# 4. 结果分析与可视化
# ==========================================
def analyze_and_plot(df):
    """
    对模拟结果进行统计分析和绘图
    """
    print("\n=== Analysis Report ===")
    print(df.describe())
    
    # --- 图1: TTE 分布直方图 (Probability Distribution) ---
    plt.figure(figsize=(10, 6))
    sns.histplot(data=df, x='TTE_Hours', hue='User_Profile', element="step", stat="density", common_norm=False, palette='viridis')
    plt.title('Distribution of Time-to-Empty (TTE) by User Profile', fontweight='bold')
    plt.xlabel('Time to Empty (Hours)')
    plt.ylabel('Probability Density')
    plt.grid(True, alpha=0.3)
    plt.savefig('output_q2/fig1_tte_distribution.png')
    plt.show()
    
    # --- 图2: 初始电量 vs TTE (非线性关系展示) ---
    plt.figure(figsize=(10, 6))
    # 使用散点图，颜色映射温度
    sc = plt.scatter(df['Initial_SOC']*100, df['TTE_Hours'], c=df['Temperature_C'], cmap='coolwarm', alpha=0.6, s=20)
    plt.colorbar(sc, label='Temperature (°C)')
    plt.title('Impact of Initial SOC and Temperature on TTE', fontweight='bold')
    plt.xlabel('Initial State of Charge (%)')
    plt.ylabel('Time to Empty (Hours)')
    plt.grid(True, linestyle='--', alpha=0.5)
    
    # 添加拟合曲线 (展示非线性)
    # 简单的多项式拟合用于视觉引导
    z = np.polyfit(df['Initial_SOC']*100, df['TTE_Hours'], 2)
    p = np.poly1d(z)
    x_range = np.linspace(20, 100, 100)
    plt.plot(x_range, p(x_range), 'k--', linewidth=2, label='Trend Line')
    plt.legend()
    plt.savefig('output_q2/fig2_soc_vs_tte.png')
    plt.show()
    
    # --- 图3: 全局灵敏度分析 (Tornado Plot 替代品) ---
    # 计算 Spearman 相关系数
    corr_cols = ['Initial_SOC', 'Temperature_C', 'Signal_dBm', 'Base_Power']
    corr_matrix = df[corr_cols + ['TTE_Hours']].corr(method='spearman')
    sensitivity = corr_matrix['TTE_Hours'].drop('TTE_Hours').sort_values()
    
    plt.figure(figsize=(10, 5))
    colors = ['red' if x < 0 else 'green' for x in sensitivity.values]
    sensitivity.plot(kind='barh', color=colors, alpha=0.8)
    plt.title('Sensitivity Analysis: Correlation with TTE', fontweight='bold')
    plt.xlabel('Spearman Correlation Coefficient')
    plt.axvline(0, color='black', linewidth=0.8)
    plt.grid(axis='x', linestyle='--', alpha=0.5)
    
    # 添加数值标签
    for index, value in enumerate(sensitivity):
        plt.text(value, index, f' {value:.2f}', va='center', fontsize=10, fontweight='bold')
        
    plt.tight_layout()
    plt.savefig('output_q2/fig3_sensitivity.png')
    plt.show()

    print("\nKey Insights:")
    print(f"1. Base Power Correlation: {sensitivity['Base_Power']:.2f} (Dominant negative factor)")
    print(f"2. Initial SOC Correlation: {sensitivity['Initial_SOC']:.2f} (Dominant positive factor)")
    print(f"3. Signal Strength Correlation: {sensitivity['Signal_dBm']:.2f} (Significant environmental factor)")

# ==========================================
# 5. 主程序入口
# ==========================================
if __name__ == "__main__":
    # 1. 配置环境
    configure_environment()
    
    # 2. 运行模拟
    # 模拟 3000 次以获得平滑的分布
    simulation_data = run_monte_carlo_simulation(n_simulations=3000)
    
    # 3. 分析与绘图
    analyze_and_plot(simulation_data)
    
    print("\nAll tasks completed. Results saved in 'output_q2' directory.")