import streamlit as st
import pandas_ta as ta
import pandas as pd
import numpy as np
from backtesting import Backtest, Strategy
from typing import Dict, List, Union
import itertools
from datetime import datetime, timedelta
from src.utils.common_utils import get_stock_data

class DynamicStrategy(Strategy):
    """Dynamic strategy class that can be configured through the UI"""
    
    def init(self):
        # Will be populated with indicator calculations
        self.indicators = {}
        
        # Initialize all selected indicators
        for ind_name, ind_config in self.indicator_configs.items():
            if ind_config['type'] == 'SMA':
                self.indicators[ind_name] = self.I(ta.sma, self.data.Close, length=ind_config['params']['length'])
            elif ind_config['type'] == 'EMA':
                self.indicators[ind_name] = self.I(ta.ema, self.data.Close, length=ind_config['params']['length'])
            elif ind_config['type'] == 'RSI':
                self.indicators[ind_name] = self.I(ta.rsi, self.data.Close, length=ind_config['params']['length'])
            elif ind_config['type'] == 'MACD':
                macd = ta.macd(self.data.Close, 
                             fast=ind_config['params']['fast'],
                             slow=ind_config['params']['slow'],
                             signal=ind_config['params']['signal'])
                self.indicators[f"{ind_name}_macd"] = self.I(lambda x: x['MACD_12_26_9'], macd)
                self.indicators[f"{ind_name}_signal"] = self.I(lambda x: x['MACDs_12_26_9'], macd)
            elif ind_config['type'] == 'BB':
                bb = ta.bbands(self.data.Close, 
                             length=ind_config['params']['length'],
                             std=ind_config['params']['std'])
                self.indicators[f"{ind_name}_upper"] = self.I(lambda x: x['BBU_20_2.0'], bb)
                self.indicators[f"{ind_name}_lower"] = self.I(lambda x: x['BBL_20_2.0'], bb)
                self.indicators[f"{ind_name}_middle"] = self.I(lambda x: x['BBM_20_2.0'], bb)

    def next(self):
        # Simple example strategy - should be customized based on user input
        price = self.data.Close[-1]
        
        # Example trading logic using indicators
        for ind_name, ind_config in self.indicator_configs.items():
            if ind_config['type'] == 'SMA':
                if price > self.indicators[ind_name][-1] and not self.position:
                    self.buy()
                elif price < self.indicators[ind_name][-1] and self.position:
                    self.position.close()
            
            elif ind_config['type'] == 'MACD':
                macd = self.indicators[f"{ind_name}_macd"][-1]
                signal = self.indicators[f"{ind_name}_signal"][-1]
                prev_macd = self.indicators[f"{ind_name}_macd"][-2]
                prev_signal = self.indicators[f"{ind_name}_signal"][-2]
                
                if macd > signal and prev_macd <= prev_signal and not self.position:
                    self.buy()
                elif macd < signal and prev_macd >= prev_signal and self.position:
                    self.position.close()

def get_available_indicators() -> Dict:
    """Returns available indicators and their parameters"""
    return {
        'SMA': {
            'params': ['length'],
            'defaults': {'length': 20},
            'ranges': {'length': (5, 200)}
        },
        'EMA': {
            'params': ['length'],
            'defaults': {'length': 20},
            'ranges': {'length': (5, 200)}
        },
        'RSI': {
            'params': ['length'],
            'defaults': {'length': 14},
            'ranges': {'length': (5, 30)}
        },
        'MACD': {
            'params': ['fast', 'slow', 'signal'],
            'defaults': {'fast': 12, 'slow': 26, 'signal': 9},
            'ranges': {'fast': (8, 20), 'slow': (20, 40), 'signal': (5, 15)}
        },
        'BB': {
            'params': ['length', 'std'],
            'defaults': {'length': 20, 'std': 2.0},
            'ranges': {'length': (10, 50), 'std': (1.5, 3.0)}
        }
    }

def prepare_data_for_backtest(df: pd.DataFrame) -> pd.DataFrame:
    """Prepare the dataframe for backtesting"""
    # Ensure the dataframe has the required columns
    required_columns = ['Open', 'High', 'Low', 'Close', 'Volume']
    
    # Rename columns if needed
    df = df.copy()
    df.columns = [c.capitalize() for c in df.columns]
    
    # Set the date as index if it's not already
    if 'Date' in df.columns:
        df.set_index('Date', inplace=True)
    
    # Verify all required columns exist
    missing_cols = [col for col in required_columns if col not in df.columns]
    if missing_cols:
        raise ValueError(f"Missing required columns: {missing_cols}")
    
    return df

def backtesting_page():
    st.title("Strategy Backtesting")
    
    # Sidebar controls
    st.sidebar.subheader("Backtest Settings")
    
    # Ticker selection
    ticker = st.sidebar.text_input("Enter Ticker Symbol", value="AAPL").upper()
    
    # Date range selection
    col1, col2 = st.sidebar.columns(2)
    with col1:
        start_date = st.date_input("Start Date", 
                                 value=datetime.now() - timedelta(days=365))
    with col2:
        end_date = st.date_input("End Date")
    
    # Indicator selection
    available_indicators = get_available_indicators()
    selected_indicators = st.sidebar.multiselect(
        "Select Technical Indicators",
        options=list(available_indicators.keys())
    )
    
    # Parameter input/optimization section
    optimize = st.sidebar.checkbox("Optimize Parameters")
    
    indicator_settings = {}
    for ind_name in selected_indicators:
        st.sidebar.subheader(f"{ind_name} Settings")
        ind_config = available_indicators[ind_name]
        
        if optimize:
            params = {}
            for param in ind_config['params']:
                param_range = ind_config['ranges'][param]
                col1, col2, col3 = st.sidebar.columns(3)
                with col1:
                    min_val = st.number_input(f"{param} Min", 
                                           value=float(param_range[0]))
                with col2:
                    max_val = st.number_input(f"{param} Max", 
                                           value=float(param_range[1]))
                with col3:
                    step = st.number_input(f"{param} Step", 
                                        value=1.0 if param == 'std' else 1)
                params[param] = {'min': min_val, 'max': max_val, 'step': step}
        else:
            params = {}
            for param in ind_config['params']:
                default_val = ind_config['defaults'][param]
                params[param] = st.sidebar.number_input(f"{param}", 
                                                     value=float(default_val))
        
        indicator_settings[ind_name] = {
            'type': ind_name,
            'params': params
        }
    
    if st.sidebar.button("Run Backtest"):
        with st.spinner('Running backtest...'):
            # Fetch data
            df = get_stock_data(ticker, start_date, end_date, 'daily')
            
            if df.empty:
                st.error("No data available for the selected period")
                return
            
            try:
                df = prepare_data_for_backtest(df)
                
                if optimize:
                    results = run_optimization(df, indicator_settings)
                    display_optimization_results(results)
                else:
                    results = run_single_backtest(df, indicator_settings)
                    display_backtest_results(results)
                    
            except Exception as e:
                st.error(f"Error during backtest: {str(e)}")

def run_optimization(df: pd.DataFrame, indicator_settings: Dict) -> List:
    """Run optimization with different parameter combinations"""
    param_combinations = generate_param_combinations(indicator_settings)
    results = []
    
    for params in param_combinations:
        # Configure strategy with current parameters
        DynamicStrategy.indicator_configs = params
        
        # Run backtest
        bt = Backtest(df, DynamicStrategy, cash=100000, commission=.002)
        stats = bt.run()
        
        results.append({
            'parameters': params,
            'Return [%]': stats['Return [%]'],
            'Sharpe Ratio': stats['Sharpe Ratio'],
            'Max Drawdown [%]': stats['Max Drawdown [%]'],
            'Win Rate [%]': stats['Win Rate [%]']
        })
    
    return results

def run_single_backtest(df: pd.DataFrame, indicator_settings: Dict) -> Dict:
    """Run a single backtest with fixed parameters"""
    DynamicStrategy.indicator_configs = indicator_settings
    bt = Backtest(df, DynamicStrategy, cash=100000, commission=.002)
    return bt.run()

def generate_param_combinations(settings: Dict) -> List[Dict]:
    """Generate all possible parameter combinations for optimization"""
    param_space = {}
    for ind_name, ind_config in settings.items():
        for param_name, param_range in ind_config['params'].items():
            if isinstance(param_range, dict):  # Optimization mode
                values = np.arange(
                    param_range['min'],
                    param_range['max'] + param_range['step'],
                    param_range['step']
                )
                param_space[f"{ind_name}_{param_name}"] = values
    
    # Generate all combinations
    keys = list(param_space.keys())
    values = list(param_space.values())
    combinations = list(itertools.product(*values))
    
    # Convert to list of parameter dictionaries
    result = []
    for combo in combinations:
        params = {}
        for key, value in zip(keys, combo):
            ind_name, param_name = key.rsplit('_', 1)
            if ind_name not in params:
                params[ind_name] = {'type': ind_name, 'params': {}}
            params[ind_name]['params'][param_name] = value
        result.append(params)
    
    return result

def display_optimization_results(results: List):
    """Display optimization results in a formatted table"""
    df_results = pd.DataFrame(results)
    
    # Format parameters column for better display
    df_results['parameters'] = df_results['parameters'].apply(str)
    
    st.subheader("Optimization Results")
    st.dataframe(df_results.sort_values('Return [%]', ascending=False))

def display_backtest_results(results: Dict):
    """Display single backtest results with metrics and plots"""
    st.subheader("Backtest Results")
    
    # Display key metrics
    col1, col2, col3, col4 = st.columns(4)
    with col1:
        st.metric("Return", f"{results['Return [%]']:.2f}%")
    with col2:
        st.metric("Sharpe Ratio", f"{results['Sharpe Ratio']:.2f}")
    with col3:
        st.metric("Max Drawdown", f"{results['Max Drawdown [%]']:.2f}%")
    with col4:
        st.metric("Win Rate", f"{results['Win Rate [%]']:.2f}%")
    
    # Display full results in expandable section
    with st.expander("See detailed results"):
        st.write(results)