Update app.py

app.py

@@ -1,5 +1,7 @@
-# app.py
-import os, io, math, warnings
+import os
+import io
+import math
+import warnings
 warnings.filterwarnings("ignore")
 
 from typing import List, Tuple, Dict, Optional
@@ -12,13 +14,16 @@ from PIL import Image
 import requests
 import yfinance as yf
 
+# Embeddings
+from sentence_transformers import SentenceTransformer, util
+
 # ---------------- config ----------------
 DATA_DIR = "data"
 DATASET_PATH = os.path.join(DATA_DIR, "investor_profiles.csv")
 
 MAX_TICKERS = 30
 DEFAULT_LOOKBACK_YEARS = 5
-MARKET_TICKER = "VOO"
+MARKET_TICKER = "VOO"  # will auto-fallback to SPY if VOO missing
 
 POS_COLS = ["ticker", "amount_usd", "weight_exposure", "beta"]
 SUG_COLS = ["ticker", "suggested_weight_pct"]
@@ -35,16 +40,44 @@ FRED_MAP = [
     (100, "DGS30"),
 ]
 
+# Embedding model cfg
+EMB_MODEL_NAME = "FinLang/finance-embeddings-investopedia"
+
+# ---------------- globals (runtime) ----------------
+HORIZON_YEARS = 5.0
+RF_CODE = "DGS5"
+RF_ANN = 0.03
+
+UNIVERSE: List[str] = [MARKET_TICKER, "QQQ", "XLK", "XLP", "XLE", "VNQ", "IEF", "HYG", "GLD", "EEM"]
+
+LAST_DATASET_PATH: Optional[str] = None
+LAST_UNIVERSE: Optional[List[str]] = None
+LAST_PLOT_STATE: Optional[Dict[str, float]] = None
+
+# embedding caches
+_EMB_MODEL = None
+_DS_TEXTS = None
+_DS_EMBS = None
+_DS_CACHE_KEY = None   # (csv_path, tuple(universe))
+
+
 # ---------------- helpers ----------------
 def ensure_data_dir():
     os.makedirs(DATA_DIR, exist_ok=True)
 
+
 def empty_positions_df():
     return pd.DataFrame(columns=POS_COLS)
 
+
 def empty_suggest_df():
     return pd.DataFrame(columns=SUG_COLS)
 
+
+def fmt_pct(x: float) -> str:
+    return f"{x*100:.2f}%"
+
+
 def fred_series_for_horizon(years: float) -> str:
     y = max(1.0, min(100.0, float(years)))
     for cutoff, code in FRED_MAP:
@@ -52,6 +85,7 @@ def fred_series_for_horizon(years: float) -> str:
             return code
     return "DGS30"
 
+
 def fetch_fred_yield_annual(code: str) -> float:
     url = f"https://fred.stlouisfed.org/graph/fredgraph.csv?id={code}"
     try:
@@ -63,31 +97,8 @@ def fetch_fred_yield_annual(code: str) -> float:
     except Exception:
         return 0.03
 
-def fetch_prices_monthly(tickers: List[str], years: int) -> pd.DataFrame:
-    start = pd.Timestamp.today(tz="UTC") - pd.DateOffset(years=years, days=7)
-    end = pd.Timestamp.today(tz="UTC")
-    syms = [str(t).upper().strip() for t in dict.fromkeys(tickers)]
-    df = yf.download(
-        syms, start=start.date(), end=end.date(),
-        interval="1mo", auto_adjust=True, progress=False
-    )["Close"]
-    if isinstance(df, pd.Series):
-        df = df.to_frame()
-    df = df.dropna(how="all").fillna(method="ffill")
-    # columns become single Index if single ticker
-    if isinstance(df.columns, pd.MultiIndex):
-        df.columns = [c[1] if isinstance(c, tuple) else c for c in df.columns]
-    return df
-
-def monthly_returns(prices: pd.DataFrame) -> pd.DataFrame:
-    return prices.pct_change().dropna()
-
-def annualize_mean(m):
-    return np.asarray(m, dtype=float) * 12.0
-
-def annualize_sigma(s):
-    return np.asarray(s, dtype=float) * math.sqrt(12.0)
 
+# -------- Yahoo symbol search ----------
 def yahoo_search(query: str):
     if not query or len(query.strip()) == 0:
         return []
@@ -111,30 +122,128 @@ def yahoo_search(query: str):
     except Exception:
         return [{"symbol": query.strip().upper(), "name": "typed symbol", "exchange": "n a"}]
 
+
+# --------- prices / returns ----------
+def _extract_close(df: pd.DataFrame, tickers: List[str]) -> pd.DataFrame:
+    """
+    Robustly extract a (date x ticker) Close DataFrame regardless of yf's column layout.
+    """
+    if isinstance(df.columns, pd.MultiIndex):
+        lv0 = df.columns.get_level_values(0)
+        lv1 = df.columns.get_level_values(1)
+        if "Close" in lv0:
+            close = df["Close"]
+        elif "Adj Close" in lv0:
+            close = df["Adj Close"]
+        elif "Close" in lv1:
+            close = df.xs("Close", level=1, axis=1)
+        elif "Adj Close" in lv1:
+            close = df.xs("Adj Close", level=1, axis=1)
+        else:
+            # fallback: if first level are tickers
+            # try to select 'Close' under each
+            try:
+                close = df.xs("Close", level=1, axis=1)
+            except Exception:
+                close = df.copy()
+    else:
+        # Single ticker case
+        if "Close" in df.columns:
+            s = df["Close"].copy()
+        elif "Adj Close" in df.columns:
+            s = df["Adj Close"].copy()
+        else:
+            # last resort: take any one numeric column
+            s = df.select_dtypes(include=[np.number]).iloc[:, 0]
+        # ensure column named as ticker
+        name = tickers[0] if len(tickers) else "T0"
+        close = s.to_frame(name=name)
+
+    # Reindex columns to requested order where possible
+    # If some symbols missing, they simply won't be present
+    close = close.dropna(how="all").ffill()
+    # Keep only requested tickers, in order
+    cols = [c for c in tickers if c in close.columns]
+    if not cols:  # if nothing matched, keep whatever is there
+        close = close.copy()
+    else:
+        close = close[cols]
+    return close
+
+
+def fetch_prices_monthly(tickers: List[str], years: int) -> pd.DataFrame:
+    start = pd.Timestamp.today(tz="UTC") - pd.DateOffset(years=years, days=7)
+    end = pd.Timestamp.today(tz="UTC")
+    dl = yf.download(
+        list(dict.fromkeys(tickers)),
+        start=start.date(),
+        end=end.date(),
+        interval="1mo",
+        auto_adjust=True,
+        progress=False
+    )
+    close = _extract_close(dl, tickers)
+    return close
+
+
+def monthly_returns(prices: pd.DataFrame) -> pd.DataFrame:
+    return prices.pct_change().dropna(how="all")
+
+
+def annualize_mean(m):
+    return np.asarray(m, dtype=float) * 12.0
+
+
+def annualize_sigma(s):
+    return np.asarray(s, dtype=float) * math.sqrt(12.0)
+
+
 def validate_tickers(symbols: List[str], years: int) -> List[str]:
-    ok, df = [], fetch_prices_monthly(list(set(symbols)), years)
-    for s in symbols:
-        if s in df.columns:
-            ok.append(s)
+    uniq = list(dict.fromkeys(symbols))
+    df = fetch_prices_monthly(uniq, years)
+    ok = [s for s in uniq if s in df.columns]
     return ok
 
+
 # -------------- aligned moments --------------
 def get_aligned_monthly_returns(symbols: List[str], years: int) -> pd.DataFrame:
-    uniq = [c for c in dict.fromkeys(symbols) if c != MARKET_TICKER]
-    tickers = uniq + [MARKET_TICKER]
+    uniq = [c for c in dict.fromkeys(symbols) if c]
+    tickers = uniq.copy()
+
+    # Ensure market present (try MARKET_TICKER then fallback to SPY)
+    market_ok = MARKET_TICKER in tickers
+    if not market_ok:
+        tickers.append(MARKET_TICKER)
+
     px = fetch_prices_monthly(tickers, years)
+    if MARKET_TICKER not in px.columns:
+        # fallback to SPY if VOO missing
+        if "SPY" not in tickers:
+            tickers.append("SPY")
+            px2 = fetch_prices_monthly(tickers, years)
+            if "SPY" in px2.columns:
+                px = px2
+            else:
+                pass  # keep px as-is
+
     rets = monthly_returns(px)
-    cols = [c for c in uniq if c in rets.columns] + ([MARKET_TICKER] if MARKET_TICKER in rets.columns else [])
-    R = rets[cols].dropna(how="any")
+    keep = [c for c in uniq if c in rets.columns]
+    if MARKET_TICKER in rets.columns:
+        keep += [MARKET_TICKER]
+    elif "SPY" in rets.columns:
+        keep += ["SPY"]
+    R = rets[keep].dropna(how="any")
     return R.loc[:, ~R.columns.duplicated()]
 
+
 def estimate_all_moments_aligned(symbols: List[str], years: int, rf_ann: float):
     R = get_aligned_monthly_returns(symbols, years)
-    if MARKET_TICKER not in R.columns or R.shape[0] < 3:
-        raise ValueError("Could not align data with market or not enough rows.")
-    rf_m = rf_ann / 12.0
+    mkt_col = MARKET_TICKER if MARKET_TICKER in R.columns else ("SPY" if "SPY" in R.columns else None)
+    if mkt_col is None or R.shape[0] < 3:
+        raise ValueError("Not enough aligned data including market")
 
-    m = R[MARKET_TICKER]
+    rf_m = rf_ann / 12.0
+    m = R[mkt_col]
     if isinstance(m, pd.DataFrame):
         m = m.iloc[:, 0].squeeze()
 
@@ -144,23 +253,25 @@ def estimate_all_moments_aligned(symbols: List[str], years: int, rf_ann: float):
 
     ex_m = m - rf_m
     var_m = float(np.var(ex_m.values, ddof=1))
-    var_m = max(var_m, 1e-6)
+    var_m = max(var_m, 1e-8)
 
     betas: Dict[str, float] = {}
-    for s in [c for c in R.columns if c != MARKET_TICKER]:
+    for s in [c for c in R.columns if c != mkt_col]:
         ex_s = R[s] - rf_m
         betas[s] = float(np.cov(ex_s.values, ex_m.values, ddof=1)[0, 1] / var_m)
-    betas[MARKET_TICKER] = 1.0
+    betas[mkt_col] = 1.0  # definition
 
-    asset_cols = [c for c in R.columns if c != MARKET_TICKER]
+    asset_cols = [c for c in R.columns if c != mkt_col]
     cov_m = np.cov(R[asset_cols].values.T, ddof=1) if asset_cols else np.zeros((0, 0))
     covA = pd.DataFrame(cov_m * 12.0, index=asset_cols, columns=asset_cols)
 
-    return {"betas": betas, "cov_ann": covA, "erp_ann": erp_ann, "sigma_m_ann": sigma_m_ann}
+    return {"betas": betas, "cov_ann": covA, "erp_ann": erp_ann, "sigma_m_ann": sigma_m_ann, "mkt_col": mkt_col}
+
 
 def capm_er(beta: float, rf_ann: float, erp_ann: float) -> float:
     return float(rf_ann + beta * erp_ann)
 
+
 def portfolio_stats(weights: Dict[str, float],
                     cov_ann: pd.DataFrame,
                     betas: Dict[str, float],
@@ -180,20 +291,23 @@ def portfolio_stats(weights: Dict[str, float],
     sigma_p = math.sqrt(float(max(w_expo.T @ cov @ w_expo, 0.0)))
     return beta_p, er_p, sigma_p
 
-# -------------- CML helpers + plot (percent axes) --------------
+
+# -------------- CML helpers --------------
 def efficient_same_sigma(sigma_target: float, rf_ann: float, erp_ann: float, sigma_mkt: float):
     if sigma_mkt <= 1e-12:
         return 0.0, 1.0, rf_ann
     a = sigma_target / sigma_mkt
     return a, 1.0 - a, rf_ann + a * erp_ann
 
+
 def efficient_same_return(mu_target: float, rf_ann: float, erp_ann: float, sigma_mkt: float):
     if abs(erp_ann) <= 1e-12:
         return 0.0, 1.0, rf_ann
     a = (mu_target - rf_ann) / erp_ann
     return a, 1.0 - a, abs(a) * sigma_mkt
 
-def plot_cml_percent(
+
+def plot_cml(
     rf_ann, erp_ann, sigma_mkt,
     pt_sigma, pt_mu,
     same_sigma_sigma, same_sigma_mu,
@@ -206,34 +320,29 @@ def plot_cml_percent(
         0.3,
         sigma_mkt * 2.0,
         pt_sigma * 1.4,
-        same_mu_sigma * 1.4,
-        same_sigma_sigma * 1.4,
+        (same_mu_sigma or 0.0) * 1.4,
+        (same_sigma_sigma or 0.0) * 1.4,
         (targ_sigma or 0.0) * 1.4,
     )
     xs = np.linspace(0, xmax, 160)
     slope = erp_ann / max(sigma_mkt, 1e-12)
     cml = rf_ann + slope * xs
+    plt.plot(xs, cml, label="CML via Market", linewidth=2.0)
 
-    def pct(x): return 100.0 * np.asarray(x)
+    # key points
+    plt.scatter([0.0], [rf_ann], label="Risk-free (FRED)")
+    plt.scatter([sigma_mkt], [rf_ann + erp_ann], label="Market")
+    plt.scatter([pt_sigma], [pt_mu], label="Your portfolio", marker="D")
 
-    plt.plot(pct(xs), pct(cml), label="CML through VOO")
-
-    plt.scatter([0.0], [pct(rf_ann)], label="Risk free")
-    plt.scatter([pct(sigma_mkt)], [pct(rf_ann + erp_ann)], label="Market VOO")
-    plt.scatter([pct(pt_sigma)], [pct(pt_mu)], label="Your portfolio")
-    plt.scatter([pct(same_sigma_sigma)], [pct(same_sigma_mu)], label="Efficient same sigma")
-    plt.scatter([pct(same_mu_sigma)], [pct(same_mu_mu)], label="Efficient same return")
+    if same_sigma_sigma is not None and same_sigma_mu is not None:
+        plt.scatter([same_sigma_sigma], [same_sigma_mu], label="Efficient same sigma", marker="o")
+    if same_mu_sigma is not None and same_mu_mu is not None:
+        plt.scatter([same_mu_sigma], [same_mu_mu], label="Efficient same return", marker="o")
     if targ_sigma is not None and targ_mu is not None:
-        plt.scatter([pct(targ_sigma)], [pct(targ_mu)], label="Target suggestion")
-
-    # Guides (keep simple)
-    plt.plot([pct(pt_sigma), pct(same_sigma_sigma)], [pct(pt_mu), pct(same_sigma_mu)],
-             linestyle="--", linewidth=1.0, alpha=0.7, color="gray")
-    plt.plot([pct(pt_sigma), pct(same_mu_sigma)], [pct(pt_mu), pct(same_mu_mu)],
-             linestyle="--", linewidth=1.0, alpha=0.7, color="gray")
+        plt.scatter([targ_sigma], [targ_mu], label="Suggestion", marker="X", s=70)
 
-    plt.xlabel("Standard deviation (%)")
-    plt.ylabel("Expected return (%)")
+    plt.xlabel("σ (annualized)")
+    plt.ylabel("Expected return (annual)")
     plt.legend(loc="best")
     plt.tight_layout()
 
@@ -243,128 +352,193 @@ def plot_cml_percent(
     buf.seek(0)
     return Image.open(buf)
 
-# -------------- dataset over *current* tickers --------------
-def dirichlet_mixture(n: int, k: int, allow_shorts: bool, rng: np.random.Generator) -> np.ndarray:
-    """Return n weight vectors (exposures) across k assets; sum |w| = 1."""
-    out = []
-    n1 = int(n * 0.6)  # diversified
-    n2 = n - n1        # concentrated
-    for _ in range(n1):
-        w = rng.dirichlet(np.ones(k))
-        if allow_shorts:
-            signs = rng.choice([-1.0, 1.0], size=k, p=[0.25, 0.75])
-            w = w * signs
-        out.append(w)
-    for _ in range(n2):
-        hot = rng.integers(0, k)
-        alpha = np.ones(k) * 0.3
-        alpha[hot] = 3.0
-        w = rng.dirichlet(alpha)
-        if allow_shorts:
-            signs = rng.choice([-1.0, 1.0], size=k, p=[0.35, 0.65])
-            w = w * signs
-        out.append(w)
-    W = np.vstack(out)
-    # normalize to exposure space (sum |w| = 1)
-    denom = np.sum(np.abs(W), axis=1, keepdims=True)
-    denom[denom == 0] = 1.0
-    return W / denom
-
-def build_fixed_universe_dataset(
-    symbols: List[str], years: int, rf_ann: float, erp_ann: float,
-    covA: pd.DataFrame, betas: Dict[str, float],
-    allow_shorts: bool, n_rows: int = 1000
-) -> pd.DataFrame:
-    rng = np.random.default_rng(12345)
-    k = len(symbols)
-    W = dirichlet_mixture(n_rows, k, allow_shorts, rng)
 
-    rows = []
-    for i in range(W.shape[0]):
-        w = W[i]
-        wmap = {symbols[j]: float(w[j]) for j in range(k)}
-        beta_p, er_p, sigma_p = portfolio_stats(wmap, covA, betas, rf_ann, erp_ann)
+def _overlay_plot_with_suggestion(sigma_sugg: Optional[float], mu_sugg: Optional[float]) -> Optional[Image.Image]:
+    if LAST_PLOT_STATE is None:
+        return None
+    s = LAST_PLOT_STATE
+    return plot_cml(
+        s["rf_ann"], s["erp_ann"], s["sigma_mkt"],
+        s["pt_sigma"], s["pt_mu"],
+        s["pt_sigma"], s["mu_eff_sigma"],
+        s["sigma_eff_mu"], s["pt_mu"],
+        targ_sigma=sigma_sugg, targ_mu=mu_sugg
+    )
+
+
+# -------------- synthetic dataset (1,000 rows over *current* universe) --------------
+def build_synthetic_dataset(universe: List[str], years: int, rf_ann: float, erp_ann: float) -> pd.DataFrame:
+    # Always include market column used in cov/beta (if present)
+    symbols = list(sorted(set([s for s in universe if s] )))[:MAX_TICKERS]
+    moms = estimate_all_moments_aligned(symbols, years, rf_ann)
+    covA, betas = moms["cov_ann"], moms["betas"]
+
+    rows, rng = [], np.random.default_rng(123)
+    n = 1000
+    for i in range(n):
+        k = rng.integers(low=min(2, len(symbols)), high=min(8, len(symbols)) + 1)
+        picks = list(rng.choice(symbols, size=k, replace=False))
+        signs = rng.choice([-1.0, 1.0], size=k, p=[0.20, 0.80])
+        raw = rng.dirichlet(np.ones(k))
+        gross = 1.0 + float(rng.gamma(2.0, 0.5))
+        w = gross * signs * raw
+        # compute stats from CAPM + cov
+        beta_p, er_p, sigma_p = portfolio_stats({picks[j]: w[j] for j in range(k)}, covA, betas, rf_ann, erp_ann)
         rows.append({
             "id": i,
-            "tickers": ",".join(symbols),
+            "tickers": ",".join(picks),
             "weights": ",".join(f"{x:.6f}" for x in w),
-            "beta_p": beta_p,
             "er_p": er_p,
-            "sigma_p": sigma_p
+            "sigma_p": sigma_p,
+            "beta_p": beta_p
         })
+
     return pd.DataFrame(rows)
 
-def save_dataset_csv(df: pd.DataFrame, path: str = DATASET_PATH):
+
+def save_synth_csv(df: pd.DataFrame, path: str = DATASET_PATH):
     os.makedirs(os.path.dirname(path), exist_ok=True)
     df.to_csv(path, index=False)
 
-def _row_to_exposures(row: pd.Series, universe: List[str]) -> Optional[np.ndarray]:
-    try:
-        ts = [t.strip().upper() for t in str(row["tickers"]).split(",")]
-        ws = [float(x) for x in str(row["weights"]).split(",")]
-        wmap = {t: ws[i] for i, t in enumerate(ts) if i < len(ws)}
-        x = np.array([wmap.get(t, 0.0) for t in universe], dtype=float)
-        g = float(np.sum(np.abs(x)))
+
+# ---------------- Embeddings helpers ----------------
+def _get_emb_model():
+    global _EMB_MODEL
+    if _EMB_MODEL is None:
+        _EMB_MODEL = SentenceTransformer(EMB_MODEL_NAME)
+    return _EMB_MODEL
+
+
+def _weights_top_phrase(universe, w, top=4):
+    pairs = sorted([(universe[i], abs(float(w[i]))) for i in range(len(universe))],
+                   key=lambda t: -t[1])[:top]
+    parts = [f"{t} {p*100:.1f}%" for t, p in pairs if p > 1e-4]
+    return ", ".join(parts)
+
+
+def portfolio_to_sentence(universe, w, er, sigma, beta):
+    return (f"portfolio with volatility {sigma*100:.2f} percent, "
+            f"expected return {er*100:.2f} percent, beta {beta:.2f}, "
+            f"weights mostly in {_weights_top_phrase(universe, w)}")
+
+
+def build_ds_embeddings(csv_path: str, universe: list):
+    global _DS_TEXTS, _DS_EMBS, _DS_CACHE_KEY
+    cache_key = (csv_path, tuple(universe))
+    if _DS_EMBS is not None and _DS_CACHE_KEY == cache_key:
+        return _DS_TEXTS, _DS_EMBS
+
+    df = pd.read_csv(csv_path)
+    texts = []
+    rows = []
+    for _, r in df.iterrows():
+        ws = np.array([float(x) for x in str(r["weights"]).split(",")], dtype=float)
+        ts = [t.strip().upper() for t in str(r["tickers"]).split(",")]
+        wmap = {ts[i]: ws[i] for i in range(min(len(ts), len(ws)))}
+        w = np.array([wmap.get(t, 0.0) for t in universe], dtype=float)
+        g = np.sum(np.abs(w))
         if g <= 1e-12:
-            return None
-        return x / g
-    except Exception:
-        return None
+            continue
+        w = w / g
+        er = float(r["er_p"]); sigma = float(r["sigma_p"]); beta = float(r["beta_p"])
+        txt = portfolio_to_sentence(universe, w, er, sigma, beta)
+        texts.append(txt); rows.append((w, er, sigma, beta))
+
+    model = _get_emb_model()
+    embs = model.encode(texts, normalize_embeddings=True, show_progress_bar=False)
+    _DS_TEXTS, _DS_EMBS, _DS_CACHE_KEY = (rows, embs, cache_key)
+    return _DS_TEXTS, _DS_EMBS
+
 
 def pick_low_med_high(csv_path: str, universe: List[str]):
     df = pd.read_csv(csv_path)
     rows = []
     for _, r in df.iterrows():
-        x = _row_to_exposures(r, universe)
-        if x is None:
+        ws = [float(x) for x in str(r["weights"]).split(",")]
+        ts = [t.strip().upper() for t in str(r["tickers"]).split(",")]
+        wmap = {ts[i]: ws[i] for i in range(min(len(ts), len(ws)))}
+        x = np.array([wmap.get(t, 0.0) for t in universe], dtype=float)
+        g = float(np.sum(np.abs(x)))
+        if g <= 1e-12:
             continue
+        x = x / g
         rows.append((x, float(r["er_p"]), float(r["sigma_p"]), float(r["beta_p"])))
     if not rows:
         return None
     rows_sorted = sorted(rows, key=lambda t: t[2])  # by sigma
-    lo = rows_sorted[0]
-    hi = rows_sorted[-1]
-    med = rows_sorted[len(rows_sorted)//2]
-    return {"low": lo, "medium": med, "high": hi}
+    return rows_sorted
 
-# -------------- summary builder --------------
-def fmt_pct(x: float) -> str:
-    return f"{x*100:.2f}%"
 
+def _band_indices(n, level):
+    if level == "low":
+        return range(0, max(1, int(0.25 * n)))
+    if level == "medium":
+        a, b = int(0.375 * n), int(0.625 * n)
+        return range(max(0, a), min(n, b))
+    return range(max(0, int(0.75 * n)), n)  # high
+
+
+def suggest_level(level: str):
+    if not LAST_DATASET_PATH or not os.path.exists(LAST_DATASET_PATH) or not LAST_UNIVERSE:
+        return empty_suggest_df(), "Run Compute first.", None
+    rows_texts, embs = build_ds_embeddings(LAST_DATASET_PATH, LAST_UNIVERSE)
+    if not rows_texts:
+        return empty_suggest_df(), "No dataset rows.", None
+
+    n = len(rows_texts)
+    band = list(_band_indices(n, level))
+    if not band:
+        return empty_suggest_df(), "No rows in band.", None
+
+    prompts = {
+        "low": "conservative low-risk portfolio with low volatility and low beta",
+        "medium": "balanced moderate-risk portfolio with moderate volatility and beta around 1",
+        "high": "aggressive high-risk growth portfolio with high volatility and beta above 1",
+    }
+    q = prompts.get(level, "balanced portfolio")
+    model = _get_emb_model()
+    q_emb = model.encode([q], normalize_embeddings=True)
+    band_embs = embs[band]
+    sims = util.cos_sim(q_emb, band_embs).cpu().numpy()[0]
+    best_idx_in_band = int(np.argmax(sims))
+    x, er_p, sig_p, beta_p = rows_texts[band[best_idx_in_band]]
+
+    rows_df = [{"ticker": LAST_UNIVERSE[i], "suggested_weight_pct": float(x[i]) * 100.0}
+               for i in range(len(LAST_UNIVERSE))]
+    df = pd.DataFrame(rows_df, columns=SUG_COLS).sort_values("suggested_weight_pct", ascending=False)
+    msg = f"{level.capitalize()} risk (embedding-ranked) → ER {fmt_pct(er_p)}, Sigma {fmt_pct(sig_p)}, Beta {beta_p:.2f}"
+    img = _overlay_plot_with_suggestion(sig_p, er_p)
+    return df, msg, img
+
+
+# -------------- summary builder --------------
 def build_summary_md(lookback, horizon, rf, rf_code, erp, sigma_mkt,
                      beta_p, er_p, sigma_p,
                      a_sigma, b_sigma, mu_eff_sigma,
-                     a_mu, b_mu, sigma_eff_mu,
-                     ds_info: str) -> str:
+                     a_mu, b_mu, sigma_eff_mu) -> str:
     lines = []
     lines.append("### Inputs")
-    lines.append(f"- Lookback years: {lookback}")
-    lines.append(f"- Horizon years: {int(round(horizon))}")
-    lines.append(f"- Risk free: {fmt_pct(rf)} from {rf_code}")
-    lines.append(f"- Market ERP: {fmt_pct(erp)}")
-    lines.append(f"- Market sigma: {fmt_pct(sigma_mkt)}")
+    lines.append(f"- Lookback years **{lookback}**")
+    lines.append(f"- Horizon years **{int(round(horizon))}**")
+    lines.append(f"- Risk free **{fmt_pct(rf)}** from **{rf_code}**")
+    lines.append(f"- Market ERP **{fmt_pct(erp)}**")
+    lines.append(f"- Market σ **{fmt_pct(sigma_mkt)}**")
     lines.append("")
-    lines.append("### Your portfolio")
-    lines.append(f"- Beta: {beta_p:.2f}")
-    lines.append(f"- Sigma: {fmt_pct(sigma_p)}")
-    lines.append(f"- Expected return: {fmt_pct(er_p)}")
+    lines.append("### Your portfolio (CAPM expectations)")
+    lines.append(f"- Beta **{beta_p:.2f}**")
+    lines.append(f"- σ **{fmt_pct(sigma_p)}**")
+    lines.append(f"- Expected return **{fmt_pct(er_p)}**")
     lines.append("")
     lines.append("### Efficient alternatives on CML")
-    lines.append(f"- Same sigma ⇒ Market {a_sigma:.2f} , Bills {b_sigma:.2f} , ER {fmt_pct(mu_eff_sigma)}")
-    lines.append(f"- Same return ⇒ Market {a_mu:.2f} , Bills {b_mu:.2f} , Sigma {fmt_pct(sigma_eff_mu)}")
+    lines.append("**Same σ as your portfolio**")
+    lines.append(f"- Market weight **{a_sigma:.2f}**, Bills weight **{b_sigma:.2f}**")
+    lines.append(f"- Expected return **{fmt_pct(mu_eff_sigma)}**")
     lines.append("")
-    lines.append("### Dataset for risk suggestions")
-    lines.append(ds_info)
+    lines.append("**Same expected return as your portfolio**")
+    lines.append(f"- Market weight **{a_mu:.2f}**, Bills weight **{b_mu:.2f}**")
+    lines.append(f"- σ **{fmt_pct(sigma_eff_mu)}**")
     return "\n".join(lines)
 
-# -------------- globals to carry session state --------------
-LAST_MOMS = None
-LAST_BASE = None
-LAST_UNIVERSE = []
-LAST_DATASET_PATH = None
-HORIZON_YEARS = 5.0
-RF_CODE = fred_series_for_horizon(HORIZON_YEARS)
-RF_ANN = fetch_fred_yield_annual(RF_CODE)
 
 # -------------- gradio callbacks --------------
 def search_tickers_cb(q: str):
@@ -374,6 +548,7 @@ def search_tickers_cb(q: str):
     opts = [f"{h['symbol']}  |  {h['name']}  |  {h['exchange']}" for h in hits]
     return "Select a symbol and click Add", opts
 
+
 def add_symbol(selection: str, table: pd.DataFrame):
     if not selection:
         return table, "Pick a row from Matches first"
@@ -395,6 +570,7 @@ def add_symbol(selection: str, table: pd.DataFrame):
         msg = f"Reached max of {MAX_TICKERS}"
     return new_table, msg
 
+
 def lock_ticker_column(tb: pd.DataFrame):
     if tb is None or len(tb) == 0:
         return pd.DataFrame(columns=["ticker", "amount_usd"])
@@ -405,6 +581,7 @@ def lock_ticker_column(tb: pd.DataFrame):
     amounts = amounts[:len(tickers)] + [0.0] * max(0, len(tickers) - len(amounts))
     return pd.DataFrame({"ticker": tickers, "amount_usd": amounts})
 
+
 def set_horizon(years: float):
     y = max(1.0, min(100.0, float(years)))
     code = fred_series_for_horizon(y)
@@ -413,67 +590,63 @@ def set_horizon(years: float):
     HORIZON_YEARS = y
     RF_CODE = code
     RF_ANN = rf
-    return f"Risk free {fmt_pct(rf)} from {code}. Will be used on Compute."
+    return f"Risk free series {code}. Latest annual rate {rf:.2%}. Will be used for CAPM and CML."
+
 
 def compute(years_lookback: int, table: pd.DataFrame):
+    if table is None or len(table) == 0:
+        return None, "Add at least one ticker", "Universe empty", empty_positions_df(), empty_suggest_df(), None
+
     df = table.dropna()
     df["ticker"] = df["ticker"].astype(str).str.upper().str.strip()
     df["amount_usd"] = pd.to_numeric(df["amount_usd"], errors="coerce").fillna(0.0)
 
     symbols = [t for t in df["ticker"].tolist() if t]
     if len(symbols) == 0:
-        return None, "Add at least one ticker", "Universe empty", empty_positions_df(), empty_suggest_df(), "", None
+        return None, "Add at least one ticker", "Universe empty", empty_positions_df(), empty_suggest_df(), None
 
     symbols = validate_tickers(symbols, years_lookback)
     if len(symbols) == 0:
-        return None, "Could not validate any tickers", "Universe invalid", empty_positions_df(), empty_suggest_df(), "", None
+        return None, "Could not validate any tickers", "Universe invalid", empty_positions_df(), empty_suggest_df(), None
+
+    global UNIVERSE
+    UNIVERSE = list(sorted(set(symbols)))[:MAX_TICKERS]
 
     df = df[df["ticker"].isin(symbols)].copy()
     amounts = {r["ticker"]: float(r["amount_usd"]) for _, r in df.iterrows()}
-    allow_shorts = any(v < 0 for v in amounts.values())
     rf_ann = RF_ANN
 
-    # moments
-    moms = estimate_all_moments_aligned(symbols, years_lookback, rf_ann)
+    moms = estimate_all_moments_aligned(UNIVERSE, years_lookback, rf_ann)
     betas, covA, erp_ann, sigma_mkt = moms["betas"], moms["cov_ann"], moms["erp_ann"], moms["sigma_m_ann"]
 
     gross = sum(abs(v) for v in amounts.values())
     if gross == 0:
-        return None, "All amounts are zero", "Universe ok", empty_positions_df(), empty_suggest_df(), "", None
+        return None, "All amounts are zero", "Universe ok", empty_positions_df(), empty_suggest_df(), None
     weights = {k: v / gross for k, v in amounts.items()}
+
     beta_p, er_p, sigma_p = portfolio_stats(weights, covA, betas, rf_ann, erp_ann)
 
     a_sigma, b_sigma, mu_eff_sigma = efficient_same_sigma(sigma_p, rf_ann, erp_ann, sigma_mkt)
     a_mu, b_mu, sigma_eff_mu = efficient_same_return(er_p, rf_ann, erp_ann, sigma_mkt)
 
-    # dataset strictly over *these* symbols
-    ensure_data_dir()
-    ds = build_fixed_universe_dataset(
-        symbols=symbols, years=years_lookback, rf_ann=rf_ann, erp_ann=erp_ann,
-        covA=covA.loc[symbols, symbols], betas=betas, allow_shorts=allow_shorts, n_rows=1000
-    )
-    save_dataset_csv(ds, DATASET_PATH)
-    ds_info = f"- Built {len(ds)} simulated mixes over current tickers ({'shorts allowed' if allow_shorts else 'long-only'})."
-
-    # plot + summary
-    img = plot_cml_percent(
+    img = plot_cml(
         rf_ann, erp_ann, sigma_mkt,
         sigma_p, er_p,
         sigma_p, mu_eff_sigma,
         sigma_eff_mu, er_p,
         targ_sigma=None, targ_mu=None
     )
+
     info = build_summary_md(
         years_lookback, HORIZON_YEARS, rf_ann, RF_CODE, erp_ann, sigma_mkt,
         beta_p, er_p, sigma_p,
         a_sigma, b_sigma, mu_eff_sigma,
-        a_mu, b_mu, sigma_eff_mu,
-        ds_info=ds_info
+        a_mu, b_mu, sigma_eff_mu
     )
 
     rows = []
-    for t in symbols:
-        beta_val = 1.0 if t == MARKET_TICKER else betas.get(t, np.nan)
+    for t in UNIVERSE:
+        beta_val = 1.0 if abs(betas.get(t, 0.0) - 1.0) < 1e-6 else betas.get(t, np.nan)
         rows.append({
             "ticker": t,
             "amount_usd": amounts.get(t, 0.0),
@@ -481,78 +654,38 @@ def compute(years_lookback: int, table: pd.DataFrame):
             "beta": beta_val,
         })
     pos_table = pd.DataFrame(rows, columns=POS_COLS)
-    pos_table["weight_exposure"] = pos_table["weight_exposure"].astype(float)
-
-    uni_msg = f"Universe set to {', '.join(symbols)}"
-    # store globals for Suggest buttons
-    global LAST_MOMS, LAST_BASE, LAST_UNIVERSE, LAST_DATASET_PATH
-    LAST_MOMS = {"betas": betas, "covA": covA, "erp_ann": erp_ann, "sigma_mkt": sigma_mkt}
-    LAST_BASE = {"rf_ann": rf_ann, "er_p": er_p, "sigma_p": sigma_p}
-    LAST_UNIVERSE = list(symbols)
-    LAST_DATASET_PATH = DATASET_PATH
 
-    return img, info, uni_msg, pos_table, empty_suggest_df(), ds_info, DATASET_PATH
+    # build 1,000-row dataset over CURRENT universe
+    synth_df = build_synthetic_dataset(UNIVERSE, years=DEFAULT_LOOKBACK_YEARS, rf_ann=rf_ann, erp_ann=erp_ann)
+    save_synth_csv(synth_df, DATASET_PATH)
 
-def _overlay_plot_with_suggestion(sigma_s, er_s):
-    if not LAST_MOMS or not LAST_BASE:
-        return None
-    rf_ann = LAST_BASE["rf_ann"]
-    erp_ann = LAST_MOMS["erp_ann"]
-    sigma_mkt = LAST_MOMS["sigma_mkt"]
-    sigma_p = LAST_BASE["sigma_p"]
-    er_p = LAST_BASE["er_p"]
-    a_sigma, b_sigma, mu_eff_sigma = efficient_same_sigma(sigma_p, rf_ann, erp_ann, sigma_mkt)
-    a_mu, b_mu, sigma_eff_mu = efficient_same_return(er_p, rf_ann, erp_ann, sigma_mkt)
-    return plot_cml_percent(
-        rf_ann, erp_ann, sigma_mkt,
-        sigma_p, er_p,
-        sigma_p, mu_eff_sigma,
-        sigma_eff_mu, er_p,
-        targ_sigma=sigma_s, targ_mu=er_s
-    )
+    # update globals for suggestion buttons
+    global LAST_DATASET_PATH, LAST_UNIVERSE, LAST_PLOT_STATE
+    LAST_DATASET_PATH = DATASET_PATH
+    LAST_UNIVERSE = UNIVERSE.copy()
+    LAST_PLOT_STATE = {
+        "rf_ann": rf_ann, "erp_ann": erp_ann, "sigma_mkt": sigma_mkt,
+        "pt_sigma": sigma_p, "pt_mu": er_p,
+        "mu_eff_sigma": mu_eff_sigma, "sigma_eff_mu": sigma_eff_mu
+    }
 
-def suggest_level(level: str):
-    if not LAST_DATASET_PATH or not os.path.exists(LAST_DATASET_PATH) or not LAST_UNIVERSE:
-        return empty_suggest_df(), "Run Compute first.", None
-    picks = pick_low_med_high(LAST_DATASET_PATH, LAST_UNIVERSE)
-    if picks is None or level not in picks:
-        return empty_suggest_df(), "No suggestion available.", None
-    x, er_p, sig_p, beta_p = picks[level]
-    # build table in percent
-    rows = [{"ticker": LAST_UNIVERSE[i], "suggested_weight_pct": float(x[i]) * 100.0} for i in range(len(LAST_UNIVERSE))]
-    df = pd.DataFrame(rows, columns=SUG_COLS)
-    msg = f"{level.capitalize()} risk → ER {fmt_pct(er_p)}, Sigma {fmt_pct(sig_p)}, Beta {beta_p:.2f}"
-    img = _overlay_plot_with_suggestion(sig_p, er_p)
-    return df, msg, img
+    uni_msg = f"Universe set to: {', '.join(UNIVERSE)} — dataset generated with 1,000 mixes."
+    return img, info, uni_msg, pos_table, empty_suggest_df(), DATASET_PATH
 
-def apply_suggestion_to_amounts(level: str, table: pd.DataFrame):
-    if table is None or len(table) == 0:
-        return table
-    df_sug, _, _ = suggest_level(level)
-    if df_sug is None or len(df_sug) == 0:
-        return table
-    # compute gross dollars (use total |amounts|; if zero, default to 10,000)
-    t = table.copy()
-    t["ticker"] = t["ticker"].astype(str).str.upper().str.strip()
-    t["amount_usd"] = pd.to_numeric(t["amount_usd"], errors="coerce").fillna(0.0)
-    gross = float(np.sum(np.abs(t["amount_usd"].values)))
-    if gross <= 1e-9:
-        gross = 10000.0
-    w = {r["ticker"]: float(r["suggested_weight_pct"]) / 100.0 for _, r in df_sug.iterrows()}
-    # map to amounts using current ticker order; missing → 0
-    new_amounts = [gross * w.get(sym.upper(), 0.0) for sym in t["ticker"].tolist()]
-    t["amount_usd"] = new_amounts
-    return t
 
 # -------------- UI --------------
 ensure_data_dir()
+# initial RF based on default horizon
+HORIZON_YEARS = 5.0
+RF_CODE = fred_series_for_horizon(HORIZON_YEARS)
+RF_ANN = fetch_fred_yield_annual(RF_CODE)
 
 with gr.Blocks(title="Efficient Portfolio Advisor") as demo:
     gr.Markdown(
         "## Efficient Portfolio Advisor\n"
         "Search symbols, enter dollar amounts, set your horizon. "
-        "Prices: Yahoo Finance. Risk free: FRED. "
-        "Suggestions (Low/Medium/High) come **only** from the 1,000-portfolio dataset built over your tickers."
+        "Prices from Yahoo Finance. Risk-free from FRED. "
+        "Low/Medium/High suggestions use embeddings over a 1,000-mix dataset generated from your current universe."
     )
 
     with gr.Row():
@@ -560,11 +693,10 @@ with gr.Blocks(title="Efficient Portfolio Advisor") as demo:
             q = gr.Textbox(label="Search symbol")
             search_note = gr.Markdown()
             matches = gr.Dropdown(choices=[], label="Matches")
-            with gr.Row():
-                search_btn = gr.Button("Search")
-                add_btn = gr.Button("Add selected to portfolio")
+            search_btn = gr.Button("Search")
+            add_btn = gr.Button("Add selected to portfolio")
 
-            gr.Markdown("### Portfolio positions (type dollar amounts, negatives allowed for shorts)")
+            gr.Markdown("### Portfolio positions — type dollar amounts (negatives allowed for shorts)")
             table = gr.Dataframe(
                 headers=["ticker", "amount_usd"],
                 datatype=["str", "number"],
@@ -573,26 +705,20 @@ with gr.Blocks(title="Efficient Portfolio Advisor") as demo:
             )
 
             horizon = gr.Number(label="Horizon in years (1–100)", value=5, precision=0)
-            lookback = gr.Slider(1, 10, value=DEFAULT_LOOKBACK_YEARS, step=1, label="Lookback years for beta & sigma")
+            lookback = gr.Slider(1, 10, value=DEFAULT_LOOKBACK_YEARS, step=1, label="Lookback years for beta and sigma")
 
-            with gr.Row():
-                run_btn = gr.Button("Compute (build dataset)", variant="primary")
+            run_btn = gr.Button("Compute (build dataset & plot)")
 
-            gr.Markdown("### Risk tolerance suggestions (dataset-based only)")
+            gr.Markdown("### Suggestions (dataset + embeddings)")
             with gr.Row():
-                btn_low = gr.Button("Low risk")
-                btn_med = gr.Button("Medium risk")
-                btn_high = gr.Button("High risk")
-            with gr.Row():
-                apply_low = gr.Button("Apply Low → $")
-                apply_med = gr.Button("Apply Medium → $")
-                apply_high = gr.Button("Apply High → $")
+                btn_low = gr.Button("Suggest LOW risk")
+                btn_med = gr.Button("Suggest MEDIUM risk")
+                btn_high = gr.Button("Suggest HIGH risk")
 
         with gr.Column(scale=1):
-            plot = gr.Image(label="Capital Market Line", type="pil")
+            plot = gr.Image(label="Capital Market Line (CML)", type="pil")
             summary = gr.Markdown(label="Summary")
-            universe_msg = gr.Textbox(label="Universe status", interactive=False)
-            dataset_info = gr.Markdown(label="Dataset info", value="")
+            universe_msg = gr.Textbox(label="Status", interactive=False)
             positions = gr.Dataframe(
                 label="Computed positions",
                 headers=POS_COLS,
@@ -602,15 +728,15 @@ with gr.Blocks(title="Efficient Portfolio Advisor") as demo:
                 interactive=False
             )
             suggestions = gr.Dataframe(
-                label="Suggested weights (percent of exposure)",
+                label="Suggested portfolio (weights as % exposures)",
                 headers=SUG_COLS,
                 datatype=["str", "number"],
                 col_count=(len(SUG_COLS), "fixed"),
                 value=empty_suggest_df(),
                 interactive=False
             )
-            sugg_msg = gr.Markdown("")
-            dl = gr.File(label="Session dataset CSV", value=None, visible=True)
+            sugg_msg = gr.Textbox(label="Suggestion detail", interactive=False)
+            dl = gr.File(label="Generated dataset (CSV)", value=None, visible=True)
 
     # wiring
     def do_search(query):
@@ -625,23 +751,24 @@ with gr.Blocks(title="Efficient Portfolio Advisor") as demo:
     run_btn.click(
         fn=compute,
         inputs=[lookback, table],
-        outputs=[plot, summary, universe_msg, positions, suggestions, dataset_info, dl]
+        outputs=[plot, summary, universe_msg, positions, suggestions, dl]
     )
 
-    # suggest buttons
-    def wrap_suggest(level):
-        df, msg, img = suggest_level(level)
-        img_out = img if img is not None else gr.update()
-        return df, msg, img_out
+    def do_low():
+        df, msg, img = suggest_level("low")
+        return df, msg, (img if img is not None else gr.update())
+
+    def do_med():
+        df, msg, img = suggest_level("medium")
+        return df, msg, (img if img is not None else gr.update())
 
-    btn_low.click(lambda: wrap_suggest("low"), outputs=[suggestions, sugg_msg, plot])
-    btn_med.click(lambda: wrap_suggest("medium"), outputs=[suggestions, sugg_msg, plot])
-    btn_high.click(lambda: wrap_suggest("high"), outputs=[suggestions, sugg_msg, plot])
+    def do_high():
+        df, msg, img = suggest_level("high")
+        return df, msg, (img if img is not None else gr.update())
 
-    # apply buttons (only updates the table; user can hit Compute again)
-    apply_low.click(lambda tb: apply_suggestion_to_amounts("low", tb), inputs=table, outputs=table)
-    apply_med.click(lambda tb: apply_suggestion_to_amounts("medium", tb), inputs=table, outputs=table)
-    apply_high.click(lambda tb: apply_suggestion_to_amounts("high", tb), inputs=table, outputs=table)
+    btn_low.click(fn=do_low, inputs=None, outputs=[suggestions, sugg_msg, plot])
+    btn_med.click(fn=do_med, inputs=None, outputs=[suggestions, sugg_msg, plot])
+    btn_high.click(fn=do_high, inputs=None, outputs=[suggestions, sugg_msg, plot])
 
 if __name__ == "__main__":
     demo.launch()