import streamlit as st
import pandas as pd
import plotly.express as px
from statsmodels.tsa.seasonal import seasonal_decompose
from statsmodels.tsa.arima.model import ARIMA
from prophet import Prophet

class TimeSeriesAnalyzer:
    def analyze(self, df):
        date_columns = df.select_dtypes(include=['datetime64']).columns
        if len(date_columns) > 0:
            date_column = st.selectbox("Select date column", date_columns)
            value_column = st.selectbox("Select value column", df.columns)
            
            df[date_column] = pd.to_datetime(df[date_column])
            df = df.sort_values(date_column)
            
            st.subheader("Time Series Plot")
            fig = px.line(df, x=date_column, y=value_column)
            st.plotly_chart(fig)
            
            analysis_type = st.selectbox("Select analysis type", ["Decomposition", "ARIMA Forecasting", "Prophet Forecasting"])
            
            if analysis_type == "Decomposition":
                self.perform_decomposition(df, date_column, value_column)
            elif analysis_type == "ARIMA Forecasting":
                self.perform_arima_forecast(df, date_column, value_column)
            elif analysis_type == "Prophet Forecasting":
                self.perform_prophet_forecast(df, date_column, value_column)
        else:
            st.write("No datetime columns found in the dataset.")

    def perform_decomposition(self, df, date_column, value_column):
        df_temp = df.set_index(date_column)
        result = seasonal_decompose(df_temp[value_column], model='additive', period=30)
        
        st.subheader("Time Series Decomposition")
        fig = px.line(x=result.seasonal.index, y=result.seasonal, title="Seasonal")
        st.plotly_chart(fig)
        fig = px.line(x=result.trend.index, y=result.trend, title="Trend")
        st.plotly_chart(fig)
        fig = px.line(x=result.resid.index, y=result.resid, title="Residual")
        st.plotly_chart(fig)

    def perform_arima_forecast(self, df, date_column, value_column):
        df_temp = df.set_index(date_column)
        model = ARIMA(df_temp[value_column], order=(1,1,1))
        results = model.fit()
        
        forecast_steps = st.slider("Select number of steps to forecast", min_value=1, max_value=365, value=30)
        forecast = results.forecast(steps=forecast_steps)
        
        st.subheader("ARIMA Forecast")
        fig = px.line(x=df_temp.index, y=df_temp[value_column], title="Original Data with Forecast")
        fig.add_scatter(x=forecast.index, y=forecast, mode='lines', name='Forecast')
        st.plotly_chart(fig)

    def perform_prophet_forecast(self, df, date_column, value_column):
        df_prophet = df[[date_column, value_column]].rename(columns={date_column: 'ds', value_column: 'y'})
        model = Prophet()
        model.fit(df_prophet)
        
        future_dates = st.slider("Select number of days to forecast", min_value=1, max_value=365, value=30)
        future = model.make_future_dataframe(periods=future_dates)
        forecast = model.predict(future)
        
        st.subheader("Prophet Forecast")
        fig = px.line(x=df_prophet['ds'], y=df_prophet['y'], title="Original Data with Forecast")
        fig.add_scatter(x=forecast['ds'], y=forecast['yhat'], mode='lines', name='Forecast')
        fig.add_scatter(x=forecast['ds'], y=forecast['yhat_lower'], mode='lines', name='Lower Bound')
        fig.add_scatter(x=forecast['ds'], y=forecast['yhat_upper'], mode='lines', name='Upper Bound')
        st.plotly_chart(fig)