Update rag_search.py

rag_search.py

@@ -157,7 +157,8 @@ def plot_line(df, x_column, y_columns, figsize=(12, 10), color='orange', title=N
 
     return fig
 
-def plot_bar(df, x_column, y_column, figsize=(12, 10), color='orange', title=None):
+
+def plot_bar(df, x_column, y_column, figsize=(12, 10), color='orange', title=None, rotation=45):
     fig, ax = plt.subplots(figsize=figsize)
 
     sns.barplot(data=df, x=x_column, y=y_column, color=color, ax=ax)
@@ -169,6 +170,8 @@ def plot_bar(df, x_column, y_column, figsize=(12, 10), color='orange', title=Non
     ax.tick_params(axis='x', colors=color)
     ax.tick_params(axis='y', colors=color)
 
+    plt.xticks(rotation=rotation)
+
     # Remove background
     fig.patch.set_alpha(0)
     ax.patch.set_alpha(0)
@@ -311,6 +314,45 @@ def filter_dataframe(df: pd.DataFrame) -> pd.DataFrame:
                 if len(user_date_input) == 2:
                     user_date_input = tuple(map(pd.to_datetime, user_date_input))
                     start_date, end_date = user_date_input
+                    
+                    # Determine the most appropriate time unit for plot
+                    time_units = {
+                        'year': df_[column].dt.year,
+                        'month': df_[column].dt.to_period('M'),
+                        'day': df_[column].dt.date
+                    }
+                    unique_counts = {unit: col.nunique() for unit, col in time_units.items()}
+                    closest_to_36 = min(unique_counts, key=lambda k: abs(unique_counts[k] - 36))
+
+                    # Group by the most appropriate time unit and count occurrences
+                    grouped = df_.groupby(time_units[closest_to_36]).size().reset_index(name='count')
+                    grouped.columns = [column, 'count']
+
+                    # Create a complete date range
+                    if closest_to_36 == 'year':
+                        date_range = pd.date_range(start=f"{start_date.year}-01-01", end=f"{end_date.year}-12-31", freq='YS')
+                    elif closest_to_36 == 'month':
+                        date_range = pd.date_range(start=start_date.replace(day=1), end=end_date + pd.offsets.MonthEnd(0), freq='MS')
+                    else:  # day
+                        date_range = pd.date_range(start=start_date, end=end_date, freq='D')
+
+                    # Create a DataFrame with the complete date range
+                    complete_range = pd.DataFrame({column: date_range})
+
+                    # Convert the date column to the appropriate format based on closest_to_36
+                    if closest_to_36 == 'year':
+                        complete_range[column] = complete_range[column].dt.year
+                    elif closest_to_36 == 'month':
+                        complete_range[column] = complete_range[column].dt.to_period('M')
+
+                    # Merge the complete range with the grouped data
+                    final_data = pd.merge(complete_range, grouped, on=column, how='left').fillna(0)
+
+                    with st.status(f"Date Distributions: {column}", expanded=False) as stat:
+                        try:
+                            st.pyplot(plot_bar(final_data, column, 'count'))
+                        except Exception as e:
+                            st.error(f"Error plotting bar chart: {e}")
                     df_ = df_.loc[df_[column].between(start_date, end_date)]
 
                 date_column = column