modules/performance.py

"""
Phase 3 — Performance Engine
Reads all history JSON files, calculates BUY signal returns vs. Nifty 50,
and computes a composite AI Score.
"""

import os
import json
import re
import pandas as pd
from datetime import datetime

try:
    from modules.price_fetcher import fetch_current_prices, fetch_nifty50_current
except ImportError:
    from price_fetcher import fetch_current_prices, fetch_nifty50_current


# ─────────────────────────────────────────────────────────────────────────────
# Helpers
# ─────────────────────────────────────────────────────────────────────────────

def _parse_nifty_from_market_summary(market_summary: list) -> float | None:
    """
    Extract a numeric Nifty 50 value from the market_summary list.
    Example entry: {"label": "Nifty 50", "value": "23,050.20 (+0.25%)"}
    """
    for item in market_summary:
        label = item.get("label", "").lower()
        if "nifty 50" in label or "nifty50" in label:
            raw = item.get("value", "")
            # Strip commas, extract first numeric value
            match = re.search(r"[\d,]+\.?\d*", raw.replace(",", ""))
            if match:
                try:
                    return float(match.group())
                except ValueError:
                    pass
    return None


def _load_signal_prices_from_snapshot(data: dict) -> dict:
    """
    Build _signal_prices by cross-referencing portfolio signals with _portfolio_snapshot LTPs.
    Used for backfilling old history files that lack _signal_prices.
    Returns dict: { symbol: { signal, price_on_day, qty, nifty_on_day } }
    """
    snapshot = data.get("_portfolio_snapshot", [])
    portfolio_signals = data.get("portfolio", [])
    market_summary = data.get("market_summary", [])

    ltp_map = {}
    qty_map = {}
    for item in snapshot:
        sym = item.get("symbol")
        if sym:
            ltp_map[sym] = float(item.get("ltp") or 0)
            qty_map[sym] = float(item.get("qty") or 0)

    nifty_on_day = _parse_nifty_from_market_summary(market_summary)

    result = {}
    for entry in portfolio_signals:
        sym = entry.get("symbol")
        sig = entry.get("signal", "")
        if sym and sig == "BUY" and sym in ltp_map and ltp_map[sym] > 0:
            result[sym] = {
                "signal": sig,
                "price_on_day": ltp_map[sym],
                "qty": qty_map.get(sym, 1.0),
                "nifty_on_day": nifty_on_day,
            }
    return result


# ─────────────────────────────────────────────────────────────────────────────
# Core Loader
# ─────────────────────────────────────────────────────────────────────────────

def load_all_buy_signals(history_dir: str = "history") -> list[dict]:
    """
    Read every history JSON or database record and collect all BUY signals with entry prices.

    Returns list of dicts:
        {
          date, symbol, signal, price_on_day, qty, nifty_on_day
        }
    """
    # Try DB load first
    try:
        from modules.db import load_all_buy_signals_from_db
        db_rows = load_all_buy_signals_from_db()
        if db_rows:
            return db_rows
    except Exception as e:
        print(f"[performance] DB load failed, falling back to files: {e}")

    if not os.path.exists(history_dir):
        return []

    files = sorted([f for f in os.listdir(history_dir) if f.endswith(".json")])
    rows = []

    for filename in files:
        date_str = filename.replace(".json", "")
        path = os.path.join(history_dir, filename)
        try:
            with open(path, "r", encoding="utf-8") as fp:
                data = json.load(fp)
        except Exception as e:
            print(f"[performance] Could not load {filename}: {e}")
            continue

        # Use stored _signal_prices if present, else derive from snapshot
        signal_prices = data.get("_signal_prices")
        if not signal_prices:
            signal_prices = _load_signal_prices_from_snapshot(data)

        for sym, sp in signal_prices.items():
            if sp.get("signal") == "BUY" and sp.get("price_on_day", 0) > 0:
                rows.append({
                    "date": date_str,
                    "symbol": sym,
                    "signal": "BUY",
                    "price_on_day": float(sp["price_on_day"]),
                    "qty": float(sp.get("qty", 1.0)),
                    "nifty_on_day": sp.get("nifty_on_day"),   # may be None
                })

    return rows


# ─────────────────────────────────────────────────────────────────────────────
# Main Computation
# ─────────────────────────────────────────────────────────────────────────────

def compute_performance(history_dir: str = "history") -> tuple[pd.DataFrame, dict]:
    """
    Full performance computation.

    Returns:
        (signals_df, summary_dict)

        signals_df columns:
            date, symbol, signal, price_on_day, qty, nifty_on_day,
            current_price, current_nifty, stock_return_pct, nifty_return_pct,
            alpha_pct, beat_market, invested_value, current_value

        summary_dict keys:
            ai_score, buy_win_rate, avg_alpha, total_signals,
            total_invested, total_current_value, total_gain_loss,
            portfolio_return_pct, nifty_return_pct
    """
    rows = load_all_buy_signals(history_dir)
    if not rows:
        return pd.DataFrame(), {}

    df = pd.DataFrame(rows)

    # ── Fetch current prices ──────────────────────────────────────────────────
    all_symbols = df["symbol"].unique().tolist()
    current_prices = fetch_current_prices(all_symbols)
    current_nifty  = fetch_nifty50_current()

    df["current_price"] = df["symbol"].map(current_prices)
    df["current_nifty"] = current_nifty

    # ── Drop rows where we couldn't get live prices ───────────────────────────
    df = df.dropna(subset=["current_price"])
    df = df[df["current_price"] > 0]

    if df.empty:
        return pd.DataFrame(), {}

    # ── Return calculations ───────────────────────────────────────────────────
    df["stock_return_pct"] = ((df["current_price"] - df["price_on_day"]) / df["price_on_day"]) * 100

    # Nifty return: use stored nifty_on_day if available, else default to None
    df["nifty_return_pct"] = df.apply(
        lambda r: ((current_nifty - r["nifty_on_day"]) / r["nifty_on_day"]) * 100
        if (r["nifty_on_day"] and current_nifty and r["nifty_on_day"] > 0) else None,
        axis=1
    )

    df["alpha_pct"] = df.apply(
        lambda r: r["stock_return_pct"] - r["nifty_return_pct"]
        if r["nifty_return_pct"] is not None else None,
        axis=1
    )

    df["beat_market"] = df.apply(
        lambda r: r["alpha_pct"] > 0 if r["alpha_pct"] is not None else None,
        axis=1
    )

    # ── Portfolio value tracking (qty-weighted) ───────────────────────────────
    df["invested_value"] = df["qty"] * df["price_on_day"]
    df["current_value"]  = df["qty"] * df["current_price"]

    # ── Summary metrics ───────────────────────────────────────────────────────
    valid_alpha = df.dropna(subset=["alpha_pct"])
    beats = valid_alpha[valid_alpha["beat_market"] == True]

    buy_win_rate = (len(beats) / len(valid_alpha) * 100) if len(valid_alpha) > 0 else 0
    avg_alpha    = float(valid_alpha["alpha_pct"].mean()) if len(valid_alpha) > 0 else 0

    total_invested     = float(df["invested_value"].sum())
    total_current_val  = float(df["current_value"].sum())
    total_gain_loss    = total_current_val - total_invested
    portfolio_return   = ((total_current_val - total_invested) / total_invested * 100) if total_invested > 0 else 0

    # Nifty return — use average nifty_on_day across signals weighted by invested value
    weighted_nifty_rows = df.dropna(subset=["nifty_on_day"])
    if current_nifty and not weighted_nifty_rows.empty:
        total_w = weighted_nifty_rows["invested_value"].sum()
        weighted_nifty_base = (
            (weighted_nifty_rows["nifty_on_day"] * weighted_nifty_rows["invested_value"]).sum() / total_w
            if total_w > 0 else weighted_nifty_rows["nifty_on_day"].mean()
        )
        nifty_return = ((current_nifty - weighted_nifty_base) / weighted_nifty_base) * 100
    else:
        nifty_return = None

    # AI Score: buy win rate (60%) + avg alpha clamped (40%)
    alpha_score = min(max((avg_alpha + 5) / 10 * 100, 0), 100)   # +5% alpha → 100, -5% → 0
    ai_score    = (buy_win_rate * 0.6) + (alpha_score * 0.4)
    ai_score    = round(min(max(ai_score, 0), 100), 1)

    summary = {
        "ai_score":             ai_score,
        "buy_win_rate":         round(buy_win_rate, 1),
        "avg_alpha":            round(avg_alpha, 2),
        "total_signals":        int(len(df)),
        "signals_beating":      int(len(beats)),
        "total_invested":       round(total_invested, 2),
        "total_current_value":  round(total_current_val, 2),
        "total_gain_loss":      round(total_gain_loss, 2),
        "portfolio_return_pct": round(portfolio_return, 2),
        "nifty_return_pct":     round(float(nifty_return), 2) if nifty_return is not None else None,
        "current_nifty":        round(float(current_nifty), 2) if current_nifty else None,
        "last_updated":         datetime.now().isoformat(),
    }

    return df, summary


# ─────────────────────────────────────────────────────────────────────────────
# Cumulative Chart Data
# ─────────────────────────────────────────────────────────────────────────────

def build_cumulative_chart_data(df: pd.DataFrame, summary: dict) -> pd.DataFrame:
    """
    Build a DataFrame for the dual-line cumulative return chart.

    Each row = one unique signal date, with:
        date, portiq_cumulative_pct, nifty_cumulative_pct

    Strategy: equal-time-point comparison — for each date, compute the
    average stock return of BUYs issued on or before that date.
    """
    if df.empty:
        return pd.DataFrame()

    dates = sorted(df["date"].unique())
    chart_rows = []

    for d in dates:
        # All BUY signals issued on or before this date
        subset = df[df["date"] <= d]
        if subset.empty:
            continue

        avg_stock_ret = float(subset["stock_return_pct"].mean())
        nifty_vals    = subset.dropna(subset=["nifty_return_pct"])
        avg_nifty_ret = float(nifty_vals["nifty_return_pct"].mean()) if not nifty_vals.empty else None

        chart_rows.append({
            "date":                d,
            "portiq_return_pct":   round(avg_stock_ret, 2),
            "nifty_return_pct":    round(avg_nifty_ret, 2) if avg_nifty_ret is not None else None,
        })

    return pd.DataFrame(chart_rows)


# ─────────────────────────────────────────────────────────────────────────────
# Cache persistence
# ─────────────────────────────────────────────────────────────────────────────

def save_performance_cache(df: pd.DataFrame, summary: dict):
    """Persist performance results to PostgreSQL database for fast re-loads."""
    from modules.db import get_db_connection, get_owner_user_id, clean_nans
    from psycopg2.extras import Json
    
    payload = {
        "summary": summary,
        "signals": df.to_dict(orient="records") if not df.empty else [],
    }
    payload = clean_nans(payload)
    owner_id = get_owner_user_id()
    conn = get_db_connection()
    try:
        with conn.cursor() as cur:
            cur.execute(
                """
                INSERT INTO performance_cache (user_id, payload, updated_at)
                VALUES (%s, %s, CURRENT_TIMESTAMP)
                ON CONFLICT (user_id) DO UPDATE
                SET payload = EXCLUDED.payload,
                    updated_at = CURRENT_TIMESTAMP;
                """,
                (owner_id, Json(payload))
            )
            conn.commit()
            print("[performance] Saved performance metrics cache to database.")
    except Exception as e:
        conn.rollback()
        print(f"[performance] Failed to save DB cache: {e}")
    finally:
        conn.close()


def load_performance_cache() -> tuple[pd.DataFrame, dict]:
    """Load cached performance data from PostgreSQL if it exists and is fresh (< 4 hours old)."""
    from modules.db import get_db_connection, get_owner_user_id
    from datetime import datetime, timezone
    
    owner_id = get_owner_user_id()
    conn = get_db_connection()
    try:
        with conn.cursor() as cur:
            cur.execute(
                "SELECT payload, updated_at FROM performance_cache WHERE user_id = %s;",
                (owner_id,)
            )
            row = cur.fetchone()
            if not row:
                return pd.DataFrame(), {}
            
            payload, updated_at = row
            # Calculate cache age from timezone-aware updated_at
            age = (datetime.now(timezone.utc) - updated_at).total_seconds()
            if age > 4 * 3600:  # Stale after 4 hours
                return pd.DataFrame(), {}
                
            summary = payload.get("summary", {})
            signals = payload.get("signals", [])
            df = pd.DataFrame(signals) if signals else pd.DataFrame()
            return df, summary
    except Exception as e:
        print(f"[performance] Database cache load failed: {e}")
        return pd.DataFrame(), {}
    finally:
        conn.close()


def compute_actual_trades_performance(owner_id) -> dict:
    """
    Computes performance benchmarking for all actual trades.
    Compares trade return vs nifty return, gets comparison list of skipped AI suggestions,
    and calculates portfolio metrics.
    """
    from modules.db import load_actual_trades_from_db, get_ai_suggestions_for_date
    
    trades = load_actual_trades_from_db(owner_id)
    if not trades:
        return {
            "trades": [],
            "comparison": [],
            "summary": {
                "total_trades": 0,
                "total_invested": 0.0,
                "total_current_value": 0.0,
                "total_gain_loss": 0.0,
                "portfolio_return_pct": 0.0,
                "nifty_return_pct": 0.0,
                "alpha_pct": 0.0,
                "win_rate": 0.0,
                "ai_suggested_count": 0,
                "you_bought_count": 0,
                "coverage_pct": 0.0,
            }
        }
        
    # Get current prices for all trade symbols
    trade_symbols = list(set(t["symbol"] for t in trades))
    
    # Also find all AI suggestions on trade dates to include in comparison & coverage
    unique_dates = list(set(t["trade_date"] for t in trades))
    
    ai_suggestions_by_date = {}
    all_comparison_symbols = set(trade_symbols)
    
    for d in unique_dates:
        sugs = get_ai_suggestions_for_date(owner_id, d)
        # Keep only BUY signals
        buy_sugs = [s for s in sugs if s["signal"] == "BUY"]
        ai_suggestions_by_date[d] = buy_sugs
        for s in buy_sugs:
            all_comparison_symbols.add(s["symbol"])
            
    # Fetch live prices for all comparison symbols
    all_symbols_list = list(all_comparison_symbols)
    current_prices = fetch_current_prices(all_symbols_list)
    current_nifty = fetch_nifty50_current() or 23000.0
    
    # Calculate returns for actual trades
    processed_trades = []
    total_invested = 0.0
    total_current_value = 0.0
    winning_trades_count = 0
    trades_with_nifty_count = 0
    sum_alpha = 0.0
    
    # Map trade_date and symbol to the trade for easy checking
    trade_map = {(t["trade_date"], t["symbol"]): t for t in trades}
    
    for t in trades:
        sym = t["symbol"]
        qty = t["qty_bought"]
        buy_price = t["buy_price"]
        nifty_entry = t["nifty_on_trade_day"]
        
        curr_price = current_prices.get(sym)
        if curr_price is None or curr_price <= 0:
            curr_price = buy_price
            
        stock_ret = ((curr_price - buy_price) / buy_price) * 100
        
        if nifty_entry and nifty_entry > 0:
            nifty_ret = ((current_nifty - nifty_entry) / nifty_entry) * 100
            alpha = stock_ret - nifty_ret
            beat_market = alpha > 0
            trades_with_nifty_count += 1
            sum_alpha += alpha
        else:
            nifty_ret = None
            alpha = None
            beat_market = None
            
        invested = qty * buy_price
        curr_val = qty * curr_price
        gain_loss = curr_val - invested
        
        total_invested += invested
        total_current_value += curr_val
        if beat_market:
            winning_trades_count += 1
            
        processed_trades.append({
            "trade_date": t["trade_date"],
            "symbol": sym,
            "qty_bought": qty,
            "buy_price": buy_price,
            "current_price": curr_price,
            "nifty_on_trade_day": nifty_entry,
            "stock_return_pct": round(stock_ret, 2),
            "nifty_return_pct": round(nifty_ret, 2) if nifty_ret is not None else None,
            "alpha_pct": round(alpha, 2) if alpha is not None else None,
            "beat_market": beat_market,
            "invested": round(invested, 2),
            "current_value": round(curr_val, 2),
            "gain_loss": round(gain_loss, 2)
        })
        
    # Calculate comparison table
    comparison_list = []
    ai_suggested_count = 0
    you_bought_count = 0
    
    for d in unique_dates:
        buy_sugs = ai_suggestions_by_date.get(d, [])
        for sug in buy_sugs:
            sym = sug["symbol"]
            price_on_day = sug["price_on_day"]
            ai_suggested_count += 1
            
            trade_key = (d, sym)
            bought = trade_key in trade_map
            
            curr_price = current_prices.get(sym) or price_on_day
            
            # Fetch nifty_on_day for the date from the trade (if any trade has it)
            nifty_entry = None
            trades_on_date = [tr for tr in trades if tr["trade_date"] == d]
            if trades_on_date:
                nifty_entry = trades_on_date[0]["nifty_on_trade_day"]
                
            if bought:
                you_bought_count += 1
                tr = trade_map[trade_key]
                qty = tr["qty_bought"]
                buy_price = tr["buy_price"]
                nifty_entry = tr["nifty_on_trade_day"] or nifty_entry
                
                stock_ret = ((curr_price - buy_price) / buy_price) * 100
                invested = qty * buy_price
                curr_val = qty * curr_price
                gain_loss = curr_val - invested
            else:
                qty = 0.0
                buy_price = None
                stock_ret = ((curr_price - price_on_day) / price_on_day) * 100
                invested = 0.0
                curr_val = 0.0
                gain_loss = 0.0
                
            if nifty_entry and nifty_entry > 0:
                nifty_ret = ((current_nifty - nifty_entry) / nifty_entry) * 100
                alpha = stock_ret - nifty_ret
            else:
                nifty_ret = None
                alpha = None
                
            comparison_list.append({
                "date": d,
                "symbol": sym,
                "signal": "BUY",
                "price_on_day": price_on_day,
                "nifty_on_day": nifty_entry,
                "bought": bought,
                "qty_bought": qty,
                "buy_price": buy_price,
                "current_price": curr_price,
                "stock_return_pct": round(stock_ret, 2),
                "nifty_return_pct": round(nifty_ret, 2) if nifty_ret is not None else None,
                "alpha_pct": round(alpha, 2) if alpha is not None else None,
                "invested": round(invested, 2),
                "current_value": round(curr_val, 2),
                "gain_loss": round(gain_loss, 2)
            })
            
    total_gain_loss = total_current_value - total_invested
    portfolio_return = (total_gain_loss / total_invested * 100) if total_invested > 0 else 0.0
    
    nifty_return = None
    trades_with_nifty = [pt for pt in processed_trades if pt["nifty_on_trade_day"] is not None]
    if current_nifty and trades_with_nifty:
        total_w = sum(pt["invested"] for pt in trades_with_nifty)
        if total_w > 0:
            weighted_nifty_base = sum(pt["nifty_on_trade_day"] * pt["invested"] for pt in trades_with_nifty) / total_w
        else:
            weighted_nifty_base = sum(pt["nifty_on_trade_day"] for pt in trades_with_nifty) / len(trades_with_nifty)
        nifty_return = ((current_nifty - weighted_nifty_base) / weighted_nifty_base) * 100
        
    win_rate = (winning_trades_count / len(processed_trades) * 100) if processed_trades else 0.0
    coverage_pct = (you_bought_count / ai_suggested_count * 100) if ai_suggested_count > 0 else 0.0
    avg_alpha = (sum_alpha / trades_with_nifty_count) if trades_with_nifty_count > 0 else 0.0
    
    summary = {
        "total_trades": len(processed_trades),
        "total_invested": round(total_invested, 2),
        "total_current_value": round(total_current_value, 2),
        "total_gain_loss": round(total_gain_loss, 2),
        "portfolio_return_pct": round(portfolio_return, 2),
        "nifty_return_pct": round(nifty_return, 2) if nifty_return is not None else None,
        "alpha_pct": round(avg_alpha, 2),
        "win_rate": round(win_rate, 1),
        "ai_suggested_count": ai_suggested_count,
        "you_bought_count": you_bought_count,
        "coverage_pct": round(coverage_pct, 1),
    }
    
    return {
        "trades": processed_trades,
        "comparison": comparison_list,
        "summary": summary
    }