import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import segyio
import matplotlib.cm as cm
from sklearn.decomposition import PCA
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsRegressor
from sklearn.metrics import mean_squared_error
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, Dense
from tensorflow.keras.callbacks import EarlyStopping
from mpl_toolkits.mplot3d import Axes3D  

# -------------- 1. Загрузка и предобработка данных --------------


segy_path = 'Cube.sgy'              # SEGY файл
horizon_b_path = 'B'                # Файл горизонта B
thickness_path = 'dh_dataset.txt'   # Файл толщин горизонта

thickness_df = pd.read_csv('dh_dataset.txt', sep=r'\s+', header=None, names=['inline', 'crossline', 'thickness'])
print("Файл dh_dataset.txt загружен:")
print(thickness_df.head())

horizon_b_df = pd.read_csv('B', sep=r'\s+', header=None, names=['x', 'y', 'inline', 'crossline', 'time_ms', 'unknown'])
print("\nФайл B загружен:")
print(horizon_b_df.head())

print(f"Горизонт B загружен: {horizon_b_df.shape}")
print(horizon_b_df['crossline'])
# Приведение колонок 'inline' и 'crossline' к целочисленному типу


# Объединение DataFrame по колонкам 'inline' и 'crossline' с левой выборкой
merged_df = thickness_df.merge(horizon_b_df, on=['inline', 'crossline'], how='left')



# Проверка количества записей после объединения
print(f"После объединения осталось записей: {merged_df.shape[0]}")

# -------------- 2. Загрузка и обработка SEGY данных --------------

with segyio.open(segy_path, "r", ignore_geometry=True) as seg:
    # Извлечение inline и crossline атрибутов всей трассы
    ilines = seg.attributes(segyio.TraceField.INLINE_3D)[:]        # Извлечение inline
    xlines = seg.attributes(segyio.TraceField.CROSSLINE_3D)[:]     # Извлечение crossline

    # Получение шага дискретизации времени в микросекундах

    sample_rate_us = segyio.tools.dt(seg)  # Шаг дискретизации в микросекундах
    sample_rate_ms = sample_rate_us / 1000  # Преобразование в миллисекунды
    print(f"Шаг дискретизации: {sample_rate_us} мкс ({sample_rate_ms} мс)")


    
    first_trace = seg.trace[0]
    num_samples = len(first_trace)
    traces = np.array([tr for tr in seg.trace])
    

    print(f"Количество сэмплов на трассу: {num_samples}")
    print(f"Размер трассового массива: {traces.shape}")

# Создание словаря для быстрого поиска трасс по (inline, crossline)
il_xl_to_trace_idx = {}
for idx, (il, xl) in enumerate(zip(ilines, xlines)):
    il_xl_to_trace_idx[(il, xl)] = idx

# Инициализация списков для выбранных трасс и координат
selected_traces = []
selected_coords = []

print("Загрузка и нормализация необходимых трасс...")

for idx, row in merged_df.iterrows():
    il = row['inline']
    xl = row['crossline']
    time_ms = row['time_ms']

    # Находим индекс трассы в SEGY файле
    trace_idx = il_xl_to_trace_idx.get((il, xl))
    if trace_idx is None:
        print(f"Трасса inline={il}, crossline={xl} не найдена в SEGY файле.")
        continue

    trace = traces[trace_idx]

    # Определение окна выборки (30 дискретов по 2 мс = 60 мс)
    window_size_samples = 30
    half_window = window_size_samples // 2  
    center_sample = int(time_ms / sample_rate_ms)
    start_sample = center_sample - half_window
    end_sample = center_sample + half_window
    trace_window = trace[start_sample:end_sample]
    # Нормализация трассы по среднеквадратичной амплитуде
    rms = np.sqrt(np.mean(trace_window**2))
    trace_normalized = trace_window / rms
    selected_traces.append(trace_normalized)
    selected_coords.append((il, xl))
selected_traces = np.array(selected_traces)
selected_coords = np.array(selected_coords)
print(f"Выбранные и нормализованные трассы: {selected_traces.shape}") Задачи:
Загрузить сейсмические трассы из SEGY во временном интервале 60мс – 30 дискретов B. При загрузке трассы нормализовать на их среднеквадратичную амплитуду.
Преобразовать сейсмические трассы, используя МГК и автокодировщик. Каждую трассу из 30 дискретов требуется превратить в три значения.
Сформировать datasetы для обучения:  первый должен в себя включать три главные компоненты и значение толщины горизонта, второй – три значения от автокодировщика и значение толщины.
Выполнить обучение для двух datasetов, используя метод ближайших соседей
Рассчитать прогнозные значения толщин
Программа должна на выходе выдать два 3d кроссплота, где по осям будут значения главных компонент и значения от автокодировщика, а точки закрашены значениями толщин горизонта. Помимо кроссплотов должны быть выведены сами карты plt.pcolormesh( cmap='jet')
plt.colorbar(label='Значение').