src/charts.py

"""
Charts module: generates matplotlib charts for embedded use in PDF reports.
All functions return a BytesIO buffer containing a PNG image.
"""

import io
import numpy as np
import matplotlib
matplotlib.use('Agg')  # non-interactive backend
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib.figure import Figure
from typing import Dict, List, Optional

# Brand colours
BRAND_BLUE   = "#1F3864"
BRAND_ACCENT = "#2E75B6"
GREENS       = ["#2ECC71", "#27AE60", "#1ABC9C", "#16A085", "#52BE80"]
REDS         = ["#E74C3C", "#C0392B", "#EC7063"]
PALETTE      = [
    "#2E75B6", "#E67E22", "#2ECC71", "#E74C3C", "#9B59B6",
    "#1ABC9C", "#F39C12", "#3498DB", "#D35400", "#27AE60",
]


def _buf(fig: Figure) -> io.BytesIO:
    buf = io.BytesIO()
    fig.savefig(buf, format='png', bbox_inches='tight', dpi=150)
    buf.seek(0)
    plt.close(fig)
    return buf


def holdings_pie_chart(holdings_data: Dict[str, float], title: str = "Portfolio Allocation") -> io.BytesIO:
    """
    Pie chart of holdings by name → value.
    holdings_data: {scheme_name: current_value}
    """
    labels = list(holdings_data.keys())
    values = list(holdings_data.values())

    # Shorten long labels
    short_labels = [l.split('-')[0].strip()[:22] for l in labels]

    fig, ax = plt.subplots(figsize=(5, 4))
    wedges, texts, autotexts = ax.pie(
        values,
        labels=None,
        autopct='%1.1f%%',
        startangle=140,
        colors=PALETTE[:len(values)],
        pctdistance=0.78,
    )
    for at in autotexts:
        at.set_fontsize(7)
        at.set_color("white")

    ax.legend(wedges, short_labels, loc="center left", bbox_to_anchor=(1, 0.5),
              fontsize=7, frameon=False)
    ax.set_title(title, fontsize=10, fontweight='bold', color=BRAND_BLUE, pad=10)
    fig.tight_layout()
    return _buf(fig)


def sector_bar_chart(sector_data: Dict[str, float], title: str = "Sector Allocation (%)") -> io.BytesIO:
    """Horizontal bar chart for sector allocation."""
    if not sector_data:
        sector_data = {"Data Not Available": 100}

    sectors = list(sector_data.keys())
    values  = list(sector_data.values())

    # Sort descending
    pairs = sorted(zip(values, sectors), reverse=True)
    values, sectors = zip(*pairs)

    fig, ax = plt.subplots(figsize=(5, max(3, len(sectors) * 0.35)))
    bars = ax.barh(sectors, values, color=BRAND_ACCENT, edgecolor='white', height=0.6)

    for bar, val in zip(bars, values):
        ax.text(bar.get_width() + 0.3, bar.get_y() + bar.get_height() / 2,
                f'{val:.1f}%', va='center', fontsize=7, color='black')

    ax.set_xlabel("Allocation (%)", fontsize=8, color='gray')
    ax.set_title(title, fontsize=10, fontweight='bold', color=BRAND_BLUE)
    ax.set_xlim(0, max(values) * 1.2)
    ax.invert_yaxis()
    ax.spines[['top', 'right']].set_visible(False)
    ax.tick_params(axis='y', labelsize=8)
    fig.tight_layout()
    return _buf(fig)


def market_cap_pie(market_cap_data: Dict[str, float]) -> io.BytesIO:
    """Pie chart for Large/Mid/Small/Other market cap split."""
    default = {"Large Cap": 0, "Mid Cap": 0, "Small Cap": 0, "Others": 0}
    data = {**default, **market_cap_data}
    data = {k: v for k, v in data.items() if v > 0}

    colors = {"Large Cap": "#2E75B6", "Mid Cap": "#E67E22",
              "Small Cap": "#2ECC71", "Others": "#BDC3C7"}

    labels = list(data.keys())
    values = list(data.values())
    clrs   = [colors.get(l, "#95A5A6") for l in labels]

    fig, ax = plt.subplots(figsize=(4, 3.5))
    wedges, _, autotexts = ax.pie(
        values, labels=None, autopct='%1.1f%%',
        colors=clrs, startangle=90, pctdistance=0.75
    )
    for at in autotexts:
        at.set_fontsize(8)
        at.set_color("white")

    ax.legend(wedges, labels, loc="lower center", bbox_to_anchor=(0.5, -0.12),
              ncol=2, fontsize=8, frameon=False)
    ax.set_title("Market Cap Allocation", fontsize=10, fontweight='bold', color=BRAND_BLUE)
    fig.tight_layout()
    return _buf(fig)


def holding_vs_benchmark_chart(
    fund_name: str,
    cagr_data: Dict[str, Dict[str, Optional[float]]],
) -> io.BytesIO:
    """
    Bar chart comparing fund CAGR vs benchmark across time periods.
    cagr_data = {
        '1Y': {'fund': 12.5, 'benchmark': 14.6, 'category': 13.4},
        '3Y': {...}, '5Y': {...}, '10Y': {...}
    }
    """
    periods = list(cagr_data.keys())
    fund_vals  = [cagr_data[p].get('fund') or 0 for p in periods]
    bm_vals    = [cagr_data[p].get('benchmark') or 0 for p in periods]
    cat_vals   = [cagr_data[p].get('category') or 0 for p in periods]

    x = np.arange(len(periods))
    width = 0.25

    fig, ax = plt.subplots(figsize=(5, 3.5))
    b1 = ax.bar(x - width, fund_vals, width, label='Fund', color=BRAND_ACCENT, zorder=2)
    b2 = ax.bar(x,          bm_vals,  width, label='Benchmark', color='#E67E22', zorder=2)
    b3 = ax.bar(x + width,  cat_vals, width, label='Category', color='#BDC3C7', zorder=2)

    def label_bars(bars):
        for bar in bars:
            h = bar.get_height()
            if h:
                ax.text(bar.get_x() + bar.get_width() / 2, h + 0.2,
                        f'{h:.1f}', ha='center', va='bottom', fontsize=6.5)

    label_bars(b1); label_bars(b2); label_bars(b3)

    ax.set_xticks(x)
    ax.set_xticklabels(periods, fontsize=9)
    ax.set_ylabel("CAGR (%)", fontsize=8, color='gray')
    ax.set_title(f"{fund_name[:30]}\nReturns vs Benchmark", fontsize=9, fontweight='bold', color=BRAND_BLUE)
    ax.legend(fontsize=7, frameon=False)
    ax.spines[['top', 'right']].set_visible(False)
    ax.yaxis.grid(True, linestyle='--', alpha=0.5, zorder=0)
    ax.set_axisbelow(True)
    fig.tight_layout()
    return _buf(fig)


def quartile_analysis_grid(holdings_data: list) -> io.BytesIO:
    """
    Quartile Analysis Grid — based on the senior's handwritten sketch.

    Layout (matching sketch exactly):
      Columns  : 1Y | 3Y | 5Y | 10Y
      For each holding, show 3 rows:
        BM    : Benchmark CAGR value for each period
        Cat   : Category Average CAGR for each period
        Scheme: Fund CAGR + Quartile (Q1/Q2/Q3/Q4) — color-coded

    holdings_data: list of dicts, each with keys:
      scheme_name, rank_in_category, total_in_category,
      cagr_1y/_bm/_cat, cagr_3y/_bm/_cat, cagr_5y/_bm/_cat, cagr_10y/_bm/_cat
    """
    PERIODS     = ["1Y", "3Y", "5Y", "10Y"]
    PERIOD_KEYS = ["1y", "3y", "5y", "10y"]
    ROW_LABELS  = ["BM", "Cat", "Scheme"]

    Q_COLORS   = {1: "#90EE90", 2: "#BDD7EE", 3: "#FFD580", 4: "#FFB3B3"}
    HEADER_CLR = "#1F3864"
    BM_CLR     = "#D6E4F0"
    CAT_CLR    = "#EBF5FB"

    def get_quartile(rank, total):
        if not rank or not total or total == 0:
            return 4
        pct = rank / total
        if pct <= 0.25: return 1
        if pct <= 0.50: return 2
        if pct <= 0.75: return 3
        return 4

    def fmt(v):
        if v is None: return "–"
        try: return f"{float(v):.1f}%"
        except: return "–"

    n_holdings = len(holdings_data)
    rows_per   = 3   # BM, Cat, Scheme
    n_rows     = n_holdings * rows_per + 1   # +1 for header row
    n_cols     = 5                            # Label + 4 periods

    fig_h = max(4.5, 0.5 * n_rows + 1.5)
    fig, ax = plt.subplots(figsize=(10, fig_h))
    ax.set_xlim(0, n_cols)
    ax.set_ylim(0, n_rows)
    ax.axis('off')

    def cell(row, col, text, bg, tc="#1F3864", bold=False, fs=8):
        ax.add_patch(plt.Rectangle(
            (col, n_rows - row - 1), 1, 1,
            facecolor=bg, edgecolor="#AAAAAA", linewidth=0.5, zorder=1))
        ax.text(col + 0.5, n_rows - row - 0.5, text,
                ha='center', va='center', fontsize=fs,
                fontweight='bold' if bold else 'normal',
                color=tc, zorder=2, wrap=True)

    # Column header row
    col_widths = [1.5, 1, 1, 1, 0.8]  # proportional, but we draw on a 5-unit grid
    cell(0, 0, "Scheme / Row", HEADER_CLR, "white", bold=True, fs=7.5)
    for ci, p in enumerate(PERIODS, 1):
        cell(0, ci, p, HEADER_CLR, "white", bold=True, fs=10)

    # Data rows
    cur = 1
    for h in holdings_data:
        rank  = h.get("rank_in_category")
        total = h.get("total_in_category")
        q     = get_quartile(rank, total)
        qc    = Q_COLORS[q]
        q_lbl = f"Q{q}"
        name  = str(h.get("scheme_name", ""))[:22]

        for ri, rl in enumerate(ROW_LABELS):
            if ri == 0:
                lbl = f"{name}\n[BM]"
                bg  = BM_CLR
            elif ri == 1:
                lbl = "[Category]"
                bg  = CAT_CLR
            else:
                lbl = f"[Scheme — {q_lbl}]"
                bg  = qc

            cell(cur + ri, 0, lbl, bg, bold=(ri == 2), fs=6.5)

            for ci, pk in enumerate(PERIOD_KEYS, 1):
                if ri == 0:
                    v = fmt(h.get(f"cagr_{pk}_bm"))
                    bg_c = BM_CLR
                elif ri == 1:
                    v = fmt(h.get(f"cagr_{pk}_cat"))
                    bg_c = CAT_CLR
                else:
                    fv  = h.get(f"cagr_{pk}")
                    bmv = h.get(f"cagr_{pk}_bm")
                    v   = fmt(fv)
                    bg_c = qc
                    # Green tick if fund beats benchmark this period
                    if fv is not None and bmv is not None and float(fv) >= float(bmv):
                        ax.text(ci + 0.88, n_rows - (cur + ri) - 0.18,
                                "✓", fontsize=8, color="#006400", va='center', zorder=3)

                cell(cur + ri, ci, v, bg_c, bold=(ri == 2), fs=8)

        # Divider between schemes
        y = n_rows - (cur + rows_per) - 0.02
        ax.axhline(y=y, xmin=0, xmax=1, color="#555555", linewidth=1.0, zorder=4)
        cur += rows_per

    # Legend
    patches = [mpatches.Patch(facecolor=Q_COLORS[i], edgecolor='#AAAAAA',
               label=f"Q{i} – {['Top Quartile','Above Avg','Below Avg','Bottom Quartile'][i-1]}")
               for i in range(1, 5)]
    ax.legend(handles=patches, loc='lower center',
              bbox_to_anchor=(0.5, -0.09), ncol=4, fontsize=7.5, frameon=False)

    ax.set_title("Quartile Analysis — Scheme vs Benchmark & Category Average",
                 fontsize=10, fontweight='bold', color=HEADER_CLR, pad=10)
    fig.tight_layout()
    return _buf(fig)


def wealth_projection_chart(projection: Dict[int, float], current_value: float) -> io.BytesIO:
    """Line chart showing projected wealth growth at 12% over years."""
    years = [0] + list(projection.keys())
    values = [current_value] + list(projection.values())

    fig, ax = plt.subplots(figsize=(5, 3))
    ax.plot(years, values, marker='o', color=BRAND_ACCENT, linewidth=2, markersize=6)

    for yr, val in zip(years, values):
        ax.annotate(f'₹{val/1e5:.1f}L', (yr, val),
                    textcoords="offset points", xytext=(0, 8),
                    ha='center', fontsize=7.5, color=BRAND_BLUE)

    ax.fill_between(years, values, alpha=0.15, color=BRAND_ACCENT)
    ax.set_xticks(years)
    ax.set_xticklabels([f'Now' if y == 0 else f'{y}Y' for y in years], fontsize=8)
    ax.set_ylabel("Portfolio Value (₹)", fontsize=8, color='gray')
    ax.set_title("Wealth Projection @ 12% p.a.", fontsize=10, fontweight='bold', color=BRAND_BLUE)
    ax.spines[['top', 'right']].set_visible(False)
    ax.yaxis.grid(True, linestyle='--', alpha=0.4)
    ax.set_axisbelow(True)
    fig.tight_layout()
    return _buf(fig)