from __future__ import annotations

import argparse
from pathlib import Path

import pandas as pd


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--equity", required=True)
    ap.add_argument("--trades", required=True)
    ap.add_argument("--top", type=int, default=3)
    args = ap.parse_args()

    eq = pd.read_parquet(args.equity)
    eq = eq.copy()
    eq.index = eq.index.astype(str)
    s = eq["equity"].astype(float)

    peak = s.cummax()
    dd = s / peak - 1.0

    # find worst drawdowns by trough
    worst = dd.nsmallest(args.top)

    tr = pd.read_parquet(args.trades)
    tr = tr.copy()
    tr["trade_date"] = tr["trade_date"].astype(str)

    for d, v in worst.items():
        # drawdown start = last peak before d
        peak_date = (s.loc[:d]).idxmax()
        print("---")
        print("trough", d, "dd", float(v))
        print("peak", peak_date, "peak_equity", float(s.loc[peak_date]), "trough_equity", float(s.loc[d]))
        w = tr[(tr["trade_date"] >= peak_date) & (tr["trade_date"] <= d)]
        print("trades in window", len(w))
        if not w.empty:
            cols = [c for c in ["trade_date", "ts_code", "side", "reason", "weight_before", "weight_after", "price"] if c in w.columns]
            print(w[cols].tail(25).to_string(index=False))


if __name__ == "__main__":
    main()