import sys
sys.path.append('D:/my_gprmax/gprMax')  # Add gprMax installation path to system for module import

import os
from gprMax.gprMax import api
import numpy as np
from tools.plot_Bscan import get_output_data, mpl_plot
import matplotlib.pyplot as plt
import scipy.ndimage
from scipy.signal import tukey
from config import Forward_Model_Config as cfg
from config import Path_Config as pcfg


# Get the current script directory
path = os.getcwd()
root = pcfg.in_data_dir               # Specify the input file directory
files = os.listdir(root)              # Get the list of files in the directory
data_per_ns = cfg.data_per_ns         # Sampling rate per nanosecond
direct_wave_time = cfg.direct_wave_time*1e9   # Direct wave duration (ns)
time = cfg.Time*1e9                       # Time window (ns)
static_time = cfg.static_time**1e9        # Static correction time (ns)


for file in files:
    if file.endswith('.in'):          # Process only .in files
        filename = root + file        # Full path of the input file
        fi = filename[:-3]            # Remove .in extension while keeping the path
        
        # Run gprMax simulation (geometry_only=False means full simulation)
        api(filename, n=1, gpu={0}, geometry_only=False)
        
        filename_b = fi + '.out'
        rxnumber = 1
        rxcomponent = 'Ez'
        
        # Retrieve reflected wave data
        outputdata, dt = get_output_data(filename_b, rxnumber, rxcomponent)
        h = len(outputdata)
        # Load direct wave data for background removal
        zhidabo_data = np.loadtxt('./dataset/direct_wave_generation/out_data/000' + file[6:8] + '.txt')
        outputdata = outputdata - zhidabo_data  # Remove direct wave background
        # Resample the data to match the required time resolution
        outputdata = scipy.ndimage.zoom(outputdata, (int((time + static_time) * data_per_ns) / h), order=1)
        # Apply static correction
        static_output = outputdata[int(static_time * data_per_ns):]
        outputdata = static_output
        # Zero out the initial direct wave
        outputdata[:int(direct_wave_time * data_per_ns)] = 0
        
        # Apply dewow filter to remove low-frequency drift
        dewow = np.mean(outputdata[int(direct_wave_time * data_per_ns):])
        outputdata[int(direct_wave_time * data_per_ns):] -= dewow
        
        # Save processed data
        np.savetxt(pcfg.out_data_dir + file[:-6] + '.txt', outputdata, delimiter=' ')
        
        # Plot and save the processed data
        plt.figure()
        plt.plot(outputdata, linestyle='-', color='b')  
        plt.title("1D Data Curve")
        plt.xlabel("Index")
        plt.ylabel("Value")
        plt.grid(True)  
        plt.savefig('./dataset/out_data_img/' + file[:-6] + '.png')
        
        # Remove original files after processing
        os.remove(filename)
        os.remove(filename_b)