import os
import argparse

import numpy as np
import matplotlib.pyplot as plt

def parse_args():
    parser = argparse.ArgumentParser(description="")
    parser.add_argument("--shifts", type=int, default=4, help="相移步数")
    parser.add_argument("--img_w", type=int, default=2360, help="图片宽度")
    parser.add_argument("--img_h", type=int, default=1640, help="图片高度")
    parser.add_argument("--freq_x", type=int, default=32, help="x轴相移频率")
    parser.add_argument("--freq_y", type=int, default=1, help="y轴相移频率")
    args = parser.parse_args()
    return args

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
            pattern = np.sin(x)
            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


if __name__ == "__main__":

    args = parse_args()
    width = args.img_w
    height = args.img_h

    shifts = args.shifts  # Number of phase shifts
    frequencies = [args.freq_x, args.freq_y]  # Two different frequencies

    # Ensure the directory exists
    directory = f'imgs/patterns_{width}x{height}_{shifts}_fx{args.freq_x}_fy{args.freq_y}'
    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

        # 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')
            plt.imsave(file_path, pattern, cmap='gray')
            pattern_index += 1
    print(f"Fringe pattern with shift = {shifts}, freqs = [{args.freq_x}, {args.freq_x}] has been saved to {directory}, with number {pattern_index}.")