import numpy as np
from scipy.interpolate import interp1d
import matplotlib.pyplot as plt
from config import Network_prediction_Config as cfg
from config import Path_Config as pcfg
from utils.plot import plot_depth_permittivity_constant

def convert2depth(dielectric_img, dt=1e-9, dz_interval=0.05, max_samples=399):
    """
    Convert a time-domain permittivity map into a depth-domain map.
    
    :param dielectric_img: 2D NumPy array (time, distance), representing relative permittivity values.
    :param dt: Time sampling interval (seconds), default is 1e-9 s.
    :param dz_interval: Depth resampling interval (meters), default is 0.05 m.
    :param max_samples: Maximum time samples per profile, default is 399.
    
    Process:
      1. Compute local wave velocity using: v = c / sqrt(ε), where c=3e8 m/s.
      2. Compute cumulative depth: depth_z = cumsum(v_line * dt) / 2.
      3. Resample permittivity to a uniform depth scale using interpolation.
    
    :return: (depth_converted_img, min_depth) 
      - depth_converted_img: 2D array with depth as rows and distance as columns.
      - min_depth: Maximum depth of all profiles (for reference).
    """
    c = 3e8  # Speed of light in vacuum (m/s)
    n_time, n_x = dielectric_img.shape
    depth_img_list = []
    z_stack_length = []
    
    for i in range(n_x):
        eps_line = dielectric_img[:, i][:max_samples]  # Extract permittivity values
        v_line = c / np.sqrt(eps_line)  # Compute local wave velocity
        depth_z = np.cumsum(v_line * dt) / 2  # Compute cumulative depth
        
        # Insert initial depth and permittivity value for accurate interpolation
        depth_z = np.insert(depth_z, 0, 0)
        eps_line = np.insert(eps_line, 0, eps_line[0])
        
        max_depth = np.max(depth_z)
        dest_depth = np.arange(0, max_depth, dz_interval)  # Define uniform depth scale
        
        # Interpolate permittivity values to uniform depth scale
        interp_func = interp1d(depth_z, eps_line, kind='linear', bounds_error=False, fill_value=np.nan)
        eps_interp = interp_func(dest_depth)
        
        depth_img_list.append(eps_interp)
        z_stack_length.append(len(dest_depth))
    
    # Standardize profile length with NaN padding for consistency
    max_length = np.max(z_stack_length)
    depth_img_padded = [np.pad(v, (0, max_length - len(v)), constant_values=np.nan) for v in depth_img_list]
    
    # Convert list to NumPy array and transpose so rows = depth, columns = distance
    depth_converted_img = np.array(depth_img_padded).T
    max_depth = np.max(z_stack_length)
    
    return depth_converted_img, max_depth * dz_interval



if __name__ == "__main__":
  # Load data and apply depth conversion
  data_file = pcfg.inversion_time_result_file
  path = pcfg.inversion_depth_result_img
  data = np.delete(np.loadtxt(data_file, delimiter=',', skiprows=0), [0], axis=0) 
  depth_converted_img, max_depth = convert2depth(data, dt=cfg.dt, dz_interval=cfg.dz_interval, max_samples=cfg.max_samples)
  plot_depth_permittivity_constant(depth_converted_img, path, line_length=100, time_length=max_depth)