"""
01_download_prices.py
Download monthly + daily adjusted prices from Yahoo Finance for all S&P 500
tickers present in the semantic panel, plus the SPY market benchmark and the
risk-free proxy (^IRX, 13-week T-bill). Produces:
  data/processed/prices_monthly.parquet   (long: ticker, date, adj_close)
  data/processed/marketcap.parquet        (ticker, sharesOutstanding, marketCap snapshot)
  data/processed/spy_monthly.parquet
"""
import time, json, warnings
import pandas as pd
import numpy as np
import yfinance as yf

warnings.filterwarnings("ignore")
PANEL = "data/parquet/annual_final/annual_final.parquet"
OUT = "data/processed"

panel = pd.read_parquet(PANEL)
tickers = sorted(panel.ticker.dropna().unique().tolist())
print(f"{len(tickers)} tickers to download")

START, END = "2009-06-01", "2026-06-17"

def dl_batch(tks, **kw):
    for attempt in range(3):
        try:
            df = yf.download(tks, start=START, end=END, auto_adjust=True,
                             progress=False, group_by="ticker", threads=True, **kw)
            return df
        except Exception as e:
            print("  retry", attempt, e); time.sleep(3)
    return None

# --- monthly prices, batched ---
all_monthly = []
BATCH = 50
for i in range(0, len(tickers), BATCH):
    batch = tickers[i:i+BATCH]
    print(f"batch {i//BATCH+1}: {batch[0]}..{batch[-1]}")
    df = dl_batch(batch, interval="1mo")
    if df is None:
        continue
    for t in batch:
        try:
            sub = df[t] if isinstance(df.columns, pd.MultiIndex) else df
            s = sub["Close"].dropna()
            if len(s) == 0:
                continue
            tmp = s.reset_index()
            tmp.columns = ["date", "adj_close"]
            tmp["ticker"] = t
            all_monthly.append(tmp)
        except Exception:
            pass
    time.sleep(1)

monthly = pd.concat(all_monthly, ignore_index=True)
monthly["date"] = pd.to_datetime(monthly["date"])
monthly = monthly.sort_values(["ticker", "date"])
monthly.to_parquet(f"{OUT}/prices_monthly.parquet", index=False)
got = monthly.ticker.nunique()
print(f"monthly: {len(monthly)} rows, {got}/{len(tickers)} tickers")

# --- benchmark + risk free ---
spy = dl_batch(["SPY"], interval="1mo")
spy_close = (spy["SPY"]["Close"] if isinstance(spy.columns, pd.MultiIndex) else spy["Close"]).dropna()
spy_df = spy_close.reset_index(); spy_df.columns = ["date", "spy_adj_close"]
spy_df.to_parquet(f"{OUT}/spy_monthly.parquet", index=False)
print("spy rows:", len(spy_df))

# --- market cap snapshot (best effort) ---
rows = []
for i, t in enumerate(tickers):
    try:
        fi = yf.Ticker(t).fast_info
        rows.append({"ticker": t,
                     "market_cap": getattr(fi, "market_cap", None),
                     "shares": getattr(fi, "shares", None)})
    except Exception:
        rows.append({"ticker": t, "market_cap": None, "shares": None})
    if i % 50 == 0:
        print("  mktcap", i)
pd.DataFrame(rows).to_parquet(f"{OUT}/marketcap.parquet", index=False)
print("done")