import segyio
import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import tukey, hilbert
import mig_fk  # Custom FK migration module


def write_to_dat_append(data, x, y, file):
    """
    Append 3D point cloud data to a .dat file in the format (x, y, z, value)
    
    Parameters:
        data (2D array): Processed data to export
        x (float): X offset for the current profile
        y (float): Y offset for the current profile
        file (file object): Output file opened in write/append mode
    """
    nz, nx = data.shape
    x_coords = np.linspace(x + 8, 213, nx)
    z_coords = np.linspace(0, 85, nz)

    for i in range(nz):
        for j in range(nx):
            file.write(f"{x_coords[j]} {y + 3} {z_coords[i]} {data[i, j]}\n")


def sliding_window_mean(data, window_size):
    """
    Compute local background mean using a sliding window
    
    Parameters:
        data (2D array): Input data of shape (n_traces, n_samples)
        window_size (int): Size of the sliding window (must be odd)
    
    Returns:
        2D array: Local mean values with the same shape as input
    """
    if window_size % 2 == 0:
        raise ValueError("Window size must be an odd number.")

    half_window = window_size // 2
    padded_data = np.pad(data, ((0, 0), (half_window, half_window)), mode='reflect')
    local_means = np.empty_like(data)

    for i in range(data.shape[1]):
        local_means[:, i] = np.mean(padded_data[:, i:i + window_size], axis=1)

    return local_means


def taper_field(in_array): 
    """
    Apply a Tukey window-based taper to suppress edge noise
    
    Parameters:
        in_array (1D array): Input radar trace
    
    Returns:
        1D array: Tapered trace
    """
    window = tukey(500, 0.3)
    window[65:] = 1  # Force central part to remain unchanged
    window = window ** 2
    return in_array * window


def load(filename, static, trace_remove, x, y, sub_aver, amplify):
    """
    Load and process GPR data from a SEG-Y file
    
    Parameters:
        filename (str): Path to SEG-Y file
        static (int): Number of top samples to remove (static correction)
        trace_remove (int): Number of traces to remove from the end
        x (float): X-coordinate offset
        y (float): Y-coordinate offset
        sub_aver (int): Window size for background subtraction
        amplify (float): Amplitude scale factor for visualization
    """
    with segyio.open(filename, 'r', ignore_geometry=True) as segyfile:
        raw_traces = segyfile.trace.raw[:]
        data_meta = raw_traces.T[static:, :-trace_remove]
        data=data_meta[0:500]

        # Subtract local mean background (clutter removal)
        reference_mean = sliding_window_mean(data, sub_aver)
        data = data - reference_mean

        # Apply tapering to suppress edge effects
        data = np.apply_along_axis(taper_field, 0, data)

        # Perform FK migration (custom)
        data, twtt, mig_pos = mig_fk.fkmig(data, dt=0.02425e-9, dx=0.005, v=1.10e8)

        # Compute envelope using Hilbert transform
        data_envelop = np.abs(hilbert(data, axis=0))

    # Plot migrated image
    plt.title(filename[7:-4])
    plt.imshow(data * amplify, cmap='gray_r', aspect='auto',
               vmin=-np.max(np.abs(data)), vmax=np.max(np.abs(data)))
    plt.colorbar()
    plt.show()

    # Plot envelope
    plt.title(filename[7:-4] + " Envelope")
    plt.imshow(data_envelop * amplify, cmap='gray_r', aspect='auto',
               vmin=0, vmax=0.0040)
    plt.colorbar()
    plt.show()

    # Save envelope to .dat file
    write_to_dat_append(data_envelop, x, y, file)


# Main process for batch loading and writing
with open('3D.dat', 'w') as file:
    load('./data/x0y0.sgy', static=35, trace_remove=1, x=0, y=0, sub_aver=51, amplify=3)
    load('./data/x0y5.sgy', static=35, trace_remove=1, x=0, y=5, sub_aver=51, amplify=5)
    load('./data/x0y10.sgy', static=35, trace_remove=1, x=0, y=10, sub_aver=51, amplify=10)
    load('./data/x0y15.sgy', static=35, trace_remove=1, x=0, y=15, sub_aver=51, amplify=10)
    load('./data/x0y20.sgy', static=35, trace_remove=1, x=0, y=20, sub_aver=51, amplify=10)
    load('./data/x10y25.sgy', static=35, trace_remove=1, x=10, y=25, sub_aver=51, amplify=10)
    load('./data/x20y30.sgy', static=35, trace_remove=1, x=20, y=30, sub_aver=51, amplify=10)
    load('./data/x30y35.sgy', static=35, trace_remove=15, x=30, y=35, sub_aver=51, amplify=10)
    load('./data/x30y40.sgy', static=35, trace_remove=1, x=30, y=40, sub_aver=51, amplify=10)
    load('./data/x40y45.sgy', static=35, trace_remove=1, x=40, y=45, sub_aver=51, amplify=10)
    load('./data/x45y50.sgy', static=35, trace_remove=14, x=45, y=50, sub_aver=51, amplify=10)