"""
Synthetic Data Generator using Statistical Methods for bid/ask time series
"""
import pandas as pd
import numpy as np
from pathlib import Path
from datetime import datetime, timedelta
import warnings
warnings.filterwarnings('ignore')


class SyntheticDataGenerator:
    def __init__(self, sample_data_path='../data/sample_data.parquet'):
        """Initialize the generator with sample data"""
        self.base_path = Path(__file__).parent.parent
        self.sample_path = self.base_path / 'data' / 'sample_data.parquet'
        self.output_path = self.base_path / 'data' / 'synthetic_data'

        # Load sample data
        self.sample_data = pd.read_parquet(self.sample_path)
        print(f"Loaded {len(self.sample_data)} sample records")

        # Calculate baseline statistics
        self.bid_mean = self.sample_data['bid'].mean()
        self.bid_std = self.sample_data['bid'].std()
        self.ask_mean = self.sample_data['ask'].mean()
        self.ask_std = self.sample_data['ask'].std()

        # Calculate spread statistics
        spread = self.sample_data['ask'] - self.sample_data['bid']
        self.spread_mean = spread.mean()
        self.spread_std = spread.std()

        # Calculate returns for time series correlation
        self.bid_returns = self.sample_data['bid'].pct_change().dropna()
        self.ask_returns = self.sample_data['ask'].pct_change().dropna()

        print(f"Baseline bid: {self.bid_mean:.2f} ± {self.bid_std:.2f}")
        print(f"Baseline ask: {self.ask_mean:.2f} ± {self.ask_std:.2f}")
        print(f"Baseline spread: {self.spread_mean:.4f} ± {self.spread_std:.4f}")

    def generate_synthetic_data(self, simulation_length, volatility, window_size=24):
        """
        Generate synthetic bid/ask data using statistical methods

        Parameters:
        - simulation_length: Number of data points to generate
        - volatility: Multiplier for std dev (1.0 = same as sample, 2.0 = 2x volatility)
        """
        print(f"Generating {simulation_length} synthetic data points...")

        # Scale volatility
        scaled_bid_std = self.bid_std * volatility
        scaled_ask_std = self.ask_std * volatility
        scaled_spread_std = self.spread_std * volatility

        # Initialize arrays
        bid_prices = np.zeros(simulation_length)
        ask_prices = np.zeros(simulation_length)

        # Start from sample mean
        bid_prices[0] = self.bid_mean
        ask_prices[0] = self.ask_mean

        # Generate correlated random walk with mean reversion
        mean_reversion_strength = 0.1  # How strongly prices revert to mean

        for i in range(1, simulation_length):
            # Calculate mean reversion force
            bid_reversion = (self.bid_mean - bid_prices[i-1]) * mean_reversion_strength
            ask_reversion = (self.ask_mean - ask_prices[i-1]) * mean_reversion_strength

            # Generate correlated random shocks
            base_shock = np.random.normal(0, scaled_bid_std / np.sqrt(252))  # Daily scale
            bid_shock = base_shock + bid_reversion
            ask_shock = base_shock * 0.95 + np.random.normal(0, scaled_ask_std / np.sqrt(252)) * 0.05 + ask_reversion

            # Update prices
            bid_prices[i] = bid_prices[i-1] + bid_shock
            ask_prices[i] = ask_prices[i-1] + ask_shock

            # Ensure ask > bid (maintain positive spread)
            if ask_prices[i] <= bid_prices[i]:
                ask_prices[i] = bid_prices[i] + np.abs(np.random.normal(self.spread_mean, scaled_spread_std))

        # Add some high-frequency noise for realism
        bid_prices += np.random.normal(0, scaled_bid_std * 0.01, simulation_length)
        ask_prices += np.random.normal(0, scaled_ask_std * 0.01, simulation_length)

        # Ensure all asks are greater than bids
        for i in range(simulation_length):
            if ask_prices[i] <= bid_prices[i]:
                ask_prices[i] = bid_prices[i] + abs(np.random.normal(self.spread_mean, scaled_spread_std))

        # Create DataFrame with timestamp index
        start_time = datetime.now()
        timestamps = [start_time + timedelta(seconds=i) for i in range(simulation_length)]

        df = pd.DataFrame(
            {
                'bid': bid_prices,
                'ask': ask_prices
            },
            index=pd.DatetimeIndex(timestamps, name='timestamp')
        )

        # Save to file
        timestamp_str = start_time.strftime('%Y%m%d_%H%M%S')
        filename = f"synthetic_{simulation_length}_{volatility}x_{timestamp_str}.parquet"
        output_file = self.output_path / filename

        df.to_parquet(output_file)
        print(f"Saved to {output_file}")

        # Print statistics
        print(f"Generated bid: {df['bid'].mean():.2f} ± {df['bid'].std():.2f}")
        print(f"Generated ask: {df['ask'].mean():.2f} ± {df['ask'].std():.2f}")
        print(f"Generated spread: {(df['ask'] - df['bid']).mean():.4f} ± {(df['ask'] - df['bid']).std():.4f}")

        return df, str(output_file)


def quick_generate(simulation_length=1000, volatility=1.0):
    """Quick generation function for testing"""
    generator = SyntheticDataGenerator()
    df, filepath = generator.generate_synthetic_data(simulation_length, volatility)
    return df, filepath


if __name__ == "__main__":
    # Test the generator
    print("Testing Synthetic Data Generator...")
    df, filepath = quick_generate(simulation_length=5000, volatility=1.5)
    print(f"\nGenerated {len(df)} records")
    print(df.head())
    print(f"\nStats:\n{df.describe()}")
