"""
Nifty 50 Ensemble Backtest Script
Simulates trading with the ensemble predictions.
"""
import argparse
import pandas as pd
import numpy as np
from nifty_ensemble_v2 import NiftyEnsembleV2

def backtest(df, results, threshold=0.55, commission=0.0005, slippage=0.0001):
    """
    Simple backtest: go long when prob_up > threshold, short when < 1-threshold.
    """
    df = df.copy()
    df['proba_up'] = results['proba'][:, 1]
    df['pred'] = results['pred']
    df['returns'] = df['Close'].pct_change().shift(-1)
    
    # Positions: +1 = long, -1 = short, 0 = flat
    df['position'] = np.where(df['proba_up'] > threshold, 1,
                       np.where(df['proba_up'] < (1 - threshold), -1, 0))
    
    # Strategy returns
    df['strat_ret'] = df['position'].shift(1) * df['returns']
    
    # Transaction costs on position changes
    df['pos_change'] = df['position'].diff().abs()
    df['costs'] = df['pos_change'] * (commission + slippage)
    df['strat_ret_net'] = df['strat_ret'] - df['costs']
    
    # Cumulative
    df['cum_market'] = (1 + df['returns'].fillna(0)).cumprod()
    df['cum_strat'] = (1 + df['strat_ret_net'].fillna(0)).cumprod()
    
    # Metrics
    valid = df['strat_ret_net'].dropna()
    sharpe = valid.mean() / valid.std() * np.sqrt(252 * 375) if valid.std() > 0 else 0  # annualized for 1m
    win_rate = (valid > 0).mean()
    max_dd = ((df['cum_strat'].cummax() - df['cum_strat']) / df['cum_strat'].cummax()).max()
    
    print(f"\n{'='*50}")
    print("BACKTEST RESULTS")
    print(f"{'='*50}")
    print(f"Threshold: {threshold}")
    print(f"Trades: {df['pos_change'].sum():.0f}")
    print(f"Win Rate: {win_rate:.2%}")
    print(f"Sharpe (ann): {sharpe:.3f}")
    print(f"Max Drawdown: {max_dd:.2%}")
    print(f"Final Return (strat): {(df['cum_strat'].iloc[-1] - 1):.2%}")
    print(f"Final Return (buy-hold): {(df['cum_market'].iloc[-1] - 1):.2%}")
    
    return df

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--data', required=True)
    parser.add_argument('--model', required=True)
    parser.add_argument('--threshold', type=float, default=0.55)
    parser.add_argument('--output', default='backtest_results.csv')
    args = parser.parse_args()
    
    df = pd.read_csv(args.data, index_col='Datetime', parse_dates=True)
    pipe = NiftyEnsembleV2()
    pipe.load(args.model)
    results = pipe.predict(df)
    
    bt = backtest(results['df'], results, threshold=args.threshold)
    bt.to_csv(args.output)
    print(f"\nSaved backtest to {args.output}")

if __name__ == '__main__':
    main()