import numpy as np
from scipy.interpolate import interp1d
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
from numpy.lib.stride_tricks import sliding_window_view


def smooth_rows_reflect(arr, window):
    """
    Row-wise moving average with reflective padding.

    Parameters
    ----------
    arr : array-like, shape (N_rows, N_cols)
        Input 2D array.
    window : int
        Window length for the moving average (>=1).

    Returns
    -------
    smoothed : ndarray
        Array with the same shape as `arr`, smoothed along axis=1.
    """
    arr = np.asarray(arr)
    w = int(window)
    if w <= 1:
        return arr.copy()

    # Ensure the window does not exceed the number of columns
    w = min(w, arr.shape[1])

    # Make output length equal to input: left + right = w - 1
    left  = w // 2
    right = w - 1 - left  # for even w: (left, right) = (w//2, w//2 - 1)

    padded = np.pad(arr, ((0, 0), (left, right)), mode='reflect')

    # Slide along axis=1 and take the mean within each window
    win = sliding_window_view(padded, window_shape=w, axis=1)
    return win.mean(axis=-1)


def custom_rainbow():
    """
    Create a custom colormap with expanded red and blue regions,
    de-emphasizing the yellow/green middle.
    """
    colors = [
        (0.00, 'red'),
        (1.00, 'blue'),
    ]
    return mcolors.LinearSegmentedColormap.from_list("custom_rainbow", colors)


def plot_water_constant(data, path, line_length=400, time_length=10):
    """
    Plot estimated water content (%) image.

    Parameters
    ----------
    data : ndarray, shape (N_time, N_dist)
        2D array of water content in fraction (0–1).
    path : str
        Output image path.
    line_length : float
        Profile length in meters (x-axis extent).
    time_length : float
        Time window in ns (y-axis extent).
    """
    num_points, num_lines = data.shape

    plt.rcParams.update({'font.family': 'Times New Roman', 'font.size': 20})

    plt.figure(figsize=(12, 4))
    im = plt.imshow(
        data * 100, aspect='auto', cmap='YlGnBu',
        extent=[0, line_length, time_length, 0],
        vmin=20, vmax=40
    )

    plt.xlabel('Distance (m)', fontsize=20)
    plt.xticks([0, 20, 40, 60, 80, 100])
    plt.ylabel('Time (ns)', fontsize=20)
    plt.yticks([0, 2, 4, 6, 8, 10])

    cbar = plt.colorbar(im)
    cbar.set_label('Water content (%)', fontsize=20)

    plt.tick_params(axis='both', direction='in', width=1)
    for spine in plt.gca().spines.values():
        spine.set_linewidth(1)

    plt.grid(False)
    plt.tight_layout()
    plt.savefig(path, dpi=300)
    plt.show()


def plot_ice_constant(data, path, line_length=400, time_length=10):
    """
    Plot estimated ice content (%) image.

    Parameters
    ----------
    data : ndarray, shape (N_time, N_dist)
        2D array of ice content in fraction (0–1).
    path : str
        Output image path.
    line_length : float
        Profile length in meters (x-axis extent).
    time_length : float
        Time window in ns (y-axis extent).
    """
    num_points, num_lines = data.shape

    plt.rcParams.update({'font.family': 'Times New Roman', 'font.size': 20})

    plt.figure(figsize=(12, 4))
    im = plt.imshow(
        data * 100, aspect='auto', cmap='PuBu',
        extent=[0, line_length, time_length, 0],
        vmin=10, vmax=40
    )

    plt.xlabel('Distance (m)', fontsize=20)
    plt.xticks([0, 20, 40, 60, 80, 100])
    plt.ylabel('Time (ns)', fontsize=20)
    plt.yticks([0, 2, 4, 6, 8, 10])

    cbar = plt.colorbar(im)
    cbar.set_label('Ice content (%)', fontsize=20)

    plt.tick_params(axis='both', direction='in', width=1)
    for spine in plt.gca().spines.values():
        spine.set_linewidth(1)

    plt.grid(False)
    plt.tight_layout()
    plt.savefig(path, dpi=300)
    plt.show()


def plot_permittivity_constant(data, path, line_length=400, time_length=10):
    """
    Plot inverted relative permittivity image.

    Parameters
    ----------
    data : ndarray, shape (N_time, N_dist)
        2D array of relative permittivity.
    path : str
        Output image path.
    line_length : float
        Profile length in meters (x-axis extent).
    time_length : float
        Time window in ns (y-axis extent).
    """
    num_points, num_lines = data.shape

    plt.rcParams.update({'font.family': 'Times New Roman', 'font.size': 20})

    plt.figure(figsize=(12, 4))
    im = plt.imshow(
        data, aspect='auto', cmap='rainbow_r',
        extent=[0, line_length, time_length, 0],
        vmin=5, vmax=25
    )

    plt.xlabel('Distance (m)', fontsize=20)
    plt.xticks([0, 20, 40, 60, 80, 100])
    plt.ylabel('Time (ns)', fontsize=20)
    plt.yticks([0, 2, 4, 6, 8, 10])

    cbar = plt.colorbar(im)
    cbar.set_label('Permittivity', fontsize=20)

    plt.tick_params(axis='both', direction='in', width=1)
    for spine in plt.gca().spines.values():
        spine.set_linewidth(1)

    plt.grid(False)
    plt.tight_layout()
    plt.savefig(path, dpi=300)
    plt.show()


def estimate_water_ice(eps_map, n_porosity=0.45,
                       eps_s=4.7, eps_w=86.0, eps_i=3.15):
    """
    Estimate volumetric ice and water contents from relative permittivity.

    Parameters
    ----------
    eps_map : ndarray
        Relative permittivity (inverted) of shape (N_time, N_dist).
    n_porosity : float
        Porosity (m³/m³).
    eps_s : float
        Permittivity of the solid phase.
    eps_w : float
        Permittivity of liquid water.
    eps_i : float
        Permittivity of ice.

    Returns
    -------
    theta_i_crim : ndarray
        Ice content (m³/m³) estimated via CRIM.
    theta_w_emp : ndarray
        Water content (m³/m³) estimated via empirical relation.
    """
    # --- Ice content via CRIM ---
    sqrt_eps = np.sqrt(eps_map)
    sqrt_es = np.sqrt(eps_s)
    sqrt_ew = np.sqrt(eps_w)
    sqrt_ei = np.sqrt(eps_i)

    # θ_i = [√ε - (1-n)√ε_s - n√ε_w] / (√ε_i - √ε_w)
    theta_i_crim = (sqrt_eps - (1 - n_porosity) * sqrt_es - n_porosity * sqrt_ew) / (sqrt_ei - sqrt_ew)
    theta_i_crim = np.clip(theta_i_crim, 0.0, n_porosity)

    # Water content (active layer) via empirical relation
    theta_w_emp = 0.458 * (eps_map ** 0.26) - 0.664
    theta_w_emp = np.clip(theta_w_emp, 0.0, n_porosity)

    # If needed, CRIM water content would be: theta_w_crim = n_porosity - theta_i_crim
    return theta_i_crim, theta_w_emp


def convert2depth(dielectric_img, dt=1e-9, dz_interval=0.05, max_samples=399):
    """
    Convert a time-domain permittivity section to a depth-domain permittivity section.

    Parameters
    ----------
    dielectric_img : ndarray, shape (N_time, N_x)
        Relative permittivity in time domain.
    dt : float
        Time sampling interval in seconds.
    dz_interval : float
        Target depth sampling interval in meters.
    max_samples : int
        Use only the first `max_samples` samples per trace when integrating.

    Processing
    ----------
    For each column (trace):
      1) Compute local EM velocity v = c / sqrt(eps).
      2) Accumulate depth: depth_z = cumsum(v * dt) / 2 (two-way time to depth).
      3) Interpolate permittivity onto a uniform depth grid with spacing `dz_interval`.

    Returns
    -------
    depth_converted_img : ndarray, shape (N_depth, N_x)
        Permittivity in depth domain (rows = depth, cols = distance).
    max_depth_m : float
        Maximum depth covered across traces (in meters).
    """
    c = 3e8  # 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]

        v_line = c / np.sqrt(eps_line)
        depth_z = np.cumsum(v_line * dt) / 2

        # Anchor at depth 0 with the first permittivity value
        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)

        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))

    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
    ]

    depth_converted_img = np.array(depth_img_padded).T
    max_depth_m = np.max(z_stack_length) * dz_interval

    return depth_converted_img, max_depth_m


# ------------------------- Script usage -------------------------
data_file = './time_result_csv/DONGRONG.csv'
path = './IMG/DONGRONG_depth.png'

# Load data, drop the first row (e.g., file index), scale by 30
data = np.delete(np.loadtxt(data_file, delimiter=",", skiprows=0), [0], axis=0) * 30

# Row-wise smoothing with reflective padding
data = smooth_rows_reflect(data.T, window=10)
data = smooth_rows_reflect(data.T, window=30)

depth_converted_img, max_depth = convert2depth(data, dt=0.2e-9, dz_interval=0.01, max_samples=1000)

plot_permittivity_constant(depth_converted_img, path, line_length=100, time_length=10)

ice, water = estimate_water_ice(depth_converted_img, n_porosity=0.4, eps_s=7, eps_w=81, eps_i=3.15)
plot_ice_constant(ice, './IMG/ice.png', line_length=100, time_length=10)
plot_water_constant(water, './IMG/water.png', line_length=100, time_length=10)