#result.py

import numpy as np
import pkg_resources
import matplotlib.pyplot as plt
from numpy.fft import fftshift, fft2, ifftshift, fftfreq, ifft2


def generate_result(sin_val, n_val, wavelength_val, m, b, op_status):
    """
    生成计算结果图
    
    参数：
    sin_val (float): 传入的 sin 值
    n_val (float): 传入的折射率值
    wavelength_val (float): 传入的波长值
    m (int): 传入的 m 值，0 为相干光，1 为非相干光
    b (int): 传入的 b 值, 0 为不启用 OPC，1 为启用 OPC
    op_status (int): 0 表示第二行输出光强分布，1 表示输出图像边界
    
    返回：
    fig (matplotlib.figure.Figure): Matplotlib 图形容器
    """
    # 固定参数
    aperture_diameter = focal_length = 1.0

    # 计算入射角 theta
    theta = np.arcsin(sin_val)

    # 计算数值孔径
    NA = n_val * np.sin(theta)

    # 波长
    wavelength = wavelength_val

    # 计算透镜孔径 D
    lens_aperture = 2 * focal_length * np.tan(theta)

    # 实空间坐标架、动量空间坐标架及光学传递函数
    num_samples = 100  # 一个维度的采样数
    p = 1 # 修正离散化导致的误差
    num_samples_extended = (2 * p + 1) * num_samples

    spatial_x = np.linspace(-0.6 * aperture_diameter, 0.6 * aperture_diameter, num_samples)
    spatial_y = spatial_x.copy()
    spatial_X, spatial_Y = np.meshgrid(spatial_x, spatial_y)
    spatial_Y = spatial_Y[::-1, :]

    kx = fftshift(fftfreq(num_samples_extended,1.2*aperture_diameter*(2*p+1)/num_samples_extended))
    ky = kx.copy()
    Kx, Ky = np.meshgrid(kx, ky)
    k = np.sqrt(Kx**2 + Ky**2)
    K = 2 * np.pi * NA / wavelength
    optical_transfer_function = np.zeros((2, num_samples_extended, num_samples_extended))
    optical_transfer_function[0][0:num_samples_extended][0:num_samples_extended] = k <= K

    for i in range(num_samples_extended):
        for j in range(num_samples_extended):
            if k[i][j] < K:
                optical_transfer_function[1][i][j] = np.arccos(k[i][j] / 2 / K) - k[i][j] / 2 / K * np.sqrt(1 - (k[i][j] / 2 / K)**2)

    # 加载文件
    file_path = pkg_resources.resource_filename(__name__, "mask.npy")
    amask = np.load(file_path)
    bmask=np.zeros((num_samples_extended,num_samples_extended))

    def calculate_intensity(mask):
        bmask[p*num_samples:(p+1)*num_samples:1,p*num_samples:(p+1)*num_samples:1]=mask
        diffracted = fftshift(fft2(ifftshift(bmask)))
        transmitted=diffracted*optical_transfer_function[m]
        E=fftshift(ifft2(ifftshift(transmitted)))
        intensity=np.abs(E)**2
        output=np.zeros((num_samples,num_samples))
        output=intensity[p*num_samples:(p+1)*num_samples,p*num_samples:(p+1)*num_samples]
        return output/output.max()


    def calculate_pixel_error(T, threshold, mask):
        mp = np.zeros((num_samples, num_samples))
        mp[calculate_intensity(T) > threshold] = 1
        return np.sum(np.abs(mp - mask).reshape(num_samples**2, 1))


    def calculate_contour(T):
        res = 1
        i = 0
        temp = num_samples**2
        while i <= 35:
            t = 0.4 - 0.01 * i
            if calculate_pixel_error(T, t, T) <= temp:
                temp = calculate_pixel_error(T, t, T)
                res = t
            i = i + 1
        Im = calculate_intensity(T)
        edge = np.zeros((num_samples, num_samples))
        edge[Im > res] = 1
        return edge


       # 创建 Matplotlib 图形容器，祖传配方3:1
    fig, axs = plt.subplots(1, 5, figsize=(12, 4))

    # 绘制光强分布或者图像边界，根据op_status决定
    for i in range(5):
        if op_status == 0:
            intensity = calculate_intensity(amask[i + 5 * b])
            axs[i].imshow(intensity, cmap='viridis', extent=(spatial_x.min(), spatial_x.max(), spatial_y.min(), spatial_y.max()))
            axs[i].set_title(f'Intensity {i}')
        else:
            axs[i].imshow(calculate_contour(amask[i + 5 * b]), cmap='gray', extent=(spatial_x.min(), spatial_x.max(), spatial_y.min(), spatial_y.max()))
            axs[i].set_title(f'Contour {i}')

    # 调整子图布局
    plt.tight_layout()

    return fig  # 返回 Matplotlib 图形容器

# 计算结果图函数调用示例
#plotter = generate_result(0.9, 1, 2.0, 1, 0, 0)
#plt.show()