from AlgorithmImports import *
from scipy.optimize import curve_fit
from scipy.stats import linregress


class SmartInsiderCorporateBuybacksAlgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(self.end_date - timedelta(5*365))
        self.set_cash(100_000)
        self.set_warm_up(252 * 5, Resolution.DAILY)
        self._min_ic = 0.05
        self._min_expected_return = 0.005
        self.set_portfolio_construction(InsightWeightingPortfolioConstructionModel(timedelta(days=252)))        
        # Equities from research: MMM, GS, HON, JNJ, MSFT, TRV, V
        tickers = ["MMM", "GS", "HON", "JNJ", "MSFT", "TRV", "V"]
        self._equities = []
        self._equity_by_buyback_symbol = {}
        # Attach session tracking, buyback symbol, and buyback data storage to each security via duck typing
        for ticker in tickers:
            equity = self.add_equity(ticker)
            equity.session.size = 252 * 5  # Track 5 years of daily sessions
            buyback_symbol = self.add_data(SmartInsiderTransaction, equity).symbol
            equity.buyback_data = {}  # Store buyback data keyed by date     
            self._equities.append(equity)
            self._equity_by_buyback_symbol[buyback_symbol] = equity
        # Add a Scheduled Event to rebalance the portfolio.
        self.schedule.on(
            self.date_rules.every_day("MMM"), 
            self.time_rules.at(10, 0), 
            self._rebalance
        )
    
    def on_data(self, data):
        """Process incoming transaction data and store it."""
        for dataset_symbol, transaction in data.get(SmartInsiderTransaction).items():
            equity = self._equity_by_buyback_symbol[dataset_symbol]
            transaction_date = transaction.end_time.date()
            buyback_pct = transaction.usd_value / transaction.usd_market_cap
            # Aggregate buybacks by date.
            if transaction_date in equity.buyback_data:
                equity.buyback_data[transaction_date] += buyback_pct
            else:
                equity.buyback_data[transaction_date] = buyback_pct
    
    def _rebalance(self):
        """Check entry signals and emit insights."""
        if self.is_warming_up:
            return
        insights = []
        for equity in self._equities:
            # Get processed data for IC and expectation calculation, plus latest buyback.
            result = self._get_data(equity)
            if not result:
                continue
            df, latest_buyback = result 
            # Calculate the information coefficient.
            ic = df.corr().values[0, 1]
            if ic < self._min_ic:
                continue            
            # Calculate expected return using latest buyback transaction.
            expectation = self._get_expectation(df, latest_buyback)
            if expectation <= self._min_expected_return:
                continue
            # Emit insight for entry.
            insights.append(
                Insight.price(equity, timedelta(252), InsightDirection.UP, weight=expectation*ic)
            )
        if insights:
            self.emit_insights(insights)
    
    def _get_data(self, equity):
        """Get processed dataframe for IC and expectation calculation.        
        Args:
            equity: The equity Security to get the processed dataframe for.            
        Returns:
            Tuple of (DataFrame with 1-year forward returns and buyback percentages, latest buyback percentage),
            or None if insufficient data.
        """
        if not equity.buyback_data:
            return
        # Calculate 1-year forward returns.
        history = pd.Series(
            [x.close for x in equity.session], 
            index=[x.end_time.date() for x in equity.session]
        ).sort_index().pct_change(252).shift(-252).dropna()        
        # Convert stored buyback data to a Series.
        buyback_aggregated = pd.Series(equity.buyback_data).sort_index()
        # Extract latest buyback percentage for expectation calculation.
        latest_buyback = np.log(buyback_aggregated.iloc[-1])
        # Concatenate dataframes keeping only rows with data
        df = pd.concat([history, np.log(buyback_aggregated)], axis=1).replace([np.inf, -np.inf], np.nan).dropna()        
        if df.empty:
            return
        return (df, latest_buyback)
    
    def _exponential_curve(self, x, a, b, c):
        """Exponential decay curve function for buyback analysis."""
        return a * np.exp(-b * x) + c
    
    def _get_expectation(self, df, transaction):
        """Calculate expected return given buyback percentage.        
        Args:
            df: DataFrame with historical data for curve fitting.
            transaction: Current buyback percentage (log-scaled).            
        Returns:
            Expected return percentage, or -1 if curve fitting fails.
        """
        x = df.iloc[:, 1].values
        y = df.iloc[:, 0].values
        # Try exponential curve fitting.
        try:            
            popt, _ = curve_fit(
                self._exponential_curve,
                x, y,
                p0=[np.mean(y), 1.0, 0.0],
                bounds=([-np.inf, 0, -np.inf], [np.inf, np.inf, np.inf]),
                maxfev=10_000
            )
            return self._exponential_curve(transaction, *popt)
        except:
            pass
        # Fallback to simple linear regression.
        try:
            slope, intercept, _, _, _ = linregress(x, y)
            return slope * transaction + intercept
        except:
            return -1