import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.interpolate import interp1d
from config import Field_data_test_Config as cfg
from config import Path_Config as pcfg

def shift_data_to_end(data, n):
    """
    Shift the first n values of a 1D data array to the end.
    
    Parameters:
        data (list or numpy array): 1D data array
        n (int): Number of elements to shift
        
    Returns:
        numpy array: Shifted data
    """
    if n < 0 or n > len(data):
        raise ValueError("Shift length n must be between 0 and the data length.")
    
    return np.concatenate((data[n:], data[:n]))

# Step 1: Read 1D data from CSV or TXT file
def read_file(file_path, idx):
    if file_path.endswith('.csv'):
        data = pd.read_csv(file_path, header=None)  # Assuming no column names
        one_d_data = data.iloc[1:, idx].values  # Read data from the idx-th column
    elif file_path.endswith('.txt'):
        data = np.loadtxt(file_path)  # Assuming TXT file contains only numeric values
        one_d_data = data if data.ndim == 1 else data[:, 0]  # Ensure it's 1D
    else:
        raise ValueError("Unsupported file format. Please provide a .csv or .txt file.")
    return one_d_data

# Step 2: Normalize data by dividing by its maximum absolute value
def normalize_data(data):
    max_abs_value = np.max(np.abs(data))
    return data if max_abs_value == 0 else data / max_abs_value  # Avoid division by zero

# Step 3: Interpolate data to a length of 1000
def interpolate_data(data, target_length=1000):
    original_length = len(data)
    x_original = np.linspace(0, 1, original_length)
    x_target = np.linspace(0, 1, target_length)
    interpolator = interp1d(x_original, data, kind='cubic')  # Cubic interpolation
    return interpolator(x_target)

# Step 4: Save processed data to a CSV file
def save_to_csv(data, output_path):
    pd.DataFrame(data).to_csv(output_path, index=False, header=False)

# Step 5: Plot original and interpolated data
def plot_data(original_data, interpolated_data):
    plt.figure(figsize=(10, 5))
    plt.plot(original_data, label='Original Data', marker='o', linestyle='--', alpha=0.7)
    plt.plot(np.linspace(0, len(original_data)-1, len(interpolated_data)), interpolated_data, label='Interpolated Data', linestyle='-', linewidth=2)
    plt.legend()
    plt.xlabel('Index')
    plt.ylabel('Value')
    plt.title('Data Normalization and Interpolation')
    plt.grid(alpha=0.3)
    plt.show()

# Main execution
if __name__ == "__main__":
    input_file = pcfg.PROCESSED_TEST_FILE  # Input file path
    output_csv = pcfg.field_impulse  # Output file path   
    idx = cfg.refer_wave_idx  # Column index to read data from
    static_time=cfg.static_time
    wavelet_range=cfg.wavelet_range
    data = read_file(input_file, idx)
    normalized_data = normalize_data(data)
    
    # Zero out unwanted data regions
    normalized_data[wavelet_range[1]:] = 0  # Set values after the source wave to zero
    normalized_data[0:wavelet_range[0]] = 0  # Set values before the source wave to zero
    
    interpolated_data = interpolate_data(normalized_data,cfg.time_window_length)
    interpolated_data = shift_data_to_end(interpolated_data, static_time)
    
    save_to_csv(interpolated_data, output_csv)
    plot_data(normalized_data, interpolated_data)