import streamlit as st
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
from Utility.data_loader  import load_train_series,load_train_events,load_sample_submission,load_test_series
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder, StandardScaler
from xgboost import XGBClassifier  # or XGBRegressor depending on your task
import xgboost as xgb
import numpy as np

@st.cache_data
def load_sampled_data():
    df3 = pd.read_parquet("train_series.parquet", columns=['series_id', 'step', 'anglez', 'enmo'])
    df4 = pd.read_parquet("test_series.parquet", columns=['series_id', 'step', 'anglez', 'enmo'])
    df2 = pd.read_csv("train_events.csv")

    # Sample safely based on available data
    df3_sample = df3.sample(n=min(5_000_000, len(df3)), random_state=42)
    df4_sample = df4.sample(n=min(1_000_000, len(df4)), random_state=42)

    return df3_sample, df4_sample, df2

# Load
df3, df4, df2 = load_sampled_data()
df = pd.concat([df3, df4], axis=0, ignore_index=True)
merged_df = pd.merge(df, df2, on=['series_id', 'step'], how='inner')

# Rename timestamp columns if they exist
if 'timestamp_x' in merged_df.columns:
    merged_df.rename(columns={'timestamp_x': 'sensor_timestamp'}, inplace=True)
if 'timestamp_y' in merged_df.columns:
    merged_df.rename(columns={'timestamp_y': 'event_timestamp'}, inplace=True)

# Box plots for each numerical feature
fig, ax = plt.subplots(figsize=(2, 1))
sns.boxplot(x=df2['step'], ax=ax)
ax.set_title('Boxplot of Step')

# Show the plot in Streamlit
st.pyplot(fig)

st.write("1. Data Visualization - Scatter Plot (feature vs feature or vs target)")
# Assume merged_df is already defined or loaded
df_sample = merged_df  # or use df_sample = merged_df.sample(n=50000) to downsample

st.subheader("Scatter Plot: anglez vs enmo")

# Create the plot
fig, ax = plt.subplots(figsize=(10, 6))
sns.scatterplot(x='anglez', y='enmo', data=df_sample, ax=ax)
ax.set_title("Scatter Plot: anglez vs enmo")

# Display in Streamlit
st.pyplot(fig)

# df_sample = merged_df.sample(n=10000)  # adjust sample size for performance

# # Subheader
# st.subheader("Pair Plot of Features")

# # Create pairplot
# fig = sns.pairplot(df_sample[['anglez', 'enmo', 'step']])
# fig.fig.suptitle("Pair Plot of Features", y=1.02)

# # Display in Streamlit
# st.pyplot(fig)
# Define columns to plot
plot_columns = ['anglez', 'enmo', 'step']

# Safety check: make sure required columns exist
if all(col in merged_df.columns for col in plot_columns):
    
    # Check data size and sample accordingly
    max_rows = len(merged_df)
    sample_size = min(10000, max_rows)  # Don't exceed available rows

    df_sample = merged_df.sample(n=sample_size)

    # Subheader
    st.subheader("Pair Plot of Features")

    # Create pairplot
    fig = sns.pairplot(df_sample[plot_columns])
    fig.fig.suptitle("Pair Plot of Features", y=1.02)

    # Display in Streamlit
    st.pyplot(fig)

else:
    st.error("One or more required columns ('anglez', 'enmo', 'step') are missing in the dataset.")


  

# Define features to plot
plot_features = ['anglez', 'enmo']

# Check if the required columns exist in the DataFrame
if all(col in merged_df.columns for col in plot_features):
    total_rows = len(merged_df)
    sample_size = 10000

    # Handle small datasets
    if total_rows < sample_size:
        st.info(f"Only {total_rows} rows available — using full dataset.")
        df_sample = merged_df.copy()
    else:
        df_sample = merged_df.sample(n=sample_size)

    # Plot
    fig, axes = plt.subplots(1, 2, figsize=(14, 5))

    sns.histplot(df_sample['anglez'], kde=True, bins=50, ax=axes[0])
    axes[0].set_title("Distribution of anglez")

    sns.histplot(df_sample['enmo'], kde=True, bins=50, ax=axes[1])
    axes[1].set_title("Distribution of enmo")

    plt.tight_layout()
    st.pyplot(fig)

else:
    st.error("Required columns not found in the dataset.")


    

st.write("Multicollinearity Check - Correlation Matrix")
features = ['anglez', 'enmo', 'step', 'night']
df_subset = merged_df[features]

# Streamlit title
st.subheader("Multicollinearity Check - Correlation Matrix")

# Calculate correlation matrix
corr_matrix = df_subset.corr()

# Plot heatmap
fig, ax = plt.subplots(figsize=(6, 4))
sns.heatmap(corr_matrix, annot=True, cmap='coolwarm', ax=ax)
ax.set_title("Correlation Matrix")

# Display in Streamlit
st.pyplot(fig)

# Encode
le = LabelEncoder()
merged_df['series_id'] = le.fit_transform(merged_df['series_id'])
merged_df['event'] = le.fit_transform(merged_df['event'])

# Drop columns with string or datetime values
drop_cols = ['sensor_timestamp', 'event_timestamp', 'night', 'step', 'sleep_duration_hrs', 'series_id']
df_cleaned = merged_df.drop(columns=[col for col in drop_cols if col in merged_df.columns])

# Ensure only numeric features in X
X = df_cleaned.drop('event', axis=1).select_dtypes(include=[np.number])
y = merged_df['event']

# Split and scale
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=27)

st.write("Feature Importance")
# Create model instance
xgb_model = XGBClassifier(use_label_encoder=False, eval_metric='logloss')  # example for classification

# Fit the model
xgb_model.fit(X_train, y_train)

# Plot feature importance
fig, ax = plt.subplots()
xgb.plot_importance(xgb_model, ax=ax)
ax.set_title("XGBoost Feature Importance")

# Show in Streamlit
st.subheader("XGBoost Feature Importance")
st.pyplot(fig)