import xarray as xr
import pandas as pd
import numpy as np
from scipy.interpolate import interp1d
from scipy.signal import lombscargle
import os
import glob
from scipy.optimize import curve_fit
import math
from sklearn.decomposition import PCA
import sympy as sp
import matplotlib.pyplot as plt
import seaborn as sns
from matplotlib.patches import Patch
from matplotlib.lines import Line2D
from windrose import WindroseAxes
from scipy.signal import lombscargle, find_peaks
from itertools import permutations
from itertools import product

# 定义四舍五入函数


def round_to_nearest_multiple(value, multiple):
    return round(value / multiple) * multiple


def read_data(path):
    # 读取数据
    with open(path, "rb") as f:
        df = xr.open_dataset(f).to_dataframe().dropna(how="any")

    # 提取指定列
    columns_to_extract = ["alt", "press", "temp", "rh", "u", "v", "wspeed"]
    extracted_df = df[columns_to_extract].copy()

    # 过渡所有 NaN 的行

    # 进行单位转换
    extracted_df["alt"] = extracted_df["alt"] / 1000  # km
    extracted_df["rh"] = extracted_df["rh"] * 100  # %

    # 移除重复的高度值
    extracted_df = extracted_df.drop_duplicates(subset=["alt"])
    # Replace infinity values and then drop NaN values
    extracted_df = extracted_df.replace([np.inf, -np.inf], np.nan)
    extracted_df = extracted_df.replace(math.nan, np.nan)
    extracted_df = extracted_df.dropna(how="any")
    # if data is empty, return None
    if extracted_df.empty:
        return None

    new_height = np.arange(extracted_df["alt"].min(
    ), extracted_df["alt"].max() + 0.05, 0.05)

    #  将每个高度值转换为最接近0.05的整数倍，并转化为数组
    rounded_heights = [round_to_nearest_multiple(
        height, 0.05) for height in new_height]
    rounded_heights_np = np.array(rounded_heights)

    # 初始化一个新的 DataFrame 用于存储插值结果
    interpolated_data = pd.DataFrame({"alt": rounded_heights_np})

    # 对每个因变量进行线性插值
    for col in ["press", "temp", "rh", "u", "v", "wspeed"]:
        interp_func = interp1d(
            extracted_df["alt"], extracted_df[col], kind="linear", fill_value="extrapolate")
        interpolated_data[col] = interp_func(rounded_heights_np)

    return interpolated_data