import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import pearsonr
import warnings
warnings.filterwarnings('ignore')

# Import our utility functions
from .backshift import backshift
from .fwdshift import fwdshift
from .smartmean import smartmean
from .smartstd import smartstd
from .calculateMaxDD import calculateMaxDD
from .data_loader import load_futures_data

def main():
    """
    Python implementation of the TU_mom.m momentum trading strategy.
    """
    print("TU Momentum Trading Strategy")
    print("=" * 40)
    
    try:
        # Try to load real Treasury futures data
        print("Loading Treasury futures data...")
        data = load_futures_data('TU', '20120813')
        tday = data['tday']
        cl = data['cl'][:, 0]  # Use first contract for simplicity
        
        print(f"Loaded {len(tday)} days of Treasury futures data")
        
    except (FileNotFoundError, KeyError) as e:
        print(f"Could not load real data ({e}), using synthetic data for demonstration...")
        
        # Synthetic data for demonstration
        np.random.seed(42)
        n_days = 2000
        tday = np.arange(20090102, 20090102 + n_days)
        cl = 100 * np.cumprod(1 + np.random.normal(0, 0.01, n_days))
    
    # Correlation tests
    print("\nCorrelation Analysis:")
    print("Lookback\tHolddays\tCorrelation\tp-value")
    print("-" * 50)
    
    for lookback in [1, 5, 10, 25, 60, 120, 250]:
        for holddays in [1, 5, 10, 25, 60, 120, 250]:
            # Calculate lagged returns
            ret_lag = (cl - backshift(lookback, cl)) / backshift(lookback, cl)
            ret_fut = (fwdshift(holddays, cl) - cl) / cl
            
            # Remove bad dates (NaN values)
            bad_dates = np.isnan(ret_lag) | np.isnan(ret_fut)
            ret_lag_clean = ret_lag[~bad_dates]
            ret_fut_clean = ret_fut[~bad_dates]
            
            if len(ret_lag_clean) == 0:
                continue
            
            # Create independent set
            if lookback >= holddays:
                indep_set = np.arange(0, len(ret_lag_clean), holddays)
            else:
                indep_set = np.arange(0, len(ret_lag_clean), lookback)
            
            ret_lag_indep = ret_lag_clean[indep_set]
            ret_fut_indep = ret_fut_clean[indep_set]
            
            # Calculate correlation
            if len(ret_lag_indep) > 1:
                cc, pval = pearsonr(ret_lag_indep, ret_fut_indep)
                print(f"{lookback:3d}\t\t{holddays:3d}\t\t{cc:7.4f}\t\t{pval:6.4f}")
    
    # Trading strategy implementation
    lookback = 250
    holddays = 25
    
    print(f"\nImplementing Trading Strategy:")
    print(f"Lookback: {lookback} days, Hold: {holddays} days")
    
    # Generate trading signals
    longs = cl > backshift(lookback, cl)
    shorts = cl < backshift(lookback, cl)
    
    # Initialize positions
    pos = np.zeros(len(cl))
    
    # Build position over holding period
    for h in range(holddays):
        long_lag = backshift(h, longs.astype(float))
        long_lag = np.nan_to_num(long_lag, nan=0).astype(bool)
        
        short_lag = backshift(h, shorts.astype(float))
        short_lag = np.nan_to_num(short_lag, nan=0).astype(bool)
        
        pos[long_lag] += 1
        pos[short_lag] -= 1
    
    # Calculate returns
    ret = (backshift(1, pos) * (cl - backshift(1, cl)) / backshift(1, cl)) / holddays
    ret = np.nan_to_num(ret, nan=0)
    
    # Find start index (equivalent to finding date 20090102)
    idx = 250  # Start after sufficient data for lookback
    
    # Calculate cumulative returns
    cumret = np.cumprod(1 + ret[idx:]) - 1
    
    # Plot results
    plt.figure(figsize=(12, 6))
    plt.plot(cumret)
    plt.title('TU Momentum Strategy - Cumulative Returns')
    plt.xlabel('Days')
    plt.ylabel('Cumulative Return')
    plt.grid(True)
    plt.show()
    
    # Performance metrics
    strategy_returns = ret[idx:]
    avg_ann_ret = 252 * smartmean(strategy_returns)
    ann_volatility = np.sqrt(252) * smartstd(strategy_returns)
    sharpe_ratio = avg_ann_ret / ann_volatility if ann_volatility != 0 else 0
    apr = np.prod(1 + strategy_returns) ** (252 / len(strategy_returns)) - 1
    
    maxDD, maxDDD = calculateMaxDD(cumret)
    kelly_f = np.mean(strategy_returns) / np.var(strategy_returns) if np.var(strategy_returns) != 0 else 0
    
    print(f"\nPerformance Metrics:")
    print(f"Average Annual Return: {avg_ann_ret:7.4f}")
    print(f"Annual Volatility: {ann_volatility:7.4f}")
    print(f"Sharpe Ratio: {sharpe_ratio:4.2f}")
    print(f"APR: {apr:10.4f}")
    print(f"Max Drawdown: {maxDD:.6f}")
    print(f"Max Drawdown Duration: {int(maxDDD)} days")
    print(f"Kelly f: {kelly_f:.6f}")

if __name__ == "__main__":
    main()