import numpy as np
import matplotlib.pyplot as plt
import os

def calculate_phase(images): 
    images = np.array(images) 
    shifts = len(images) 
    phase = np.arctan2(np.sum(images * np.sin(np.linspace(0, 2 * np.pi, shifts)[:, np.newaxis, np.newaxis]), axis=0), 
                       np.sum(images * np.cos(np.linspace(0, 2 * np.pi, shifts)[:, np.newaxis, np.newaxis]), axis=0)) 
    return phase 
def unwrap_phase(phase): 
    unwrapped_phase = np.unwrap(phase, axis=0) 
    unwrapped_phase = np.unwrap(unwrapped_phase, axis=1) 
    return unwrapped_phase 

def find_subpixel_extrema(phase_map): 
    # 找到最大值和最小值的索引 
    min_pos = np.unravel_index(np.argmin(phase_map), phase_map.shape) 
    max_pos = np.unravel_index(np.argmax(phase_map), phase_map.shape) 
    # 返回亚像素坐标 
    return np.array(min_pos), np.array(max_pos)


def generate_phase_shifted_patterns(width, height, shifts, orientation='horizontal', frequency=1):
    """
    Generates multiple phase-shifted fringe patterns, either horizontal or vertical, with phase shifting along the stripes.
    :param width: image width
    :param height: image height
    :param shifts: number of phase shifts
    :param orientation: 'horizontal' or 'vertical' indicating the direction of the stripes
    :param frequency: number of full wave cycles across the dimension
    :return: array of generated fringe patterns
    """
    patterns = []
    for i in range(shifts):
        shift = (i * 2 * np.pi) / shifts  # Calculating the phase shift for each pattern
        if orientation == 'horizontal':
            x = np.linspace(0, 2 * np.pi * frequency, width) + shift  # Apply frequency scaling on horizontal stripes
            print('frequency, horizontal x = ', x, frequency)
            print('step = ', 2 * np.pi * frequency )
            pattern = np.sin(x)
            print('pattern = ', pattern)
            pattern = np.tile(pattern, (height, 1))  # Extend the pattern vertically across the height
        elif orientation == 'vertical':
            y = np.linspace(0, 2 * np.pi * frequency, height) + shift  # Apply frequency scaling on vertical stripes
            pattern = np.sin(y)
            pattern = np.tile(pattern.reshape(-1, 1), (1, width))  # Extend the pattern horizontally across the width
        patterns.append(pattern)
    return patterns

# Image size parameters
width = 1920
height = 1200
width = 800
height = 400
shifts = 4  # Number of phase shifts
frequencies = [32, 8 ,1]  # Two different frequencies
#frequencies = [32, 1]  # Two different frequencies

# Ensure the directory exists
directory = 'patterns3_honor_800_400'
os.makedirs(directory, exist_ok=True)

# Pattern naming and saving
pattern_index = 0  # Start indexing from 0

for freq in frequencies:
    # Generate horizontal patterns first, then vertical for each frequency
    horizontal_patterns = generate_phase_shifted_patterns(width, height, shifts, 'horizontal', freq)
    vertical_patterns = generate_phase_shifted_patterns(width, height, shifts, 'vertical', freq)
    all_patterns = horizontal_patterns + vertical_patterns  # Combine lists to maintain the order

    # phase = np.arctan2(
    #     np.sum(horizontal_patterns * np.sin(np.linspace(0, 2 * np.pi, shifts)[:, np.newaxis, np.newaxis]), axis=0),
    #     np.sum(horizontal_patterns * np.cos(np.linspace(0, 2 * np.pi, shifts)[:, np.newaxis, np.newaxis]), axis=0)
    #     )
    # print('截断相位 = ', phase)
    # y_coords, x_coords = np.meshgrid(np.arange(height), np.arange(width), indexing='ij')

    # # 提取所有像素坐标和截断相位值
    # coordinates = np.stack((y_coords, x_coords), axis=-1)  # 生成每个像素的坐标
    # target_mask = (phase >= -1) & (phase <= 1)
    # target_coordinates = coordinates[target_mask]  # 提取目标范围内的像素坐标

    # print("对应的像素坐标：", target_coordinates)
    # # 可视化截断相位
    # plt.imshow(phase, cmap='jet')
    # plt.colorbar(label='Phase (wrapped)')
    # plt.scatter(target_coordinates[:, 1], target_coordinates[:, 0], color='red', s=1, label='Target Pixels')
    # plt.title("Wrapped Phase with Target Pixels")
    # plt.legend()
    # plt.show()
    # 计算相位图 
    # phase_map = calculate_phase(horizontal_patterns) 
    # # 展开相位图 
    # unwrapped_phase_map = unwrap_phase(phase_map) 
    # print('unwrapped_phase_map = ', unwrapped_phase_map.shape)
    # # 找到亚像素级别的极值点坐标 
    # min_subpixel, max_subpixel = find_subpixel_extrema(unwrapped_phase_map) 
    # # 可视化相位图并标注极值点 
    # plt.imshow(unwrapped_phase_map, cmap='jet') 
    # plt.colorbar() 
    # plt.scatter([min_subpixel[1]], [min_subpixel[0]], color='red', label='Min') 
    # plt.scatter([max_subpixel[1]], [max_subpixel[0]], color='blue', label='Max') 
    # plt.legend() 
    # plt.title("Unwrapped Phase Map with Extremes") 
    # plt.show() 
    # print(f"Min subpixel position: {min_subpixel}") 
    # print(f"Max subpixel position: {max_subpixel}")

    # Save the patterns with the new naming convention in the specified directory
    for pattern in all_patterns:
        file_path = os.path.join(directory, f'pat{pattern_index:02}.bmp')
        print('pattern shape =',pattern.shape)
        print(pattern)
        plt.imsave(file_path, pattern, cmap='gray')
        pattern_index += 1