#!/usr/bin/env python3
# Plot heat & gap metrics from the cache-out SQLite DB (heat_stats table)
# Usage:
#   python3 plot_heat_stats.py --db path/to/combined_cache_session.db --outdir charts
# Optional filters:
#   --symbol BTC/USDT:USDT
#   --mode live   (or paper-api, etc.)
#
# This script creates separate PNG files for:
#   - Heat (1 - combined_gap)
#   - Combined gap
#   - Gap components: ATR, VolSurge, qv_24h, qv_1h, Momentum, Breadth
#
# Notes:
#   - Uses matplotlib only (no seaborn), one chart per figure, no explicit colors.

import argparse
import os
import sqlite3
from datetime import datetime
import matplotlib.pyplot as plt

COLUMNS = [
    'datetime_utc','mode','symbol','combined_gap','heat',
    'gap_atr','gap_volsurge','gap_qv24','gap_qv1h','gap_momentum','gap_breadth'
]

def read_heat_stats(db_path, symbol_filter=None, mode_filter=None):
    con = sqlite3.connect(db_path)
    cur = con.cursor()
    q = f"""SELECT {','.join(COLUMNS)}
             FROM heat_stats
             ORDER BY datetime_utc ASC"""
    cur.execute(q)
    rows = cur.fetchall()
    con.close()

    # Pack columns
    data = {k: [] for k in COLUMNS}
    for r in rows:
        for k, v in zip(COLUMNS, r):
            data[k].append(v)

    # Convert timestamps
    ts = []
    for t in data['datetime_utc']:
        try:
            ts.append(datetime.fromisoformat(str(t).replace('Z', '+00:00')))
        except Exception:
            # last resort
            try:
                ts.append(datetime.strptime(str(t), '%Y-%m-%dT%H:%M:%S+00:00'))
            except Exception:
                ts.append(None)
    data['ts'] = ts

    # Create index considering filters and valid timestamps
    idx_all = list(range(len(ts)))
    idx = [i for i in idx_all if data['ts'][i] is not None]
    if symbol_filter:
        idx = [i for i in idx if str(data['symbol'][i]) == symbol_filter]
    if mode_filter:
        idx = [i for i in idx if str(data['mode'][i]) == mode_filter]

    return data, idx

def ensure_outdir(path):
    os.makedirs(path, exist_ok=True)

def plot_series(ts, ys, title, path_png):
    plt.figure()
    plt.plot(ts, ys)
    plt.title(title)
    plt.xlabel('Time')
    plt.ylabel(title)
    plt.grid(True)
    plt.tight_layout()
    plt.savefig(path_png)
    print('Saved', path_png)

def main():
    ap = argparse.ArgumentParser(description='Plot heat/gap series from heat_stats table')
    ap.add_argument('--db', required=True, help='Path to cache-out sqlite db (e.g., combined_cache_session.db)')
    ap.add_argument('--outdir', default='charts', help='Output directory for PNG files')
    ap.add_argument('--symbol', default='', help='Filter by nearest symbol in heat (exact match)')
    ap.add_argument('--mode', default='', help='Filter by mode label (e.g., live, paper-api)')
    args = ap.parse_args()

    ensure_outdir(args.outdir)
    symbol_filter = args.symbol if args.symbol else None
    mode_filter = args.mode if args.mode else None

    data, idx = read_heat_stats(args.db, symbol_filter, mode_filter)

    if not idx:
        print('No rows match the filters (or table is empty).')
        return

    # Prepare series
    ts = [data['ts'][i] for i in idx]
    series = {
        'Heat (1 - combined_gap)': [float(data['heat'][i]) for i in idx],
        'Combined gap':            [float(data['combined_gap'][i]) for i in idx],
        'Gap ATR':                 [float(data['gap_atr'][i]) for i in idx],
        'Gap VolSurge':            [float(data['gap_volsurge'][i]) for i in idx],
        'Gap qv_24h':              [float(data['gap_qv24'][i]) for i in idx],
        'Gap qv_1h':               [float(data['gap_qv1h'][i]) for i in idx],
        'Gap Momentum':            [float(data['gap_momentum'][i]) for i in idx],
        'Gap Breadth':             [float(data['gap_breadth'][i]) for i in idx],
    }

    # File name suffix for filters
    suffix = []
    if symbol_filter:
        safe_sym = symbol_filter.replace('/', '_').replace(':', '_')
        suffix.append(safe_sym)
    if mode_filter:
        suffix.append(mode_filter)
    suf = ('_' + '_'.join(suffix)) if suffix else ''

    # Plot each series to a separate PNG
    for title, ys in series.items():
        fname = title.lower().replace(' ', '_').replace('(', '').replace(')', '').replace('-', '')
        out_path = os.path.join(args.outdir, f"{fname}{suf}.png")
        plot_series(ts, ys, title, out_path)

if __name__ == '__main__':
    main()