import streamlit as st
import pandas as pd
import joblib
from tensorflow import keras
from datetime import date
import os

from utils import forecast_to_annual_return, plot_forecast_vs_actual
from optimization import (
    load_processed_prices, 
    compute_mu_and_cov,
    optimize_portfolio,
    optimize_for_target,
    plot_efficient_frontier,
)
from forecast import forecast_lstm

# --- Config ---
tickers = ["TSLA", "BND", "SPY"]
sequence_length = 60
start_date_data = "2015-01-01" 

APP_DIR = os.path.dirname(os.path.abspath(__file__))
PROJECT_ROOT = os.path.dirname(APP_DIR)
proc_dir = os.path.join(PROJECT_ROOT, "src", "data", "processed")
models_dir = os.path.join(PROJECT_ROOT, "src","models")

# --- App Layout ---
st.set_page_config(layout="wide")
st.title("Portfolio Optimization with LSTM-Forecasted Returns")

st.sidebar.header("User Inputs")
investment_capital = st.sidebar.number_input("Enter Investment Capital ($)", min_value=1000, value=10000, step=1000)
target_return = st.sidebar.slider("Select Target Annual Return (%)", min_value=5.0, max_value=50.0, value=20.0, step=1.0) / 100

if st.sidebar.button("Run Optimization"):

    # --- 1. Fetch live historical data ---
    with st.spinner("Loading historical data..."):
        end_date_data = date.today().strftime("%Y-%m-%d")
        data = load_processed_prices(proc_dir, tickers)

    # --- 2. Load trained LSTM model & scaler from local models folder ---
    with st.spinner("Loading forecasting model..."):
        model_path = os.path.join(models_dir, "lstm_model.keras")
        scaler_path = os.path.join(models_dir, "scaler.joblib")

        model = keras.models.load_model(model_path)
        scaler = joblib.load(scaler_path)

    # --- 3. Prepare TSLA data & forecast ---
    with st.spinner("Generating TSLA price forecast..."):
        tsla_close = data["TSLA"][["Close"]].dropna()
        test_data_input = tsla_close.tail(2 * sequence_length) 
        tsla_forecast = forecast_lstm(model, tsla_close, test_data_input, scaler, sequence_length)
        last_price = tsla_close["Close"].iloc[-1]

    # --- 4. Convert forecast → annual expected return ---
    tsla_annual = forecast_to_annual_return(
        tsla_forecast, last_price, method="avg_daily", cap=2.0
    )

    # --- 5. Build mu & cov ---
    with st.spinner("Calculating expected returns and covariance..."):
        mu, cov = compute_mu_and_cov(data, tsla_annual)

    # --- 6. Optimize portfolios ---
    with st.spinner("Running portfolio optimizations..."):
        general_results = optimize_portfolio(mu, cov, rf=0.02)
        user_portfolio = optimize_for_target(mu, cov, target_type="return", target_value=target_return)

    # --- 7. Display Results ---
    st.header("Optimization Results")
    
    col1, col2 = st.columns(2)

    with col1:
        st.subheader("General Optimal Portfolios")
        st.write("**Max Sharpe Portfolio** (Best risk-adjusted return)")
        st.json(general_results["weights_sharpe"])
        st.write(pd.DataFrame([general_results["perf_sharpe"]]))

        st.write("**Min Volatility Portfolio** (Lowest risk)")
        st.json(general_results["weights_minvol"])
        st.write(pd.DataFrame([general_results["perf_minvol"]]))

    with col2:
        st.subheader(f"Your Custom Portfolio (Target Return: {target_return:.1%})")
        if "error" in user_portfolio["performance"]:
            st.error(f"Could not find a portfolio for this target return. Error: {user_portfolio['performance']['error']}")
        else:
            st.write("**Recommended Weights**")
            st.json(user_portfolio["weights"])
            st.write("**Expected Performance**")
            st.write(pd.DataFrame([user_portfolio["performance"]]))
            
            st.write("**Investment Allocation**")
            weights_df = pd.Series(user_portfolio["weights"])
            allocation_df = (weights_df * investment_capital).map('${:,.2f}'.format)
            st.dataframe(allocation_df)

    # --- 8. Efficient Frontier Plot ---
    st.subheader("Efficient Frontier")
    fig_frontier = plot_efficient_frontier(mu, cov, general_results)
    st.pyplot(fig_frontier)

    # --- 9. Forecast vs Actual Plot ---
    st.subheader("TSLA Forecast vs. Recent Actuals")
    actual_plot_data = tsla_close.tail(len(tsla_forecast) * 2) 
    fig_forecast = plot_forecast_vs_actual(actual_plot_data, tsla_forecast)
    st.pyplot(fig_forecast)

else:
    st.info("Adjust the inputs in the sidebar and click 'Run Optimization' to begin.")