# generate_pickles.py

import pandas as pd
import numpy as np
import pickle
import streamlit as st
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder, StandardScaler
from xgboost import XGBClassifier

st.title("📦 Data Preprocessing and Pickle Saver")

@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)

# Drop only columns that exist
required_cols = ['night', 'event', 'event_timestamp']
existing_cols = [col for col in required_cols if col in merged_df.columns]
merged_df.dropna(subset=existing_cols, inplace=True)
merged_df.reset_index(drop=True, inplace=True)

# Convert timestamps and calculate sleep duration
if 'event_timestamp' in merged_df.columns and 'sensor_timestamp' in merged_df.columns:
    merged_df['event_timestamp'] = pd.to_datetime(merged_df['event_timestamp'], errors='coerce', utc=True)
    merged_df['sensor_timestamp'] = pd.to_datetime(merged_df['sensor_timestamp'], errors='coerce', utc=True)
    merged_df['sleep_duration_hrs'] = (merged_df['sensor_timestamp'] - merged_df['event_timestamp']).dt.total_seconds() / 3600
    merged_df.dropna(subset=['sensor_timestamp', 'event_timestamp', 'sleep_duration_hrs'], inplace=True)

# 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)
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)

# Train model
model = XGBClassifier(n_estimators=100, max_depth=3, learning_rate=0.1, eval_metric='logloss')
model.fit(X_train_scaled, y_train)

# Save pickles
with open("model.pkl", "wb") as f: pickle.dump(model, f)
with open("scaler.pkl", "wb") as f: pickle.dump(scaler, f)
with open("label_encoder.pkl", "wb") as f: pickle.dump(le, f)
with open("X_test.pkl", "wb") as f: pickle.dump(X_test_scaled, f)
with open("y_test.pkl", "wb") as f: pickle.dump(y_test, f)

st.success("✅ Pickle files saved successfully.")