import os
import requests

from dotenv import load_dotenv


load_dotenv()  # load environment variables from .env file


API_AUTH_HEADERS = {"Content-Type": "application/x-www-form-urlencoded"}
API_HOST = "https://api.empiricalmarkets.com"

def get_request_headers() -> dict:
    api_key = os.getenv("EM_API_KEY")
    if api_key is None:
        raise ValueError("API key not found. Please set the EM_API_KEY environment variable.")
    r = requests.post(
        f"{API_HOST}/v1/token",
        data={"api_key": api_key},
        headers={"Content-Type": "application/x-www-form-urlencoded"}
    )
    r.raise_for_status()
    token = r.json().get("access_token")
    request_headers = {"Authorization": "Bearer {}".format(token)}
    return request_headers

request_headers = get_request_headers()
print(str(request_headers)[:50] + "...}")

import time

import requests

def get_me(request_headers: dict) -> dict:
    r = requests.get(f"{API_HOST}/v1/me", headers=request_headers)
    r.raise_for_status()
    return r.json()

me_info = get_me(request_headers)

# Use the rate limit from the /me endpoint; fall back to 1 req/s (free tier)
_rate_limit_per_second = me_info.get("api_rate_limit_per_second") or 1
_min_interval = 1.0 / _rate_limit_per_second  # e.g. 1.0s for free, 0.33s for 3 req/s

def rate_limited_get(url: str, **kwargs) -> requests.Response:
    """Drop-in replacement for requests.get that respects the tier rate limit."""
    response = requests.get(url, **kwargs)
    time.sleep(_min_interval)
    return response

print(f"Built rate limit utility with allowed rate limit: {_rate_limit_per_second} req/s. The program will sleep {_min_interval:.2f}s between API calls.")

import pandas as pd

def get_sector_sentiment(sector: str) -> dict:
    """Fetch the historical sentiment data for a given sector."""
    r = rate_limited_get(f"{API_HOST}/v1/ai/sentiment/fed-market-sectors/dimensions/{sector}", headers=request_headers)
    r.raise_for_status()
    json_data = r.json()
    value_map = json_data["value_map"]
    df = pd.DataFrame(json_data["sentiment"])
    df["score_name"] = df["score"].astype(str).map(value_map)
    df["date"] = pd.to_datetime(df["date"])
    return df

def get_all_sector_sentiments(request_headers: dict) -> pd.DataFrame:
    """Query the list of available sectors, and fetch the historical sentiment data for each."""
    r = rate_limited_get(f"{API_HOST}/v1/ai/sentiment/fed-market-sectors", headers=request_headers)
    r.raise_for_status()
    fed_sector_sentiment_info = r.json()
    all_data = []
    for sector in fed_sector_sentiment_info["dimensions"]:
        print(f"Fetching sentiment data for sector: {sector}")
        df = get_sector_sentiment(sector)
        all_data.append(df)
    all_data = pd.concat(all_data, ignore_index=True)
    all_data = all_data.sort_values("date").reset_index(drop=True)
    return all_data.set_index("date")

sentiment_df = get_all_sector_sentiments(request_headers)
sentiment_df

from datetime import datetime
from time import sleep

def get_ticker_ohlc_data(
    ticker: str,
    request_headers: dict,
    start_date: datetime = datetime(1996, 1, 1),
    chunk_size: int = 1260,
) -> pd.DataFrame:
    print(f"Fetching OHLC data for ticker: {ticker}")
    frames = []
    params = {
        "limit": chunk_size,
        "start_date": start_date.strftime("%Y-%m-%d"),
    }

    while True:
        r = rate_limited_get(
            f"{API_HOST}/v1/ohlc/tickers/{ticker}",
            params=params,
            headers=request_headers,
        )
        sleep(0.5)
        r.raise_for_status()
        payload = r.json()
        data = payload["data"]
        if not data:
            break
        frames.append(pd.DataFrame(data))
        date_cursor = payload["date_cursor"]
        if not date_cursor["has_more"]:
            break

        # next page: swap start_date for the exclusive date_cursor
        params = {
            "limit": chunk_size,
            "date_cursor": date_cursor["next_cursor"],
        }

    if not frames:
        return pd.DataFrame()

    df = pd.concat(frames, ignore_index=True)
    df["date"] = pd.to_datetime(df["date"])
    return df.set_index("date").sort_index()


ohlc_data = {
    "SPY": get_ticker_ohlc_data("SPY", request_headers),
    "XLF": get_ticker_ohlc_data("XLF", request_headers),
    "XLE": get_ticker_ohlc_data("XLE", request_headers),
}
ohlc_data["SPY"]

import matplotlib.pyplot as plt
import numpy as np

plt.style.use("dark_background")

def _ffill_price_sentiment_df(df: pd.DataFrame) -> pd.DataFrame:
    """Forward fill."""
    df = df.sort_index()
    df["score"] = df["score"].ffill()
    df["close"] = df["close"].ffill()
    return df

def plot_price_vs_sector_sentiment(
    ticker: str,
    sector_name: str,
    sentiment_df: pd.DataFrame,
    ohlc_data: dict[str, pd.DataFrame]
):
    plot_df = (
        sentiment_df[sentiment_df["dimension"] == sector_name][["score", "score_name"]]
        .merge(
            ohlc_data[ticker]["close"],
            how="outer",
            left_index=True,
            right_index=True
        )
    )
    plot_df = _ffill_price_sentiment_df(plot_df).dropna()

    fig, ax = plt.subplots(figsize=(12, 8))
    np.log(plot_df["close"]).plot(ax=ax, label=f"{ticker} (Log)")

    ax2 = ax.twinx()
    # plot_df["score"].plot(ax=ax2, label="Sentiment Score", color="orange")
    ax2.bar(
        plot_df.index,
        plot_df["score"],
        width=pd.Timedelta(days=20),  # adjust if your sentiment dates are closer/farther apart
        alpha=0.5,
        label="Sentiment Score",
        color="orange",
        align="center",
    )
    ax.set_title(f"{sector_name} Sector Sentiment vs {ticker} (Log)")
    ax.set_ylabel(f"{ticker} (Log)")
    ax2.set_ylabel(f"FOMC {sector_name} Sentiment Score")

    h1, l1 = ax.get_legend_handles_labels()
    h2, l2 = ax2.get_legend_handles_labels()
    fig.legend(
        h1 + h2,
        l1 + l2,
        loc="lower center",
        bbox_to_anchor=(0.5, 0.93),
        ncol=2,
        columnspacing=1.0,
        handlelength=1.0,
    )
    plt.show()

plot_price_vs_sector_sentiment(ticker="SPY", sector_name="Overall", sentiment_df=sentiment_df, ohlc_data=ohlc_data)
plot_price_vs_sector_sentiment(ticker="XLF", sector_name="Financials", sentiment_df=sentiment_df, ohlc_data=ohlc_data)
plot_price_vs_sector_sentiment(ticker="XLE", sector_name="Energy", sentiment_df=sentiment_df, ohlc_data=ohlc_data)

def _sector_most_bearish_sentiment(df: pd.DataFrame, sector_name: str) -> dict:
    """Identify the date with the most bearish sentiment score for a given sector."""
    sector_df = df[df["dimension"] == sector_name].reset_index()
    most_bearish_row = sector_df.loc[sector_df["score"].idxmin()]
    return most_bearish_row.to_dict()

_sector_most_bearish_sentiment(sentiment_df, "Financials")

_sector_most_bearish_sentiment(sentiment_df, "Technology")

REGIME_LABEL_MAP = {1: "risk_on", 0: "neutral", -1: "risk_off"}


def join_price_history(ticker: str, sector_name: str, sentiment_df: pd.DataFrame, ohlc_data: dict[str, pd.DataFrame]) -> pd.DataFrame:
    """Join FOMC Sentiment to OHLC Prices."""
    joined_df = ohlc_data[ticker][["close"]].merge(
        sentiment_df[sentiment_df["dimension"] == sector_name],
        how="outer",
        left_index=True,
        right_index=True
     )
    joined_df = _ffill_price_sentiment_df(joined_df).dropna()
    return joined_df

def compute_regimes(
    df_price: pd.DataFrame,
    sector_name: str,
    ema_span: int = 30,
    z_window: int = 300,
    z_threshold: float = 2.,
) -> pd.DataFrame:
    """Calculate rolling sentiment z-score and create regime labels for a given sector."""
    sector_df = df_price[df_price["dimension"] == sector_name].copy()
    sector_df = sector_df.sort_index()

    # Smooth sentiment and standardize
    sector_df["score_smooth"] = sector_df["score"].ewm(span=ema_span, adjust=False).mean()
    roll_mean = sector_df["score_smooth"].rolling(window=z_window, min_periods=max(ema_span, 63)).mean()
    roll_std = sector_df["score_smooth"].rolling(window=z_window, min_periods=max(ema_span, 63)).std(ddof=0)
    sector_df["z"] = (sector_df["score_smooth"] - roll_mean) / roll_std

    # Regime classification
    sector_df["regime"] = np.select([sector_df["z"] >= z_threshold, sector_df["z"] <= -z_threshold], [1, -1], default=0)
    sector_df["regime_label"] = sector_df["regime"].map(REGIME_LABEL_MAP)
    return sector_df

def _shade_regimes(ax, index: pd.DatetimeIndex, regime: pd.Series):
    current, start = None, None
    for t, r in regime.items():
        if current is None:
            current, start = r, t
        elif r != current:
            color = "green" if current == 1 else ("red" if current == -1 else "gray")
            alpha = 0.2 if color == "gray" else 0.5
            ax.axvspan(start, t, color=color, alpha=alpha, lw=0)
            current, start = r, t
    if current is not None and len(index):
        color = "green" if current == 1 else ("red" if current == -1 else "gray")
        alpha = 0.2 if color == "gray" else 0.5
        ax.axvspan(start, index[-1], color=color, alpha=alpha, lw=0)

def plot_price_with_regimes(sector_df: pd.DataFrame, ticker: str, sector_name: str):
    """
    sector_df: output of compute_regimes
    """
    sector_df = sector_df.dropna(subset=["z"])
    fig, ax = plt.subplots(figsize=(12, 7))
    np.log(sector_df["close"]).plot(ax=ax, label=f"${ticker.upper()} (Log Close)")
    # z-score on secondary axis
    ax2 = ax.twinx()
    sector_df["z"].clip(-3, 3).plot(ax=ax2, color="orange", alpha=0.7, label="FOMC Sentiment z-score")

    _shade_regimes(ax, sector_df.index, sector_df["regime"])

    ax.set_title(f"{sector_name} Sentiment Regimes vs ${ticker.upper()}")
    ax.set_ylabel(f"${ticker.upper()} (log)")
    ax2.set_ylabel("Sentiment z-score")

    h1, l1 = ax.get_legend_handles_labels()
    h2, l2 = ax2.get_legend_handles_labels()
    fig.legend(h1 + h2, l1 + l2, loc="lower center", bbox_to_anchor=(0.5, 0.93), ncol=2)
    plt.show()

# overall sentiment vs spy
overall = compute_regimes(
    join_price_history("SPY", "Overall", sentiment_df, ohlc_data),
    sector_name="Overall",
    ema_span=20,
    z_window=252,
    z_threshold=1.5
)
plot_price_with_regimes(overall, ticker="SPY", sector_name="Overall")

# financials sentiment vs xlf
financials = compute_regimes(
    join_price_history("XLF", "Financials", sentiment_df, ohlc_data),
    sector_name="Financials",
    ema_span=20,
    z_window=252,
    z_threshold=1.5
)
plot_price_with_regimes(financials, ticker="XLF", sector_name="Financials")


# energy sentiment vs xle
energy = compute_regimes(
    join_price_history("XLE", "Energy", sentiment_df, ohlc_data),
    sector_name="Energy",
    ema_span=20,
    z_window=252,
    z_threshold=1.5
)
plot_price_with_regimes(energy, ticker="XLE", sector_name="Energy")