!pip install pandas numpy matplotlib seaborn scipy --quiet

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import seaborn as sns
from scipy.signal import argrelextrema
from scipy.stats import pearsonr
from datetime import datetime, timezone
import warnings

warnings.filterwarnings('ignore')
%matplotlib inline
plt.rcParams['figure.figsize'] = (14, 5)
plt.rcParams['axes.spines.top'] = False
plt.rcParams['axes.spines.right'] = False
sns.set_palette('husl')
print('Imports ready.')

# ── CONFIGURATION ─────────────────────────────────────────────────────────────
DIVERGENCE_WINDOW    = 14   # days to look back when measuring divergence
CORRELATION_WINDOW   = 7    # rolling correlation window in days
DIVERGENCE_THRESHOLD = -0.4 # rolling correlation below this → divergence detected
MIN_HOLDING_DAYS     = 3    # minimum days to hold a divergence signal
SIGNAL_FILE          = 'sentiment_signal.csv'   # from Notebook 121
# ─────────────────────────────────────────────────────────────────────────────


def generate_synthetic_data(n_days: int = 120, seed: int = 42) -> tuple:
    """
    Generate synthetic price and sentiment data with built-in divergence episodes.

    Parameters
    ----------
    n_days : int
        Number of days to simulate.
    seed : int
        Random seed.

    Returns
    -------
    tuple : (price_df, sentiment_series)
        price_df        : pd.DataFrame with OHLCV data
        sentiment_series: pd.Series with daily composite sentiment [-1, +1]

    Notes
    -----
    Three divergence episodes are embedded:
    - Days 20-35: bearish divergence (price up, sentiment down)
    - Days 60-75: bullish divergence (price down, sentiment up)
    - Days 90-105: bearish divergence
    """
    np.random.seed(seed)
    dates = pd.date_range(end=datetime.now(tz=timezone.utc).date(), periods=n_days, freq='D')

    # Base price process
    log_ret = np.random.randn(n_days) * 0.02
    # Base sentiment correlated with price
    sent = np.random.randn(n_days) * 0.15

    # Inject divergence episodes: sentiment moves opposite to price
    for start, end, direction in [(20, 35, -1), (60, 75, 1), (90, 105, -1)]:
        trend = np.linspace(0, direction * 0.04 * (end - start), end - start)
        log_ret[start:end] += trend * direction         # price moves one way
        sent[start:end]    -= trend * direction * 2     # sentiment moves opposite

    close = 40000 * np.exp(np.cumsum(log_ret))
    sentiment = np.clip(np.cumsum(sent * 0.1), -1, 1)

    price_df = pd.DataFrame({
        'open':  close * (1 + np.random.randn(n_days) * 0.004),
        'high':  close * (1 + np.abs(np.random.randn(n_days) * 0.008)),
        'low':   close * (1 - np.abs(np.random.randn(n_days) * 0.008)),
        'close': close,
        'volume': np.random.exponential(1e9, n_days)
    }, index=dates)

    sentiment_series = pd.Series(sentiment, index=dates, name='sentiment')
    print(f'Generated {n_days} days of synthetic data with embedded divergence episodes.')
    return price_df, sentiment_series


def load_real_data(signal_file: str, ohlcv_file: str) -> tuple:
    """
    Load real sentiment signal and OHLCV data from CSV files.

    Parameters
    ----------
    signal_file : str
        Path to sentiment_signal.csv from Notebook 121.
    ohlcv_file : str
        Path to OHLCV CSV with at minimum a 'close' column and date index.

    Returns
    -------
    tuple : (price_df, sentiment_series)
    """
    sig = pd.read_csv(signal_file, parse_dates=['date']).set_index('date')
    sentiment_series = sig['ema_fast'].rename('sentiment')
    price_df = pd.read_csv(ohlcv_file, parse_dates=[0], index_col=0)
    common = price_df.index.intersection(sentiment_series.index)
    return price_df.loc[common], sentiment_series.loc[common]


# ── Load data ─────────────────────────────────────────────────────────────────
USE_SYNTHETIC = True
if USE_SYNTHETIC:
    price_df, sentiment_series = generate_synthetic_data(n_days=120)
else:
    price_df, sentiment_series = load_real_data(SIGNAL_FILE, 'ohlcv.csv')

def compute_rolling_correlation(
    price_df: pd.DataFrame,
    sentiment: pd.Series,
    window: int = 7
) -> pd.Series:
    """
    Compute rolling Pearson correlation between daily price returns and sentiment changes.

    Parameters
    ----------
    price_df : pd.DataFrame
        OHLCV DataFrame with at minimum a 'close' column.
    sentiment : pd.Series
        Daily composite sentiment score aligned to price_df's index.
    window : int
        Rolling window size in days.

    Returns
    -------
    pd.Series
        Rolling correlation in [-1, +1]. Values near -1 indicate
        strong divergence (price and sentiment moving oppositely).
    """
    price_ret  = price_df['close'].pct_change()
    sent_delta = sentiment.diff()
    rolling_corr = price_ret.rolling(window).corr(sent_delta)
    return rolling_corr.rename('rolling_corr')


def compute_zscore_spread(
    price_df: pd.DataFrame,
    sentiment: pd.Series,
    window: int = 20
) -> pd.Series:
    """
    Compute the z-score spread between normalized price and sentiment.

    Parameters
    ----------
    price_df : pd.DataFrame
        OHLCV DataFrame.
    sentiment : pd.Series
        Daily sentiment score.
    window : int
        Lookback window for rolling z-score normalization.

    Returns
    -------
    pd.Series
        Spread series. Large positive values → price above sentiment (bearish divergence).
        Large negative values → price below sentiment (bullish divergence).
    """
    def rolling_zscore(s, w):
        mean = s.rolling(w).mean()
        std  = s.rolling(w).std().clip(lower=1e-6)
        return (s - mean) / std

    price_z = rolling_zscore(price_df['close'], window)
    sent_z  = rolling_zscore(sentiment, window)
    return (price_z - sent_z).rename('zscore_spread')


def detect_divergence_signal(
    rolling_corr: pd.Series,
    zscore_spread: pd.Series,
    corr_threshold: float = -0.4,
    spread_threshold: float = 1.5
) -> pd.DataFrame:
    """
    Classify each day as a bullish divergence, bearish divergence, or no divergence.

    Parameters
    ----------
    rolling_corr : pd.Series
        Rolling correlation from compute_rolling_correlation().
    zscore_spread : pd.Series
        Z-score spread from compute_zscore_spread().
    corr_threshold : float
        Correlation must be below this to qualify as divergence (should be negative).
    spread_threshold : float
        Absolute z-score spread must exceed this to confirm divergence strength.

    Returns
    -------
    pd.DataFrame
        Columns: rolling_corr, zscore_spread, divergence_type, signal_numeric.
        divergence_type: 'bullish', 'bearish', or 'none'.
        signal_numeric: +1 (bullish), -1 (bearish), 0 (none).

    Notes
    -----
    Bullish divergence: correlation is low AND spread is strongly negative
    (sentiment is above price on z-score scale — crowd more optimistic than price implies).
    Bearish divergence: correlation is low AND spread is strongly positive.
    """
    div_df = pd.DataFrame({
        'rolling_corr':  rolling_corr,
        'zscore_spread': zscore_spread,
    })
    low_corr = div_df['rolling_corr'] < corr_threshold
    strong_negative_spread = div_df['zscore_spread'] < -spread_threshold
    strong_positive_spread = div_df['zscore_spread'] >  spread_threshold

    conditions = [
        low_corr & strong_negative_spread,
        low_corr & strong_positive_spread,
    ]
    div_df['divergence_type'] = np.select(conditions, ['bullish', 'bearish'], default='none')
    div_df['signal_numeric']  = np.select(conditions, [1, -1], default=0)

    bull_days = (div_df['divergence_type'] == 'bullish').sum()
    bear_days = (div_df['divergence_type'] == 'bearish').sum()
    print(f'Divergence detected: {bull_days} bullish days, {bear_days} bearish days')
    return div_df


# ── Run detection ─────────────────────────────────────────────────────────────
rolling_corr  = compute_rolling_correlation(price_df, sentiment_series, window=CORRELATION_WINDOW)
zscore_spread = compute_zscore_spread(price_df, sentiment_series, window=DIVERGENCE_WINDOW)
div_df        = detect_divergence_signal(rolling_corr, zscore_spread,
                                          DIVERGENCE_THRESHOLD, spread_threshold=1.2)

def plot_divergence_overview(
    price_df: pd.DataFrame,
    sentiment: pd.Series,
    div_df: pd.DataFrame
) -> None:
    """
    Four-panel chart: price, sentiment, rolling correlation, and divergence signal.

    Parameters
    ----------
    price_df : pd.DataFrame
        OHLCV DataFrame.
    sentiment : pd.Series
        Daily composite sentiment score.
    div_df : pd.DataFrame
        Output of detect_divergence_signal().
    """
    fig, axes = plt.subplots(4, 1, figsize=(14, 16), sharex=True)

    # Panel 1: Price
    ax = axes[0]
    ax.plot(price_df.index, price_df['close'], color='navy', linewidth=1.5)
    bull_mask = div_df['divergence_type'] == 'bullish'
    bear_mask = div_df['divergence_type'] == 'bearish'
    ax.fill_between(price_df.index, price_df['close'].min(), price_df['close'],
                    where=bull_mask, alpha=0.2, color='green', label='Bullish div.')
    ax.fill_between(price_df.index, price_df['close'].min(), price_df['close'],
                    where=bear_mask, alpha=0.2, color='red',   label='Bearish div.')
    ax.set_title('Price with Divergence Episodes Highlighted')
    ax.set_ylabel('Price (USD)')
    ax.legend(fontsize=9)

    # Panel 2: Sentiment
    ax = axes[1]
    ax.plot(sentiment.index, sentiment, color='steelblue', linewidth=1.5)
    ax.axhline(0, color='black', linewidth=0.5, linestyle='--')
    ax.set_title('Composite Sentiment Score')
    ax.set_ylabel('Score [-1, +1]')

    # Panel 3: Rolling Correlation
    ax = axes[2]
    ax.plot(div_df.index, div_df['rolling_corr'], color='purple', linewidth=1.5)
    ax.axhline(DIVERGENCE_THRESHOLD, color='red', linewidth=0.8, linestyle='--', label=f'Threshold ({DIVERGENCE_THRESHOLD})')
    ax.axhline(0, color='black', linewidth=0.5)
    ax.set_title('Rolling Price-Sentiment Correlation')
    ax.set_ylabel('Pearson r')
    ax.legend(fontsize=9)

    # Panel 4: Z-score spread
    ax = axes[3]
    colors = div_df['divergence_type'].map({'bullish': 'green', 'bearish': 'red', 'none': 'grey'})
    ax.bar(div_df.index, div_df['zscore_spread'], color=colors, alpha=0.7)
    ax.axhline(0, color='black', linewidth=0.5)
    ax.set_title('Z-Score Spread (Price − Sentiment)')
    ax.set_ylabel('Z-Score')
    bull_patch = mpatches.Patch(color='green', alpha=0.7, label='Bullish')
    bear_patch = mpatches.Patch(color='red',   alpha=0.7, label='Bearish')
    ax.legend(handles=[bull_patch, bear_patch], fontsize=9)

    plt.tight_layout()
    plt.show()


plot_divergence_overview(price_df, sentiment_series, div_df)

def backtest_divergence_signal(
    price_df: pd.DataFrame,
    div_df: pd.DataFrame,
    holding_period: int = 5
) -> pd.DataFrame:
    """
    Backtest a divergence-triggered entry strategy.

    Parameters
    ----------
    price_df : pd.DataFrame
        OHLCV DataFrame.
    div_df : pd.DataFrame
        Output of detect_divergence_signal().
    holding_period : int
        Number of days to hold a position after a divergence signal fires.

    Returns
    -------
    pd.DataFrame
        Trade log with entry_date, exit_date, direction, entry_price, exit_price, return_pct.

    Notes
    -----
    Entry rule: on the first day of a divergence episode.
    Exit rule:  after holding_period days OR when divergence ends (whichever comes first).
    Direction:  +1 for bullish divergence entry, -1 for bearish.
    """
    close = price_df['close']
    signal = div_df['signal_numeric']

    trades = []
    in_trade = False
    entry_idx = None
    direction = 0

    for i, (date, sig) in enumerate(signal.items()):
        if not in_trade and sig != 0:
            in_trade   = True
            entry_idx  = i
            direction  = sig
            entry_price = close.iloc[i]
        elif in_trade:
            days_held = i - entry_idx
            if days_held >= holding_period or sig == 0:
                exit_price = close.iloc[i]
                ret = direction * (exit_price / entry_price - 1)
                trades.append({
                    'entry_date':  close.index[entry_idx],
                    'exit_date':   date,
                    'direction':   'long' if direction == 1 else 'short',
                    'entry_price': round(entry_price, 2),
                    'exit_price':  round(exit_price, 2),
                    'return_pct':  round(ret * 100, 2),
                })
                in_trade = False

    trade_df = pd.DataFrame(trades)
    if len(trade_df):
        win_rate = (trade_df['return_pct'] > 0).mean()
        avg_ret  = trade_df['return_pct'].mean()
        print(f'Trades: {len(trade_df)} | Win rate: {win_rate:.0%} | Avg return: {avg_ret:.2f}%')
        print(trade_df.to_string(index=False))
    else:
        print('No trades generated. Try lowering divergence thresholds.')
    return trade_df


trade_log = backtest_divergence_signal(price_df, div_df, holding_period=MIN_HOLDING_DAYS)

def export_divergence_results(
    div_df: pd.DataFrame,
    trade_log: pd.DataFrame,
    prefix: str = 'sentiment_divergence'
) -> None:
    """
    Export divergence signal and trade log to CSV files.

    Parameters
    ----------
    div_df : pd.DataFrame
        Daily divergence signal DataFrame.
    trade_log : pd.DataFrame
        Trade log from backtest_divergence_signal().
    prefix : str
        Filename prefix for output files.
    """
    signal_out = f'{prefix}_signal.csv'
    trades_out = f'{prefix}_trades.csv'
    div_df.to_csv(signal_out)
    if len(trade_log):
        trade_log.to_csv(trades_out, index=False)
    print(f'Exported: {signal_out}')
    print(f'Exported: {trades_out}')


export_divergence_results(div_df, trade_log)