"""
MIMO GPR — coherent time-domain BackProjection
=================================================

Эта ячейка полностью независима от верхнего эллипсного алгоритма:
  1. загружает те же S21-данные;
  2. строит oversampled A-сканы через тот же частотный сдвиг перед IFFT;
  3. для каждой точки (x,z) вычисляет tau_ij = (Rtx + Rrx) / v;
  4. интерполирует комплексный A-скан в этой задержке;
  5. когерентно суммирует комплексные вклады всех Tx/Rx-пар с компенсацией geo·pattern.

Это coherent BP: суммируются комплексные h_ij(tau), затем строится |sum h_ij|.
"""

import numpy as np
import matplotlib.pyplot as plt
from scipy.ndimage import gaussian_filter, label
from pathlib import Path

# ══════════════════════════════════════════════════════
# 0.1 ИЗМЕНЯЕМЫЕ ПАРАМЕТРЫ
# ══════════════════════════════════════════════════════
INPUT_IDX = [0, 1, 2, 3]
OUTPUT_IDX = [2, 3]

# Частотный диапазон и глубинный gate оставлены в прошлой логике.
F_START = 27 * 1e8
F_STOP = 6*1e9
MIN_DEPTH = 2.7
MAX_DEPTH = 15.0

# Данные и Вычитание среднего фона - как обычно 
BG_SUBTRACT = True
BG_PATH = Path('/Users/ivan_root/Desktop/GPR data/20260422_for_Vanya/20260422_night/20260422_evening_scene_balvanka_z/preprocessed')

DATA_PATH = Path('/Users/ivan_root/Desktop/GPR data/20260422_for_Vanya/20260422_night/20260422_evening_scene_balvanka_z/preprocessed/0002_id1_ns5952943855654')

# Убрать паразитные боковые лепестки с 2D карты
BP_REMOVE_SIDELOBE_OBJECTS = True  # True: убрать SL-кандидаты из финальной таблицы и разметки

# Компенсация затухания в этой версии разделена на 2 части:
#   range: геометрическое расхождение 1/(Rtx*Rrx)
#   angle: диаграмма направленности cos_tx^2 * cos_rx^2
COMP_RANGE_POWER = 0.28 # Можно менять свободно, текущее значение подобрано экспериментально 
COMP_ANGLE_POWER = 0.10 # Можно менять но лучше ставить не больше 0.25. Текущее значение подобрано нормально

# ══════════════════════════════════════════════════════
# 0.2 КОНФИГИ АНТЕНН
# ══════════════════════════════════════════════════════

# Физические координаты антенн по их реальным индексам, [м] - текущий конфиг
TX_POSITIONS = {
    2: -74.5 * 0.01,
    3:  75.0 * 0.01,
}

RX_POSITIONS = {
    0:  19.0 * 0.01,
    1:  44.5 * 0.01,
    2: -40.0 * 0.01,
    3: -19.0 * 0.01,
}
# Вариант для теста на старых конфигурациях
# TX_POSITIONS = {0:  90.5*0.01,   # Tx с индексом 0 → x = +0.9 м
#                 3: -90.5*0.01}   # Tx с индексом 3 → x = -0.9 м

# RX_POSITIONS = {0: -18*0.01,
#                 1:  48.5*0.01,
#                 2: -49.0*0.01,
#                 3:  18.5*0.01}

# ══════════════════════════════════════════════════════
# 0.3 НЕИЗМЕНЯЕМЫЕ ПАРАМЕТРЫ (ЛУЧШЕ НЕ ТРОГАТЬ)
# ══════════════════════════════════════════════════════

eps_r = 1.0
v = 3e8 / np.sqrt(eps_r)

x_tx = np.array([TX_POSITIONS[i] for i in sorted(TX_POSITIONS)])
x_rx = np.array([RX_POSITIONS[j] for j in sorted(RX_POSITIONS)])


# Сетка BP-карты
x_min, x_max = x_tx.min() - 2.0, x_tx.max() + 2.0
z_min, z_max = 0.1, MAX_DEPTH
NX_BP = 300
NZ_BP = 300

# Oversampling A-сканов: искусственно повышает плотность точек по t, но не физическое разрешение.
BP_OVERSAMPLE = 8

# Компенсация ослабления разделена на две части:
#   range: геометрическое расхождение 1/(Rtx*Rrx)
#   angle: диаграмма направленности cos_tx^2 * cos_rx^2
# Компенсация нормируется на точку под виртуальным центром пары на глубине COMP_REF_DEPTH.
COMP_RANGE_WEIGHT_MAX = 5.0
COMP_ANGLE_WEIGHT_MAX = 2.0
COMP_WEIGHT_MAX = 8.0          # общий финальный потолок после перемножения range*angle
COMP_REF_DEPTH = 3.0

# Сглаживание только для удобства поиска/визуализации максимума.
BP_SMOOTH_SIGMA = 1.5

# Параметры поиска объектов на BP-карте.
MAX_OBJECTS = 10
BP_OBJECT_MIN_FRAC = 0.35       # остановка: пик ниже этой доли от глобального максимума
BP_REGION_THRESH_FRAC = 0.75    # область объекта: связная область выше этой доли от локального пика
BP_SUPPRESS_THRESH_FRAC = 0.2  # подавление: более широкая связная область вокруг найденного пика
BP_SUPPRESS_USE_WINDOW = True # False: подавляыть всю связанную область; True: ограничить окно вокруг пика
BP_SUPPRESS_RX_CM = 80.0        # используется только если BP_SUPPRESS_USE_WINDOW=True
BP_SUPPRESS_RZ_CM = 40.0       # используется только если BP_SUPPRESS_USE_WINDOW=True
BP_MIN_REGION_AREA_CM2 = 10.0    # отсечение совсем мелких шумовых пятен

# Компактный центроид вокруг локального максимума, устойчивее центроида всей вытянутой области.
BP_CENTER_USE_COMPACT = True
BP_CENTER_RX_CM = 60.0
BP_CENTER_RZ_CM = 25.0
BP_CENTER_THRESH_FRAC = 0.88
BP_CENTER_WEIGHT_POWER = 2.0

# Диагностика боковых лепестков coherent BP.
BP_SIDELOBE_DETECT = True
BP_SIDELOBE_RANGE_RMS_TOL_CM = 20.0  # RMS-разница бистатических глубин по всем парам
BP_SIDELOBE_MIN_DX_CM = 35.0         # боковой лепесток должен быть заметно смещен по X
BP_SIDELOBE_MAX_DZ_CM = 70.0         # но находиться примерно на той же глубине
BP_SIDELOBE_MAX_REL_PEAK = 0.85      # кандидат должен быть слабее родительского максимума




# ══════════════════════════════════════════════════════
# 1. ЗАГРУЗКА ДАННЫХ
# ══════════════════════════════════════════════════════

def load_mimo_data(data_path, input_idx, output_idx):
    data_path = Path(data_path)
    s21_data = {}
    freq_data = {}

    for f in data_path.glob('i*_o*_s21.npy'):
        name = f.stem
        parts = name.split('_')
        i_tx_phys = int(parts[1][1:])
        i_rx_phys = int(parts[0][1:])

        if i_tx_phys not in output_idx or i_rx_phys not in input_idx:
            continue

        i_tx = sorted(output_idx).index(i_tx_phys)
        i_rx = sorted(input_idx).index(i_rx_phys)

        s21_data[(i_tx, i_rx)] = np.load(f)
        freq_file = data_path / f'i{i_rx_phys}_o{i_tx_phys}_freq.npy'
        freq_data[(i_tx, i_rx)] = np.load(freq_file)

    return s21_data, freq_data


def compute_background(bg_path, input_idx, output_idx):
    bg_path = Path(bg_path)
    snapshots = [s for s in sorted(bg_path.glob('*/')) if s.is_dir()]

    if len(snapshots) == 0:
        print('Фоновые снимки не найдены, BG_SUBTRACT отключен.')
        return None

    print(f'Вычисление фона по {len(snapshots)} снимкам...', end=' ', flush=True)
    bg_sum = {}
    bg_count = {}

    for snap_dir in snapshots:
        for f in snap_dir.glob('i*_o*_s21.npy'):
            name = f.stem
            parts = name.split('_')
            i_tx_phys = int(parts[1][1:])
            i_rx_phys = int(parts[0][1:])

            if i_tx_phys not in output_idx or i_rx_phys not in input_idx:
                continue

            i_tx = sorted(output_idx).index(i_tx_phys)
            i_rx = sorted(input_idx).index(i_rx_phys)
            key = (i_tx, i_rx)

            s21 = np.load(f)
            if key not in bg_sum:
                bg_sum[key] = np.zeros_like(s21, dtype=np.complex128)
                bg_count[key] = 0
            bg_sum[key] += s21
            bg_count[key] += 1

    bg = {key: bg_sum[key] / bg_count[key] for key in bg_sum}
    print(f'готово. Пар: {len(bg)}')
    return bg


s21_data, freq_data = load_mimo_data(DATA_PATH, INPUT_IDX, OUTPUT_IDX)
N_tx = len(OUTPUT_IDX)
N_rx = len(INPUT_IDX)
N_pairs = len(s21_data)

assert len(x_tx) == N_tx, f'x_tx должен содержать {N_tx} элементов'
assert len(x_rx) == N_rx, f'x_rx должен содержать {N_rx} элементов'
assert N_pairs > 0, 'Не найдено ни одной Tx/Rx-пары. Проверьте DATA_PATH.'

if BG_SUBTRACT:
    background = compute_background(BG_PATH, INPUT_IDX, OUTPUT_IDX)
    if background is None:
        BG_SUBTRACT = False
else:
    background = None
    print('BG_SUBTRACT = False, вычитание фона отключено.')

first_key = list(freq_data.keys())[0]
freqs = freq_data[first_key]
freq_mask = (freqs >= F_START) & (freqs <= F_STOP)
freqs_bp = freqs[freq_mask]

if len(freqs_bp) < 2:
    raise ValueError('В выбранном частотном диапазоне меньше двух точек.')

f_min = float(freqs_bp[0])
f_max = float(freqs_bp[-1])
BW = f_max - f_min
df_values = np.diff(freqs_bp)
df_median = float(np.median(df_values))
df_min = float(np.min(df_values))
df_max = float(np.max(df_values))
df_rel_spread = (df_max - df_min) / (df_median + 1e-30)
range_resolution = v / (2 * BW)
unambiguous_depth = v / (2 * df_median)
unambiguous_total_path = v / df_median

# print(f'Загружено пар Tx/Rx: {N_pairs}')
# print(f'Частотный диапазон BP: {f_min/1e9:.3f} - {f_max/1e9:.3f} ГГц')
# print(f'Точек частоты в BP-диапазоне: {len(freqs_bp)}')
# print(f'Шаг частоты df: median={df_median/1e6:.3f} МГц, '
#       f'min={df_min/1e6:.3f} МГц, max={df_max/1e6:.3f} МГц')
# print(f'Неравномерность df: {(df_rel_spread*100):.3f}% от median')
# print(f'Полоса B: {BW/1e9:.3f} ГГц')
# print(f'Теоретический предел разрешения по глубине deltaZ = {range_resolution*100:.2f} см')
# print(f'Unambiguous range по глубине = {unambiguous_depth:.2f} м '
#       f'(max total path = {unambiguous_total_path:.2f} м)')
# print(f'BP_OVERSAMPLE = {BP_OVERSAMPLE}')


# ══════════════════════════════════════════════════════
# 2. OVERSAMPLED A-СКАНЫ
# ══════════════════════════════════════════════════════

def compute_ascan_bp(s21, freq, f_start, f_stop, window=True, oversample=8):
    """
    S21(f) -> A-скан с правильным частотным сдвигом и oversampling.

    Частотный шаг df остается тем же, а n_fft увеличивается в oversample раз.
    Поэтому временная сетка становится плотнее: dt = 1 / (n_fft * df).
    """
    mask = (freq >= f_start) & (freq <= f_stop)
    freq_cut = freq[mask]
    s21_cut = s21[mask]

    if len(freq_cut) < 2:
        raise ValueError('После обрезки по частоте осталось меньше двух точек.')

    df = float(np.median(np.diff(freq_cut)))
    n = len(freq_cut)
    k0 = int(round(freq_cut[0] / df))

    min_len = 2 * (k0 + n - 1)
    n_fft_base = 1 << int(np.ceil(np.log2(min_len)))
    n_fft = int(n_fft_base * oversample)

    dt = 1.0 / (n_fft * df)
    t_sec = np.arange(n_fft, dtype=float) * dt

    s = s21_cut * np.hanning(n) if window else s21_cut.copy()

    H = np.zeros(n_fft, dtype=np.complex128)
    H[k0:k0 + n] = s

    h_complex = np.fft.ifft(H)
    a_abs = np.abs(h_complex)

    return t_sec, a_abs, h_complex, n_fft_base, n_fft


# print('Вычисление oversampled A-сканов...', end=' ', flush=True)
A_bp = {}
H_bp = {}
T_bp = {}
Z_bp = {}
n_fft_info = {}

for (i, j), s21 in s21_data.items():
    s21_proc = s21.copy()
    if BG_SUBTRACT and background is not None and (i, j) in background:
        s21_proc = s21_proc - background[(i, j)]

    t_pair, a_pair, h_pair, n_fft_base, n_fft = compute_ascan_bp(
        s21_proc,
        freq_data[(i, j)],
        f_start=F_START,
        f_stop=F_STOP,
        oversample=BP_OVERSAMPLE,
    )
    T_bp[(i, j)] = t_pair
    Z_bp[(i, j)] = t_pair * v / 2
    A_bp[(i, j)] = a_pair
    H_bp[(i, j)] = h_pair
    n_fft_info[(i, j)] = (n_fft_base, n_fft)

z_h_bp = Z_bp[first_key]
t_h_bp = T_bp[first_key]
# print('готово.')
# print(f'dt = {(t_h_bp[1] - t_h_bp[0])*1e12:.2f} пс')
# print(f'dz_sample = {(z_h_bp[1] - z_h_bp[0])*100:.3f} см')


# ══════════════════════════════════════════════════════
# 3. TIME-DOMAIN INCOHERENT BACKPROJECTION
# ══════════════════════════════════════════════════════

x_grid_bp = np.linspace(x_min, x_max, NX_BP)
z_grid_bp = np.linspace(z_min, z_max, NZ_BP)
XX_bp, ZZ_bp = np.meshgrid(x_grid_bp, z_grid_bp)
depth_gate = (ZZ_bp >= MIN_DEPTH) & (ZZ_bp <= MAX_DEPTH)


def attenuation_components_map(i_tx, i_rx, XX, ZZ):
    Rtx = np.sqrt((XX - x_tx[i_tx])**2 + ZZ**2)
    Rrx = np.sqrt((XX - x_rx[i_rx])**2 + ZZ**2)
    geo = 1.0 / (Rtx * Rrx + 1e-12)
    angle = (ZZ / (Rtx + 1e-12))**2 * (ZZ / (Rrx + 1e-12))**2
    return geo + 1e-30, angle + 1e-30


def attenuation_components_at_ref_depth(i_tx, i_rx, z_ref):
    xc = (x_tx[i_tx] + x_rx[i_rx]) / 2.0
    Rtx = np.sqrt((xc - x_tx[i_tx])**2 + z_ref**2)
    Rrx = np.sqrt((xc - x_rx[i_rx])**2 + z_ref**2)
    geo = 1.0 / (Rtx * Rrx + 1e-12)
    angle = (z_ref / (Rtx + 1e-12))**2 * (z_ref / (Rrx + 1e-12))**2
    return geo + 1e-30, angle + 1e-30


def bp_compensation_weight(i_tx, i_rx, XX, ZZ):
    geo, angle = attenuation_components_map(i_tx, i_rx, XX, ZZ)
    geo_ref, angle_ref = attenuation_components_at_ref_depth(i_tx, i_rx, COMP_REF_DEPTH)

    geo_norm = geo / geo_ref
    angle_norm = angle / angle_ref

    range_weight = 1.0 / (geo_norm ** COMP_RANGE_POWER + 1e-12)
    angle_weight = 1.0 / (angle_norm ** COMP_ANGLE_POWER + 1e-12)

    range_weight = np.clip(range_weight, 0.0, COMP_RANGE_WEIGHT_MAX)
    angle_weight = np.clip(angle_weight, 0.0, COMP_ANGLE_WEIGHT_MAX)

    weight = range_weight * angle_weight
    return np.clip(weight, 0.0, COMP_WEIGHT_MAX)


def interpolate_ascan_amplitude(tau, t_axis, a_axis):
    return np.interp(tau.ravel(), t_axis, a_axis, left=0.0, right=0.0).reshape(tau.shape)


def interpolate_ascan_complex(tau, t_axis, h_axis):
    h_real = np.interp(tau.ravel(), t_axis, h_axis.real, left=0.0, right=0.0)
    h_imag = np.interp(tau.ravel(), t_axis, h_axis.imag, left=0.0, right=0.0)
    return (h_real + 1j * h_imag).reshape(tau.shape)


def backproject_coherent(H, T, compensate=True):
    bp_complex = np.zeros_like(XX_bp, dtype=np.complex128)
    contribution_count = np.zeros_like(XX_bp, dtype=float)

    for i in range(N_tx):
        for j in range(N_rx):
            key = (i, j)
            if key not in H:
                continue

            Rtx = np.sqrt((XX_bp - x_tx[i])**2 + ZZ_bp**2)
            Rrx = np.sqrt((XX_bp - x_rx[j])**2 + ZZ_bp**2)
            tau = (Rtx + Rrx) / v
            valid = depth_gate & (tau >= T[key][0]) & (tau <= T[key][-1])

            h_tau = interpolate_ascan_complex(tau, T[key], H[key])
            h_tau = np.where(valid, h_tau, 0.0 + 0.0j)

            if compensate:
                w = bp_compensation_weight(i, j, XX_bp, ZZ_bp)
                w = np.where(valid, w, 0.0)
            else:
                w = np.where(valid, 1.0, 0.0)

            bp_complex += h_tau * w
            contribution_count += valid.astype(float)

    bp_complex = bp_complex / (contribution_count + 1e-12)
    bp_complex = np.where(depth_gate, bp_complex, 0.0 + 0.0j)
    bp_abs = np.abs(bp_complex)
    return bp_abs, bp_complex

def component_containing_peak(image, iz, ix, threshold, window_mask=None):
    mask = image >= threshold
    if window_mask is not None:
        mask &= window_mask

    labels, n_labels = label(mask, structure=np.ones((3, 3), dtype=int))
    if n_labels == 0 or labels[iz, ix] == 0:
        fallback = np.zeros_like(image, dtype=bool)
        fallback[iz, ix] = True
        return fallback

    return labels == labels[iz, ix]


def weighted_centroid(image, region_mask, threshold=0.0):
    values = image[region_mask]
    weights = np.clip(values - threshold, 0.0, None)
    if weights.sum() <= 1e-15:
        weights = values.copy()
    if weights.sum() <= 1e-15:
        iz, ix = np.argwhere(region_mask)[0]
        return x_grid_bp[ix], z_grid_bp[iz]

    x_vals = XX_bp[region_mask]
    z_vals = ZZ_bp[region_mask]
    return (x_vals * weights).sum() / weights.sum(), (z_vals * weights).sum() / weights.sum()



def compact_peak_centroid(image, iz, ix, peak):
    x0 = x_grid_bp[ix]
    z0 = z_grid_bp[iz]
    window_mask = (
        (np.abs(XX_bp - x0) <= BP_CENTER_RX_CM / 100.0) &
        (np.abs(ZZ_bp - z0) <= BP_CENTER_RZ_CM / 100.0)
    )
    threshold = BP_CENTER_THRESH_FRAC * peak
    center_mask = window_mask & (image >= threshold)

    if center_mask.sum() == 0:
        center_mask = window_mask.copy()
        center_mask[iz, ix] = True

    values = image[center_mask]
    weights = np.clip(values - threshold, 0.0, None) ** BP_CENTER_WEIGHT_POWER
    if weights.sum() <= 1e-15:
        weights = values.copy()
    if weights.sum() <= 1e-15:
        return x0, z0, center_mask

    x_vals = XX_bp[center_mask]
    z_vals = ZZ_bp[center_mask]
    x_c = (x_vals * weights).sum() / weights.sum()
    z_c = (z_vals * weights).sum() / weights.sum()
    return x_c, z_c, center_mask

def find_bp_objects(bp_image):
    """
    CLEAN-подобный поиск объектов на BP-карте.

    Для каждого шага берется максимум рабочей карты, вокруг него выделяется
    связная область выше BP_REGION_THRESH_FRAC от локального пика, затем
    считается взвешенный центроид этой области. После этого более широкая
    область вокруг той же цели подавляется на рабочей карте.
    """
    work = bp_image.copy()
    objects = []
    global_peak = float(work.max())
    stop_level = BP_OBJECT_MIN_FRAC * global_peak
    dx_cm = abs(x_grid_bp[1] - x_grid_bp[0]) * 100
    dz_cm = abs(z_grid_bp[1] - z_grid_bp[0]) * 100
    pixel_area_cm2 = dx_cm * dz_cm

    for step in range(MAX_OBJECTS):
        peak = float(work.max())
        if peak <= stop_level or peak <= 0:
            break

        iz, ix = np.unravel_index(np.argmax(work), work.shape)
        x_peak_local = x_grid_bp[ix]
        z_peak_local = z_grid_bp[iz]

        region_threshold = BP_REGION_THRESH_FRAC * peak
        region_mask = component_containing_peak(work, iz, ix, region_threshold)
        region_area_cm2 = float(region_mask.sum() * pixel_area_cm2)

        if region_area_cm2 < BP_MIN_REGION_AREA_CM2:
            work[iz, ix] = 0.0
            continue

        x_region_centroid, z_region_centroid = weighted_centroid(work, region_mask, threshold=region_threshold)
        if BP_CENTER_USE_COMPACT:
            x_centroid, z_centroid, center_mask = compact_peak_centroid(work, iz, ix, peak)
        else:
            x_centroid, z_centroid = x_region_centroid, z_region_centroid
            center_mask = region_mask.copy()

        center_area_cm2 = float(center_mask.sum() * pixel_area_cm2)
        region_values = work[region_mask]
        objects.append({
            'index': len(objects) + 1,
            'x_peak': float(x_peak_local),
            'z_peak': float(z_peak_local),
            'x': float(x_centroid),
            'z': float(z_centroid),
            'x_region': float(x_region_centroid),
            'z_region': float(z_region_centroid),
            'peak': peak,
            'area_cm2': region_area_cm2,
            'center_area_cm2': center_area_cm2,
            'region_mask': region_mask.copy(),
            'center_mask': center_mask.copy(),
            'mean_value': float(region_values.mean()),
            'sum_value': float(region_values.sum()),
        })

        if BP_SUPPRESS_USE_WINDOW:
            suppress_window = (
                (np.abs(XX_bp - x_peak_local) <= BP_SUPPRESS_RX_CM / 100.0) &
                (np.abs(ZZ_bp - z_peak_local) <= BP_SUPPRESS_RZ_CM / 100.0)
            )
        else:
            suppress_window = None

        suppress_threshold = BP_SUPPRESS_THRESH_FRAC * peak
        suppress_mask = component_containing_peak(
            work, iz, ix, suppress_threshold, window_mask=suppress_window
        )

        # Если широкий порог дал слишком маленькую область, подавляем хотя бы область детекции.
        if suppress_mask.sum() < region_mask.sum():
            suppress_mask = region_mask

        work[suppress_mask] = 0.0

    return objects, work


def bistatic_depth_signature(x_obj, z_obj):
    signature = []
    for i in range(N_tx):
        for j in range(N_rx):
            Rtx = np.sqrt((x_obj - x_tx[i])**2 + z_obj**2)
            Rrx = np.sqrt((x_obj - x_rx[j])**2 + z_obj**2)
            signature.append(0.5 * (Rtx + Rrx))
    return np.array(signature, dtype=float)


def mark_sidelobe_candidates(objects):
    for obj in objects:
        obj['sidelobe_candidate'] = False
        obj['sidelobe_parent'] = None
        obj['sidelobe_range_rms_cm'] = np.nan
        obj['sidelobe_dx_cm'] = np.nan
        obj['sidelobe_dz_cm'] = np.nan

    if not BP_SIDELOBE_DETECT:
        return objects

    signatures = [bistatic_depth_signature(obj['x'], obj['z']) for obj in objects]

    for k, obj in enumerate(objects):
        best_parent = None
        best_rms_cm = np.inf
        best_dx_cm = np.nan
        best_dz_cm = np.nan

        for p in range(k):
            parent = objects[p]
            rel_peak = obj['peak'] / (parent['peak'] + 1e-12)
            dx_cm = abs(obj['x'] - parent['x']) * 100.0
            dz_cm = abs(obj['z'] - parent['z']) * 100.0
            rms_cm = float(np.sqrt(np.mean((signatures[k] - signatures[p])**2)) * 100.0)

            is_candidate = (
                rel_peak <= BP_SIDELOBE_MAX_REL_PEAK and
                dx_cm >= BP_SIDELOBE_MIN_DX_CM and
                dz_cm <= BP_SIDELOBE_MAX_DZ_CM and
                rms_cm <= BP_SIDELOBE_RANGE_RMS_TOL_CM
            )

            if is_candidate and rms_cm < best_rms_cm:
                best_parent = parent
                best_rms_cm = rms_cm
                best_dx_cm = dx_cm
                best_dz_cm = dz_cm

        if best_parent is not None:
            obj['sidelobe_candidate'] = True
            obj['sidelobe_parent'] = best_parent['index']
            obj['sidelobe_range_rms_cm'] = best_rms_cm
            obj['sidelobe_dx_cm'] = best_dx_cm
            obj['sidelobe_dz_cm'] = best_dz_cm

    return objects


#print('Расчет coherent time-domain BP...', end=' ', flush=True)
bp_raw, bp_complex = backproject_coherent(H_bp, T_bp, compensate=True)
bp_map = bp_raw / (bp_raw.max() + 1e-12)
bp_map_s = gaussian_filter(bp_map, sigma=BP_SMOOTH_SIGMA)
bp_map_s = np.where(depth_gate, bp_map_s, 0.0)
bp_map_s = bp_map_s / (bp_map_s.max() + 1e-12)
#print('готово.')


bp_objects_all, bp_residual = find_bp_objects(bp_map_s)
# bp_objects_all = add_phase_metrics(bp_objects_all, bp_complex)
bp_objects_all = mark_sidelobe_candidates(bp_objects_all)

bp_map_display = bp_map_s.copy()
if BP_REMOVE_SIDELOBE_OBJECTS:
    for obj in bp_objects_all:
        if obj['sidelobe_candidate']:
            bp_map_display[obj['region_mask']] = 0.0
    bp_objects = [obj for obj in bp_objects_all if not obj['sidelobe_candidate']]
else:
    bp_objects = bp_objects_all

iz_max, ix_max = np.unravel_index(np.argmax(bp_map_s), bp_map_s.shape)
x_peak = x_grid_bp[ix_max]
z_peak = z_grid_bp[iz_max]
peak_value = bp_map_s[iz_max, ix_max]
main_obj = bp_objects[0] if bp_objects else None

# print('\n' + '=' * 72)
# print('  COHERENT TIME-DOMAIN BP: максимум карты и центроид области')
# print('=' * 72)
# print(f'  argmax:   x = {x_peak*100:+.1f} см, z = {z_peak*100:.1f} см, BP = {peak_value:.3f}')
# if main_obj is not None:
#     print(f"  centroid: x = {main_obj['x']*100:+.1f} см, z = {main_obj['z']*100:.1f} см, "
#           f"area = {main_obj['area_cm2']:.1f} см^2")
# print('=' * 72)

# print('\n' + '=' * 72)
# n_sl_all = sum(obj['sidelobe_candidate'] for obj in bp_objects_all)
# if BP_REMOVE_SIDELOBE_OBJECTS:
#     print(f'  НАЙДЕННЫЕ ОБЪЕКТЫ НА COHERENT BP-КАРТЕ, max {MAX_OBJECTS} '
#           f'(SL скрыты: {n_sl_all})')
# else:
#     print(f'  НАЙДЕННЫЕ ОБЪЕКТЫ НА COHERENT BP-КАРТЕ, max {MAX_OBJECTS} '
#           f'(SL показаны: {n_sl_all})')
# print('=' * 126)
# print(f"  {'#':<4} {'Xc [см]':>10} {'Zc [см]':>10} {'Xmax [см]':>11} {'Zmax [см]':>11} "
#       f"{'Peak':>8} {'Area [см2]':>11} {'PhCoh':>7} {'PhVar':>7} {'SL?':>5} {'Parent':>6} {'RMSr [см]':>10}")
# print('-' * 126)
# for obj in bp_objects:
#     sl_label = 'yes' if obj['sidelobe_candidate'] else 'no'
#     parent_label = '-' if obj['sidelobe_parent'] is None else str(obj['sidelobe_parent'])
#     rms_label = '-' if np.isnan(obj['sidelobe_range_rms_cm']) else f"{obj['sidelobe_range_rms_cm']:.1f}"
#     print(f"  {obj['index']:<4} {obj['x']*100:>+10.1f} {obj['z']*100:>10.1f} "
#           f"{obj['x_peak']*100:>+11.1f} {obj['z_peak']*100:>11.1f} "
#           f"{obj['peak']:>8.3f} {obj['area_cm2']:>11.1f} "
#           #f"{obj['phase_coherence']:>7.3f} {obj['phase_circular_variance']:>7.3f} "
#           f"{sl_label:>5} {parent_label:>6} {rms_label:>10}")
# print('=' * 126)


# ══════════════════════════════════════════════════════
# 4. ГРАФИКИ ДЛЯ СРАВНЕНИЯ
# ══════════════════════════════════════════════════════

# BP-карта
fig, ax = plt.subplots(figsize=(12, 7))
im = ax.imshow(
    bp_map_display,
    extent=[x_grid_bp[0]*100, x_grid_bp[-1]*100, z_grid_bp[-1]*100, z_grid_bp[0]*100],
    aspect='auto',
    cmap='jet',
    vmin=0.25,  # Было 0.0
    vmax=0.95, # Было 1.0
)
plt.colorbar(im, ax=ax, label='Нормированная |coherent BP|')
ax.plot(x_tx * 100, np.zeros(N_tx), 'r^', ms=12, label='Tx', zorder=5)
ax.plot(x_rx * 100, np.zeros(N_rx), 'bv', ms=12, label='Rx', zorder=5)
for obj in bp_objects:
    ax.contour(
        x_grid_bp * 100,
        z_grid_bp * 100,
        obj['region_mask'].astype(float),
        levels=[0.5],
        colors='white',
        linewidths=0.9,
        alpha=0.75,
    )
    ax.contour(
        x_grid_bp * 100,
        z_grid_bp * 100,
        obj['center_mask'].astype(float),
        levels=[0.5],
        colors='cyan',
        linewidths=0.8,
        alpha=0.85,
    )
    if obj['sidelobe_candidate']:
        ax.plot(obj['x'] * 100, obj['z'] * 100, 'x', color='yellow', ms=9,
                mew=2.0, zorder=7)
        label_text = f"{obj['index']} SL"
        text_color = 'yellow'
    else:
        ax.plot(obj['x'] * 100, obj['z'] * 100, 'wo', ms=7,
                markeredgecolor='k', mew=0.8, zorder=6)
        label_text = str(obj['index'])
        text_color = 'white'
    ax.text(obj['x'] * 100 + 3, obj['z'] * 100, label_text,
            color=text_color, fontsize=9, weight='bold', zorder=7)

ax.plot(x_peak * 100, z_peak * 100, 'w*', ms=18, markeredgecolor='k', mew=0.9,
        label='Глобальный максимум BP', zorder=6)
ax.plot([], [], 'wo', ms=7, markeredgecolor='k', mew=0.8, label='Центроид области')
if not BP_REMOVE_SIDELOBE_OBJECTS:
    ax.plot([], [], 'x', color='yellow', ms=9, mew=2.0, label='Sidelobe candidate')

ax.axhline(MIN_DEPTH * 100, color='white', lw=1.0, ls='--', alpha=0.75)
ax.set_xlabel('X [см]')
ax.set_ylabel('Глубина Z [см]')
map_title = 'Time-domain coherent BackProjection с geo·pattern компенсацией'
if BP_REMOVE_SIDELOBE_OBJECTS:
    map_title += ' (SL области скрыты на карте)'
ax.set_title(map_title)
ax.set_xlim(x_grid_bp[0] * 100, x_grid_bp[-1] * 100)
ax.set_ylim(z_grid_bp[-1] * 100, 0)
ax.legend(loc='lower right', fontsize=9)
ax.grid(alpha=0.22)
ax.invert_yaxis()
plt.tight_layout()
plt.show()