"""Earnings Model v1.0 — Earnings Event Intelligence
Detects earnings proximity, estimates implied move from options, 
analyzes post-earnings drift patterns, and adjusts position sizing.
Based on: Ball & Brown (1968), Frazzini & Lamont (2007) PEAD
"""
import yfinance as yf
import numpy as np
import pandas as pd
from datetime import datetime, timedelta
from typing import Dict, Optional, Tuple

# Typical quarterly earnings windows by month (approximate)
EARNINGS_CALENDAR = {
    # Q1 (Jan-Mar): reports Apr-May
    'Q1':  {'months': [4, 5], 'label': 'Q1'},
    # Q2 (Apr-Jun): reports Jul-Aug
    'Q2':  {'months': [7, 8], 'label': 'Q2'},
    # Q3 (Jul-Sep): reports Oct-Nov
    'Q3':  {'months': [10, 11], 'label': 'Q3'},
    # Q4 (Oct-Dec): reports Jan-Feb
    'Q4':  {'months': [1, 2], 'label': 'Q4'},
}

# Sector-specific typical implied moves (based on historical options data)
SECTOR_IMPLIED_MOVES = {
    'Technology':        0.045,
    'Healthcare':        0.040,
    'Financials':        0.030,
    'Energy':            0.055,
    'Consumer Discretionary': 0.050,
    'Consumer Staples': 0.025,
    'Industrials':       0.035,
    'Communication':     0.045,
    'Utilities':         0.020,
    'Materials':         0.045,
    'Real Estate':       0.030,
    'default':           0.040,
}


class EarningsModel:
    """Earnings event intelligence for quant trading."""

    def __init__(self):
        self._cache = {}

    def estimate_earnings_date(self, ticker: str) -> Optional[datetime.date]:
        """Estimate next earnings date from yfinance calendar."""
        try:
            t = yf.Ticker(ticker)
            cal = t.calendar
            if cal is not None and not cal.empty:
                # yfinance calendar returns next earnings date
                if hasattr(cal, 'index') and 'Earnings Date' in cal.index:
                    date_str = cal.loc['Earnings Date'].values[0]
                    if isinstance(date_str, str):
                        return datetime.strptime(date_str, '%Y-%m-%d').date()
                elif 'Earnings Date' in cal.columns:
                    date_str = cal['Earnings Date'].iloc[0]
                    if isinstance(date_str, str):
                        return datetime.strptime(date_str, '%Y-%m-%d').date()
            # Fallback: estimate from current month
            today = datetime.now()
            for q, data in EARNINGS_CALENDAR.items():
                if today.month in data['months']:
                    # Next likely report window: mid-month
                    return datetime(today.year, today.month, 15).date()
            return None
        except Exception:
            return None

    def days_to_earnings(self, ticker: str) -> Tuple[Optional[int], Optional[datetime.date]]:
        """Return (days_until, date)."""
        ed = self.estimate_earnings_date(ticker)
        if ed is None:
            return None, None
        today = datetime.now().date()
        delta = (ed - today).days
        return delta, ed

    def implied_move(self, ticker: str, default_move: float = None) -> Dict:
        """Estimate implied earnings move from options chain if available.
        Fallback to sector average or default.
        """
        move = default_move or 0.04  # 4% default
        source = 'default'

        try:
            t = yf.Ticker(ticker)
            # Try to get implied move from near-the-money straddle
            # Find the closest expiration before/after estimated earnings
            expiry = t.options
            if expiry and len(expiry) > 0:
                # Get first expiration
                opt = t.option_chain(expiry[0])
                calls = opt.calls
                puts = opt.puts

                if len(calls) > 0 and len(puts) > 0:
                    # ATM strike
                    spot = calls['strike'].iloc[len(calls)//2]
                    atm_call = calls[calls['strike'].sub(spot).abs().idxmin()]
                    atm_put = puts[puts['strike'].sub(spot).abs().idxmin()]

                    # Straddle price
                    call_p = atm_call['lastPrice'] if 'lastPrice' in atm_call else atm_call['ask']
                    put_p = atm_put['lastPrice'] if 'lastPrice' in atm_put else atm_put['ask']
                    straddle = float(call_p) + float(put_p)

                    # Implied move as % of stock price
                    current_price = yf.Ticker(ticker).history(period='1d')['Close'].iloc[-1]
                    move = straddle / current_price
                    source = 'options_chain'

                    # Annualize: if 1 month until expiry
                    days = 30  # approximate
                    move_annual = move * np.sqrt(365 / max(1, days))

                    return {
                        'implied_move_pct': round(move * 100, 2),
                        'annualized_pct': round(move_annual * 100, 2),
                        'straddle_price': round(float(straddle), 2),
                        'source': source,
                        'expiry': expiry[0],
                        'atm_strike': float(spot),
                    }
        except Exception as e:
            pass

        # Try sector-based
        try:
            info = yf.Ticker(ticker).info
            sector = info.get('sector', '')
            for sector_name, sector_move in SECTOR_IMPLIED_MOVES.items():
                if sector_name.lower() in sector.lower():
                    move = sector_move
                    source = f'sector_{sector_name}'
                    break
        except:
            pass

        return {
            'implied_move_pct': round(move * 100, 2),
            'annualized_pct': round(move * np.sqrt(12) * 100, 2),
            'source': source,
        }

    def historical_pead(self, ticker: str, lookback_years: int = 3) -> Dict:
        """Post-Earnings Announcement Drift analysis.
        Measures how much stock drifts in direction of surprise after earnings.
        Returns drift score: positive = bullish drift tendency, negative = bearish.
        """
        try:
            # Fetch enough history
            df = yf.Ticker(ticker).history(period=f"{lookback_years}y")
            if len(df) < 100:
                return {'drift_score': 0, 'confidence': 0, 'n_events': 0}

            # Detect earnings dates (volume spikes on quarterly frequency)
            df['volume_z'] = (df['Volume'] - df['Volume'].rolling(20).mean()) / df['Volume'].rolling(20).std()
            # Find volume spike days
            spike_days = df[df['volume_z'] > 2.5].index

            if len(spike_days) < 4:
                return {'drift_score': 0, 'confidence': 0, 'n_events': 0}

            # Analyze returns post-spike
            drifts = []
            for spike in spike_days:
                try:
                    idx = df.index.get_loc(spike)
                    if idx + 5 >= len(df): continue
                    # Day +1 to +5 return
                    post_ret = (df['Close'].iloc[idx+5] / df['Close'].iloc[idx]) - 1
                    # Day 0 overnight surprise (gap)
                    overnight_gap = (df['Open'].iloc[idx] / df['Close'].iloc[idx-1]) - 1 if idx > 0 else 0
                    # PEAD: continuation of surprise
                    drift = np.sign(overnight_gap) * post_ret
                    drifts.append(drift)
                except:
                    continue

            if len(drifts) < 3:
                return {'drift_score': 0, 'confidence': 0, 'n_events': len(spike_days)}

            drift_score = np.mean(drifts)
            confidence = min(1.0, len(drifts) / 12)  # More events = more confidence

            return {
                'drift_score': round(float(drift_score), 4),
                'confidence': round(float(confidence), 2),
                'n_events': len(drifts),
                'avg_overnight_gap': round(float(np.mean([d for d in drifts])), 4) if drifts else 0,
                'interpretation': (
                    'Positive PEAD: earnings surprises tend to continue' if drift_score > 0.02 else
                    'Negative PEAD: post-earnings reversals typical' if drift_score < -0.02 else
                    'Weak PEAD pattern'
                ),
            }

        except Exception:
            return {'drift_score': 0, 'confidence': 0, 'n_events': 0}

    def earnings_position_size(self, base_size: float, days_to_earnings: int,
                              implied_move: float = 0.04) -> Dict:
        """Adjust position size for earnings proximity.
        
        Strategy:
        - D-30 to D-7: Full size (can position for earnings)
        - D-7 to D-3: Reduce to 50% (avoid theta decay, lock profits)
        - D-3 to D-1: Reduce to 25% (extreme event risk)
        - D-Day: 0% (do not hold into earnings)
        - D+1 to D+5: Can re-enter at 50% (capture PEAD)
        """
        if days_to_earnings is None:
            return {
                'adjusted_size': base_size,
                'reduction': 0.0,
                'strategy': 'No earnings date detected — normal sizing',
            }

        if days_to_earnings > 7:
            adj = base_size
            strategy = 'Pre-earnings positioning window — full size'
        elif days_to_earnings > 3:
            adj = base_size * 0.5
            strategy = 'Reduce to 50% — earnings week risk building'
        elif days_to_earnings > 0:
            adj = base_size * 0.25
            strategy = 'Reduce to 25% — imminent earnings, extreme theta'
        elif days_to_earnings == 0:
            adj = 0.0
            strategy = 'DO NOT HOLD INTO EARNINGS — day-of closure'
        elif days_to_earnings >= -1:
            adj = base_size * 0.5
            strategy = 'Post-earnings PEAD window — 50% re-entry'
        elif days_to_earnings >= -5:
            adj = base_size * 0.5
            strategy = 'PEAD continuation — 50% size'
        else:
            adj = base_size
            strategy = 'Post-earnings quiet period — normal sizing'

        # Adjust for implied move magnitude
        if implied_move > 0.08:  # > 8% expected move
            adj *= 0.7
            strategy += ' | High implied move (>8%), further reduced'
        elif implied_move > 0.06:
            adj *= 0.85
            strategy += ' | Elevated implied move, slight reduction'

        return {
            'days_to_earnings': days_to_earnings,
            'implied_move_pct': round(implied_move * 100, 2),
            'base_size': round(base_size, 4),
            'adjusted_size': round(adj, 4),
            'reduction': round(1 - (adj / (base_size + 1e-10)), 2),
            'strategy': strategy,
        }

    def full_analysis(self, ticker: str, base_size: float = 1.0) -> Dict:
        """Complete earnings intelligence for a ticker."""
        days, ed = self.days_to_earnings(ticker)
        implied = self.implied_move(ticker)
        pead = self.historical_pead(ticker)

        sizing = self.earnings_position_size(
            base_size, days,
            implied.get('implied_move_pct', 4) / 100
        )

        return {
            'ticker': ticker,
            'estimated_earnings_date': ed.strftime('%Y-%m-%d') if ed else 'unknown',
            'days_to_earnings': days,
            'implied_move': implied,
            'pead_analysis': pead,
            'position_sizing': sizing,
            'recommendation': sizing['strategy'],
        }


if __name__ == '__main__':
    model = EarningsModel()
    result = model.full_analysis('AAPL', base_size=1.0)
    print(f"Estimated Earnings: {result['estimated_earnings_date']}")
    print(f"Days Until: {result['days_to_earnings']}")
    print(f"Implied Move: {result['implied_move'].get('implied_move_pct', 'N/A')}%")
    print(f"PEAD Score: {result['pead_analysis'].get('drift_score', 0):.4f}")
    print(f"Position Sizing: {result['position_sizing']['adjusted_size']*100:.0f}% ({result['position_sizing']['strategy']})")