import streamlit as st
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics import mean_squared_error
from datetime import timedelta
import torch
import torch.nn as nn
from PIL import Image
from pathlib import Path

# --- Configuration de la page ---
# Doit être la première commande Streamlit du script
st.set_page_config(
    page_title="Prédiction Boursière GRU",
    page_icon="📈",
    layout="wide"
)

script = Path(__file__)
script_dir = script.parent


# --- Définition des modèles et fonctions (partie non visible dans l'UI) ---

# === Définition de l'architecture du modèle GRU ===
class GRUModel(nn.Module):
    def __init__(self, input_size=1, hidden_layer_size=50, num_layers=2, output_size=1):
        super().__init__()
        self.hidden_layer_size = hidden_layer_size
        self.gru = nn.GRU(input_size, hidden_layer_size, num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_layer_size, output_size)

    def forward(self, input_seq):
        gru_out, _ = self.gru(input_seq)
        predictions = self.fc(gru_out[:, -1])
        return predictions

# === Fonctions de chargement et de traitement (avec cache pour la performance) ===
@st.cache_data
def load_data(csv_path= script_dir/"action_amd.csv"):
    """Charge les données depuis le fichier CSV et les formate correctement."""
    try:
        df = pd.read_csv(csv_path)
    except FileNotFoundError:
        st.error(f"Erreur : Le fichier '{csv_path}' est introuvable. "
                 "Assurez-vous qu'il se trouve dans le même dossier que votre script.")
        st.stop()
    if 'Date' not in df.columns or 'Close' not in df.columns:
        st.error("Le fichier CSV doit contenir les colonnes 'Date' et 'Close'.")
        st.stop()
    df_filtered = df[['Date', 'Close']].copy()
    df_filtered['Date'] = pd.to_datetime(df_filtered['Date'])
    df_renamed = df_filtered.rename(columns={'Date': 'ds', 'Close': 'y'})
    return df_renamed.sort_values(by='ds')

@st.cache_resource
def load_gru_model(path, model_class):
    """Charge le modèle GRU pré-entraîné."""
    model = model_class()
    model.load_state_dict(torch.load(path, map_location=torch.device('cpu')))
    model.eval()
    return model

def predict_gru(model, df, forecast_days, window_size=20):
    """Effectue des prédictions sur les N prochains jours."""
    data_values = df['y'].values
    predictions = []
    input_seq_np = data_values[-window_size:]
    
    for _ in range(forecast_days):
        input_seq_torch = torch.from_numpy(input_seq_np).float().view(1, window_size, 1)
        with torch.no_grad():
            pred = model(input_seq_torch).item()
        predictions.append(pred)
        input_seq_np = np.append(input_seq_np[1:], pred)
        
    last_date = df['ds'].max()
    future_dates = []
    current_date = last_date
    while len(future_dates) < forecast_days:
        current_date += timedelta(days=1)
        if current_date.weekday() < 5:
            future_dates.append(current_date)
            
    return pd.DataFrame({'ds': future_dates, 'yhat': predictions})

def calculate_rmse(y_true, y_pred):
    return np.sqrt(mean_squared_error(y_true, y_pred))


# --- Définition des pages de l'application ---

def page_accueil():
    """Affiche la page d'accueil."""
    st.title("Projet de Prédiction de Séries Temporelles avec GRU")
    st.markdown("---")
    
    col1, col2 = st.columns([1, 3])
    with col1:
        try:
            logo_keyce = Image.open(script_dir/"Keyce_Logo.jpg")
            st.image(logo_keyce, width=150)
        except FileNotFoundError:
            st.warning("Logo Keyce Keyce_Logo.jpg non trouvé.")

    with col2:
        st.header("KEYCE INFORMATIQUE - Master II IA")
        st.subheader("Session Normale de Réseaux de Neurones Récurrents (RNN)")
        
    st.markdown("---")
    
    st.header("Présentation de l'étudiant")
    st.markdown("### **Nom :** TATSA TCHINDA Colince")
    
    st.info("Utilisez le menu de navigation à gauche pour accéder à la page de prédiction.")

def page_prediction():
    """Affiche la page de prédiction et ses résultats."""
    st.title("📈 Prédiction du Cours de l'Action AMD")

    # --- Étape 1: Chargement des données ---
    st.header("Étape 1 : Chargement et Visualisation des Données")
    with st.spinner("Chargement des données historiques..."):
        data = load_data()
    st.success("Données chargées avec succès !")
    
    fig, ax = plt.subplots(figsize=(12, 5))
    ax.plot(data['ds'], data['y'], label="Historique", color='black')
    ax.set_xlabel("Date")
    ax.set_ylabel("Prix de clôture ($)")
    ax.set_title("Cours historique de l'action AMD")
    ax.grid(True)
    ax.legend()
    st.pyplot(fig)

    # --- Étape 2: Chargement du modèle ---
    st.header("Étape 2 : Chargement du Modèle GRU")
    with st.spinner("Chargement du modèle pré-entraîné..."):
        try:
            gru_model = load_gru_model(script_dir/"model_gru.pth", GRUModel)
            st.success("Modèle GRU chargé avec succès !")
        except FileNotFoundError:
            st.error("Erreur : Le fichier 'model_gru.pth' est introuvable.")
            st.stop()

    # --- Étape 3: Prédictions ---
    st.header("Étape 3 : Génération des Prédictions")
    WINDOW_SIZE = 20
    FORECAST_DAYS = 21
    
    with st.spinner(f"Calcul des prédictions pour les {FORECAST_DAYS} prochains jours..."):
        gru_forecast = predict_gru(gru_model, data, FORECAST_DAYS, window_size=WINDOW_SIZE)

    st.subheader("Prédictions du modèle GRU vs Historique récent")
    fig, ax = plt.subplots(figsize=(12, 5))
    ax.plot(data['ds'].tail(100), data['y'].tail(100), label="Historique récent", color='black')
    ax.plot(gru_forecast['ds'], gru_forecast['yhat'], label=f"Prédiction GRU ({FORECAST_DAYS} jours)", color='orange', linestyle='--')
    ax.set_xlabel("Date")
    ax.set_ylabel("Prix de clôture ($)")
    ax.set_title("Prédictions GRU vs Historique")
    ax.grid(True)
    ax.legend()
    st.pyplot(fig)

    # --- Étape 4: Évaluation et Résultats ---
    st.header("Étape 4 : Évaluation et Résultats")
    
    st.subheader("Performance du Modèle (Backtesting)")
    true_values = data['y'].values[-WINDOW_SIZE:]
    # ... (Le reste de votre logique de backtesting est correct)
    input_for_backtest_np = data['y'].values[-(WINDOW_SIZE * 2):-WINDOW_SIZE]
    backtest_preds = []
    input_seq_for_backtest = input_for_backtest_np.copy()
    for _ in range(WINDOW_SIZE):
        input_tensor = torch.from_numpy(input_seq_for_backtest[-WINDOW_SIZE:]).float().view(1, WINDOW_SIZE, 1)
        with torch.no_grad():
            pred = gru_model(input_tensor).item()
        backtest_preds.append(pred)
        input_seq_for_backtest = np.append(input_seq_for_backtest, pred)
    gru_rmse = calculate_rmse(true_values, backtest_preds)
    
    st.metric(label="RMSE (Backtest sur 20 jours)", value=f"{gru_rmse:.4f}")
    st.info("Le RMSE évalue l'écart moyen entre les valeurs prédites et les valeurs réelles sur la période de test.")
    
    st.subheader("Tableau des Prédictions")
    styled_df = gru_forecast.style.format({'yhat': '{:.2f}'})
    st.dataframe(styled_df, use_container_width=True)
    
    csv = gru_forecast.to_csv(index=False).encode('utf-8')
    st.download_button(
        label="📥 Télécharger les prédictions (.csv)",
        data=csv,
        file_name="predictions_gru_amd.csv",
        mime="text/csv",
    )

# --- Barre latérale et Navigation ---

st.sidebar.header("Navigation")
try:
    logo_theme = Image.open(script_dir/"Theme_Logo.jpg")
    st.sidebar.image(logo_theme,  use_container_width=True)
except FileNotFoundError:
    st.sidebar.warning("Logo thème Theme_Logo.jpg non trouvé.")

page = st.sidebar.selectbox(
    "Choisissez une page",
    ["Accueil", "Prédiction"]
)

# --- Affichage de la page sélectionnée ---

if page == "Accueil":
    page_accueil()
elif page == "Prédiction":
    page_prediction()