# %% [code]
# %%
!pip install pymap3d

import pandas as pd  # 添加这行，导入 pandas 库并使用别名 pd
import matplotlib.pyplot as plt  # 如果需要绘制图表，也需要导入 matplotlib
import numpy as np
import pandas as pd
import pymap3d as pm
import pymap3d.vincenty as pmv
import matplotlib.pyplot as plt
import glob as gl
import scipy.optimize
from tqdm.auto import tqdm
from scipy.interpolate import InterpolatedUnivariateSpline

from scipy.spatial import distance

from scipy.interpolate import InterpolatedUnivariateSpline
from scipy.spatial import distance
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense, Dropout
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import mean_squared_error



# %% [markdown]
# # 0 Constants

# %%
# Constants
CLIGHT = 299_792_458   # speed of light (m/s)
RE_WGS84 = 6_378_137   # earth semimajor axis (WGS84) (m)
OMGE = 7.2921151467E-5  # earth angular velocity (IS-GPS) (rad/s)

# %% [markdown]
# # 1 Satellite Selection
# satellite_selection(df, column)

# %%
# Satellite selection using carrier frequency error, elevation angle, and C/N0
def satellite_selection(df, column):
    """
    Args:
        df : DataFrame from device_gnss.csv
        column : Column name
    Returns:
        df: DataFrame with eliminated satellite signals
    """
    idx = df[column].notnull()
    idx &= df['CarrierErrorHz'] < 2.0e6  # carrier frequency error (Hz)
    idx &= df['SvElevationDegrees'] > 10.0  # elevation angle (deg)
    idx &= df['Cn0DbHz'] > 15.0  # C/N0 (dB-Hz)
    idx &= df['MultipathIndicator'] == 0 # Multipath flag

    return df[idx]

# %% [markdown]
# # 2 Pseudorange/Doppler Residuals and Jacobian
#  los_vector(xusr, xsat);
#  jac_pr_residuals(x, xsat, pr, W);
#  pr_residuals(x, xsat, pr, W);
#  jac_prr_residuals(v, vsat, prr, x, xsat, W);
#  prr_residuals(v, vsat, prr, x, xsat, W);
#  

# %%
# Compute line-of-sight vector from user to satellite
def los_vector(xusr, xsat):
    """
    Args:
        xusr : user position in ECEF (m)
        xsat : satellite position in ECEF (m)
    Returns:
        u: unit line-of-sight vector in ECEF (m)
        rng: distance between user and satellite (m)
    """
    u = xsat - xusr
    rng = np.linalg.norm(u, axis=1).reshape(-1, 1)
    u /= rng
    
    return u, rng.reshape(-1)


# Compute Jacobian matrix
def jac_pr_residuals(x, xsat, pr, W):
    """
    Args:
        x : current position in ECEF (m)
        xsat : satellite position in ECEF (m)
        pr : pseudorange (m)
        W : weight matrix
    Returns:
        W*J : Jacobian matrix
    """
    u, _ = los_vector(x[:3], xsat)
    J = np.hstack([-u, np.ones([len(pr), 1])])  # J = [-ux -uy -uz 1]

    return W @ J


# Compute pseudorange residuals
def pr_residuals(x, xsat, pr, W):
    """
    Args:
        x : current position in ECEF (m)
        xsat : satellite position in ECEF (m)
        pr : pseudorange (m)
        W : weight matrix
    Returns:
        residuals*W : pseudorange residuals
    """
    u, rng = los_vector(x[:3], xsat)

    # Approximate correction of the earth rotation (Sagnac effect) often used in GNSS positioning
    rng += OMGE * (xsat[:, 0] * x[1] - xsat[:, 1] * x[0]) / CLIGHT

    # Add GPS L1 clock offset
    residuals = rng - (pr - x[3])

    return residuals @ W


# Compute Jacobian matrix
def jac_prr_residuals(v, vsat, prr, x, xsat, W):
    """
    Args:
        v : current velocity in ECEF (m/s)
        vsat : satellite velocity in ECEF (m/s)
        prr : pseudorange rate (m/s)
        x : current position in ECEF (m)
        xsat : satellite position in ECEF (m)
        W : weight matrix
    Returns:
        W*J : Jacobian matrix
    """
    u, _ = los_vector(x[:3], xsat)
    J = np.hstack([-u, np.ones([len(prr), 1])])

    return np.dot(W, J)


# Compute pseudorange rate residuals
def prr_residuals(v, vsat, prr, x, xsat, W):
    """
    Args:
        v : current velocity in ECEF (m/s)
        vsat : satellite velocity in ECEF (m/s)
        prr : pseudorange rate (m/s)
        x : current position in ECEF (m)
        xsat : satellite position in ECEF (m)
        W : weight matrix
    Returns:
        residuals*W : pseudorange rate residuals
    """
    u, rng = los_vector(x[:3], xsat)
    rate = np.sum((vsat-v[:3])*u, axis=1) \
          + OMGE / CLIGHT * (vsat[:, 1] * x[0] + xsat[:, 1] * v[0]
                           - vsat[:, 0] * x[1] - xsat[:, 0] * v[1])

    residuals = rate - (prr - v[3])

    return residuals @ W

# %% [markdown]
# 
# # 3 Carrier Smoothing
# carrier_smoothing(gnss_df):

# %%
# Carrier smoothing of pseudarange
def carrier_smoothing(gnss_df):
    """
    Args:
        df : DataFrame from device_gnss.csv
    Returns:
        df: DataFrame with carrier-smoothing pseudorange 'pr_smooth'
    """
    carr_th = 1.0 # carrier phase jump threshold [m] 2->1.5 (best)->1.0
    pr_th =  15.0 # pseudorange jump threshold [m] 20->15

    prsmooth = np.full_like(gnss_df['RawPseudorangeMeters'], np.nan)
    # Loop for each signal
    for (i, (svid_sigtype, df)) in enumerate((gnss_df.groupby(['Svid', 'SignalType']))):
        df = df.replace(
            {'AccumulatedDeltaRangeMeters': {0: np.nan}})  # 0 to NaN

        # Compare time difference between pseudorange/carrier with Doppler
        drng1 = df['AccumulatedDeltaRangeMeters'].diff() - df['PseudorangeRateMetersPerSecond']
        drng2 = df['RawPseudorangeMeters'].diff() - df['PseudorangeRateMetersPerSecond']

        # Check cycle-slip
        slip1 = (df['AccumulatedDeltaRangeState'].to_numpy() & 2**1) != 0  # reset flag
        slip2 = (df['AccumulatedDeltaRangeState'].to_numpy() & 2**2) != 0  # cycle-slip flag
        slip3 = np.fabs(drng1.to_numpy()) > carr_th # Carrier phase jump
        slip4 = np.fabs(drng2.to_numpy()) > pr_th # Pseudorange jump

        idx_slip = slip1 | slip2 | slip3 | slip4
        idx_slip[0] = True

        # groups with continuous carrier phase tracking
        df['group_slip'] = np.cumsum(idx_slip)

        # Psudorange - carrier phase
        df['dpc'] = df['RawPseudorangeMeters'] - df['AccumulatedDeltaRangeMeters']

        # Absolute distance bias of carrier phase
        meandpc = df.groupby('group_slip')['dpc'].mean()
        df = df.merge(meandpc, on='group_slip', suffixes=('', '_Mean'))

        # Index of original gnss_df
        idx = (gnss_df['Svid'] == svid_sigtype[0]) & (
            gnss_df['SignalType'] == svid_sigtype[1])

        # Carrier phase + bias
        prsmooth[idx] = df['AccumulatedDeltaRangeMeters'] + df['dpc_Mean']

    # If carrier smoothing is not possible, use original pseudorange
    idx_nan = np.isnan(prsmooth)
    prsmooth[idx_nan] = gnss_df['RawPseudorangeMeters'][idx_nan]
    gnss_df['pr_smooth'] = prsmooth

    return gnss_df

# %% [markdown]
# # 4 score
# vincenty_distance(llh1, llh2):
# calc_score(llh, llh_gt)-vincenty_distance(llh, llh_gt)

# %%
# Compute distance by Vincenty's formulae
def vincenty_distance(llh1, llh2):
    """
    Args:
        llh1 : [latitude,longitude] (deg)
        llh2 : [latitude,longitude] (deg)
    Returns:
        d : distance between llh1 and llh2 (m)
    """
    d, az = np.array(pmv.vdist(llh1[:, 0], llh1[:, 1], llh2[:, 0], llh2[:, 1]))

    return d


# Compute score
def calc_score(llh, llh_gt):
    """
    Args:
        llh : [latitude,longitude] (deg)
        llh_gt : [latitude,longitude] (deg)
    Returns:
        score : (m)
    """
    d = vincenty_distance(llh, llh_gt)
    score = np.mean([np.quantile(d, 0.50), np.quantile(d, 0.95)])

    return score

# %% [markdown]
# # 5 Robust WLS
# point_positioning(gnss_df)

# %%
# GNSS single point positioning using pseudorange

def point_positioning(gnss_df):
    # Add nominal frequency to each signal
    # Note: GLONASS is an FDMA signal, so each satellite has a different frequency
    CarrierFrequencyHzRef = gnss_df.groupby(['Svid', 'SignalType'])[
        'CarrierFrequencyHz'].median()
    gnss_df = gnss_df.merge(CarrierFrequencyHzRef, how='left', on=[
                            'Svid', 'SignalType'], suffixes=('', 'Ref'))
    gnss_df['CarrierErrorHz'] = np.abs(
        (gnss_df['CarrierFrequencyHz'] - gnss_df['CarrierFrequencyHzRef']))

    # Carrier smoothing
    gnss_df = carrier_smoothing(gnss_df)

    # GNSS single point positioning
    utcTimeMillis = gnss_df['utcTimeMillis'].unique()
    nepoch = len(utcTimeMillis)
    x0 = np.zeros(4)  # [x,y,z,tGPSL1]
    v0 = np.zeros(4)  # [vx,vy,vz,dtGPSL1]
    x_wls = np.full([nepoch, 3], np.nan)  # For saving position
    v_wls = np.full([nepoch, 3], np.nan)  # For saving velocity

 # Store additional features for each epoch!!!!!!!!!
    features = []

    # Loop for epochs
    for i, (t_utc, df) in enumerate(tqdm(gnss_df.groupby('utcTimeMillis'), total=nepoch)):
        # Valid satellite selection
        df_pr = satellite_selection(df, 'pr_smooth')
        df_prr = satellite_selection(df, 'PseudorangeRateMetersPerSecond')

        # Corrected pseudorange/pseudorange rate
        pr = (df_pr['pr_smooth'] + df_pr['SvClockBiasMeters'] - df_pr['IsrbMeters'] -
              df_pr['IonosphericDelayMeters'] - df_pr['TroposphericDelayMeters']).to_numpy()
        prr = (df_prr['PseudorangeRateMetersPerSecond'] +
               df_prr['SvClockDriftMetersPerSecond']).to_numpy()

        # Satellite position/velocity
        xsat_pr = df_pr[['SvPositionXEcefMeters', 'SvPositionYEcefMeters',
                         'SvPositionZEcefMeters']].to_numpy()
        xsat_prr = df_prr[['SvPositionXEcefMeters', 'SvPositionYEcefMeters',
                           'SvPositionZEcefMeters']].to_numpy()
        vsat = df_prr[['SvVelocityXEcefMetersPerSecond', 'SvVelocityYEcefMetersPerSecond',
                       'SvVelocityZEcefMetersPerSecond']].to_numpy()

        # Weight matrix for peseudorange/pseudorange rate
        Wx = np.diag(1 / df_pr['RawPseudorangeUncertaintyMeters'].to_numpy())
        Wv = np.diag(1 / df_prr['PseudorangeRateUncertaintyMetersPerSecond'].to_numpy())
#####
# Compute additional features for this epoch
        num_sats = len(df_pr)
        avg_elevation = df_pr['SvElevationDegrees'].mean()
        avg_cn0 = df_pr['Cn0DbHz'].mean()
        max_cn0 = df_pr['Cn0DbHz'].max()
        min_cn0 = df_pr['Cn0DbHz'].min()
        std_cn0 = df_pr['Cn0DbHz'].std()

         # 初始化GDOP数组
        gdop_values = []

# 在循环中保存GDOP值
        if len(df_pr) >= 4:
    # 计算GDOP# Compute GDOP (Geometric Dilution of Precision)
            u, _ = los_vector(np.zeros(3), xsat_pr)
             G = np.hstack([-u, np.ones([len(pr), 1])])
              Q = np.linalg.inv(G.T @ G)
             gdop = np.sqrt(np.trace(Q))
        else:
             gdop = np.nan

        gdop_values.append(gdop)  # 保存GDOP值

# 在函数结束前，将GDOP值作为额外返回值
        
  
        # Store features
        features.append([num_sats, avg_elevation, avg_cn0, max_cn0, min_cn0, std_cn0, gdop])

#######

        # Robust WLS requires accurate initial values for convergence,
        # so perform normal WLS for the first time
        if len(df_pr) >= 4:
            # Normal WLS
            if np.all(x0 == 0):
                opt = scipy.optimize.least_squares(
                    pr_residuals, x0, jac_pr_residuals, args=(xsat_pr, pr, Wx))
                x0 = opt.x 
            # Robust WLS for position estimation
            opt = scipy.optimize.least_squares(
                 pr_residuals, x0, jac_pr_residuals, args=(xsat_pr, pr, Wx), loss='soft_l1')
            if opt.status < 1 or opt.status == 2:
                 print(f'i = {i} position lsq status = {opt.status}')
            else:
                 x_wls[i, :] = opt.x[:3]
                 x0 = opt.x
                 
        # Velocity estimation
        if len(df_prr) >= 4:
            if np.all(v0 == 0): # Normal WLS
                opt = scipy.optimize.least_squares(
                    prr_residuals, v0, jac_prr_residuals, args=(vsat, prr, x0, xsat_prr, Wv))
                v0 = opt.x
            # Robust WLS for velocity estimation
            opt = scipy.optimize.least_squares(
                prr_residuals, v0, jac_prr_residuals, args=(vsat, prr, x0, xsat_prr, Wv), loss='soft_l1')
            if opt.status < 1:
                print(f'i = {i} velocity lsq status = {opt.status}')
            else:
                v_wls[i, :] = opt.x[:3]
                v0 = opt.x
 # Convert features to numpy array
    features = np.array(features)
    
    return utcTimeMillis, x_wls, v_wls,features, np.array(gdop_values)


# %% [markdown]
# # 6 outlier detection
# exclude_interpolate_outlier(x_wls, v_wls)

# %%
# Simple outlier detection and interpolation
def exclude_interpolate_outlier(x_wls, v_wls):
    # Up velocity threshold
    v_up_th = 2.0 # m/s

    # Coordinate conversion
    x_llh = np.array(pm.ecef2geodetic(x_wls[:, 0], x_wls[:, 1], x_wls[:, 2])).T
    v_enu = np.array(pm.ecef2enuv(
        v_wls[:, 0], v_wls[:, 1], v_wls[:, 2], x_llh[0, 0], x_llh[0, 1])).T

    # Up velocity jump detection
    # Cars don't jump suddenly!
    idx_v_out = np.abs(v_enu[:, 2]) > v_up_th
    v_wls[idx_v_out, :] = np.nan
    
    # Interpolate NaNs at beginning and end of array
    x_df = pd.DataFrame({'x': x_wls[:, 0], 'y': x_wls[:, 1], 'z': x_wls[:, 2]})
    x_df = x_df.interpolate(limit_area='outside', limit_direction='both')
    
    # Interpolate all NaN data
    v_df = pd.DataFrame({'x': v_wls[:, 0], 'y': v_wls[:, 1], 'z': v_wls[:, 2]})
    v_df = v_df.interpolate(limit_area='outside', limit_direction='both')
    v_df = v_df.interpolate('spline', order=3)

    return x_df.to_numpy(), v_df.to_numpy()

# %% [markdown]
# # 7 Kalman 
# Kalman_filter(zs, us, phone)
# Kalman_smoothing(x_wls, v_wls, phone):
# Application of Kalman smoother; integration of velocity and position obtained by WLS. The covariance matrix is a fixed value.

# %%
#Kalman filter
def Kalman_filter(zs, us, phone):
    # Parameters
    # I don't know why only XiaomiMi8 seems to be inaccurate ... 
    sigma_v = 0.6 if phone == 'XiaomiMi8' else 0.1 # velocity SD m/s
    sigma_x = 5.0  # position SD m
    sigma_mahalanobis = 30.0 # Mahalanobis distance for rejecting innovation
    
    n, dim_x = zs.shape
    F = np.eye(3)  # Transition matrix
    Q = sigma_v**2 * np.eye(3)  # Process noise

    H = np.eye(3)  # Measurement function
    R = sigma_x**2 * np.eye(3)  # Measurement noise

    # Initial state and covariance
    x = zs[0, :3].T  # State
    P = sigma_x**2 * np.eye(3)  # State covariance
    I = np.eye(dim_x)

    x_kf = np.zeros([n, dim_x])
    P_kf = np.zeros([n, dim_x, dim_x])

    # Kalman filtering
    for i, (u, z) in enumerate(zip(us, zs)):
        # First step
        if i == 0:
            x_kf[i] = x.T
            P_kf[i] = P
            continue

        # Prediction step
        x = F @ x + u.T
        P = (F @ P) @ F.T + Q

        # Check outliers for observation
        d = distance.mahalanobis(z, H @ x, np.linalg.pinv(P))

        # Update step
        if d < sigma_mahalanobis:
            y = z.T - H @ x
            S = (H @ P) @ H.T + R
            K = (P @ H.T) @ np.linalg.inv(S)
            x = x + K @ y
            P = (I - (K @ H)) @ P
        else:
            # If no observation update is available, increase covariance
            P += 10**2*Q

        x_kf[i] = x.T
        P_kf[i] = P

    return x_kf, P_kf


# Forward + backward Kalman filter and smoothing
def Kalman_smoothing(x_wls, v_wls, phone):
    n, dim_x = x_wls.shape

    # Forward
    v = np.vstack([np.zeros([1, 3]), (v_wls[:-1, :] + v_wls[1:, :])/2])
    x_f, P_f = Kalman_filter(x_wls, v, phone)

    # Backward
    v = -np.flipud(v_wls)
    v = np.vstack([np.zeros([1, 3]), (v[:-1, :] + v[1:, :])/2])
    x_b, P_b = Kalman_filter(np.flipud(x_wls), v, phone)

    # Smoothing
    x_fb = np.zeros_like(x_f)
    P_fb = np.zeros_like(P_f)
    for (f, b) in zip(range(n), range(n-1, -1, -1)):
        P_fi = np.linalg.inv(P_f[f])
        P_bi = np.linalg.inv(P_b[b])

        P_fb[f] = np.linalg.inv(P_fi + P_bi)
        x_fb[f] = P_fb[f] @ (P_fi @ x_f[f] + P_bi @ x_b[b])

    return x_fb, x_f, np.flipud(x_b)

# %% [markdown]
# # 8 Kalman 
# enhanced_lstm_prediction(x_kf, v_wls, features, sequence_length=10)
# evaluate_model(llh_pred, llh_gt, method_name):

# %%
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense, Dropout
import numpy as np
def enhanced_lstm_prediction(x_kf, v_wls, features, sequence_length=10):
    """
    Train an LSTM model with additional features to predict future positions.
    
    Args:
        x_kf: ECEF positions after Kalman filtering
        v_wls: ECEF velocities
        features: Additional features (satellite count, elevation, C/N0, etc.)
        sequence_length: Number of previous positions to use for prediction
        
    Returns:
        x_lstm: Predicted positions using LSTM
    """
    # Combine all features
    all_features = np.hstack([x_kf, v_wls, features])
    
    # Normalize features
    scaler = StandardScaler()
    normalized_features = scaler.fit_transform(all_features)
    
    # Prepare data for LSTM
    n_samples = len(normalized_features) - sequence_length
    n_features = normalized_features.shape[1]
    
    X = np.zeros((n_samples, sequence_length, n_features))
    y = np.zeros((n_samples, 3))  # We only predict position (3 dimensions)
    
    for i in range(n_samples):
        X[i] = normalized_features[i:i+sequence_length]
        y[i] = x_kf[i+sequence_length]  # Target is the next position
    
    # Build enhanced LSTM model
    model = Sequential()
    model.add(LSTM(128, activation='relu', input_shape=(sequence_length, n_features), return_sequences=True))
    model.add(Dropout(0.2))
    model.add(LSTM(64, activation='relu'))
    model.add(Dropout(0.2))
    model.add(Dense(3))  # Output layer for 3D position
    model.compile(optimizer='adam', loss='mse')
    
    # Train the model
    print("Training enhanced LSTM model...")
    model.fit(X, y, epochs=50, batch_size=32, verbose=1)
    
    # Make predictions
    x_lstm = np.copy(x_kf)
    for i in range(sequence_length, len(x_kf)):
        input_features = normalized_features[i-sequence_length:i].reshape((1, sequence_length, n_features))
        x_lstm[i] = model.predict(input_features, verbose=0)
    
    return x_lstm

# %%
# Evaluate model performance
def evaluate_model(llh_pred, llh_gt, method_name):
    """
    Evaluate the performance of a positioning method.
    
    Args:
        llh_pred: Predicted positions (latitude, longitude)
        llh_gt: Ground truth positions (latitude, longitude)
        method_name: Name of the positioning method
        
    Returns:
        score: Competition score (mean of 50th and 95th percentile errors)
        rmse: Root mean squared error
        errors: Individual position errors
    """
    # Calculate distance errors
    errors = vincenty_distance(llh_pred, llh_gt)
    
    # Calculate metrics
    score = calc_score(llh_pred, llh_gt)
    rmse = np.sqrt(np.mean(errors**2))
    percentile_50 = np.percentile(errors, 50)
    percentile_95 = np.percentile(errors, 95)
    
    # Print results
    print(f"\n{method_name} Performance:")
    print(f"  Score (mean of 50th and 95th percentile): {score:.4f} m")
    print(f"  RMSE: {rmse:.4f} m")
    print(f"  50th percentile error: {percentile_50:.4f} m")
    print(f"  95th percentile error: {percentile_95:.4f} m")
    
    return score, rmse, errors


# %% [markdown]
# # Train Data

# %%
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense, Dropout
import numpy as np
import pandas as pd  # 添加这行，导入 pandas 库并使用别名 pd
import matplotlib.pyplot as plt  # 如果需要绘制图表，也需要导入 matplotlib

# Target course/phone换回去
########path = '/kaggle/input/smartphone-decimeter-2023/sdc2023/train/2023-09-07-22-48-us-ca-routebc2/pixel4xl'
########drive, phone = path.split('/')[-2:]

path = r'G:\Machine Learning\sdc2023\train\2023-09-07-22-48-us-ca-routebc2\pixel4xl'
drive, phone = path.split('\\')[-2:]
# Read data
gnss_df = pd.read_csv(f'{path}/device_gnss.csv')  # GNSS data
gt_df = pd.read_csv(f'{path}/ground_truth.csv')  # ground truth

# Point positioning
#utc, x_wls, v_wls = point_positioning(gnss_df)
utc, x_wls, v_wls, features = point_positioning(gnss_df)

 ##### Extract GDOP values from features (7th column)
gdop_values = features[:, 6]
    
   ##### # 可视化 GDOP 值
plt.figure(figsize=(12, 6))
plt.plot(utc, gdop_values, label='GDOP')
plt.axhline(y=5, color='r', linestyle='--', label='Poor GDOP Threshold (5)')
plt.axhline(y=3, color='y', linestyle='--', label='Fair GDOP Threshold (3)')
plt.axhline(y=2, color='g', linestyle='--', label='Good GDOP Threshold (2)')
plt.title('Geometric Dilution of Precision (GDOP) over Time')
plt.xlabel('Time (UTC)')
plt.ylabel('GDOP')
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.show()
    
# Exclude velocity outliers
#x_wls, v_wls = exclude_interpolate_outlier(x_wls, v_wls)
x_wls, v_wls, features = exclude_interpolate_outlier(x_wls, v_wls, features)

# Kalman smoothing
x_kf, _, _ = Kalman_smoothing(x_wls, v_wls, phone)

# LSTM prediction！！！！！！
x_lstm = enhanced_lstm_prediction(x_kf, v_wls, features)


# Convert to latitude and longitude！！！（3排）
llh_wls = np.array(pm.ecef2geodetic(x_wls[:, 0], x_wls[:, 1], x_wls[:, 2])).T
llh_kf = np.array(pm.ecef2geodetic(x_kf[:, 0], x_kf[:, 1], x_kf[:, 2])).T
llh_lstm = np.array(pm.ecef2geodetic(x_lstm[:, 0], x_lstm[:, 1], x_lstm[:, 2])).T

# Baseline
x_bl = gnss_df.groupby('TimeNanos')[
    ['WlsPositionXEcefMeters', 'WlsPositionYEcefMeters', 'WlsPositionZEcefMeters']].mean().to_numpy()
llh_bl = np.array(pm.ecef2geodetic(x_bl[:, 0], x_bl[:, 1], x_bl[:, 2])).T

# Ground truth
llh_gt = gt_df[['LatitudeDegrees', 'LongitudeDegrees']].to_numpy()

# Distance from ground truth！！！！（4排）
vd_bl = vincenty_distance(llh_bl, llh_gt)
vd_wls = vincenty_distance(llh_wls, llh_gt)
vd_kf = vincenty_distance(llh_kf, llh_gt)
vd_lstm = vincenty_distance(llh_lstm, llh_gt)

# Score！！！！（4排）
score_bl = calc_score(llh_bl, llh_gt)
score_wls = calc_score(llh_wls, llh_gt)
score_kf = calc_score(llh_kf[:-1, :], llh_gt[:-1, :])
score_lstm = calc_score(llh_lstm[:-1, :], llh_gt[:-1, :])

#！！！！（4排）
print(f'Score Baseline   {score_bl:.4f} [m]')
print(f'Score Robust WLS {score_wls:.4f} [m]')
print(f'Score KF         {score_kf:.4f} [m]')
print(f'Score LSTM       {score_lstm:.4f} [m]')

# Plot distance error！！！！（7排）
plt.figure()
plt.title('Distance error')
plt.ylabel('Distance error [m]')
plt.plot(vd_bl, label=f'Baseline, Score: {score_bl:.4f} m')
plt.plot(vd_wls, label=f'Robust WLS, Score: {score_wls:.4f} m')
plt.plot(vd_kf, label=f'Robust WLS + KF, Score: {score_kf:.4f} m')
plt.plot(vd_lstm, label=f'Robust WLS + KF + LSTM, Score: {score_lstm:.4f} m')
plt.legend()
plt.grid()
plt.ylim([0, 30])

# Compute velocity error
speed_wls = np.linalg.norm(v_wls[:, :3], axis=1)
speed_gt = gt_df['SpeedMps'].to_numpy()
speed_rmse = np.sqrt(np.sum((speed_wls-speed_gt)**2)/len(speed_gt))

# Plot velocity error
plt.figure()
plt.title('Speed error')
plt.ylabel('Speed Error [m/s]')
plt.plot(speed_wls - speed_gt, label=f'Speed RMSE: {speed_rmse:.4f} m')
plt.legend()
plt.grid()

# %% [markdown]
# # Test and submission

# %%
#######path = '/kaggle/input/smartphone-decimeter-2023/sdc2023'


# Target course/phone
path = r'G:\Machine Learning\sdc2023'


sample_df = pd.read_csv(f'{path}/sample_submission.csv')
test_dfs = []

# Loop for each trip
for i, dirname in enumerate(tqdm(sorted(gl.glob(f'{path}/test/*/*/')))):
    #drive, phone = dirname.split('/')[-3:-1]
    drive, phone = path.split('\\')[-3:-1]
    tripID = f'{drive}/{phone}'
    print(tripID)

    # Read data
    gnss_df = pd.read_csv(f'{dirname}/device_gnss.csv')

    # Point positioning
    utc, x_wls, v_wls = point_positioning(gnss_df)

    # Exclude velocity outliers
    x_wls, v_wls = exclude_interpolate_outlier(x_wls, v_wls)

    # Kalman smoothing
    x_kf, _, _ = Kalman_smoothing(x_wls, v_wls, phone)

    # LSTM prediction！！！！
    x_lstm = lstm_prediction(x_kf)

    # Convert to latitude and longitude
    #llh_kf = np.array(pm.ecef2geodetic(x_kf[:, 0], x_kf[:, 1], x_kf[:, 2])).T
    llh_lstm = np.array(pm.ecef2geodetic(x_lstm[:, 0], x_lstm[:, 1], x_lstm[:, 2])).T


    # Interpolation for submission
    UnixTimeMillis = sample_df[sample_df['tripId'] == tripID]['UnixTimeMillis'].to_numpy()
    lat = InterpolatedUnivariateSpline(utc, llh_kf[:,0], ext=3)(UnixTimeMillis)+0.00001
    lng = InterpolatedUnivariateSpline(utc, llh_kf[:,1], ext=3)(UnixTimeMillis)-0.000001
    trip_df = pd.DataFrame({
        'tripId' : tripID,
        'UnixTimeMillis': UnixTimeMillis,
        'LatitudeDegrees': lat,
        'LongitudeDegrees': lng
        })

    test_dfs.append(trip_df)

# Write submission.csv
test_df = pd.concat(test_dfs)
#test_df.to_csv('submission.csv', index=False)！！！！
test_df.to_csv(r'G:\Machine Learning\submission\submission.csv', index=False)

# %%