import streamlit as st
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from amazon_chronos import ChronosModel, TimeSeriesDataSet
import tempfile
import os

class TimeSeriesForecaster:
    def __init__(self):
        self.model = None
        self.dataset = None
        self.original_series = None

    def preprocess_data(self, df, date_column, value_column, context_length=30, prediction_length=7):
        """
        Persiapkan data time series dari DataFrame
        """
        # Pastikan data terurut berdasarkan tanggal
        df = df.sort_values(by=date_column)
        
        # Konversi kolom tanggal ke datetime
        df[date_column] = pd.to_datetime(df[date_column])
        
        # Set index ke tanggal
        df.set_index(date_column, inplace=True)
        
        # Ekstrak series numerik
        time_series = df[value_column].values
        
        # Buat dataset Chronos
        self.original_series = time_series
        self.dataset = TimeSeriesDataSet.from_series(
            time_series, 
            context_length=context_length, 
            prediction_length=prediction_length
        )
        
        return self.dataset

    def train_model(self, model_id='chronos-t5-small'):
        """
        Latih model Chronos
        """
        self.model = ChronosModel.from_pretrained(model_id)
        self.model.fit(self.dataset)
        return self.model

    def forecast(self, n_samples=100):
        """
        Lakukan prediksi
        """
        if not self.model or not self.dataset:
            raise ValueError("Model belum dilatih. Latih model terlebih dahulu.")
        
        forecasts = self.model.predict(self.dataset, num_samples=n_samples)
        return forecasts

    def visualize_forecast(self, forecasts):
        """
        Buat visualisasi prediksi
        """
        plt.figure(figsize=(12, 6))
        
        # Plot series asli
        plt.plot(self.original_series, label='Data Historis', color='blue')
        
        # Plot prediksi
        forecast_mean = forecasts.mean(axis=0)
        forecast_lower = np.percentile(forecasts, 10, axis=0)
        forecast_upper = np.percentile(forecasts, 90, axis=0)
        
        forecast_start = len(self.original_series)
        plt.plot(
            range(forecast_start, forecast_start + len(forecast_mean)), 
            forecast_mean, 
            label='Prediksi Rata-rata', 
            color='red'
        )
        plt.fill_between(
            range(forecast_start, forecast_start + len(forecast_mean)), 
            forecast_lower, 
            forecast_upper, 
            alpha=0.3, 
            color='red'
        )
        
        plt.title('Peramalan Time Series dengan Amazon Chronos')
        plt.xlabel('Indeks Waktu')
        plt.ylabel('Nilai')
        plt.legend()
        
        return plt

def main():
    st.title('🕰️ Time Series Forecasting dengan Amazon Chronos')
    
    # Sidebar untuk upload dan konfigurasi
    st.sidebar.header('Pengaturan Prediksi')
    
    # Upload file CSV
    uploaded_file = st.sidebar.file_uploader(
        "Unggah File CSV", 
        type=['csv'], 
        help="Pastikan file CSV memiliki kolom tanggal dan nilai numerik"
    )
    
    # Pilihan kolom
    if uploaded_file is not None:
        # Baca CSV
        df = pd.read_csv(uploaded_file)
        
        # Pilih kolom
        date_column = st.sidebar.selectbox(
            'Pilih Kolom Tanggal', 
            options=df.columns
        )
        value_column = st.sidebar.selectbox(
            'Pilih Kolom Nilai', 
            options=[col for col in df.columns if col != date_column]
        )
        
        # Parameter prediksi
        context_length = st.sidebar.slider(
            'Panjang Konteks', 
            min_value=10, 
            max_value=100, 
            value=30
        )
        prediction_length = st.sidebar.slider(
            'Panjang Prediksi', 
            min_value=1, 
            max_value=30, 
            value=7
        )
        
        # Tombol proses
        if st.sidebar.button('Lakukan Prediksi'):
            try:
                # Inisiasi forecaster
                forecaster = TimeSeriesForecaster()
                
                # Preprocessing
                dataset = forecaster.preprocess_data(
                    df, 
                    date_column, 
                    value_column, 
                    context_length, 
                    prediction_length
                )
                
                # Latih model
                model = forecaster.train_model()
                
                # Lakukan prediksi
                forecasts = forecaster.forecast()
                
                # Tampilkan hasil
                st.subheader('Visualisasi Prediksi')
                plt = forecaster.visualize_forecast(forecasts)
                st.pyplot(plt)
                
                # Tampilkan detail prediksi
                forecast_mean = forecasts.mean(axis=0)
                forecast_lower = np.percentile(forecasts, 10, axis=0)
                forecast_upper = np.percentile(forecasts, 90, axis=0)
                
                prediction_df = pd.DataFrame({
                    'Prediksi Rata-rata': forecast_mean,
                    'Batas Bawah (10%)': forecast_lower,
                    'Batas Atas (90%)': forecast_upper
                })
                
                st.subheader('Detail Prediksi')
                st.dataframe(prediction_df)
                
            except Exception as e:
                st.error(f"Terjadi kesalahan: {str(e)}")

if __name__ == '__main__':
    main()