import numpy as np
import matplotlib.pyplot as plt
from scipy.special import erfc

# 设置参数
rho = 0.7798
z0 = 4.25  # 确保使用大写 Z0，与定义一致

def calculate_tau(z, z_prime, rho, s_values):
    return np.log(rho * z * z_prime * s_values / 2)

# 定义 chi_c（Conformal Kernel）
def chi_c(z, z_prime, rho, tau):
    return np.exp((1.0 - rho) * tau) * np.exp(-np.log(z / z_prime) ** 2.0 / (rho * tau)) / np.sqrt(tau)

# 定义 F 函数
def F(z, z_prime, rho, tau):
    eta = (-np.log(z * z_prime / (z0 ** 2.0)) + rho * tau) / np.sqrt(rho * tau)
    return 1.0 - 2.0 * np.sqrt(rho * np.pi * tau) * np.exp(eta ** 2.0) * erfc(eta)

# 定义 chi_mod（Hard-Wall Kernel）
def chi_mod(z, z_prime, rho, tau):
    return chi_c(z, z_prime, rho, tau) + F(z, z_prime, rho, tau) * chi_c(z, z0 ** 2.0 / z_prime, rho, tau)

# 固定 z' 值
z_prime = 4.0
z_values_list = [0.1, 0.2, 0.4, 0.8, 1.0]  # 不同的 z 值

# 创建 s 的值
s_values = np.logspace(1, 5, 1000)  # 从 10 到 10^7

# 创建绘图
plt.figure(figsize=(6, 4))

for z in z_values_list:
    tau_values = calculate_tau(z, z_prime, rho, s_values)  # 根据 z 计算 tau
    chi_values = chi_c(z, z_prime, rho, tau_values)
    chi_mod_values = chi_mod(z, z_prime, rho, tau_values)
    # 避免除以零
    r = np.divide(chi_mod_values, chi_values, where=chi_values != 0)
    # 绘制不同 z 的曲线
    if z == 0.1:
        plt.plot(s_values, r, label=f'z = {z} (GeV$^{{-1}}$)', linestyle='-', color='red')
    elif z == 0.2:
        plt.plot(s_values, r, label=f'z = {z} (GeV$^{{-1}}$)', linestyle='--', color='green')
    elif z == 0.4:
        plt.plot(s_values, r, label=f'z = {z} (GeV$^{{-1}}$)', linestyle=':', color='coral')
    elif z == 0.8:
        plt.plot(s_values, r, label=f'z = {z} (GeV$^{{-1}}$)', linestyle='-.', color='blue')
# 添加图例与标签
plt.legend(loc='lower left', bbox_to_anchor=(0.01, 0.01))
plt.text(0.45, 0.92, r"$z_0 = 4.25$ (GeV$^{-1})$", fontsize=12, ha='center', va='center', transform=plt.gca().transAxes)
plt.xlabel(r"$s$", fontsize=12)
plt.ylabel(r"$r(s, z, 4.0)$", fontsize=12)
plt.xscale('log')
plt.legend()
plt.xlim(1e1, 1e5)
plt.ylim(0, 1.4)  # 根据需要调整 y 轴范围
plt.show()