import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import torch
from scipy.interpolate import interp1d
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def load_spectra_from_QM9(data_dir):
    df = pd.read_csv(data_dir)
    filtered_data = df.dropna().values 
    print(filtered_data[0])
    def normalize_data(data):
        min_val = np.min(data, axis=0, keepdims=True)
        max_val = np.max(data, axis=0, keepdims=True)
        return (data - min_val) / (max_val - min_val)
    if filtered_data.shape[1] > 3500:
        filtered_data = filtered_data[:, 1:3501]
    filtered_data = np.array([normalize_data(row) for row in filtered_data])
    return filtered_data

def gpu_normalize_safe(data_tensor):
        nan_mask = torch.isnan(data_tensor)
        inf_mask = torch.isinf(data_tensor)
        
        if torch.any(nan_mask) or torch.any(inf_mask):
            data_tensor = torch.nan_to_num(data_tensor, nan=0.0, posinf=1.0, neginf=0.0)
        min_vals = torch.min(data_tensor, dim=1, keepdim=True)[0]
        max_vals = torch.max(data_tensor, dim=1, keepdim=True)[0]
        range_vals = max_vals - min_vals
        range_vals[range_vals == 0] = 1.0
        normalized = (data_tensor - min_vals) / range_vals
        return normalized

def load_spectra_from_QMe14S(data_dir):
    df = pd.read_csv(data_dir)
    spectral_columns = [col for col in df.columns if col.startswith('wavenumber_')]
    
    if len(spectral_columns) == 0:
        raise ValueError("did not find spectral data columns starting with 'wavenumber_'")
    
    spectral_data = df[spectral_columns].copy()
    spectral_data = spectral_data.apply(pd.to_numeric, errors='coerce')
    spectral_data = spectral_data.dropna()
    
    if len(spectral_data) == 0:
        raise ValueError("all rows contain NaN values after conversion")
    # Ensure writable contiguous memory before converting to torch tensor.
    filtered_data = spectral_data.to_numpy(dtype=np.float32, copy=True)
    filtered_data = np.ascontiguousarray(filtered_data)
    data_tensor = torch.from_numpy(filtered_data).to(device)
    print(f"data has transferred to device: {data_tensor.device}")
    if data_tensor.shape[1] > 3500:
        data_tensor = data_tensor[:, :3500] 
    data_tensor = gpu_normalize_safe(data_tensor)
    filtered_data = data_tensor.cpu().numpy()
    dedimed_spectrum_list = []
    for i, spectrum in enumerate(filtered_data):
        original_wave_min = 500
        original_wave_max = 4000
        original_wavenumbers = np.linspace(original_wave_min, original_wave_max, len(spectrum))
        target_wave_min = 0
        target_wave_max = 3500
        target_wavenumbers = np.linspace(target_wave_min, target_wave_max, 3500)
        interp_func = interp1d(
            original_wavenumbers, 
            spectrum, 
            kind='linear', 
            bounds_error=False, 
            fill_value=0.0, 
            assume_sorted=True
        )
        interpolated_spectrum = interp_func(target_wavenumbers)
        mask_low = target_wavenumbers < 500
        interpolated_spectrum[mask_low] = 1e-10 
        interpolated_spectrum = np.nan_to_num(interpolated_spectrum, nan=1e-10)
        dedimed_spectrum_list.append(interpolated_spectrum)
    dedimed_spectrum_list = np.array(dedimed_spectrum_list)
    return dedimed_spectrum_list

def load_real_data(data_path, labels_path=None, wavenumbers_path=None, normalize=True):
    try:
        spectra = np.load(data_path)
        if np.isnan(spectra).any():
            spectra = np.nan_to_num(spectra, nan=0.0)

        if normalize:
            for i in range(spectra.shape[0]):
                spectrum = spectra[i]
                min_val = np.min(spectrum)
                max_val = np.max(spectrum)
                if max_val > min_val:
                    spectra[i] = (spectrum - min_val) / (max_val - min_val)
        labels = None
        if labels_path:
            try:
                labels = np.load(labels_path)
                if labels.shape[0] != spectra.shape[0]:
                    print(f"warning：label_num({labels.shape[0]})doesn't match spectra_num({spectra.shape[0]})!")
            except Exception as e:
                print(f"fail to load labels: {e}")
        wavenumbers = None
        if wavenumbers_path is not None:
            try:
                wavenumbers = np.load(wavenumbers_path)
                if wavenumbers.shape[0] != spectra.shape[1]:
                    print(f"warning：wavenumber_length({wavenumbers.shape[0]})doesn't match spectra_length({spectra.shape[1]})!")
            except Exception as e:
                print(f"fail to load wavenumbers: {e}")

        return spectra, labels, wavenumbers

    except Exception as e:
        print(f"fail to load spectra data: {e}")
        return None, None, None