from scipy import integrate
import math
import numpy as np
import os
# import sys
# import pandas as pd
import numpy as np
import matplotlib.pyplot as plt



def w1(lam, z, d, x):
    return - math.sqrt(2/(lam*z)) * (d + x) / 2


def w2(lam, z, d, x):
    return math.sqrt(2/(lam*z)) * (d - x) / 2


def Fc(w):
#     val, err = integrate.quad(lambda tau : math.cos(math.pi * tau**2 / 2), 0, w)
    val = integrate.quad(lambda tau : math.cos(math.pi * tau**2 / 2), 0, w)[0]
    return val


def Fs(w):
#     val, err = integrate.quad(lambda tau : math.sin(math.pi * tau**2 / 2), 0, w)
    val = integrate.quad(lambda tau : math.sin(math.pi * tau**2 / 2), 0, w)[0]
    return val


# 光强 ( 中心强度 I0, 波长 lambda, 等效距离 z, 狭缝宽度(非2倍) d, 观察位置 x )
def Intensity(I0, lam, z, d, x):
    Fcw1 = Fc(w1(lam, z, d, x))
    Fcw2 = Fc(w2(lam, z, d, x))
    Fsw1 = Fs(w1(lam, z, d, x))
    Fsw2 = Fs(w2(lam, z, d, x))
    return I0/2 * ((Fcw2 - Fcw1)**2 + (Fsw2 - Fsw1)**2)


# 返回值为图片名称
# 中心强度 I0, 波长 lambda(nm), 等效距离 z(mm), 狭缝宽度(非2倍) d(mm), 观察位置x(mm)
def plotIntensity(I0, lam, z, d, xRangeDiv, folder):
    z = round(z, 8)
    d = round(d, 8)
    figYmax = 1.8 * I0  # 自动调整高度
    # 生成存储空间
    x = np.arange(xRangeDiv[0], xRangeDiv[1], xRangeDiv[2])  # 此时仍然为mm单位
    y = np.zeros(len(x))
    # 进行计算
    for i in range(len(x)):
        y[i] = Intensity(I0, lam*1e-9, z*1e-3, d*1e-3, x[i]*1e-3)
        if y[i] > figYmax:
            print("At position", x[i], ", intensity is", y[i], ">", figYmax)
    # 绘图
    plt.figure(figsize=(8, 4), dpi=256)
    plt.xlabel('Position $x$ ($\mathrm{mm}$)')
    plt.ylabel('Relative Light Intensity ($\mathrm{arb}$)', rotation=90)
    plt.title('Fresnel Diffraction of a Single Slit\n'
              + 'wave length $\lambda = $' + str(lam) + '$\mathrm{nm}$' + ' / '
              + 'equivalent distance $z = $' + str(z) + '$\mathrm{nm}$' + ' / '
              + 'slit width $d = $' + str(d) + '$\mathrm{mm}$')
    plt.plot(x, y, linewidth=1);
    plt.ylim(-figYmax/40, figYmax)
    plt.grid();
    plt.show()
    figName = "lam" + str(lam) + "_z" + str(z) + "_d" + str(d) + ".png";  # 保存文件
#     plt.savefig(folder + "/" + figName);  # 保存文件
    plt.close()
    return figName


workFolder = "/home/nqhq/Python/PhyExpr/0607fresnel"
dataFolder = "/home/nqhq/Python/PhyExpr/0607data"
figFolder = "/home/nqhq/Python/PhyExpr/0607fig"
os.chdir(workFolder)
# myDatas = os.listdir(os.getcwd())
# print(myDatas)


if __name__ == '__main__':
    plotIntensity(17.5, 632.8, 236, 1.2, [-2,2,0.01], workFolder)
    print("Programme Exits.")