engine/friction_matrix.py

import logging
from typing import Dict, Any, List
import pandas as pd
import numpy as np
import vectorbt as vbt

logger = logging.getLogger("bqe.engine.friction_matrix")

class IndianFrictionMatrix:
    """
    Implements statutory transaction fee calculations for the Indian markets (NSE/BSE).
    
    Supports:
    - Brokerage
    - Securities Transaction Tax (STT)
    - Stamp Duty
    - Exchange Transaction Charges
    - SEBI Turnover Fees
    - Goods and Services Tax (GST)
    - Execution Slippage (Market Impact)
    """
    
    def __init__(self, slippage_pct: float = 0.05):
        """
        Initializes the friction matrix.
        
        Args:
            slippage_pct: Constant execution slippage rate per leg (default 0.05%).
        """
        self.slippage_pct = slippage_pct / 100.0 # convert to fraction

    def calculate_delivery_friction(self, trade_records: pd.DataFrame) -> pd.DataFrame:
        """
        Parses a DataFrame of individual trade executions and computes itemized charges.
        
        Statutory parameters (Equity Delivery):
        - Brokerage: min(20, 0.0003 * trade_value) per leg
        - STT: 0.1% of trade value on both entry and exit legs
        - Stamp Duty: 0.015% of trade value on the entry (buy) leg only
        - Exchange Charges: 0.00322% of trade value per leg
        - SEBI Fee: 0.00001% of trade value per leg
        - GST: 18% applied on (Brokerage + Exchange Charges + SEBI Fee)
        - Slippage: Constant execution slippage per leg
        
        Args:
            trade_records: DataFrame from portfolio.trades.records_readable.
            
        Returns:
            pd.DataFrame: Copy of DataFrame with added itemized friction columns.
        """
        if trade_records.empty:
            logger.warning("Empty trade records provided to IndianFrictionMatrix.")
            return trade_records.copy()
            
        df = trade_records.copy()
        
        # Resolve column names dynamically (handling vectorbt version naming variations)
        entry_price_col = next((c for c in ['Avg Entry Price', 'Entry Price'] if c in df.columns), 'Price')
        exit_price_col = next((c for c in ['Avg Exit Price', 'Exit Price'] if c in df.columns), entry_price_col)
        
        size = df['Size']
        entry_price = df[entry_price_col]
        exit_price = df[exit_price_col].fillna(entry_price)
        
        entry_value = size * entry_price
        exit_value = size * exit_price
        
        # Calculate entry leg costs
        entry_brokerage = np.minimum(20.0, 0.0003 * entry_value)
        entry_stt = 0.001 * entry_value
        entry_stamp_duty = 0.00015 * entry_value
        entry_exchange = 0.0000322 * entry_value
        entry_sebi = 0.0000001 * entry_value
        entry_gst = 0.18 * (entry_brokerage + entry_exchange + entry_sebi)
        entry_slippage = self.slippage_pct * entry_value
        entry_friction = (
            entry_brokerage + entry_stt + entry_stamp_duty + 
            entry_exchange + entry_sebi + entry_gst + entry_slippage
        )
        
        # Calculate exit leg costs
        exit_brokerage = np.minimum(20.0, 0.0003 * exit_value)
        exit_stt = 0.001 * exit_value
        exit_stamp_duty = 0.0  # Stamp duty is only charged on buying (entry)
        exit_exchange = 0.0000322 * exit_value
        exit_sebi = 0.0000001 * exit_value
        exit_gst = 0.18 * (exit_brokerage + exit_exchange + exit_sebi)
        exit_slippage = self.slippage_pct * exit_value
        exit_friction = (
            exit_brokerage + exit_stt + exit_stamp_duty + 
            exit_exchange + exit_sebi + exit_gst + exit_slippage
        )
        
        # Add entry columns to DataFrame
        df['entry_brokerage'] = entry_brokerage
        df['entry_stt'] = entry_stt
        df['entry_stamp_duty'] = entry_stamp_duty
        df['entry_exchange_charges'] = entry_exchange
        df['entry_sebi_fee'] = entry_sebi
        df['entry_gst'] = entry_gst
        df['entry_slippage'] = entry_slippage
        df['entry_friction'] = entry_friction
        
        # Add exit columns to DataFrame
        df['exit_brokerage'] = exit_brokerage
        df['exit_stt'] = exit_stt
        df['exit_stamp_duty'] = exit_stamp_duty
        df['exit_exchange_charges'] = exit_exchange
        df['exit_sebi_fee'] = exit_sebi
        df['exit_gst'] = exit_gst
        df['exit_slippage'] = exit_slippage
        df['exit_friction'] = exit_friction
        
        # Calculate totals
        df['total_friction'] = entry_friction + exit_friction
        df['net_pnl'] = df['PnL'] - df['total_friction']
        
        return df

    def apply_friction_to_portfolio(
        self, 
        portfolio: vbt.Portfolio, 
        close_matrix: pd.DataFrame
    ) -> Dict[str, Any]:
        """
        Extracts trades, applies friction calculations, and recursively adjusts daily NAV.
        
        Args:
            portfolio: The executed vectorbt Portfolio object.
            close_matrix: The Close price matrix used in the backtest.
            
        Returns:
            Dict: Net performance metrics payload.
        """
        logger.info("Extracting portfolio trades for friction calculation...")
        raw_trades_df = portfolio.trades.records_readable
        
        net_returns = {}
        net_sharpes = {}
        
        # If no trades occurred, net returns match raw returns (zero trades)
        if raw_trades_df.empty:
            logger.info("No trades occurred in portfolio. Friction adjustment bypassed.")
            for symbol in portfolio.wrapper.columns:
                net_returns[symbol] = float(portfolio.total_return()[symbol] * 100)
                net_sharpes[symbol] = float(portfolio.sharpe_ratio()[symbol])
            return {
                "net_total_return": {
                    "portfolio_average": sum(net_returns.values()) / len(net_returns),
                    "per_symbol": net_returns
                },
                "net_sharpe_ratio": {
                    "portfolio_average": sum(net_sharpes.values()) / len(net_sharpes) if net_sharpes else 0.0,
                    "per_symbol": net_sharpes
                },
                "friction_details": pd.DataFrame()
            }
            
        # Calculate statutory friction on all trades
        friction_trades_df = self.calculate_delivery_friction(raw_trades_df)
        
        col_key = next((k for k in ['Column', 'Symbol', 'column', 'symbol'] if k in friction_trades_df.columns), None)
        entry_time_col = next((c for c in ['Entry Timestamp', 'Entry Date'] if c in friction_trades_df.columns), 'Entry Timestamp')
        exit_time_col = next((c for c in ['Exit Timestamp', 'Exit Date'] if c in friction_trades_df.columns), 'Exit Timestamp')
        
        init_cash = portfolio.init_cash
        
        # Adjust daily values per symbol based on transaction dates
        for symbol in portfolio.wrapper.columns:
            raw_val_series = portfolio.value()[symbol]
            
            # Filter trades for this symbol
            if col_key:
                symbol_trades = friction_trades_df[friction_trades_df[col_key] == symbol]
            else:
                symbol_trades = pd.DataFrame()
                
            # Create a daily friction series
            daily_friction = pd.Series(0.0, index=close_matrix.index)
            
            if not symbol_trades.empty:
                for _, trade in symbol_trades.iterrows():
                    entry_t = trade[entry_time_col]
                    exit_t = trade[exit_time_col]
                    
                    # Add entry cost on entry date
                    if pd.notna(entry_t):
                        entry_idx = pd.to_datetime(entry_t)
                        if entry_idx in daily_friction.index:
                            daily_friction.loc[entry_idx] += trade['entry_friction']
                        else:
                            closest = daily_friction.index[daily_friction.index.normalize() == entry_idx.normalize()]
                            if not closest.empty:
                                daily_friction.loc[closest[0]] += trade['entry_friction']
                                
                    # Add exit cost on exit date
                    if pd.notna(exit_t):
                        exit_idx = pd.to_datetime(exit_t)
                        if exit_idx in daily_friction.index:
                            daily_friction.loc[exit_idx] += trade['exit_friction']
                        else:
                            closest = daily_friction.index[daily_friction.index.normalize() == exit_idx.normalize()]
                            if not closest.empty:
                                daily_friction.loc[closest[0]] += trade['exit_friction']

            # Resolve initial capital for the asset
            symbol_init_cash = init_cash[symbol] if isinstance(init_cash, (pd.Series, dict)) else init_cash
            
            # Reconstruct net NAV path-dependently
            net_val = symbol_init_cash
            net_val_history = []
            prev_raw_val = symbol_init_cash
            
            for date in close_matrix.index:
                raw_val = raw_val_series.loc[date]
                daily_ratio = raw_val / prev_raw_val if prev_raw_val > 0.0 else 1.0
                
                # Subtract friction incurred on this date from the adjusted NAV
                net_val = (net_val * daily_ratio) - daily_friction.loc[date]
                net_val = max(0.0, net_val) # floor at zero to avoid negative values
                net_val_history.append(net_val)
                prev_raw_val = raw_val
                
            net_val_series = pd.Series(net_val_history, index=close_matrix.index)
            
            # Compute net total return (%)
            net_returns[symbol] = float((net_val_series.iloc[-1] / symbol_init_cash - 1) * 100)
            
            # Compute net daily returns to calculate Sharpe
            net_daily_returns = net_val_series.pct_change().fillna(0.0)
            std_ret = net_daily_returns.std()
            
            # Resolve annualization factor (defaults to 365 daily calendars in BQE backtests)
            ann_factor = getattr(portfolio.wrapper, 'frequency_ann_factor', 365.0)
            
            net_sharpe = float((net_daily_returns.mean() / std_ret * (ann_factor ** 0.5)) if std_ret > 0.0 else 0.0)
            net_sharpes[symbol] = net_sharpe
            
        portfolio_avg_return = sum(net_returns.values()) / len(net_returns) if net_returns else 0.0
        portfolio_avg_sharpe = sum(net_sharpes.values()) / len(net_sharpes) if net_sharpes else 0.0
        
        return {
            "net_total_return": {
                "portfolio_average": portfolio_avg_return,
                "per_symbol": net_returns
            },
            "net_sharpe_ratio": {
                "portfolio_average": portfolio_avg_sharpe,
                "per_symbol": net_sharpes
            },
            "friction_details": friction_trades_df
        }