import streamlit as st
import pandas as pd
import numpy as np
import plotly.graph_objects as go
from tvDatafeed import TvDatafeed, Interval
import timesfm
import io
import math

st.set_page_config(
    page_title="חיזוי מניות AI",
    layout="wide",
    page_icon="📈"
)

# =========================
# עיצוב בהיר מקצועי (אכיפת RTL)
# =========================
st.markdown("""
<style>
@import url('https://fonts.googleapis.com/css2?family=Assistant:wght@300;400;600;700&display=swap');

html, body, [class*="css"] {
    font-family: 'Assistant', sans-serif;
    direction: rtl;
    text-align: right;
}

div[data-testid="stMarkdownContainer"], div[data-testid="stAlert"] {
    direction: rtl;
    text-align: right;
}

.stApp { background-color: #f4f6f9; }

.main-title {
    text-align: right;
    font-size: 2.2rem;
    font-weight: 700;
    margin-bottom: 0.3rem;
}

.warning-box {
    background: #fff3cd;
    border: 1px solid #ffeeba;
    padding: 0.8rem;
    border-radius: 8px;
    margin-bottom: 1rem;
    font-size: 0.9rem;
    text-align: right;
    direction: rtl;
}

.table-header {
    font-weight: bold;
    color: #475569;
    padding-bottom: 10px;
    border-bottom: 2px solid #cbd5e1;
    margin-bottom: 10px;
}
</style>
""", unsafe_allow_html=True)

st.markdown("<div class='main-title'>📈 חיזוי מניות ומדדים (Google TimesFM 2.5)</div>", unsafe_allow_html=True)

st.markdown("""
<div class="warning-box">
⚠️ המערכת עובדת באופן בלעדי על מודלים של תשואות (Current/Prev - 1) לצורך דיוק מקסימלי. אינו מהווה ייעוץ השקעות.
</div>
""", unsafe_allow_html=True)

# =========================
# טעינת מודל AI (נשמר בזיכרון)
# =========================
@st.cache_resource(show_spinner=False)
def load_model():
    model = timesfm.TimesFM_2p5_200M_torch.from_pretrained(
        "google/timesfm-2.5-200m-pytorch", 
        torch_compile=False 
    )
    model.compile(
        timesfm.ForecastConfig(
            max_context=2048, 
            max_horizon=128,
            normalize_inputs=True,
            use_continuous_quantile_head=True,
            force_flip_invariance=True,
            infer_is_positive=False, 
            fix_quantile_crossing=True,
        )
    )
    return model

# =========================
# נכסים לבחירה וקישורי Yahoo
# =========================
ASSETS = {
    "לאומי": ("LUMI", "TASE"), "פועלים": ("POLI", "TASE"), "דיסקונט": ("DSCT", "TASE"),
    "מזרחי טפחות": ("MZTF", "TASE"), "אלביט מערכות": ("ESLT", "TASE"), "טבע": ("TEVA", "TASE"),
    "נייס": ("NICE", "TASE"), "בזק": ("BEZQ", "TASE"), "דלק קבוצה": ("DLEKG", "TASE"),
    "מדד ת\"א 35": ("TA35", "TASE"), "S&P 500 ETF": ("SPY", "AMEX"), 
    'נאסד"ק 100 ETF': ("QQQ", "NASDAQ"), "USD/ILS (דולר-שקל)": ("USDILS", "FX_IDC")
}

YAHOO_LINKS = {
    "לאומי": "https://finance.yahoo.com/quote/LUMI.TA",
    "פועלים": "https://finance.yahoo.com/quote/POLI.TA",
    "דיסקונט": "https://finance.yahoo.com/quote/DSCT.TA",
    "מזרחי טפחות": "https://finance.yahoo.com/quote/MZTF.TA",
    "אלביט מערכות": "https://finance.yahoo.com/quote/ESLT.TA",
    "טבע": "https://finance.yahoo.com/quote/TEVA.TA",
    "נייס": "https://finance.yahoo.com/quote/NICE.TA",
    "בזק": "https://finance.yahoo.com/quote/BEZQ.TA",
    "דלק קבוצה": "https://finance.yahoo.com/quote/DLEKG.TA",
    "מדד ת\"א 35": "https://finance.yahoo.com/quote/^TA35",
    "S&P 500 ETF": "https://finance.yahoo.com/quote/SPY",
    'נאסד"ק 100 ETF': "https://finance.yahoo.com/quote/QQQ",
    "USD/ILS (דולר-שקל)": "https://finance.yahoo.com/quote/ILS=X"
}

# =========================
# הגדרות ממשק משתמש
# =========================
col1, col2 = st.columns(2)

with col1:
    stock = st.selectbox("בחר נכס פיננסי", list(ASSETS.keys()))

with col2:
    mode = st.radio(
        "סוג ניתוח",
        ["חיזוי רגיל (עתיד + מבחני עבר)", "חיזוי רב-שכבתי כפול (Multi-Timeframe)", "בדיקת אסטרטגיה (Matrix)"],
        horizontal=False
    )

interval_choice = "1d"
data_source = "שערי סגירה"

if mode == "חיזוי רגיל (עתיד + מבחני עבר)":
    c_res, c_meth = st.columns(2)
    with c_res:
        int_map = {"5 דקות": "5m", "15 דקות": "15m", "30 דקות": "30m", "שעתי (60m)": "60m", "יומי (1d)": "1d", "שבועי (1W)": "1W"}
        resolution_label = st.selectbox("רזולוציית זמן:", list(int_map.keys()), index=4)
        interval_choice = int_map[resolution_label]
    with c_meth:
        data_source = st.radio("מקור הנתונים (יומרו לתשואות באופן אוטומטי):", ["שערי סגירה", "מחיר משוקלל נפח (VWAP 20)"])
elif mode == "חיזוי רב-שכבתי כפול (Multi-Timeframe)":
    st.info("🧬 **מצב מחקר מתקדם:** המערכת תריץ במקביל: שערי סגירה ומחיר משוקלל נפח VWAP על כל רזולוציית זמן (החיזוי עצמו מבוצע על בסיס תשואות).")
else:
    st.info("🧮 **בדיקת אסטרטגיה (מטריצה RMSE):** המערכת מבצעת עד 40 בדיקות לכל תא ומחשבת שגיאה ריבועית. החישוב הפנימי הוא בלעדית על בסיס תשואות (Current/Prev - 1).")

# =========================
# פונקציות ליבה (תאריכים, משיכה, וחיזוי מבוסס תשואות בלבד)
# =========================
def generate_israel_trading_dates(start_date, periods, tf):
    dates = []
    curr = start_date
    if tf == "60m": step = pd.Timedelta(hours=1)
    elif tf == "30m": step = pd.Timedelta(minutes=30)
    elif tf == "15m": step = pd.Timedelta(minutes=15)
    elif tf == "5m": step = pd.Timedelta(minutes=5)
    elif tf == "1W": step = pd.Timedelta(weeks=1)
    else: step = pd.Timedelta(days=1)
    
    while len(dates) < periods:
        curr += step
        if tf == "1W":
            dates.append(curr)
            continue
        weekday = curr.weekday()
        if tf == "1d":
            if weekday in [0, 1, 2, 3, 4]: dates.append(curr)
        else:
            if weekday in [0, 1, 2, 3]:
                if 10 <= curr.hour < 17: dates.append(curr)
            elif weekday == 4:
                if 10 <= curr.hour < 14: dates.append(curr)
    return dates

@st.cache_data(ttl=600, show_spinner=False)
def fetch_data(symbol, interval_str, n_bars=5000):
    tv = TvDatafeed()
    tv_intervals = {"5m": Interval.in_5_minute, "15m": Interval.in_15_minute, "30m": Interval.in_30_minute, "60m": Interval.in_1_hour, "1d": Interval.in_daily, "1W": Interval.in_weekly}
    inter = tv_intervals.get(interval_str, Interval.in_daily)
    df = tv.get_hist(symbol=symbol[0], exchange=symbol[1], interval=inter, n_bars=n_bars)
    
    if df is None or df.empty: return pd.DataFrame()
    if df.index.tz is None: df.index = df.index.tz_localize("UTC").tz_convert("Asia/Jerusalem")
    else: df.index = df.index.tz_convert("Asia/Jerusalem")
    df.index = df.index.tz_localize(None) 
    
    window = 20
    if 'volume' in df.columns and not df['volume'].empty:
        df['vwap'] = (df['close'] * df['volume']).rolling(window=window).sum() / df['volume'].rolling(window=window).sum()
        df['vwap'] = df['vwap'].fillna(df['close']) 
    else:
        df['vwap'] = df['close'] 
        
    return df[['close', 'vwap', 'volume']]

def get_forecast(model, ctx_prices, horizon=128):
    # =================================================================
    # שינוי קריטי 2: המערכת עובדת תמיד אך ורק על תשואות (Current/Prev - 1)
    # =================================================================
    returns = (ctx_prices[1:] / ctx_prices[:-1]) - 1
    returns = np.nan_to_num(returns)
    
    point_fcst, quant_fcst = model.forecast(
        horizon=horizon, 
        inputs=[returns]
    )
    
    fcst_ret = quant_fcst[0, :, 5]
    lower_ret = quant_fcst[0, :, 1]
    upper_ret = quant_fcst[0, :, -1]
    
    # המרה חזרה למחיר אך ורק לשם הצגת הגרף הויזואלי (המודל עצמו עבד על תשואות)
    last_price = ctx_prices[-1]
    fcst_prices = last_price * np.cumprod(1 + fcst_ret)
    fcst_lower = last_price * np.cumprod(1 + lower_ret)
    fcst_upper = last_price * np.cumprod(1 + upper_ret)
    
    return fcst_prices, fcst_lower, fcst_upper

def create_forecast_figure(data_dict):
    ctx_dates, ctx_prices = data_dict['ctx_dates'], data_dict['ctx_prices']
    actual_dates, actual_prices = data_dict['actual_dates'], data_dict['actual_prices']
    fcst_dates, fcst_prices = data_dict['fcst_dates'], data_dict['fcst_prices']
    fcst_lower, fcst_upper = data_dict['fcst_lower'], data_dict['fcst_upper']
    c_val = data_dict['c_val']
    
    hist_len = min(200, len(ctx_prices))
    x_hist_int = list(range(-hist_len + 1, 1))
    x_fcst_int = list(range(1, len(fcst_dates) + 1))
    x_conn_int = [0] + x_fcst_int
    
    custom_hist = [[f"T{x}", d.strftime("%Y-%m-%d %H:%M")] for x, d in zip(x_hist_int, ctx_dates[-hist_len:])]
    custom_hist[-1][0] = "T=0" 
    
    custom_fcst = [[f"T+{x}", d.strftime("%Y-%m-%d %H:%M")] for x, d in zip(x_fcst_int, fcst_dates)]
    custom_conn = [custom_hist[-1]] + custom_fcst

    last_price = ctx_prices[-1]
    conn_fcst = [last_price] + list(fcst_prices)
    conn_lower = [last_price] + list(fcst_lower)
    conn_upper = [last_price] + list(fcst_upper)
    
    fig = go.Figure()
    fig.add_trace(go.Scatter(
        x=x_hist_int, y=ctx_prices[-hist_len:], 
        mode="lines", name="היסטוריה (בסיס)", 
        line=dict(color='#2563eb', width=2),
        customdata=custom_hist,
        hovertemplate="<b>%{customdata[0]}</b> | %{customdata[1]}<br>מחיר: %{y:.2f}<extra></extra>"
    ))
    
    fig.add_trace(go.Scatter(
        x=x_conn_int, y=conn_upper, 
        mode="lines", line=dict(width=0), 
        name="גבול עליון", showlegend=False,
        customdata=custom_conn,
        hovertemplate="<b>%{customdata[0]}</b> | %{customdata[1]}<br>גבול עליון: %{y:.2f}<extra></extra>"
    ))
    
    fig.add_trace(go.Scatter(
        x=x_conn_int, y=conn_lower, 
        mode="lines", fill="tonexty", fillcolor="rgba(245, 158, 11, 0.2)", 
        line=dict(width=0), name="טווח הסתברות",
        customdata=custom_conn,
        hovertemplate="<b>%{customdata[0]}</b> | %{customdata[1]}<br>גבול תחתון: %{y:.2f}<extra></extra>"
    ))
    
    fig.add_trace(go.Scatter(
        x=x_conn_int, y=conn_fcst, 
        mode="lines", name="תחזית AI", 
        line=dict(color='#f59e0b', width=2.5, dash="dash"),
        customdata=custom_conn,
        hovertemplate="<b>%{customdata[0]}</b> | %{customdata[1]}<br>תחזית AI: %{y:.2f}<extra></extra>"
    ))

    if c_val > 0:
        x_act_int = list(range(0, len(actual_dates) + 1))
        custom_act = [custom_hist[-1]] + [[f"T+{x}", d.strftime("%Y-%m-%d %H:%M")] for x, d in zip(range(1, len(actual_dates)+1), actual_dates)]
        conn_act_prices = [last_price] + list(actual_prices)
        
        fig.add_trace(go.Scatter(
            x=x_act_int, y=conn_act_prices, 
            mode="lines", name="מציאות בפועל", 
            line=dict(color='#10b981', width=3),
            customdata=custom_act,
            hovertemplate="<b>%{customdata[0]}</b> | %{customdata[1]}<br>מציאות בפועל: %{y:.2f}<extra></extra>"
        ))
        fig.add_vline(x=0, line_width=2, line_dash="dot", line_color="#94a3b8")
        fig.add_annotation(x=0, y=1.05, yref="paper", text="נקודת עיוורון", showarrow=False, font=dict(color="#94a3b8", size=12), xanchor="center")
    else:
        fig.add_vline(x=0, line_width=2, line_dash="dot", line_color="#94a3b8")
        fig.add_annotation(x=0, y=1.05, yref="paper", text="הווה (T=0)", showarrow=False, font=dict(color="#94a3b8", size=12), xanchor="center")

    fig.update_layout(template="plotly_white", hovermode="x unified", legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1), margin=dict(l=10, r=10, t=40, b=80))
    
    min_x = min(x_hist_int)
    max_x = max(x_fcst_int)
    tick_vals = list(range((min_x // 10) * 10, max_x + 1, 10))
    tick_texts = [f"T+{v}" if v > 0 else f"T{v}" if v < 0 else "T=0" for v in tick_vals]
    
    fig.update_xaxes(tickvals=tick_vals, ticktext=tick_texts, tickangle=-45, automargin=True, title="ציר זמן (מספר נרות ביחס להווה)")
    return fig

@st.dialog("📊 גרף מפורט - חיזוי מול מציאות", width="large")
def show_chart_dialog(c_idx):
    data = st.session_state['backtest_data'][c_idx]
    fig = create_forecast_figure(data)
    st.plotly_chart(fig, use_container_width=True)

def generate_excel(data_dict, stock_name):
    output = io.BytesIO()
    with pd.ExcelWriter(output, engine='openpyxl') as writer:
        link_df = pd.DataFrame({"נכס פיננסי": [stock_name], "קישור לאימות (Yahoo Finance)": [YAHOO_LINKS.get(stock_name, "אין נתון")]})
        link_df.to_excel(writer, index=False, sheet_name="מידע וקישורים")
        
        for sheet_name, df in data_dict.items():
            export_df = df.copy()
            export_df.reset_index(inplace=True)
            cols = list(export_df.columns)
            
            if 'vwap' in cols and 'volume' in cols:
                export_df = export_df[[cols[0], 'close', 'vwap', 'volume']]
                export_df.columns = ["תאריך ושעה", "שער סגירה", "מחיר משוקלל נפח (VWAP)", "נפח מסחר"]
            else:
                export_df = export_df[[cols[0], 'close']]
                export_df.columns = ["תאריך ושעה", "שער סגירה"]
                
            export_df.to_excel(writer, index=False, sheet_name=sheet_name[:31])
    return output.getvalue()

def generate_context_excel(sent_data_dict, stock_name):
    output = io.BytesIO()
    with pd.ExcelWriter(output, engine='openpyxl') as writer:
        # עבודה על סדרות עקב אורכים משתנים, הקובץ מציג בדיוק את התשואות
        df_export = pd.DataFrame(dict([ (k, pd.Series(v)) for k, v in sent_data_dict.items() ]))
        df_export.to_excel(writer, index=False, sheet_name="מחרוזות תשואה שנשלחו")
    return output.getvalue()

@st.dialog("📊 פירוט מלא של הבדיקות", width="large")
def show_cell_details(row_name, col_name):
    key = f"{row_name}_{col_name}"
    if key in st.session_state.get('matrix_details', {}):
        details = st.session_state['matrix_details'][key]
        st.markdown(f"**תא נבחר:** שורה `{row_name}` | עמודה `{col_name}`")
        df_details = pd.DataFrame(details)
        st.dataframe(df_details, use_container_width=True)
    else:
        st.warning("לא נמצאו נתונים עבור התא המבוקש.")

# =========================
# הפעלת הלולאה והחישובים
# =========================
if st.button("🚀 הפעל ניתוח AI מקיף", type="primary", use_container_width=True):

    with st.spinner("טוען מודל ומושך נתונים מ-TradingView..."):
        model = load_model()
    
    st.session_state['selected_stock'] = stock
    st.session_state['raw_data_export'] = {}
        
    # ===============
    # מצב בדיקת אסטרטגיה (Matrix)
    # ===============
    if mode == "בדיקת אסטרטגיה (Matrix)":
        
        ROWS_DEF = [
            ("לפי 5 דקות", 5/60, "5m"), ("לפי 15 דקות", 15/60, "15m"), ("לפי שעה", 1, "1h"),
            ("לפי שעתיים", 2, "1h"), ("לפי 3 שעות", 3, "1h"), ("לפי 4 שעות", 4, "1h"),
            ("לפי 5 שעות", 5, "1h"), ("לפי 6 שעות", 6, "1h"), ("לפי 7 שעות", 7, "1h"),
            ("לפי 8 שעות", 8, "1h"), ("לפי 9 שעות", 9, "1h"), ("לפי 10 שעות מסחר", 10, "1h"),
            ("לפי 11 שעות מסחר", 11, "1h"), ("לפי 12 שעות מסחר", 12, "1h"), ("לפי 13 שעות מסחר", 13, "1h")
        ]
        
        COLS_DEF = [
            ("שעה קדימה", 1), ("שעתיים קדימה", 2), ("5 שעות", 5),
            ("יום", 7.5), ("יומיים", 15), ("3 ימים", 22.5),
            ("4 ימים", 30), ("5 ימים", 37.5), ("6 ימים", 45), ("7 ימים", 52.5)
        ]

        matrix_df = pd.DataFrame(index=[r[0] for r in ROWS_DEF], columns=[c[0] for c in COLS_DEF])
        details_dict = {}
        sent_data_dict = {}
        
        # סנכרון זמנים לחיתוך אחיד!
        df_5m = fetch_data(ASSETS[stock], "5m", 5000)
        df_15m = fetch_data(ASSETS[stock], "15m", 5000)
        df_1h = fetch_data(ASSETS[stock], "60m", 5000)
        
        if not df_5m.empty and not df_15m.empty and not df_1h.empty:
            snapshot_time = min(df_5m.index[-1], df_15m.index[-1], df_1h.index[-1])
            df_5m = df_5m[df_5m.index <= snapshot_time]
            df_15m = df_15m[df_15m.index <= snapshot_time]
            df_1h = df_1h[df_1h.index <= snapshot_time]
        
        progress_bar = st.progress(0)
        status_text = st.empty()
        
        total_cells = len(ROWS_DEF) * len(COLS_DEF)
        current_cell = 0
        
        for row_idx, (row_name, row_hours, base_tf) in enumerate(ROWS_DEF):
            if base_tf == "5m" and not df_5m.empty:
                prices = df_5m['close'].values
                dates = df_5m.index.values
            elif base_tf == "15m" and not df_15m.empty:
                prices = df_15m['close'].values
                dates = df_15m.index.values
            elif base_tf == "1h" and not df_1h.empty:
                if row_hours > 1:
                    chunk_size = int(row_hours)
                    grouped = df_1h.groupby(np.arange(len(df_1h)) // chunk_size)
                    prices = grouped['close'].last().values
                    dates = grouped.apply(lambda x: x.index[-1]).values
                else:
                    prices = df_1h['close'].values
                    dates = df_1h.index.values
            else:
                prices = np.array([])
                dates = np.array([])
                
            if len(prices) > 1:
                # שמירת הנתונים לאקסל: עכשיו שומרים נטו את התשואות!
                ret_array = (prices[1:] / prices[:-1]) - 1
                sent_data_dict[row_name] = ret_array[-1024:]
                
            for col_idx, (col_name, col_hours) in enumerate(COLS_DEF):
                status_text.text(f"מחשב אסטרטגיה ריבועית: {row_name} -> {col_name}...")
                current_cell += 1
                
                H_candles = int(round(col_hours / row_hours))
                
                if len(prices) < 100 or H_candles < 1:
                    matrix_df.loc[row_name, col_name] = "---"
                    progress_bar.progress(current_cell / total_cells)
                    continue
                
                contexts = []
                actual_rets = []
                test_dates = []
                tests_run = 0
                
                # שינוי קריטי 3: 40 בדיקות בכל תא במקום 30!
                for i in range(40):
                    end_idx = len(prices) - 1 - H_candles - (i * 20)
                    start_idx = end_idx - 512 
                    
                    if start_idx < 0:
                        start_idx = max(0, end_idx - 128) 
                        if end_idx - start_idx < 32: break 
                            
                    ctx_prices = prices[start_idx:end_idx]
                    fut_prices = prices[end_idx:end_idx+H_candles]
                    
                    if len(ctx_prices) > 1 and len(fut_prices) == H_candles:
                        # שינוי קריטי 2: ממירים את המחרוזת לתשואות בלבד
                        ctx_ret = (ctx_prices[1:] / ctx_prices[:-1]) - 1
                        contexts.append(np.nan_to_num(ctx_ret))
                        
                        # חישוב תשואה בפועל של התקופה (Current / Prev - 1)
                        act = (fut_prices[-1] / prices[end_idx-1]) - 1
                        actual_rets.append(act)
                        test_dates.append(pd.to_datetime(dates[end_idx-1]).strftime('%Y-%m-%d %H:%M'))
                        tests_run += 1
                
                if tests_run == 0:
                    matrix_df.loc[row_name, col_name] = "---"
                else:
                    try:
                        _, quant_fcst = model.forecast(horizon=H_candles, inputs=contexts)
                        cell_results = []
                        squared_errors = []
                        
                        for b in range(tests_run):
                            fcst_ret = quant_fcst[b, :, 5] 
                            pred = np.prod(1 + fcst_ret) - 1
                            diff = actual_rets[b] - pred
                            sq_err = diff ** 2
                            squared_errors.append(sq_err)
                            
                            cell_results.append({
                                'תאריך T=0': test_dates[b],
                                'תשואה בפועל': f"{actual_rets[b]*100:+.2f}%",
                                'תשואת המודל': f"{pred*100:+.2f}%",
                                'פער מוחלט': f"{diff*100:+.2f}%",
                                'שגיאה ריבועית': f"{sq_err*10000:.4f}"
                            })
                            
                        rmse = math.sqrt(sum(squared_errors) / tests_run) * 100 
                        
                        # שינוי חכם: שמירת השגיאה יחד עם כמות הבדיקות להצגה חלקה
                        matrix_df.loc[row_name, col_name] = f"{rmse}|{tests_run}"
                        details_dict[f"{row_name}_{col_name}"] = cell_results
                        
                    except Exception as e:
                        matrix_df.loc[row_name, col_name] = "---"
                        
                progress_bar.progress(current_cell / total_cells)

        status_text.empty()
        progress_bar.empty()
        
        st.session_state['matrix_df'] = matrix_df
        st.session_state['matrix_details'] = details_dict
        st.session_state['sent_contexts'] = sent_data_dict
        st.session_state['run_done'] = True
        st.session_state['run_mode'] = mode

    # ===============
    # מצב חיזוי רב-שכבתי
    # ===============
    elif mode == "חיזוי רב-שכבתי כפול (Multi-Timeframe)":
        tfs = {"1d": ("יומי", "#f59e0b"), "60m": ("שעתי", "#8b5cf6"), "15m": ("15 דקות", "#ef4444")}
        data_sources = ["שערי סגירה", "VWAP"] # החישוב עצמו הוא תמיד לפי תשואות
        
        fig_mtf = go.Figure()
        
        progress_bar = st.progress(0)
        status_text = st.empty()
        
        bg_df = fetch_data(ASSETS[stock], "60m")
        if not bg_df.empty:
            hist_len = 150
            bg_dates_str = [d.strftime("%Y-%m-%d %H:%M") for d in bg_df.index[-hist_len:]]
            bg_labels = [[f"T-{hist_len - i}", d] for i, d in enumerate(bg_dates_str)]
            
            fig_mtf.add_trace(go.Scatter(
                x=bg_df.index[-hist_len:], y=bg_df['close'].tail(hist_len), mode="lines", 
                name="היסטוריה קרובה (שעתי)", line=dict(color='#cbd5e1', width=1.5),
                customdata=bg_labels,
                hovertemplate="<b>%{customdata[0]}</b> | %{customdata[1]}<br>מחיר: %{y:.2f}<extra></extra>"
            ))

        total_steps = len(tfs) * len(data_sources)
        current_step = 0

        for tf, (name, color) in tfs.items():
            df = fetch_data(ASSETS[stock], tf)
            if df.empty or len(df) < 512: 
                current_step += 2
                continue
            
            st.session_state['raw_data_export'][f"נתוני_{name}"] = df
            last_date = df.index[-1]
            
            if tf == "1d": draw_periods = 25
            elif tf == "60m": draw_periods = 80
            else: draw_periods = 128
            
            fcst_dates = generate_israel_trading_dates(last_date, draw_periods, tf)
            conn_dates = [last_date] + list(fcst_dates)
            
            for source in data_sources:
                status_text.text(f"מנתח שכבת זמן: {name} | מקור: {source}...")
                
                prices_full = df['vwap'].values if source == "VWAP" else df['close'].values
                ctx_prices = prices_full[-1024:] if len(prices_full) > 1024 else prices_full
                last_price = ctx_prices[-1]
                
                try:
                    fcst_prices, _, _ = get_forecast(model, ctx_prices, horizon=draw_periods)
                    conn_fcst = [last_price] + list(fcst_prices)
                    
                    conn_dates_str = [d.strftime("%Y-%m-%d %H:%M") for d in conn_dates]
                    mtf_labels = [["T=0", conn_dates_str[0]]] + [[f"T+{i+1} ({name})", conn_dates_str[i+1]] for i in range(len(fcst_prices))]
                    
                    if source == "שערי סגירה": dash_style = "solid"; opac = 1.0
                    else: dash_style = "dashdot"; opac = 0.9
                    
                    fig_mtf.add_trace(go.Scatter(
                        x=conn_dates, y=conn_fcst, mode="lines", 
                        name=f"תחזית {name} ({source})", 
                        line=dict(color=color, width=2.5, dash=dash_style),
                        opacity=opac,
                        customdata=mtf_labels,
                        hovertemplate="<b>%{customdata[0]}</b> | %{customdata[1]}<br>תחזית: %{y:.2f}<extra></extra>"
                    ))
                except Exception as e: pass
                
                current_step += 1
                progress_bar.progress(current_step / total_steps)
            
        status_text.empty()
        progress_bar.empty()
        
        fig_mtf.update_layout(
            template="plotly_white", hovermode="x unified", title_x=0.5,
            title=f"תצוגה רב-שכבתית (מחושבת מבוסס תשואות): שערים מול משוקלל נפח ({stock})",
            legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1), 
            margin=dict(l=10, r=10, t=40, b=80) 
        )
        fig_mtf.update_xaxes(nticks=40, tickangle=-45, automargin=True)
        
        st.markdown("### 🧬 תרשים רב-שכבתי (Multi-Timeframe)")
        st.plotly_chart(fig_mtf, use_container_width=True)
        st.session_state['run_done'] = True
        st.session_state['run_mode'] = mode

    # ===============
    # מצב חיזוי רגיל
    # ===============
    else:
        df = fetch_data(ASSETS[stock], interval_choice)
        
        if df.empty or len(df) < 1200:
            st.error("❌ אין מספיק נתונים עבור הנכס הזה. נסה רזולוציית זמן קצרה יותר.")
            st.stop()
            
        st.session_state['raw_data_export']["נתונים_גולמיים"] = df

        if interval_choice == "1d":
            unit = "ימי מסחר"
            test_cutoffs = [0, 5, 10, 15, 21, 42, 63, 84, 105, 126]
            labels_dict = {
                5: "שבוע (5 ימים) אחורה", 
                10: "שבועיים (10 ימים) אחורה", 
                21: "חודש (21 ימים) אחורה", 
                42: "חודשיים (42 ימים) אחורה", 
                63: "3 חודשים (63 ימים) אחורה", 
                126: "חצי שנה (126 ימים) אחורה"
            }
            test_labels = [labels_dict.get(c, f"{c} {unit} אחורה") if c > 0 else "חיזוי עתידי אמיתי (היום והלאה)" for c in test_cutoffs]
        else:
            unit = "נרות"
            test_cutoffs = [0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100, 110, 120]
            test_labels = ["חיזוי עתידי אמיתי (היום והלאה)"] + [f"{c} {unit} אחורה ({c} כפול {resolution_label})" for c in test_cutoffs[1:]]

        st.session_state['test_cutoffs'] = test_cutoffs
        st.session_state['backtest_data'] = {}
        results_list = []

        prices_full = df['vwap'].values if "VWAP" in data_source else df['close'].values
        dates_full = df.index

        progress_bar = st.progress(0)
        status_text = st.empty()

        for i, (c, label) in enumerate(zip(test_cutoffs, test_labels)):
            status_text.text(f"מחשב מודל (מקור: {data_source}) עבור: {label}...")
            
            if len(prices_full) - c >= 1024:
                if c > 0:
                    ctx_prices = prices_full[:-c]
                    ctx_dates = dates_full[:-c]
                    actual_prices = prices_full[-c:]
                    actual_dates = dates_full[-c:]
                else:
                    ctx_prices = prices_full
                    ctx_dates = dates_full
                    actual_prices = []
                    actual_dates = []

                last_date = ctx_dates[-1]
                last_price = ctx_prices[-1]

                try:
                    fcst_prices, fcst_lower, fcst_upper = get_forecast(model, ctx_prices, horizon=128)
                    fcst_dates = generate_israel_trading_dates(last_date, 128, interval_choice)

                    if c > 0:
                        pred_for_actual = fcst_prices[:c]
                        mape = np.mean(np.abs((actual_prices - pred_for_actual) / actual_prices)) * 100
                        act_dir = actual_prices[-1] - last_price
                        pred_dir = pred_for_actual[-1] - last_price
                        is_correct = (act_dir > 0 and pred_dir > 0) or (act_dir < 0 and pred_dir < 0)
                        trend_str = "✅ קלע לכיוון" if is_correct else "❌ טעה בכיוון"
                        mape_str = f"{mape:.2f}%"
                    else:
                        trend_str = "🔮 עתיד"
                        mape_str = "---"
                        is_correct = None

                    if c > 0:
                        results_list.append({
                            "label": label,
                            "mape": mape_str,
                            "trend": trend_str,
                            "_c_val": c,
                            "_is_correct": is_correct
                        })

                    st.session_state['backtest_data'][c] = {
                        'ctx_dates': ctx_dates, 'ctx_prices': ctx_prices,
                        'actual_dates': actual_dates, 'actual_prices': actual_prices,
                        'fcst_dates': fcst_dates, 'fcst_prices': fcst_prices,
                        'fcst_lower': fcst_lower, 'fcst_upper': fcst_upper,
                        'c_val': c, 'label': label
                    }
                except Exception as e: pass 
                    
            progress_bar.progress((i + 1) / len(test_cutoffs))

        status_text.empty()
        progress_bar.empty()

        if results_list or mode == "חיזוי רגיל (עתיד + מבחני עבר)":
            st.session_state['results_df'] = pd.DataFrame(results_list)
            st.session_state['run_done'] = True
            st.session_state['run_mode'] = mode


# =========================
# תצוגת התוצאות (לפי מצב הריצה)
# =========================
if st.session_state.get('run_done'):
    
    if st.session_state.get('run_mode') == "בדיקת אסטרטגיה (Matrix)":
        st.markdown("### 🧮 מפת חום לפי סטיית תקן ריבועית (RMSE Backtesting)")
        st.markdown("""
        **מפת החום מייצגת שגיאה ריבועית המבוססת על תשואות (לא שערים):**<br>
        המספר בכל תא הוא מדד ה-RMSE באחוזים (הסטייה הכוללת של המודל תוך קנס על שגיאות קיצוניות).<br>
        אם התבצעו פחות מ-40 בדיקות באותו תא, הכמות תוצג בסוגריים. **ירוק חזק** = קרוב ל-0 (דיוק מקסימלי).
        """, unsafe_allow_html=True)
        
        # פונקציית פיצול מחרוזות חכמה לצביעה
        def color_rmse(val):
            if pd.isna(val) or val == "---": return 'color: transparent;'
            try:
                num = float(str(val).split('|')[0])
                if num <= 0.5: color = '#86efac' 
                elif num <= 1.0: color = '#dcfce3' 
                elif num <= 2.0: color = '#fee2e2' 
                else: color = '#fca5a5' 
                return f'background-color: {color}; color: black; font-weight: 600;'
            except: return ''
            
        def format_rmse(val):
            if pd.isna(val) or val == "---": return "---"
            try:
                parts = str(val).split('|')
                rmse = float(parts[0])
                tests = int(parts[1])
                if tests < 40:
                    return f"{rmse:.2f}% ({tests})"
                return f"{rmse:.2f}%"
            except: return str(val)

        styled_df = st.session_state['matrix_df'].style.applymap(color_rmse).format(format_rmse)
        
        selection_event = st.dataframe(
            styled_df, 
            use_container_width=True, 
            height=600,
            on_select="rerun", 
            selection_mode="single-cell"
        )
        
        st.divider()
        
        st.markdown("#### 🔍 פירוט בדיקות היסטוריות לתא נבחר")
        
        # שינוי קריטי 1: מנגנון Try-Except למניעת קריסת הלחיצה על התא
        if len(selection_event.selection.cells) > 0:
            try:
                cell = selection_event.selection.cells[0]
                # תמיכה בכל גרסאות Streamlit
                if isinstance(cell, dict) and "row" in cell:
                    selected_row_idx = cell["row"]
                    selected_col_idx = cell["column"]
                elif isinstance(cell, tuple):
                    selected_row_idx = cell[0]
                    selected_col_idx = cell[1]
                else:
                    selected_row_idx, selected_col_idx = 0, 0
                    
                r_name = st.session_state['matrix_df'].index[selected_row_idx]
                c_name = st.session_state['matrix_df'].columns[selected_col_idx]
                
                if st.button(f"👁️ פתח פירוט מלא עבור התא המסומן ({r_name} -> {c_name})", type="primary"):
                    show_cell_details(r_name, c_name)
            except Exception as e:
                st.error("הייתה בעיה בזיהוי התא. אנא לחץ שוב, פעם אחת בלבד.")
        else:
            st.info("👆 לחץ (לחיצה בודדת) על תא כלשהו בטבלה למעלה, ואז יופיע כאן כפתור לפתיחת פירוט ה-40 בדיקות שלו.")
            
        if 'sent_contexts' in st.session_state:
            st.divider()
            st.markdown("#### 📥 שקיפות מלאה: תשואות המודל")
            st.write("באפשרותך להוריד קובץ אקסל שמכיל בכל עמודה את **רצף התשואות העשרוני** המדויק שנשלח למודל.")
            excel_context = generate_context_excel(st.session_state['sent_contexts'], st.session_state['selected_stock'])
            st.download_button(
                label="💾 הורד מחרוזות תשואה (Excel)",
                data=excel_context,
                file_name=f"{st.session_state['selected_stock']}_ModelReturnsContext.xlsx",
                mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
                use_container_width=True
            )

    elif st.session_state.get('run_mode') == "חיזוי רגיל (עתיד + מבחני עבר)":
        st.markdown("### 📈 תחזית עתידית (מהיום והלאה)")
        future_data = st.session_state['backtest_data'][0]
        fig_future = create_forecast_figure(future_data)
        st.plotly_chart(fig_future, use_container_width=True)
        
        st.divider()
        
        df_res = st.session_state.get('results_df', pd.DataFrame())

        if not df_res.empty:
            correct_count = sum(1 for x in df_res['_is_correct'] if x == True)
            total_tests = sum(1 for x in df_res['_is_correct'] if x is not None)
            win_rate = (correct_count / total_tests) * 100 if total_tests > 0 else 0

            st.markdown("### 🔬 מבחני אמינות אוטומטיים למודל")
            st.info("💡 המערכת חזרה אחורה בזמן ובדקה אם התחזיות שלה אכן התממשו במציאות. **לחץ על לחצן 'הצג' בכל שורה כדי לראות את הגרף!**")

            col_h1, col_h2, col_h3, col_h4 = st.columns([2, 2, 2, 1])
            col_h1.markdown("<div class='table-header'>נקודת התחלה (בדיקת עבר)</div>", unsafe_allow_html=True)
            col_h2.markdown("<div class='table-header'>סטייה מהמציאות (MAPE)</div>", unsafe_allow_html=True)
            col_h3.markdown("<div class='table-header'>זיהוי כיוון מגמה</div>", unsafe_allow_html=True)
            col_h4.markdown("<div class='table-header'>פעולה</div>", unsafe_allow_html=True)
            
            for index, row in df_res.iterrows():
                c1, c2, c3, c4 = st.columns([2, 2, 2, 1])
                c1.write(row['label'])
                c2.write(row['mape'])
                
                trend = row['trend']
                if "✅" in trend: c3.markdown(f"<span style='color: #047857; font-weight: bold;'>{trend}</span>", unsafe_allow_html=True)
                else: c3.markdown(f"<span style='color: #b91c1c; font-weight: bold;'>{trend}</span>", unsafe_allow_html=True)
                
                if c4.button("📊 הצג", key=f"btn_show_{row['_c_val']}"):
                    show_chart_dialog(row['_c_val'])
                    
                st.markdown("<hr style='margin: 0.2rem 0; opacity: 0.2;'>", unsafe_allow_html=True)

            if total_tests > 1:
                if win_rate >= 60:
                    st.success(f"🏆 **ציון אמינות כללי:** {win_rate:.0f}% הצלחה בזיהוי המגמה. (נחשב למודל יציב ואמין עבור הנכס הזה)")
                elif win_rate <= 40:
                    st.error(f"⚠️ **ציון אמינות כללי:** {win_rate:.0f}% הצלחה בזיהוי המגמה. (המודל מתקשה לקרוא את הנכס הזה, לא מומלץ להסתמך עליו כאן)")
                else:
                    st.warning(f"⚖️ **ציון אמינות כללי:** {win_rate:.0f}% הצלחה בזיהוי המגמה. (תוצאה בינונית - כדאי לשלב כלים נוספים בהחלטה)")

            with st.expander("❓ איך מחושבת 'הסטייה מהמציאות' (MAPE)?"):
                st.markdown("""
                **MAPE (Mean Absolute Percentage Error)** הוא מדד סטטיסטי שמראה בכמה אחוזים המודל "פספס" בממוצע.
                """)

    if 'raw_data_export' in st.session_state and st.session_state['raw_data_export']:
        st.divider()
        st.markdown("### 📥 הורדת נתונים גולמיים")
        excel_file = generate_excel(st.session_state['raw_data_export'], st.session_state['selected_stock'])
        st.download_button(
            label="💾 הורד קובץ נתונים מ-TradingView (Excel)",
            data=excel_file,
            file_name=f"{st.session_state['selected_stock']}_RawData.xlsx",
            mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
            use_container_width=True
        )

st.divider()
st.markdown("""
<div style='text-align: center; color: #64748b; font-size: 0.85rem; padding-top: 1rem; padding-bottom: 2rem; direction: rtl;'>
    מודל החיזוי מופעל באמצעות Google TimesFM 2.5. האתר לצורכי מחקר, ועל אחריות המשתמש.<br>
    לשיתופי פעולה ניתן לפנות ליוצר במייל: <a href="mailto:147590@gmail.com" style="color: #3b82f6; text-decoration: none;" dir="ltr">147590@gmail.com</a>
</div>
""", unsafe_allow_html=True)