#!/usr/bin/env python3
"""Summary statistics for a bars file produced by fetch_bars.py (one or several tickers).

  analyze.py aapl.csv [--plot aapl.png] [--periods-per-year 252]

Prints, per ticker: rows, range, CAGR, annualised volatility, Sharpe (rf=0), max drawdown (+ dates),
best/worst bar, skew, kurtosis, number of gaps > 3 calendar days, monthly returns table.
"""
from __future__ import annotations

import argparse
import sys

import numpy as np
import pandas as pd


def infer_periods_per_year(dt: pd.Series) -> int:
    if len(dt) < 3:
        return 252
    step = dt.diff().dropna().median()
    if step >= pd.Timedelta(days=1):
        return 252 if dt.dt.dayofweek.max() <= 4 else 365
    per_day = int(pd.Timedelta(hours=6.5) / step) if step < pd.Timedelta(hours=1) else int(pd.Timedelta(hours=24) / step)
    return max(per_day, 1) * 252


def stats(df: pd.DataFrame, ppy: int | None) -> dict:
    df = df.sort_values("datetime").reset_index(drop=True)
    ppy = ppy or infer_periods_per_year(df["datetime"])
    close = df["close"].astype(float)
    ret = close.pct_change().dropna()
    years = max((df["datetime"].iloc[-1] - df["datetime"].iloc[0]).days / 365.25, 1e-9)
    equity = (1 + ret).cumprod()
    dd = equity / equity.cummax() - 1
    trough = dd.idxmin() if len(dd) else None
    peak = equity.loc[:trough].idxmax() if trough is not None else None
    gaps = (df["datetime"].diff() > pd.Timedelta(days=3)).sum()
    out = {
        "rows": len(df),
        "first": str(df["datetime"].iloc[0]),
        "last": str(df["datetime"].iloc[-1]),
        "periods_per_year": ppy,
        "total_return": close.iloc[-1] / close.iloc[0] - 1,
        "cagr": (close.iloc[-1] / close.iloc[0]) ** (1 / years) - 1,
        "ann_vol": ret.std() * np.sqrt(ppy),
        "sharpe_rf0": (ret.mean() / ret.std() * np.sqrt(ppy)) if ret.std() > 0 else np.nan,
        "max_drawdown": dd.min() if len(dd) else np.nan,
        "dd_peak": str(df["datetime"].iloc[peak]) if peak is not None else None,
        "dd_trough": str(df["datetime"].iloc[trough]) if trough is not None else None,
        "best_bar": ret.max(),
        "worst_bar": ret.min(),
        "skew": ret.skew(),
        "kurtosis": ret.kurt(),
        "gaps_gt_3d": int(gaps),
        "avg_volume": float(df["volume"].mean()) if "volume" in df else None,
    }
    return out


def monthly_table(df: pd.DataFrame) -> pd.DataFrame:
    s = df.set_index("datetime")["close"].astype(float).resample("ME").last().pct_change().dropna()
    t = s.to_frame("r")
    t["year"], t["month"] = t.index.year, t.index.month
    return (t.pivot(index="year", columns="month", values="r") * 100).round(1)


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("file")
    ap.add_argument("--plot", help="write a price + drawdown PNG (needs matplotlib)")
    ap.add_argument("--periods-per-year", type=int)
    a = ap.parse_args()

    df = pd.read_parquet(a.file) if a.file.endswith(".parquet") else pd.read_csv(a.file)
    df["datetime"] = pd.to_datetime(df["datetime"])
    tickers = df["ticker"].unique() if "ticker" in df else ["?"]
    for t in tickers:
        sub = df[df["ticker"] == t] if "ticker" in df else df
        s = stats(sub, a.periods_per_year)
        print(f"\n=== {t} ===")
        for k, v in s.items():
            if isinstance(v, float):
                print(f"{k:>18}: {v:,.4f}" if abs(v) < 1000 else f"{k:>18}: {v:,.0f}")
            else:
                print(f"{k:>18}: {v}")
        if s["periods_per_year"] <= 365 and len(sub) > 40:
            print("\nmonthly returns (%):")
            print(monthly_table(sub).to_string())

    if a.plot:
        try:
            import matplotlib
            matplotlib.use("Agg")
            import matplotlib.pyplot as plt
        except ImportError:
            print("matplotlib not installed: pip install matplotlib", file=sys.stderr)
            return 0
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(11, 6.5), sharex=True, gridspec_kw={"height_ratios": [3, 1]})
        for t in tickers:
            sub = (df[df["ticker"] == t] if "ticker" in df else df).sort_values("datetime")
            close = sub["close"].astype(float)
            norm = close / close.iloc[0] * 100 if len(tickers) > 1 else close
            ax1.plot(sub["datetime"], norm, lw=1.1, label=t)
            eq = (1 + close.pct_change().fillna(0)).cumprod()
            ax2.fill_between(sub["datetime"], (eq / eq.cummax() - 1) * 100, 0, alpha=0.35)
        ax1.set_title(f"{', '.join(tickers)} — {'rebased to 100' if len(tickers) > 1 else 'close'}")
        ax1.grid(alpha=0.25)
        ax1.legend(loc="upper left")
        ax2.set_ylabel("drawdown %")
        ax2.grid(alpha=0.25)
        fig.tight_layout()
        fig.savefig(a.plot, dpi=130)
        print(f"\nwrote {a.plot}")
    return 0


if __name__ == "__main__":
    sys.exit(main())