import numpy as np
from scipy.constants import h, c, k
from scipy.integrate import simps
import matplotlib.pyplot as plt

from org.q1 import emissivity


# ------------------------------------------------
# 1. 引用第一题的 emissivity() 函数
# ------------------------------------------------
# (假设已运行第一题代码)
# emissivity(lambda_um, thickness_um)


# ------------------------------------------------
# 2. 黑体辐射谱
# ------------------------------------------------
def planck_lambda(lambda_um, T):
    lambda_m = lambda_um * 1e-6
    return (2*h*c**2)/(lambda_m**5) / (np.exp(h*c/(lambda_m*k*T)) - 1)

# ------------------------------------------------
# 3. 大气透过率（示例，可替换 MODTRAN 数据）
# ------------------------------------------------
def atmosphere_tau(lambda_um):
    tau = np.ones_like(lambda_um)
    window = (lambda_um >= 8) & (lambda_um <= 13)
    tau[window] = 0.8
    tau[~window] = 0.1
    return tau

# ------------------------------------------------
# 4. 太阳辐照谱（简化 AM1.5）
# ------------------------------------------------
def solar_spectrum(lambda_um):
    I0 = 1.5e3  # simplified scaling
    return I0 * np.exp(-(lambda_um - 0.5)**2 / 0.4)

# ------------------------------------------------
# 5. 能量项计算
# ------------------------------------------------
def P_rad(Ts, lambda_um, eps):
    return simps(eps * planck_lambda(lambda_um, Ts), lambda_um)

def P_atm(Ta, lambda_um, eps):
    return simps(eps * planck_lambda(lambda_um, Ta) * (1 - atmosphere_tau(lambda_um)), lambda_um)

def P_solar(lambda_um, eps):
    A = eps  # approximate absorption
    return simps(A * solar_spectrum(lambda_um), lambda_um)

def P_conv(Ts, Ta, h=5):
    return h * (Ts - Ta)

# ------------------------------------------------
# 6. 净冷却功率
# ------------------------------------------------
def P_net(Ts, Ta, lambda_um, eps, h=5):
    return P_rad(Ts, lambda_um, eps) - P_atm(Ta, lambda_um, eps) - P_solar(lambda_um, eps) - P_conv(Ts, Ta, h)

# ------------------------------------------------
# 7. 求稳态温度（解 Pnet=0）
# ------------------------------------------------
def solve_temperature(Ta, lambda_um, eps):
    Ts = Ta
    for _ in range(1000):
        f = P_net(Ts, Ta, lambda_um, eps)
        Ts -= 0.1 * f  # simple iteration
    return Ts

# ------------------------------------------------
# 8. 运行示例
# ------------------------------------------------
lambda_range = np.linspace(0.35, 20, 2000)
d_pdms = 10
eps = emissivity(lambda_range, d_pdms)

Ta = 300  # 环境温度
Ts = solve_temperature(Ta, lambda_range, eps)

print("Steady-state temperature:", Ts)

# 曲线可视化
T_list = np.linspace(250, 330, 200)
P_list = [P_net(T, Ta, lambda_range, eps) for T in T_list]

plt.figure(figsize=(10,5))
plt.plot(T_list, P_list)
plt.axhline(0, color='r')
plt.xlabel("Temperature (K)")
plt.ylabel("Net Cooling Power")
plt.title("Cooling Power vs Temperature")
plt.grid(True)
plt.show()