import unittest
import pandas as pd
import numpy as np
import os
import sys
from unittest.mock import patch
import tempfile
import shutil

# Import the necessary modules from the core package
from core.data import load_data
from core.strategies import load_strategies, register_strategy, get_strategy, list_strategies
from core.spd import backtest_dynamic_dca


class TestSubmitStrategies(unittest.TestCase):
    """Test submitted strategies in the submit_strategies directory"""

    @classmethod
    def setUpClass(cls):
        """Load strategies and test data once for all tests"""
        # Load all strategies
        load_strategies()
        
        # Load sample data for testing
        cls.test_data = load_data()
        
        # Store original strategies list
        cls.original_strategies = list(list_strategies().keys())
        
        # Identify strategies from submit_strategies directory
        cls.submit_strategies = []
        for strategy_name in cls.original_strategies:
            # Try to check if this is from submit_strategies
            try:
                # Test by importing - if it's in submit_strategies, this will succeed
                __import__(f"submit_strategies.{strategy_name}", fromlist=[strategy_name])
                cls.submit_strategies.append(strategy_name)
            except ImportError:
                # If import fails, it's a core strategy
                continue

    def test_submit_strategies_loaded(self):
        """Test that strategies from submit_strategies are properly loaded"""
        # Verify at least one strategy from submit_strategies was loaded
        self.assertGreater(len(self.submit_strategies), 0, 
                         "No strategies from submit_strategies were loaded")
        
        # Print the loaded strategies
        print(f"Successfully loaded {len(self.submit_strategies)} strategies from submit_strategies:")
        for strategy in self.submit_strategies:
            print(f"  - {strategy}")

    def test_submit_strategies_execution(self):
        """Test that each strategy from submit_strategies can execute"""
        for strategy_name in self.submit_strategies:
            with self.subTest(strategy=strategy_name):
                try:
                    strategy_fn = get_strategy(strategy_name)
                    # Call the strategy function and check that it returns weights
                    weights = strategy_fn(self.test_data)
                    
                    # Assert that weights is a Series 
                    self.assertIsInstance(weights, pd.Series)
                    
                    # Assert that weights has values (not empty)
                    self.assertGreater(len(weights), 0)
                    
                    # Assert that all weights are non-negative
                    self.assertTrue((weights >= 0).all())
                    
                    print(f"Strategy {strategy_name} executed successfully")
                except Exception as e:
                    self.fail(f"Strategy {strategy_name} execution failed: {str(e)}")

    def test_submit_strategies_backtest(self):
        """Test that each strategy from submit_strategies can be backtested"""
        for strategy_name in self.submit_strategies:
            with self.subTest(strategy=strategy_name):
                try:
                    # Run a backtest with the strategy
                    backtest_results = backtest_dynamic_dca(
                        self.test_data, 
                        strategy_name=strategy_name, 
                        show_plots=False
                    )
                    
                    # Assert that backtest returns a DataFrame with results
                    self.assertIsInstance(backtest_results, pd.DataFrame)
                    
                    # Assert that key columns are present
                    self.assertIn('dynamic_spd', backtest_results.columns)
                    self.assertIn('dynamic_pct', backtest_results.columns)
                    
                    print(f"Backtest for {strategy_name} completed successfully")
                    
                    # Print summary statistics
                    print(f"  Mean SPD: {backtest_results['dynamic_spd'].mean():.4f}")
                    print(f"  Mean Excess %: {backtest_results['excess_pct'].mean():.2f}%")
                except Exception as e:
                    self.fail(f"Backtest for {strategy_name} failed: {str(e)}")

    def test_submit_strategies_comparison(self):
        """Test that submit_strategies perform differently than core strategies"""
        if not self.submit_strategies:
            self.skipTest("No strategies from submit_strategies available for comparison")
        
        # Choose one submit strategy to compare with core strategies
        submit_strategy = self.submit_strategies[0]
        
        # Find a core strategy to compare with
        core_strategies = [s for s in self.original_strategies if s not in self.submit_strategies]
        if not core_strategies:
            self.skipTest("No core strategies available for comparison")
            
        core_strategy = core_strategies[0]
        
        # Run backtests for both
        submit_results = backtest_dynamic_dca(
            self.test_data, 
            strategy_name=submit_strategy, 
            show_plots=False
        )
        
        core_results = backtest_dynamic_dca(
            self.test_data, 
            strategy_name=core_strategy, 
            show_plots=False
        )
        
        # Print comparison
        submit_mean = submit_results['dynamic_spd'].mean()
        core_mean = core_results['dynamic_spd'].mean()
        
        print(f"Strategy performance comparison:")
        print(f"  {submit_strategy} (submit): {submit_mean:.4f}")
        print(f"  {core_strategy} (core): {core_mean:.4f}")
        print(f"  Difference: {abs(submit_mean - core_mean):.4f}")
        
        # We don't assert which is better, just that they're different
        # Allow a small tolerance for floating point comparison
        self.assertNotAlmostEqual(
            submit_mean, 
            core_mean, 
            places=4, 
            msg=f"Submit strategy {submit_strategy} performs identically to core strategy {core_strategy}"
        )

    def test_results_output(self):
        """Test that strategies from submit_strategies produce output files"""
        if not self.submit_strategies:
            self.skipTest("No strategies from submit_strategies available for testing output")
        
        # Create a temporary directory for results
        with tempfile.TemporaryDirectory() as temp_dir:
            # Choose a strategy to test
            strategy_name = self.submit_strategies[0]
            
            # Get the function to run backtest
            from core.spd import backtest_dynamic_dca
            
            # Run backtest with the specific strategy
            results = backtest_dynamic_dca(
                self.test_data, 
                strategy_name=strategy_name, 
                show_plots=False
            )
            
            # Add strategy name to results for proper CSV output
            results['strategy'] = strategy_name
            
            # Create summary dictionary
            summary = {
                'strategy': [strategy_name],
                'min_spd': [results['dynamic_spd'].min()],
                'max_spd': [results['dynamic_spd'].max()],
                'mean_spd': [results['dynamic_spd'].mean()],
                'median_spd': [results['dynamic_spd'].median()],
                'min_pct': [results['dynamic_pct'].min()],
                'max_pct': [results['dynamic_pct'].max()],
                'mean_pct': [results['dynamic_pct'].mean()],
                'median_pct': [results['dynamic_pct'].median()],
                'avg_excess_pct': [results['excess_pct'].mean()]
            }
            
            # Convert summary to DataFrame
            summary_df = pd.DataFrame(summary)
            
            # Create output directory if needed
            os.makedirs(temp_dir, exist_ok=True)
            
            # Save to CSV files
            spd_csv_path = os.path.join(temp_dir, 'spd_by_cycle.csv')
            summary_csv_path = os.path.join(temp_dir, 'strategy_summary.csv')
            
            results.reset_index().to_csv(spd_csv_path, index=False)
            summary_df.to_csv(summary_csv_path, index=False)
            
            # Check that output files exist
            self.assertTrue(os.path.exists(spd_csv_path))
            self.assertTrue(os.path.exists(summary_csv_path))
            
            # Check that our strategy is in the summary
            loaded_summary_df = pd.read_csv(summary_csv_path)
            strategies_in_summary = loaded_summary_df['strategy'].tolist()
            
            self.assertIn(
                strategy_name, 
                strategies_in_summary, 
                f"Strategy {strategy_name} not found in summary results"
            )
            
            print(f"Output files generated successfully and include {strategy_name}")


class TestDynamicSubmitStrategies(unittest.TestCase):
    """Test creating and running strategies on the fly in submit_strategies"""
    
    def setUp(self):
        """Set up for each test"""
        # Create a temporary strategy with a safe name (no invalid Python identifiers)
        self.temp_strategy_name = f"temp_test_strategy_{abs(hash(self)) % 10000}"
        self.temp_strategy_file = f"submit_strategies/{self.temp_strategy_name}.py"
        
        # Load test data
        self.test_data = load_data()
        
        # Write a simple strategy file
        with open(self.temp_strategy_file, 'w') as f:
            f.write("""
import pandas as pd
import numpy as np
from core.config import BACKTEST_START, BACKTEST_END
from core.strategies import register_strategy
from core.strategies.base_strategy import StrategyTemplate

class TempTestStrategy(StrategyTemplate):
    \"\"\"
    A simple test strategy that allocates more weight when price is below 50-day MA.
    \"\"\"
    
    @staticmethod
    def construct_features(df):
        df = df.copy()
        # Calculate 50-day moving average
        df['ma_50'] = df['btc_close'].rolling(window=50).mean()
        df['below_ma'] = (df['btc_close'] < df['ma_50']).astype(int)
        return df
    
    @staticmethod
    def compute_weights(df):
        df_backtest = df.loc[BACKTEST_START:BACKTEST_END]
        weights = pd.Series(index=df_backtest.index, data=1.0)
        # More weight when below MA
        weights[df_backtest['below_ma'] == 1] = 1.5
        
        # Normalize within each cycle
        start_year = pd.to_datetime(BACKTEST_START).year
        cycle_labels = df_backtest.index.to_series().apply(
            lambda dt: (dt.year - start_year) // 4
        )
        
        for cycle, group in weights.groupby(cycle_labels):
            cycle_sum = group.sum()
            if cycle_sum > 0:
                weights.loc[group.index] = weights.loc[group.index] / cycle_sum
                
        return weights

@register_strategy("{0}")
def {0}(df):
    \"\"\"Simple test strategy\"\"\"
    return TempTestStrategy.get_strategy_function()(df)
""".format(self.temp_strategy_name))

        # Reload strategies to pick up the new one
        load_strategies()
        
    def tearDown(self):
        """Clean up after each test"""
        # Remove the temporary strategy file
        if os.path.exists(self.temp_strategy_file):
            os.remove(self.temp_strategy_file)
            
        # Also remove the cached .pyc file if it exists
        pyc_dir = "submit_strategies/__pycache__"
        if os.path.exists(pyc_dir):
            pyc_files = [f for f in os.listdir(pyc_dir) 
                       if f.startswith(f"{self.temp_strategy_name}.cpython-")]
            for f in pyc_files:
                os.remove(os.path.join(pyc_dir, f))

    def test_dynamic_strategy_loading(self):
        """Test that dynamically created strategies can be loaded"""
        # Check that our strategy is in the list
        strategies = list_strategies()
        self.assertIn(self.temp_strategy_name, strategies)
        
        print(f"Dynamic strategy {self.temp_strategy_name} loaded successfully")

    def test_dynamic_strategy_execution(self):
        """Test that dynamically created strategies can be executed"""
        # Get our strategy
        strategy_fn = get_strategy(self.temp_strategy_name)
        
        # Execute it
        weights = strategy_fn(self.test_data)
        
        # Check the results
        self.assertIsInstance(weights, pd.Series)
        self.assertTrue(len(weights) > 0)
        
        print(f"Dynamic strategy {self.temp_strategy_name} executed successfully")

    def test_dynamic_strategy_backtest(self):
        """Test that dynamically created strategies can be backtested"""
        # Run a backtest
        results = backtest_dynamic_dca(
            self.test_data,
            strategy_name=self.temp_strategy_name,
            show_plots=False
        )
        
        # Check the results
        self.assertIsInstance(results, pd.DataFrame)
        self.assertIn('dynamic_spd', results.columns)
        
        # Print performance stats
        print(f"Dynamic strategy {self.temp_strategy_name} backtest results:")
        print(f"  Mean SPD: {results['dynamic_spd'].mean():.4f}")
        print(f"  Mean Excess %: {results['excess_pct'].mean():.2f}%")


if __name__ == '__main__':
    unittest.main() 