import numpy as np
import matplotlib.pyplot as plt
import math
from PIL import Image

# --- 1. 全局参数设置 ---

# 图片尺寸 (像素)
IMG_WIDTH = 2400
IMG_HEIGHT = 300

# K9玻璃赛尔麦耶(Sellmeier)色散公式常数
A1 = 0.567853792
A2 = 0.702269689
A3 = 1.08102455
C1 = 0.00238581579
C2 = 0.0135348332
C3 = 107.748228

# 物理常数 (源自您的公式)
PRISM_CONSTANT_RAD = 0.261799

# A型星的有效温度 (开尔文)
STAR_TEMPERATURE = 29000# K

# --- 2. 可调整的谱线参数 ---
# 您可以在这里轻松添加、删除或修改谱线

#Type B absorption lines
ABSORPTION_LINES = [
    {'name': 'He I ', 'lambda_nm': 447.15, 'depth': 0.25, 'width_px': 10},
    {'name': 'He I',   'lambda_nm': 402.62,  'depth': 0.15, 'width_px': 7},
    {'name': 'He I', 'lambda_nm': 438.79, 'depth': 0.25, 'width_px': 11},
    {'name': 'Si III',  'lambda_nm': 455.4, 'depth': 0.1, 'width_px': 6},
    {'name': 'N II', 'lambda_nm': 409.73, 'depth': 0.15, 'width_px': 6},

    {'name': 'H-alpha', 'lambda_nm': 656.28, 'depth': 0.1, 'width_px': 6},
    {'name': 'H-beta',  'lambda_nm': 486.13, 'depth': 0.1, 'width_px': 8},
    {'name': 'H-gamma', 'lambda_nm': 434.05, 'depth': 0.1, 'width_px': 8},
    {'name': 'H-delta', 'lambda_nm': 410.17, 'depth': 0.1, 'width_px': 8},
    {'name': 'H-epsilon','lambda_nm': 397.01, 'depth': 0.15, 'width_px': 9},
    {'name': 'H-zeta',   'lambda_nm': 388.91, 'depth': 0.1, 'width_px': 9},
    {'name': 'H-eta',    'lambda_nm': 383.54, 'depth': 0.1, 'width_px': 9},

    {'name': 'Atmospheric O2 A-band', 'lambda_nm': 760.4, 'depth': 0.60, 'width_px': 7},
    {'name': 'Atmospheric O2 Z-band', 'lambda_nm': 822.7, 'depth': 0.20, 'width_px': 7},
    {'name': 'Atmospheric H2O ρστ','lambda_nm': 928, 'depth': 0.5, 'width_px':5}
]
"""
ABSORPTION_LINES = [O   
    {'name': 'He II', 'lambda_nm': 468.57, 'depth': 0.35, 'width_px': 10},
    {'name': 'He I ', 'lambda_nm': 447.15, 'depth': 0.2, 'width_px': 7},
    {'name': 'He II',   'lambda_nm': 420.0,  'depth': 0.2, 'width_px': 7},
    {'name': 'Si IV',  'lambda_nm': 411.61, 'depth': 0.15, 'width_px': 6},
    {'name': 'He I', 'lambda_nm': 438.79, 'depth': 0.3, 'width_px': 10},
    {'name': 'H-alpha', 'lambda_nm': 656.28, 'depth': 0.1, 'width_px': 4},
    {'name': 'H-beta',  'lambda_nm': 486.13, 'depth': 0.15, 'width_px': 6},
    {'name': 'O III', 'lambda_nm': 433.35, 'depth': 0.15, 'width_px': 6},
    {'name': 'C III 'lambda_nm': 465.14, 'depth': 0.1, 'width_px': 6},
    {'name': 'N III', 'lambda_nm': 409.73, 'depth': 0.15, 'width_px': 6},

    {'name': 'Atmospheric O2 A-band', 'lambda_nm': 760.4, 'depth': 0.55, 'width_px': 7},
    {'name': 'Atmospheric O2 Z-band', 'lambda_nm': 822.7, 'depth': 0.20, 'width_px': 7},
    {'name': 'Atmospheric H2O ρστ','lambda_nm': 928, 'depth': 0.5, 'width_px':5}
]
"""
"""
ABSORPTION_LINES = [B
    {'name': 'He I ', 'lambda_nm': 447.15, 'depth': 0.25, 'width_px': 10},
    {'name': 'He I',   'lambda_nm': 402.62,  'depth': 0.15, 'width_px': 7},
    {'name': 'He I', 'lambda_nm': 438.79, 'depth': 0.25, 'width_px': 11},
    {'name': 'Si III',  'lambda_nm': 455.4, 'depth': 0.1, 'width_px': 6},
    {'name': 'N II', 'lambda_nm': 409.73, 'depth': 0.15, 'width_px': 6},

    {'name': 'H-alpha', 'lambda_nm': 656.28, 'depth': 0.1, 'width_px': 6},
    {'name': 'H-beta',  'lambda_nm': 486.13, 'depth': 0.1, 'width_px': 8},
    {'name': 'H-gamma', 'lambda_nm': 434.05, 'depth': 0.1, 'width_px': 8},
    {'name': 'H-delta', 'lambda_nm': 410.17, 'depth': 0.1, 'width_px': 8},
    {'name': 'H-epsilon','lambda_nm': 397.01, 'depth': 0.15, 'width_px': 9},
    {'name': 'H-zeta',   'lambda_nm': 388.91, 'depth': 0.1, 'width_px': 9},
    {'name': 'H-eta',    'lambda_nm': 383.54, 'depth': 0.1, 'width_px': 9},

    {'name': 'Atmospheric O2 A-band', 'lambda_nm': 760.4, 'depth': 0.60, 'width_px': 7},
    {'name': 'Atmospheric O2 Z-band', 'lambda_nm': 822.7, 'depth': 0.20, 'width_px': 7},
    {'name': 'Atmospheric H2O ρστ','lambda_nm': 928, 'depth': 0.5, 'width_px':5}
]
"""
"""
ABSORPTION_LINES = [A
    {'name': 'H-alpha', 'lambda_nm': 656.28, 'depth': 0.30, 'width_px': 8},
    {'name': 'H-beta',  'lambda_nm': 486.13, 'depth': 0.45, 'width_px': 11},
    {'name': 'H-gamma', 'lambda_nm': 434.05, 'depth': 0.45, 'width_px': 11},
    {'name': 'H-delta', 'lambda_nm': 410.17, 'depth': 0.46, 'width_px': 11},
    {'name': 'H-epsilon','lambda_nm': 397.01, 'depth': 0.45, 'width_px': 12},
    {'name': 'H-zeta',   'lambda_nm': 388.91, 'depth': 0.44, 'width_px': 12},
    {'name': 'H-eta',    'lambda_nm': 383.54, 'depth': 0.42, 'width_px': 13},

    {'name': 'Ca II K', 'lambda_nm': 393.37, 'depth': 0.1, 'width_px': 5},
    {'name': 'Ca II H', 'lambda_nm': 396.85, 'depth': 0.1, 'width_px': 5},
    {'name': 'CH',   'lambda_nm': 448.1,  'depth': 0.081, 'width_px': 6},
    {'name': 'Na I D',  'lambda_nm': 589.29, 'depth': 0.08, 'width_px': 4},

    {'name': 'Atmospheric O2 A-band', 'lambda_nm': 760.4, 'depth': 0.60, 'width_px': 7},
    {'name': 'Atmospheric O2 Z-band', 'lambda_nm': 822.7, 'depth': 0.20, 'width_px': 7},
    {'name': 'Atmospheric H2O ρστ','lambda_nm': 928, 'depth': 0.5, 'width_px':5}
]
"""
"""
ABSORPTION_LINES = [F
    {'name': 'H-alpha', 'lambda_nm': 656.28, 'depth': 0.1, 'width_px': 6},
    {'name': 'H-beta',  'lambda_nm': 486.13, 'depth': 0.25, 'width_px': 8},
    {'name': 'H-gamma', 'lambda_nm': 434.05, 'depth': 0.25, 'width_px': 8},
    {'name': 'H-delta', 'lambda_nm': 410.17, 'depth': 0.25, 'width_px': 8},
    {'name': 'H-epsilon','lambda_nm': 397.01, 'depth': 0.3, 'width_px': 9},
    {'name': 'H-zeta',   'lambda_nm': 388.91, 'depth': 0.25, 'width_px': 9},
    {'name': 'H-eta',    'lambda_nm': 383.54, 'depth': 0.25, 'width_px': 9},

    {'name': 'Ca II K', 'lambda_nm': 393.37, 'depth': 0.36, 'width_px': 9},
    {'name': 'Ca II H', 'lambda_nm': 396.85, 'depth': 0.36, 'width_px': 9},
    {'name': 'Ca II', 'lambda_nm': 854.2, 'depth': 0.15, 'width_px': 6},

    {'name': 'Ti II',   'lambda_nm': 416.36,  'depth': 0.1, 'width_px': 4},
    {'name': 'Fe II',  'lambda_nm': 417.34, 'depth': 0.08, 'width_px': 4},
    {'name': 'Ti II',   'lambda_nm': 444.38,  'depth': 0.08, 'width_px': 4},
    {'name': 'Fe II',  'lambda_nm': 442.4, 'depth': 0.08, 'width_px': 4},
    {'name': 'Fe I',  'lambda_nm': 720.74, 'depth': 0.1, 'width_px': 10},
    {'name': 'Fe I',  'lambda_nm': 925.6, 'depth': 0.1, 'width_px': 6},
    {'name': 'Fe II',  'lambda_nm': 685.56, 'depth': 0.1, 'width_px': 8},
    {'name': 'Fe I',  'lambda_nm': 679.60, 'depth': 0.15, 'width_px': 8},
    {'name': 'Fe I',  'lambda_nm': 673.02, 'depth': 0.1, 'width_px': 8},    

    {'name': 'Atmospheric O2 A-band', 'lambda_nm': 760.4, 'depth': 0.60, 'width_px': 7},
    {'name': 'Atmospheric O2 Z-band', 'lambda_nm': 822.7, 'depth': 0.20, 'width_px': 7},
    {'name': 'Atmospheric H2O ρστ','lambda_nm': 928, 'depth': 0.5, 'width_px':5}
]
"""
"""
ABSORPTION_LINES = [G
 {'name': 'Ca II H&K', 'lambda_nm': 393.37, 'depth': 0.65, 'width_px': 8},  # 钙H线
    {'name': 'Ca II H&K', 'lambda_nm': 396.85, 'depth': 0.60, 'width_px': 8},  # 钙K线
    {'name': 'G-band (CH)', 'lambda_nm': 430.0, 'depth': 0.55, 'width_px': 10},  # 分子带
    {'name': 'Mg I', 'lambda_nm': 517.27, 'depth': 0.50, 'width_px': 8},  # 镁三重线
    {'name': 'Mg I', 'lambda_nm': 518.36, 'depth': 0.50, 'width_px': 8},
    {'name': 'Na I D', 'lambda_nm': 589.0, 'depth': 0.65, 'width_px': 12},  # 钠双线
    {'name': 'Na I D', 'lambda_nm': 589.6, 'depth': 0.65, 'width_px': 12},
    
    # 铁线 (众多)
    {'name': 'Fe I', 'lambda_nm': 438.35, 'depth': 0.40, 'width_px': 5},
    {'name': 'Fe I', 'lambda_nm': 527.03, 'depth': 0.35, 'width_px': 5},
    {'name': 'Fe I', 'lambda_nm': 532.80, 'depth': 0.35, 'width_px': 5},
    
    # 钙线
    {'name': 'Ca I', 'lambda_nm': 422.67, 'depth': 0.45, 'width_px': 6},
    
    # 较弱的氢线
    {'name': 'H-alpha', 'lambda_nm': 656.28, 'depth': 0.25, 'width_px': 7},
    {'name': 'H-beta', 'lambda_nm': 486.13, 'depth': 0.20, 'width_px': 6},
    
    # 大气吸收线
    {'name': 'Atmospheric O2 A-band', 'lambda_nm': 760.4, 'depth': 0.60, 'width_px': 7}
]

"""
"""
ABSORPTION_LINES = [K
    # === 极强金属线 ===
    {'name': 'Ca II K', 'lambda_nm': 393.37, 'depth': 0.80, 'width_px': 12},
    {'name': 'Ca II H', 'lambda_nm': 396.85, 'depth': 0.75, 'width_px': 11},
    
    # === 突出的钠双线 ===
    {'name': 'Na I D1', 'lambda_nm': 589.59, 'depth': 0.90, 'width_px': 15},
    {'name': 'Na I D2', 'lambda_nm': 589.00, 'depth': 0.95, 'width_px': 15},
    
    # === 钛/钒等重元素线 ===
    {'name': 'Ti I', 'lambda_nm': 499.11, 'depth': 0.60, 'width_px': 8},
    {'name': 'Ti I', 'lambda_nm': 521.04, 'depth': 0.55, 'width_px': 7},
    {'name': 'V I', 'lambda_nm': 572.75, 'depth': 0.50, 'width_px': 6},
    
    # === 分子带主导 === (K 型星关键特征)
    # TiO 分子带
    {'name': 'TiO α-band', 'lambda_nm': 495.5, 'depth': 0.35, 'width_px': 40},
    {'name': 'TiO γ-band', 'lambda_nm': 544.8, 'depth': 0.30, 'width_px': 35},
    {'name': 'TiO ε-band', 'lambda_nm': 615.0, 'depth': 0.25, 'width_px': 30},
    
    # VO 分子带
    {'name': 'VO γ-band', 'lambda_nm': 740.0, 'depth': 0.20, 'width_px': 50},
    {'name': 'VO δ-band', 'lambda_nm': 790.0, 'depth': 0.18, 'width_px': 45},
    
    # === 极弱氢线 ===
    {'name': 'H-alpha', 'lambda_nm': 656.28, 'depth': 0.15, 'width_px': 5},
    
    # === 大气吸收 ===
    {'name': 'H₂O band', 'lambda_nm': 720, 'depth': 0.40, 'width_px': 60},  # 更宽的水汽带
    {'name': 'O₂ A-band', 'lambda_nm': 760.4, 'depth': 0.60, 'width_px': 10},
    
    # === 众多铁线 === (400-500nm 区域密集)
    {'name': 'Fe I', 'lambda_nm': 404.58, 'depth': 0.55, 'width_px': 6},
    {'name': 'Fe I', 'lambda_nm': 406.36, 'depth': 0.52, 'width_px': 5},
    # 添加更多铁线...
]
"""

"""
ABSORPTION_LINES = [M    
    {'name': 'Ca II K', 'lambda_nm': 393.37, 'depth': 0.4, 'width_px': 5},
    {'name': 'Ca II H', 'lambda_nm': 396.85, 'depth': 0.4, 'width_px': 5},
    {'name': 'Mg II',   'lambda_nm': 430.0,  'depth': 0.4, 'width_px': 6},
    {'name': 'H-beta',  'lambda_nm': 486.1, 'depth': 0.4, 'width_px': 4},
    {'name': 'H-gamma', 'lambda_nm': 434.05, 'depth': 0.4, 'width_px': 4},
    {'name': 'H-p17', 'lambda_nm': 846.7, 'depth': 0.15, 'width_px': 4},
    {'name': 'H-p16', 'lambda_nm': 850.2, 'depth': 0.15, 'width_px': 4},
    {'name': 'H-p15', 'lambda_nm': 854.5, 'depth': 0.15, 'width_px': 4},
    {'name': 'H-p14', 'lambda_nm': 859.8, 'depth': 0.15, 'width_px': 4},
    {'name': 'H-p13', 'lambda_nm': 866.5, 'depth': 0.15, 'width_px': 4},
    {'name': 'H-p12', 'lambda_nm': 875.0, 'depth': 0.15, 'width_px': 4},
    {'name': 'H-p11', 'lambda_nm': 886.3, 'depth': 0.15, 'width_px': 4},
    {'name': 'H-p10', 'lambda_nm': 901.5, 'depth': 0.15, 'width_px': 4},
    {'name': 'H-p19', 'lambda_nm': 922.9, 'depth': 0.15, 'width_px': 4},

    {'name': 'Ca-I1', 'lambda_nm': 849.8, 'depth': 0.2, 'width_px': 4},
    {'name': 'Ca-I2', 'lambda_nm': 854.2, 'depth': 0.2, 'width_px': 4},
    {'name': 'Ca-I3', 'lambda_nm': 866.2, 'depth': 0.2, 'width_px': 4},

    {'name': 'Mg I', 'lambda_nm': 908.7, 'depth': 0.15, 'width_px': 4},

    {'name': 'Ti-I', 'lambda_nm': 911.5, 'depth': 0.15, 'width_px': 4},

    
    {'name': 'Fe-1',  'lambda_nm': 372, 'depth': 0.3, 'width_px': 4},
    {'name': 'Fe-2',  'lambda_nm': 386, 'depth': 0.3, 'width_px': 4},
    {'name': 'Fe-3',  'lambda_nm': 404.6, 'depth': 0.3, 'width_px': 4},
    {'name': 'Fe-4',  'lambda_nm': 458, 'depth': 0.3, 'width_px': 4},
    {'name': 'Fe-5',  'lambda_nm': 492, 'depth': 0.3, 'width_px': 4},
    {'name': 'Fe-6',  'lambda_nm': 517, 'depth': 0.3, 'width_px': 4},
    {'name': 'Fe-7',  'lambda_nm': 438, 'depth': 0.4, 'width_px': 4},
    {'name': 'Fe-7',  'lambda_nm': 427, 'depth': 0.4, 'width_px': 4},
    {'name': 'Fe-8',  'lambda_nm': 527, 'depth': 0.4, 'width_px': 4},     
    {'name': 'Cr-2',  'lambda_nm': 520.8, 'depth': 0.3, 'width_px': 4},
    {'name': 'Cr-3',  'lambda_nm': 425.4, 'depth': 0.3, 'width_px': 4},
    {'name': 'Cr-4',  'lambda_nm': 427.5, 'depth': 0.3, 'width_px': 4},
    {'name': 'Cr-5',  'lambda_nm': 429, 'depth': 0.3, 'width_px': 4},
    {'name': 'Cr-6',  'lambda_nm': 456, 'depth': 0.3, 'width_px': 4},
    {'name': 'Cr-7',  'lambda_nm': 459, 'depth': 0.3, 'width_px': 4},
    {'name': 'Cr-1',  'lambda_nm': 482, 'depth': 0.3, 'width_px': 4},

    {'name': 'Atmospheric O2 A-band', 'lambda_nm': 760.4, 'depth': 0.60, 'width_px': 7},
    {'name': 'Atmospheric O2 Z-band', 'lambda_nm': 822.7, 'depth': 0.20, 'width_px': 7},
    {'name': 'Atmospheric H2O ρστ','lambda_nm': 928, 'depth': 0.5, 'width_px':5}
]
"""

# --- 3. 核心计算函数 ---

def sellmeier(lam_um):
#    lam_um = np.asarray(lam_um)
    lam_sq = lam_um**2
    sum_of_terms = (A1 * lam_sq) / (lam_sq - C1) + \
                   (A2 * lam_sq) / (lam_sq - C2) + \
                   (A3 * lam_sq) / (lam_sq - C3)
    return np.sqrt(sum_of_terms + 1)    


def wavelength_to_pixel_raw(lam_um):
    """
    根据最终确认的棱镜色散公式，将波长(μm)转换为【原始的、未缩放的】像素x坐标。
    """

    lam_um = np.asarray(lam_um)
    """
    lam_sq = lam_um**2
    sum_of_terms = (A1 * lam_sq) / (lam_sq - C1) + \
                   (A2 * lam_sq) / (lam_sq - C2) + \
                   (A3 * lam_sq) / (lam_sq - C3)
    n = np.sqrt(sum_of_terms + 1)
    """
    n = sellmeier(lam_um)
    
    arcsin_arg = np.clip(n * np.sin(PRISM_CONSTANT_RAD), -1.0, 1.0)
    delta_d = 735 * (np.arcsin(arcsin_arg) - PRISM_CONSTANT_RAD)
    x = (108.725 - delta_d) * 317.3
    return x

def planck_law(lam_m, T):
    """
    普朗克定律计算黑体辐射强度。
    """
    h = 6.62607015e-34; c = 2.99792458e8; k_B = 1.380649e-23
    exponent = h * c / (lam_m * k_B * T)
    intensity = (2 * h * c**2) / (lam_m**5) * 1 / (np.exp(np.clip(exponent, -700, 700)) - 1)
    return intensity

def extinction_modifier(lam_um):
    """
    根据给定的消光公式计算光线透过率因子 (0.0 - 1.0)。
    """
#    s = -1.667 * (lam_um**3) + 0.4881 * (lam_um**2) + 1.923 * lam_um - 0.532
    s = -4.1667 * (lam_um**5) + 15.4356 * (lam_um**4) - 17.3201 * (lam_um**3) + \
        0.0611 * (lam_um**2) + 8.685 * lam_um - 2.385
#    s = 17.5 * (lam_um**5) - 66.9508 * (lam_um**4) +105.2595 * (lam_um**3) - \
#        88.8088 * (lam_um**2) + 39.9017 * lam_um - 6.5481
#    s_min, s_max = np.min(s), np.max(s)
#    return (s - s_min) / (s_max - s_min) if s_max > s_min else np.ones_like(lam_um)
    s[s<0.04]=0.04
    step = 1.2/30000
    k = np.absolute((sellmeier(lam_um)-sellmeier(lam_um + step))/step)
    
    return s/k

# --- 4. 主程序 ---

def generate_spectrum_image():
    """
    生成并保存光谱图片的主函数。
    """
    print("开始进行光谱模拟...")

    # --- Part A: 计算一个宽广波长范围内的光谱数据 ---
    print("  - 步骤1: 计算高分辨率光谱数据...")
    # 创建一个覆盖范围足够广、分辨率足够高的波长数组
    source_wavelengths_um = np.linspace(0.3, 1.5, 30000) # 300nm to 1200nm
    source_wavelengths_m = source_wavelengths_um * 1e-6

    # --- Part B: 计算每个波长的最终强度 ---
    continuum_intensity = planck_law(source_wavelengths_m, STAR_TEMPERATURE)
    extinction_effect = extinction_modifier(source_wavelengths_um)
    #extinction_effect = 1
    continuum_after_extinction = continuum_intensity * extinction_effect
    
    absorption_mask = np.ones_like(source_wavelengths_um)
    
    # 计算局部色散 dx/dλ (在原始坐标系中)
    dx_raw_dlambda = np.gradient(wavelength_to_pixel_raw(source_wavelengths_um), source_wavelengths_um)

    for line in ABSORPTION_LINES:
        lam_c_um = line['lambda_nm'] / 1000.0
        depth = line['depth']
        width_px = line['width_px']
        idx = np.argmin(np.abs(source_wavelengths_um - lam_c_um))
        local_dispersion_raw = dx_raw_dlambda[idx]
        width_um = (width_px / abs(local_dispersion_raw)) / 2.355 if abs(local_dispersion_raw) > 0 else 0.001
        gaussian_dip = depth * np.exp(-((source_wavelengths_um - lam_c_um)**2) / (2 * width_um**2))
        absorption_mask -= gaussian_dip

    spectrum_intensity = continuum_after_extinction * absorption_mask
    
    # 归一化
    final_source_intensities = np.clip(spectrum_intensity, 0, None)
    max_continuum = np.max(continuum_after_extinction)
    final_source_intensities /= max_continuum if max_continuum > 0 else 1.0

    # --- Part C: 将光谱数据投影到图像像素网格上 ---
    print("  - 步骤2: 将光谱投影到图像...")
    
    # 1. 计算源光谱中每个点对应的原始像素坐标
    raw_pixel_coords = wavelength_to_pixel_raw(source_wavelengths_um)

    # 2. 创建一个空的图像强度数组
    image_intensities = np.zeros(IMG_WIDTH)

    # 3. 使用插值将源数据“绘制”到图像网格上
    #    我们只关心落在[0, 2399]范围内的坐标
    
    # 筛选出有效的坐标和对应的强度
    valid_mask = (raw_pixel_coords >= 0) & (raw_pixel_coords < IMG_WIDTH)
    valid_coords = raw_pixel_coords[valid_mask]
    valid_intensities = final_source_intensities[valid_mask]
    
    # 确保坐标是单调递增的，以供插值函数使用
    sort_indices = np.argsort(valid_coords)
    sorted_coords = valid_coords[sort_indices]
    sorted_intensities = valid_intensities[sort_indices]
    
    # 创建最终的像素网格
    pixel_grid = np.arange(IMG_WIDTH)
    
    # 进行插值
    # np.interp(x, xp, fp) -> x:目标点, xp:数据点x, fp:数据点y
    # left=0, right=0 表示在插值范围之外的区域强度为0（黑色）
    image_intensities = np.interp(pixel_grid, sorted_coords, sorted_intensities, left=0.0, right=0.0)

    # --- Part D: 创建并渲染图像 ---
    print("  - 步骤3: 生成并保存图像...")
    
    image_data = np.tile(image_intensities, (IMG_HEIGHT, 1))
    image_data_uint8 = (image_data * 185).astype(np.uint8)+70
    
    img = Image.fromarray(image_data_uint8, 'L')
    output_filename_raw = 'B_Type_star_spectrogram_raw.png'
    img.save(output_filename_raw)
    print(f"\n纯光谱灰度图已保存为: {output_filename_raw}")

    # --- Part E: 创建带坐标轴的预览图 ---
    
    dpi = 72
    fig, ax = plt.subplots(figsize=(IMG_WIDTH / dpi, (IMG_HEIGHT + 50) / dpi), dpi=dpi)

    ax.imshow(image_data_uint8, cmap='gray', aspect='auto', vmin=0, vmax=255)
    
    ax.set_yticks([]); ax.set_frame_on(False)
    
    # 设置X轴刻度
    tick_wavelengths_nm = np.arange(300, 1201, 50) # 检查一个更宽的刻度范围
    tick_wavelengths_um = tick_wavelengths_nm / 1000.0
    tick_pixel_positions = wavelength_to_pixel_raw(tick_wavelengths_um)
    
    # 筛选出在当前图像范围内的刻度
    valid_ticks_mask = (tick_pixel_positions >= 0) & (tick_pixel_positions < IMG_WIDTH)
    
    ax.set_xticks(tick_pixel_positions[valid_ticks_mask])
    ax.set_xticklabels([f"{wl}" for wl in tick_wavelengths_nm[valid_ticks_mask]])
    ax.tick_params(axis='x', colors='black', bottom=True, top=False, labelbottom=True)
    
    ax.set_xlabel("Wavelength (nm)", color='black', labelpad=10)
    
    plt.tight_layout(pad=0)
    
    output_filename_annotated = 'B_type_star_spectrogram_final_annotated.png'
    plt.savefig(output_filename_annotated)
    print(f"带坐标轴的光谱图已保存为: {output_filename_annotated}")
    
    plt.show()

if __name__ == '__main__':
    generate_spectrum_image()