#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
绘制黑体辐射曲线
本脚本绘制温度为25000K和35000K的理论黑体辐射曲线
"""

import numpy as np
import matplotlib.pyplot as plt
from scipy import constants as const
import matplotlib.font_manager as fm
from matplotlib.ticker import ScalarFormatter

# 检查是否有中文字体
try:
    chinese_font = fm.FontProperties(fname='/usr/share/fonts/truetype/droid/DroidSansFallbackFull.ttf')
except:
    # 如果找不到指定中文字体，尝试使用系统默认字体
    chinese_font = fm.FontProperties()

# 物理常数
h = const.h  # 普朗克常数，J·s
c = const.c  # 光速，m/s
k = const.k  # 玻尔兹曼常数，J/K

# 普朗克函数计算黑体辐射
def planck(wavelength, T):
    """
    计算黑体辐射的辐射强度
    
    参数:
    wavelength: 波长，单位：米
    T: 温度，单位：开尔文
    
    返回:
    B_lambda: 黑体辐射强度，单位：W·m^-2·steradian^-1·m^-1
    """
    a = 2.0 * h * c**2
    b = h * c / (wavelength * k * T)
    B_lambda = a / (wavelength**5 * (np.exp(b) - 1.0))
    return B_lambda

# 维恩位移定律：λ_max * T = b，其中b ≈ 2.8978e-3 m·K
def wien_peak(T):
    """计算维恩位移定律预测的峰值波长（米）"""
    b = 2.8978e-3  # 维恩常数，m·K
    return b / T

# 设置温度
T1 = 25000  # 25000K
T2 = 35000  # 35000K

# 计算峰值波长
peak_wavelength_T1 = wien_peak(T1)
peak_wavelength_T2 = wien_peak(T2)

# 转换为纳米
peak_wavelength_T1_nm = peak_wavelength_T1 * 1e9
peak_wavelength_T2_nm = peak_wavelength_T2 * 1e9

# 设置波长范围（纳米）- 调整范围以更好地显示峰值
wavelength_nm = np.linspace(10, 1000, 1000)
wavelength_m = wavelength_nm * 1e-9  # 转换为米

# 计算不同温度的黑体辐射强度
intensity_T1 = planck(wavelength_m, T1)
intensity_T2 = planck(wavelength_m, T2)

# 绘图设置
plt.figure(figsize=(12, 8))

# 创建主图和次坐标轴（线性刻度）
plt.subplot(211)  # 第一个子图：线性刻度
plt.plot(wavelength_nm, intensity_T1, 'r-', linewidth=2, label=f'T = {T1}K')
plt.plot(wavelength_nm, intensity_T2, 'b-', linewidth=2, label=f'T = {T2}K')

# 标记维恩峰值
plt.axvline(x=peak_wavelength_T1_nm, color='r', linestyle='--', alpha=0.7)
plt.axvline(x=peak_wavelength_T2_nm, color='b', linestyle='--', alpha=0.7)

plt.annotate(f'峰值：{peak_wavelength_T1_nm:.1f}nm',
             xy=(peak_wavelength_T1_nm, planck(peak_wavelength_T1*1e-9, T1)),
             xytext=(peak_wavelength_T1_nm+50, planck(peak_wavelength_T1*1e-9, T1)*0.8),
             arrowprops=dict(arrowstyle='->'), fontproperties=chinese_font)

plt.annotate(f'峰值：{peak_wavelength_T2_nm:.1f}nm',
             xy=(peak_wavelength_T2_nm, planck(peak_wavelength_T2*1e-9, T2)),
             xytext=(peak_wavelength_T2_nm+50, planck(peak_wavelength_T2*1e-9, T2)*0.8),
             arrowprops=dict(arrowstyle='->'), fontproperties=chinese_font)

plt.xlabel('波长 (nm)', fontproperties=chinese_font)
plt.ylabel('辐射强度 (W·m⁻²·sr⁻¹·m⁻¹)', fontproperties=chinese_font)
plt.title('黑体辐射曲线 - 线性刻度', fontproperties=chinese_font)
plt.legend(prop=chinese_font)
plt.grid(True, alpha=0.3)
plt.xlim(0, 500)  # 限制X轴范围以更好地显示峰值

# 第二个子图：对数刻度
plt.subplot(212)
plt.loglog(wavelength_nm, intensity_T1, 'r-', linewidth=2, label=f'T = {T1}K')
plt.loglog(wavelength_nm, intensity_T2, 'b-', linewidth=2, label=f'T = {T2}K')

# 标记维恩峰值（对数刻度）
plt.axvline(x=peak_wavelength_T1_nm, color='r', linestyle='--', alpha=0.7)
plt.axvline(x=peak_wavelength_T2_nm, color='b', linestyle='--', alpha=0.7)

plt.annotate(f'峰值：{peak_wavelength_T1_nm:.1f}nm',
             xy=(peak_wavelength_T1_nm, planck(peak_wavelength_T1*1e-9, T1)),
             xytext=(peak_wavelength_T1_nm*2, planck(peak_wavelength_T1*1e-9, T1)/5),
             arrowprops=dict(arrowstyle='->'), fontproperties=chinese_font)

plt.annotate(f'峰值：{peak_wavelength_T2_nm:.1f}nm',
             xy=(peak_wavelength_T2_nm, planck(peak_wavelength_T2*1e-9, T2)),
             xytext=(peak_wavelength_T2_nm*2, planck(peak_wavelength_T2*1e-9, T2)/5),
             arrowprops=dict(arrowstyle='->'), fontproperties=chinese_font)

plt.xlabel('波长 (nm)', fontproperties=chinese_font)
plt.ylabel('辐射强度 (W·m⁻²·sr⁻¹·m⁻¹)', fontproperties=chinese_font)
plt.title('黑体辐射曲线 - 对数刻度', fontproperties=chinese_font)
plt.legend(prop=chinese_font)
plt.grid(True, alpha=0.3, which='both')

# 添加总标题
plt.suptitle(f'温度为{T1}K和{T2}K的黑体辐射曲线对比', fontsize=16, fontproperties=chinese_font)

# 显示图形
plt.tight_layout()
plt.subplots_adjust(top=0.9)  # 为总标题留出空间
plt.show()

# 打印维恩峰值信息
print(f"维恩位移定律预测的峰值波长:")
print(f"T = {T1}K: {peak_wavelength_T1_nm:.2f}纳米")
print(f"T = {T2}K: {peak_wavelength_T2_nm:.2f}纳米")