visuals.py

# visuals.py
"""
Visualization helpers for the Portfolio Manager.

Design:
- Canonical positions come from Static_Portfolio.owned_stocks()
- Live prices & high-level metrics come from Market_Data
- This module focuses on matplotlib-based plots for:
    * sector & allocation visuals
    * unrealized / realized P&L
    * portfolio value over time & drawdowns
    * portfolio vs S&P 500 comparison
    * correlation heatmaps

Most functions are interactive (they call plt.show()) and are intended for
CLI / notebook usage. The helper `_get_portfolio_value_series()` is also
used by the Flask app to generate PNGs.
"""

from __future__ import annotations

from typing import List

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import yfinance as yf

from Static_Portfolio import owned_stocks
from Market_Data import (
    get_current_prices,
    unrealized_pl,
    portfolio_market_value,
    portfolio_performance_summary,
    Sector_Allocation,
    compare_to_sp500,
)


# --------------------------------------------------------------------
# SMALL INTERNAL HELPERS
# --------------------------------------------------------------------

def _ensure_nonempty_df(df: pd.DataFrame | None, what: str) -> bool:
    """Utility: print a message & return False if df is empty/None."""
    if df is None or df.empty:
        print(f"[visuals] No data available for {what}.")
        return False
    return True


def _tz_naive_series(s: pd.Series) -> pd.Series:
    """
    Return a copy of a Series with a tz-naive DatetimeIndex (drop timezone).
    If index is not datetime-like or already tz-naive, returns unchanged.
    """
    idx = s.index
    if isinstance(idx, pd.DatetimeIndex) and idx.tz is not None:
        s = s.copy()
        s.index = idx.tz_localize(None)
    return s


# --------------------------------------------------------------------
# PORTFOLIO VALUE SERIES (CORE HELPER)
# --------------------------------------------------------------------

def _get_portfolio_value_series(period: str = "1y") -> pd.Series:
    """
    Compute daily portfolio value over time using yfinance history.

    - Reads positions from Static_Portfolio.owned_stocks()
    - Downloads prices for all tickers via yfinance
    - Multiplies each price series by its quantity and sums across tickers

    Returns:
        pandas Series indexed by tz-naive dates, dtype float.
        Empty Series if there are no positions or no price data.
    """
    df_pos, _ = owned_stocks()
    if df_pos.empty:
        return pd.Series(dtype=float)

    tickers: List[str] = df_pos["Ticker"].astype(str).str.upper().tolist()
    quantities = df_pos.set_index("Ticker")["Quantity"]

    # Download prices (do not auto-adjust so we can pick Adj Close or Close)
    raw = yf.download(tickers, period=period, auto_adjust=False)

    if raw.empty:
        return pd.Series(dtype=float)

    # Handle MultiIndex vs single-index columns
    if isinstance(raw.columns, pd.MultiIndex):
        # First level is price field: 'Adj Close', 'Close', etc.
        level0 = raw.columns.get_level_values(0)
        price_field = "Adj Close" if "Adj Close" in level0 else "Close"
        data = raw[price_field]
    else:
        # Single-ticker case: plain columns
        if "Adj Close" in raw.columns:
            data = raw[["Adj Close"]]
        elif "Close" in raw.columns:
            data = raw[["Close"]]
        else:
            # No usable price field
            return pd.Series(dtype=float)

        # Ensure column name(s) match our tickers list
        data.columns = tickers[: len(data.columns)]

    # If it's a Series (odd single-ticker case), convert to DataFrame
    if isinstance(data, pd.Series):
        data = data.to_frame(name=tickers[0])

    # Align quantities to data columns
    aligned_quantities = quantities.reindex(data.columns).fillna(0.0)

    # Multiply each column by its quantity, then sum across tickers
    portfolio_values = data.mul(aligned_quantities, axis=1).sum(axis=1)

    # Ensure tz-naive index for downstream concatenation / plotting
    portfolio_values = _tz_naive_series(portfolio_values)

    return portfolio_values

def get_drawdown_series(period: str = "1y") -> pd.Series:
    """
    Compute portfolio drawdown series over time:

        drawdown = (value - running_max) / running_max

    Returns:
        pandas Series indexed by date, values in [-1, 0],
        or empty Series if no data.
    """
    series = _get_portfolio_value_series(period=period)
    if series.empty:
        return series

    running_max = series.cummax()
    drawdown = (series - running_max) / running_max
    return drawdown


def get_portfolio_vs_sp500_returns(period: str = "1y") -> pd.DataFrame:
    """
    Build a combined DataFrame of cumulative returns for:
        - Portfolio
        - S&P 500 (^GSPC)

    Both series are normalized to 0 at the start (returns, not prices):

        return_t = price_t / price_0 - 1

    Returns:
        DataFrame with columns ["Portfolio", "SP500"] and a DatetimeIndex,
        or empty DataFrame if data cannot be fetched.
    """
    # Portfolio value series
    port_series = _get_portfolio_value_series(period=period)
    if port_series.empty:
        return pd.DataFrame(columns=["Portfolio", "SP500"])

    # S&P 500
    sp = yf.Ticker("^GSPC")
    sp_hist = sp.history(period=period)
    if sp_hist.empty:
        return pd.DataFrame(columns=["Portfolio", "SP500"])

    sp_series = sp_hist["Close"]

    # Make both tz-naive
    port_series = _tz_naive_index(port_series)
    sp_series = _tz_naive_index(sp_series)

    # Convert to cumulative returns (starting at 0)
    port_ret = port_series / port_series.iloc[0] - 1.0
    sp_ret = sp_series / sp_series.iloc[0] - 1.0

    combined = pd.concat(
        [port_ret.rename("Portfolio"), sp_ret.rename("SP500")],
        axis=1,
    ).dropna()

    return combined

def get_drawdown_series(period: str = "1y") -> pd.Series:
    """
    Compute portfolio drawdown series over time:

        drawdown = (value - running_max) / running_max

    Returns:
        pandas Series indexed by date, values in [-1, 0],
        or empty Series if no data.
    """
    series = _get_portfolio_value_series(period=period)
    if series.empty:
        return series

    running_max = series.cummax()
    drawdown = (series - running_max) / running_max
    return drawdown


def get_portfolio_vs_sp500_returns(period: str = "1y") -> pd.DataFrame:
    """
    Build a combined DataFrame of cumulative returns for:
        - Portfolio
        - S&P 500 (^GSPC)

    Both series are normalized to 0 at the start (returns, not prices):

        return_t = price_t / price_0 - 1

    Returns:
        DataFrame with columns ["Portfolio", "SP500"] and a DatetimeIndex,
        or empty DataFrame if data cannot be fetched.
    """
    # Portfolio value series
    port_series = _get_portfolio_value_series(period=period)
    if port_series.empty:
        return pd.DataFrame(columns=["Portfolio", "SP500"])

    # S&P 500
    sp = yf.Ticker("^GSPC")
    sp_hist = sp.history(period=period)
    if sp_hist.empty:
        return pd.DataFrame(columns=["Portfolio", "SP500"])

    sp_series = sp_hist["Close"]

    # Make both tz-naive
    port_series = _tz_naive_series(port_series)
    sp_series = _tz_naive_series(sp_series)

    # Convert to cumulative returns (starting at 0)
    port_ret = port_series / port_series.iloc[0] - 1.0
    sp_ret = sp_series / sp_series.iloc[0] - 1.0

    combined = pd.concat(
        [port_ret.rename("Portfolio"), sp_ret.rename("SP500")],
        axis=1,
    ).dropna()

    return combined


# --------------------------------------------------------------------
# SECTOR & ALLOCATION VISUALS
# --------------------------------------------------------------------

def plot_sector_allocation(by: str = "cost") -> None:
    """
    Plot sector allocation as a bar chart.

    Args:
        by:
            "cost"   -> uses cost basis from Static_Portfolio
            "market" -> uses current market value based on live prices
    """
    df = Sector_Allocation(by=by)
    if not _ensure_nonempty_df(df, "sector allocation"):
        return

    plt.figure(figsize=(8, 5))
    plt.bar(df["Sector"], df["Weight"])
    plt.xticks(rotation=45, ha="right")
    plt.ylabel("Weight")
    plt.title(f"Sector Allocation by {by.capitalize()}")
    plt.tight_layout()
    plt.show()


def plot_allocation_by_ticker(by: str = "market") -> None:
    """
    Plot allocation by ticker as a bar chart.

    Args:
        by:
            "market" (default) uses Quantity * current price
            "cost"             uses Quantity * Buy_Price
    """
    df, _ = owned_stocks()
    if not _ensure_nonempty_df(df, "owned stocks"):
        return

    df = df.copy()

    if by.lower() == "market":
        prices = get_current_prices()
        df["Value"] = df.apply(
            lambda r: float(r["Quantity"]) * float(
                prices.get(str(r["Ticker"]).strip().upper()) or 0.0
            ),
            axis=1,
        )
        title = "Allocation by Ticker (Market Value)"
    else:
        df["Value"] = df["Quantity"] * df["Buy_Price"]
        title = "Allocation by Ticker (Cost Basis)"

    df = df.sort_values("Value", ascending=False)

    plt.figure(figsize=(8, 5))
    plt.bar(df["Ticker"], df["Value"])
    plt.xticks(rotation=45, ha="right")
    plt.ylabel("Value ($)")
    plt.title(title)
    plt.tight_layout()
    plt.show()


# --------------------------------------------------------------------
# P/L VISUALS
# --------------------------------------------------------------------

def plot_unrealized_pl_bar() -> None:
    """
    Bar chart of unrealized P/L ($) by ticker.
    """
    df = unrealized_pl()
    if not _ensure_nonempty_df(df, "unrealized P/L"):
        return

    df = df.sort_values("PL_$", ascending=False)

    plt.figure(figsize=(8, 5))
    plt.bar(df["Ticker"], df["PL_$"])
    plt.xticks(rotation=45, ha="right")
    plt.axhline(0, linestyle="--")
    plt.ylabel("Unrealized P/L ($)")
    plt.title("Unrealized Profit / Loss by Ticker")
    plt.tight_layout()
    plt.show()


def plot_realized_vs_unrealized() -> None:
    """
    Compare total realized P/L vs total unrealized P/L.
    """
    positions_df, _ = owned_stocks()
    if not _ensure_nonempty_df(positions_df, "positions for realized P/L"):
        return

    unreal_df = unrealized_pl()
    realized_total = positions_df["Realized_PL"].sum()
    unrealized_total = unreal_df["PL_$"].sum() if not unreal_df.empty else 0.0

    labels = ["Realized", "Unrealized"]
    values = [realized_total, unrealized_total]

    plt.figure(figsize=(6, 4))
    plt.bar(labels, values)
    plt.axhline(0, linestyle="--")
    plt.ylabel("P/L ($)")
    plt.title("Realized vs Unrealized Profit / Loss")
    plt.tight_layout()
    plt.show()


# --------------------------------------------------------------------
# TIME-SERIES VISUALS
# --------------------------------------------------------------------

def plot_portfolio_value_over_time(period: str = "1y") -> None:
    """
    Plot portfolio total market value over time.
    """
    series = _get_portfolio_value_series(period=period)
    if series.empty:
        print("[visuals] No data for portfolio value over time.")
        return

    plt.figure(figsize=(9, 5))
    plt.plot(series.index, series.values)
    plt.xlabel("Date")
    plt.ylabel("Portfolio Value ($)")
    plt.title(f"Portfolio Value Over Time ({period})")
    plt.tight_layout()
    plt.show()


def plot_ticker_history(ticker: str | None = None, period: str = "1y") -> None:
    """
    Plot price history for a single ticker in the portfolio.

    If ticker is None, uses the first ticker from owned_stocks().
    """
    df_pos, _ = owned_stocks()
    if df_pos.empty:
        print("[visuals] No positions available.")
        return

    if ticker is None:
        ticker = str(df_pos["Ticker"].iloc[0]).strip().upper()
    else:
        ticker = str(ticker).strip().upper()

    hist = yf.Ticker(ticker).history(period=period)
    if hist.empty:
        print(f"[visuals] No history data for {ticker}.")
        return

    plt.figure(figsize=(9, 5))
    plt.plot(hist.index, hist["Close"])
    plt.xlabel("Date")
    plt.ylabel("Price ($)")
    plt.title(f"{ticker} Price History ({period})")
    plt.tight_layout()
    plt.show()


def plot_drawdown(period: str = "1y") -> None:
    """
    Plot portfolio drawdown curve over time:
        drawdown = (value - running_max) / running_max
    """
    drawdown = get_drawdown_series(period=period)
    if drawdown.empty:
        print("[visuals] No data for drawdown.")
        return

    plt.figure(figsize=(9, 5))
    plt.plot(drawdown.index, drawdown.values)
    plt.xlabel("Date")
    plt.ylabel("Drawdown")
    plt.title(f"Portfolio Drawdown Over Time ({period})")
    plt.axhline(0, linestyle="--")
    plt.tight_layout()
    plt.show()


def plot_benchmark_comparison(period: str = "1y") -> None:
    """
    Plot portfolio vs S&P 500 normalized to 1 at the start.
    """
    df = get_portfolio_vs_sp500_returns(period=period)
    if df.empty:
        print("[visuals] No data for benchmark comparison.")
        return

    port_ret = df["Portfolio"]
    sp_ret = df["SP500"]

    plt.figure(figsize=(9, 5))
    plt.plot(port_ret.index, 1.0 + port_ret.values, label="Portfolio")
    plt.plot(sp_ret.index, 1.0 + sp_ret.values, label="S&P 500")
    plt.xlabel("Date")
    plt.ylabel("Normalized Value (start = 1.0)")
    plt.title(f"Portfolio vs S&P 500 ({period})")
    plt.legend()
    plt.tight_layout()
    plt.show()


# --------------------------------------------------------------------
# CORRELATION / RISK VISUALS
# --------------------------------------------------------------------

def plot_correlation_heatmap(period: str = "1y") -> None:
    """
    Plot a correlation heatmap between daily returns of portfolio tickers.
    """
    df_pos, _ = owned_stocks()
    if df_pos.empty:
        print("[visuals] No positions for correlation heatmap.")
        return

    tickers = df_pos["Ticker"].astype(str).str.upper().tolist()
    raw = yf.download(tickers, period=period, auto_adjust=False)

    if raw.empty:
        print("[visuals] No price data for correlation heatmap.")
        return

    if isinstance(raw.columns, pd.MultiIndex):
        level0 = raw.columns.get_level_values(0)
        price_field = "Adj Close" if "Adj Close" in level0 else "Close"
        data = raw[price_field]
    else:
        if "Adj Close" in raw.columns:
            data = raw[["Adj Close"]]
        elif "Close" in raw.columns:
            data = raw[["Close"]]
        else:
            print("[visuals] No Adj Close / Close field for correlation heatmap.")
            return
        data.columns = tickers[: len(data.columns)]

    if isinstance(data, pd.Series) or data.shape[1] <= 1:
        print("[visuals] Only one ticker; correlation heatmap not meaningful.")
        return

    returns = data.pct_change().dropna()
    corr = returns.corr()

    plt.figure(figsize=(8, 6))
    plt.imshow(corr, interpolation="nearest")
    plt.colorbar(label="Correlation")
    plt.xticks(range(len(corr.columns)), corr.columns, rotation=45, ha="right")
    plt.yticks(range(len(corr.index)), corr.index)
    plt.title(f"Correlation Heatmap ({period})")
    plt.tight_layout()
    plt.show()