import gradio as gr
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import joblib
from tensorflow import keras


try:
    scaler = joblib.load("scaler_minmax.joblib")
    lstm_loaded = keras.models.load_model("lstm_stock_model.keras")
    print("Loaded OK!")
except Exception as e:
    print(f"Load error: {e} — Check files in folder.")

WINDOW_SIZE = 60

def get_last_known(close_prices_list):
    mean_close = np.mean(close_prices_list)
    return {
        "High": mean_close * 1.05, "Low": mean_close * 0.95, "Open": mean_close,
        "Log_Volume": np.log(30000000 + 1)
    }

def predict_multi_days(csv_file, days_to_predict):
    try:
        if csv_file is None:
            fig, ax = plt.subplots(figsize=(10,5))
            ax.text(0.5, 0.5, 'Upload CSV with Close column', ha='center', va='center', transform=ax.transAxes, fontsize=14)
            ax.axis('off')
            return None, fig, pd.DataFrame({"Error": ["Upload CSV"]})

        df = pd.read_csv(csv_file)
        close_col = next((col for col in df.columns if 'close' in col.lower()), None)
        if close_col is None:
            raise ValueError("No 'Close' column found.")
        close_prices_series = df[close_col].dropna()
        close_prices_list = close_prices_series.tolist()
        close_prices = np.array(close_prices_list, dtype=float).reshape(-1, 1)

        if len(close_prices) < WINDOW_SIZE:
            raise ValueError(f"Need 60+ prices (got {len(close_prices)}).")

        last_60 = close_prices[-WINDOW_SIZE:]
        last_known = get_last_known(close_prices_list)
        other_features = np.tile(list(last_known.values()), (WINDOW_SIZE, 1))
        input_data = np.hstack([last_60, other_features])
        scaled_input = scaler.transform(input_data)
        window_3d = np.expand_dims(scaled_input, axis=0)

        predictions = []
        current_window = window_3d.copy()
        for _ in range(int(days_to_predict)):
            pred_scaled = lstm_loaded.predict(current_window, verbose=0)
            pred = scaler.inverse_transform(pred_scaled)
            next_close = float(pred[0, 0])
            predictions.append(next_close)
            new_row = np.append([next_close], list(last_known.values()))
            current_window = np.roll(current_window, -1, axis=1)
            current_window[0, -1, :] = new_row

        fig, ax = plt.subplots(figsize=(12,6))
        hist_x = range(1, WINDOW_SIZE + 1)
        fut_x = range(WINDOW_SIZE + 1, WINDOW_SIZE + len(predictions) + 1)
        ax.plot(hist_x, last_60, label="Last 60 Days", color='blue', lw=2)
        ax.plot(fut_x, predictions, label=f"Predicted {days_to_predict} Days", color='red', ls='--', lw=2, marker='o')
        ax.set_title(f"AAPL Close Prediction ({days_to_predict} Days)")
        ax.set_xlabel("Days")
        ax.set_ylabel("Close Price ($)")
        ax.legend()
        ax.grid(True, alpha=0.3)
        plt.tight_layout()

        pred_df = pd.DataFrame({
            "Day": [f"Day {i+1}" for i in range(len(predictions))],
            "Predicted Close": predictions
        })

        return predictions[0] if predictions else None, fig, pred_df

    except Exception as e:
        print(f"Error: {e}")
        fig, ax = plt.subplots(figsize=(10,5))
        ax.text(0.5, 0.5, f'Error: {str(e)}', ha='center', va='center', fontsize=14)
        ax.axis('off')
        plt.tight_layout()
        return None, fig, pd.DataFrame({"Error": [str(e)]})

with gr.Blocks(title="AAPL Multi-Day Predictor") as demo:
    gr.Markdown("# AAPL Multi-Day Stock Predictor")
    gr.Markdown("Upload CSV with 'Close' column, choose days—get plot + predictions.")

    with gr.Row():
        csv_input = gr.File(label="Upload CSV", file_types=[".csv"])
        days_slider = gr.Slider(1, 90, value=30, step=1, label="Days to Predict")

    submit_btn = gr.Button("Predict", variant="primary")

    with gr.Row():
        pred_num = gr.Number(label="First Day Prediction")

    plot_out = gr.Plot(label="Historical + Future")
    table_out = gr.Dataframe(label="Predictions")

    submit_btn.click(
        fn=predict_multi_days,
        inputs=[csv_input, days_slider],
        outputs=[pred_num, plot_out, table_out]
    )

demo.launch(share=True, quiet=True)