import json
import numpy as np
import pandas as pd
from scipy.interpolate import interp1d
from statsmodels.nonparametric.smoothers_lowess import lowess

def get_trags(trajs):
    trajs = [json.loads(a) for a in trajs]
    result_trajs = []
    for index in range(3):
        trajs_n = [[b[index] for b in a] for a in trajs]
        data = {'group': [], 'x': [], 'y': []}
        for i, traj in enumerate(trajs_n, start=1):
            data['group'].extend([i] * len(traj))
            data['x'].extend(np.linspace(0, 1, num=len(traj)))
            data['y'].extend(traj)
        df = pd.DataFrame(data)
        x_points = np.linspace(0, 1, num=50)

        df_complete = pd.concat(
            [df[df['group'] == g].set_index('x').reindex(x_points).reset_index().assign(group=g) for g in df['group'].unique()],
            ignore_index=True
        )

        def interpolate(group):
            interpolator = interp1d(group.dropna().x, group.dropna().y, kind='linear', bounds_error=False, fill_value="extrapolate")
            group['y'] = interpolator(group['x'])
            return group

        df_interpolated = df_complete.groupby('group').apply(interpolate).reset_index(drop=True)

        df_interpolated_filtered = df_interpolated[df_interpolated['x'].isin(x_points)]

        def smooth(group):
            smoothed_values = lowess(group['mean'], group['x'], frac=0.05, return_sorted=False)
            group['y_smooth'] = smoothed_values
            smoothed_stds = lowess(group['sd'], group['x'], frac=0.05, return_sorted=False)
            group['std_smooth'] = smoothed_stds
            return group

        df_smoothed = df_interpolated_filtered.groupby('x').agg(mean=('y', 'mean'), sd=('y', 'std')).reset_index()
        print(df_smoothed)
        df_smoothed = smooth(df_smoothed)
        
        result_trajs.append(df_smoothed)
    
    result = pd.concat(result_trajs, keys=["(1)", "(2)", "(3)"])
    result = result.reset_index(level=0).rename(columns={'level_0': 'source'})
    return result