import yfinance as yf
import pandas as pd
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from datetime import datetime, timedelta
import warnings
warnings.filterwarnings('ignore')

# ============================================
# DEFINE EVENT TYPES FOR PREDICTION
# ============================================
EVENT_TYPES = {
    0: "MARKET_CRASH",      # > 5% single-day drop
    1: "VOLATILITY_SPIKE",  # VIX > 30
    2: "BULL_RUN",          # > 15% gain over 5 days
    3: "LIQUIDITY_CRISIS",  # Volume drop + spread widening
    4: "GEOPOLITICAL",      # Oil/Vix spike + currency volatility
    5: "NATURAL_DISASTER",  # Commodity disruption + sector impact
    6: "POLICY_SHOCK",      # Interest rate/Sovereign yield spike
    7: "TECH_BUBBLE",       # Sector overvaluation + put/call ratio
    8: "CURRENCY_CRISIS",   # FX volatility > 5% daily
    9: "NORMAL"            # Baseline - no event
}

# ============================================
# GLOBAL MARKET INDICES (by country for map visualization)
# ============================================
GLOBAL_INDICES = {
    "USA": ["^GSPC", "^DJI", "^IXIC"],
    "UK": ["^FTSE"],
    "Japan": ["^N225"],
    "Germany": ["^GDAXI"],
    "France": ["^FCHI"],
    "Hong_Kong": ["^HSI"],
    "India": ["^BSESN"],
    "China": ["000001.SS"],
    "South_Korea": ["^KS11"],
    "Brazil": ["^BVSP"],
    "Canada": ["^GSPTSE"],
    "Australia": ["^AXJO"]
}

# Country coordinates for map plotting
COUNTRY_COORDS = {
    "USA": (37.0902, -95.7129),
    "UK": (51.5074, -0.1278),
    "Japan": (36.2048, 138.2529),
    "Germany": (51.1657, 10.4515),
    "France": (46.2276, 2.2137),
    "Hong_Kong": (22.3193, 114.1694),
    "India": (20.5937, 78.9629),
    "China": (35.8617, 104.1954),
    "South_Korea": (35.9078, 127.7669),
    "Brazil": (-14.2350, -51.9253),
    "Canada": (56.1304, -106.3468),
    "Australia": (-25.2744, 133.7751)
}


# ============================================
# DATA COLLECTION CLASS
# ============================================
class GlobalEventDataCollector:
    """Collect global financial data and event labels"""
    
    def __init__(self, start_date='2020-01-01', end_date=None):
        self.start_date = start_date
        self.end_date = end_date or datetime.now().strftime('%Y-%m-%d')
        self.data = {}
        self.events = []
        
    def fetch_global_indices(self):
        """Fetch all global indices from Yahoo Finance"""
        print("Fetching global market data...")
        all_indices = []
        for country, tickers in GLOBAL_INDICES.items():
            for ticker in tickers:
                try:
                    data = yf.download(ticker, start=self.start_date, end=self.end_date, progress=False)
                    data['Country'] = country
                    data['Ticker'] = ticker
                    data['Index_Name'] = ticker
                    all_indices.append(data)
                    print(f"  ✓ Fetched {ticker} ({country})")
                except Exception as e:
                    print(f"  ✗ Failed {ticker}: {e}")
        
        # Merge all indices on date
        self.market_data = pd.concat(all_indices, keys=[i for i in range(len(all_indices))])
        return self.market_data
    
    def fetch_macro_indicators(self):
        """Fetch macroeconomic indicators (VIX, yields, commodities)"""
        print("\nFetching macro indicators...")
        macro_tickers = {
            'VIX': '^VIX',           # Volatility index
            'TNX': '^TNX',           # 10-year Treasury yield
            'CL=F': 'CL=F',          # Crude Oil
            'GC=F': 'GC=F',          # Gold
            'DX-Y.NYB': 'DX-Y.NYB'   # US Dollar Index
        }
        
        macro_data = {}
        for name, ticker in macro_tickers.items():
            try:
                data = yf.download(ticker, start=self.start_date, end=self.end_date, progress=False)
                if not data.empty:
                    macro_data[name] = data['Close'].squeeze()
                    print(f"  ✓ Fetched {name}")
                else:
                    print(f"  ✗ No data for {name}")
            except Exception as e:
                print(f"  ✗ Failed {name}: {e}")
        
        self.macro_data = pd.DataFrame(macro_data)
        return self.macro_data
    
    def detect_events(self, window=5):
        """
        Detect events from market data using conditional logic
        Returns event labels for each date and country
        """
        print("\nDetecting events from market data...")
        event_labels = {}
        
        for country in GLOBAL_INDICES.keys():
            country_data = self.market_data[self.market_data['Country'] == country]
            if len(country_data) == 0:
                continue
                
            # Get closing prices and volumes
            if isinstance(self.market_data.index, pd.MultiIndex):
                prices = country_data.groupby(level=1)['Close'].mean()
                volumes = country_data.groupby(level=1)['Volume'].mean() if 'Volume' in country_data.columns.levels[0] else None
            else:
                prices = country_data['Close']
                volumes = country_data.get('Volume')
            
            # Ensure they are Series
            if isinstance(prices, pd.DataFrame): prices = prices.iloc[:, 0]
            if isinstance(volumes, pd.DataFrame): volumes = volumes.iloc[:, 0]
            
            # Calculate returns
            returns = prices.pct_change()
            volatility = returns.rolling(window).std() * np.sqrt(252)
            
            # Initialize event array
            events = np.zeros(len(prices))
            
            for i in range(len(prices)):
                if i < window:
                    continue
                    
                # Event detection logic
                daily_return = returns.iloc[i]
                vol = volatility.iloc[i]
                price_slice = prices.iloc[max(0, i-5):i+1]
                
                # MARKET_CRASH: > 5% drop
                if daily_return < -0.05:
                    events[i] = 0
                # BULL_RUN: > 15% gain over 5 days
                elif (prices.iloc[i] / prices.iloc[max(0, i-5)] - 1) > 0.15:
                    events[i] = 2
                # VOLATILITY_SPIKE: VIX-like > 0.4
                elif vol > 0.4:
                    events[i] = 1
                # LIQUIDITY_CRISIS: volume drop
                elif i > 0 and volumes is not None:
                    vol_drop = volumes.iloc[i] / (volumes.iloc[max(0, i-5):i].mean() + 1)
                    if vol_drop < 0.3:
                        events[i] = 3
                else:
                    events[i] = 9  # NORMAL
            
            event_labels[country] = pd.Series(events, index=prices.index, name='event_type')
        
        self.event_labels = event_labels
        return event_labels


# ============================================
# PYTORCH DATASET AND TRANSFORMER MODEL
# ============================================
class EventDataset(Dataset):
    """PyTorch Dataset for time series event prediction"""
    
    def __init__(self, market_data, macro_data, event_labels, seq_length=30):
        self.seq_length = seq_length
        self.features = []
        self.labels = []
        
        # Align dates across data sources
        print(f"Market Data Columns: {market_data.columns.tolist()[:10]}")
        if 'Country' in market_data.columns:
            print(f"Country unique values: {market_data['Country'].unique()}")
        else:
            # Maybe it's MultiIndex column
            country_col = [c for c in market_data.columns if c[0] == 'Country']
            if country_col:
                print(f"Country column found: {country_col}")
                print(f"Country unique values: {market_data[country_col[0]].unique()}")

        common_dates = market_data.index.get_level_values(1) if isinstance(market_data.index, pd.MultiIndex) else market_data.index
        common_dates = pd.Index(sorted(list(set(common_dates).intersection(macro_data.index))))
        
        if len(common_dates) <= seq_length:
            print(f"Warning: Not enough data points ({len(common_dates)}) for sequence length {seq_length}")
            return

        # Pre-process market data into a simpler format for fast access
        # We'll use one representative ticker per country as in detect_events
        market_series = {}
        for country in GLOBAL_INDICES.keys():
            country_data = market_data[market_data['Country'] == country]
            if not country_data.empty:
                if isinstance(market_data.index, pd.MultiIndex):
                    prices = country_data.groupby(level=1)['Close'].mean()
                else:
                    prices = country_data['Close']
                
                if isinstance(prices, pd.DataFrame):
                    prices = prices.iloc[:, 0]
                series = prices.reindex(common_dates).ffill().bfill()
                market_series[country] = series
                print(f"  {country}: {series.notna().sum()} valid points")

        # Aligned macro data
        macro_aligned = macro_data.reindex(common_dates).ffill().bfill()
        
        print(f"Aligning data: common_dates={len(common_dates)}, seq_length={seq_length}")
        
        for date_idx in range(seq_length, len(common_dates)):
            date = common_dates[date_idx]
            feature_vector = []
            
            # Market features
            for country in GLOBAL_INDICES.keys():
                if country in market_series:
                    prices = market_series[country].iloc[date_idx-seq_length:date_idx]
                    rets = prices.pct_change().fillna(0)
                    feature_vector.extend([
                        prices.iloc[-1],      # current price
                        rets.mean(),          # avg return
                        rets.std(),           # volatility
                        (prices.iloc[-1] / prices.iloc[0] - 1) if prices.iloc[0] != 0 else 0  # cumulative return
                    ])
                else:
                    feature_vector.extend([0, 0, 0, 0])
            
            # Macro features
            macro_slice = macro_aligned.iloc[date_idx-seq_length:date_idx]
            for col in macro_slice.columns:
                feature_vector.extend([
                    macro_slice[col].iloc[-1],
                    macro_slice[col].pct_change().mean() if macro_slice[col].iloc[0] != 0 else 0,
                    macro_slice[col].std()
                ])
            
            # Get event label
            # us_event = event_labels.get("USA", pd.Series()).get(date, 9)
            # Use majority or specific logic? Original used USA event
            us_events = event_labels.get("USA", pd.Series())
            label = us_events.get(date, 9)
            
            if not np.isnan(feature_vector).any():
                self.features.append(feature_vector)
                self.labels.append(int(label))
            else:
                if date_idx == seq_length:
                    nan_indices = [i for i, x in enumerate(feature_vector) if np.isnan(x)]
                    print(f"NaNs found in feature_vector at date {date}: {nan_indices}")
                    print(f"Feature vector length: {len(feature_vector)}")
        
        print(f"Dataset created with {len(self.features)} samples")
        self.features = np.array(self.features)
        self.labels = np.array(self.labels)
        
    def __len__(self):
        return len(self.features)
    
    def __getitem__(self, idx):
        return torch.FloatTensor(self.features[idx]), torch.LongTensor([self.labels[idx]])[0]


class TimeSeriesEventTransformer(nn.Module):
    """
    Transformer model for event prediction
    Combines encoder with classification head
    """
    
    def __init__(self, input_dim, d_model=128, nhead=8, num_layers=3, num_classes=10, dropout=0.1):
        super().__init__()
        self.input_projection = nn.Linear(input_dim, d_model)
        self.pos_encoder = nn.Parameter(torch.randn(1, 1, d_model))
        
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, 
            nhead=nhead, 
            dim_feedforward=d_model*4,
            dropout=dropout,
            batch_first=True
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        
        self.classifier = nn.Sequential(
            nn.Linear(d_model, 64),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(64, num_classes)
        )
        
    def forward(self, x):
        # x shape: (batch, features) - we treat as sequence of length 1
        # Project to d_model
        x = self.input_projection(x).unsqueeze(1)  # (batch, 1, d_model)
        x = x + self.pos_encoder[:, :x.size(1), :]
        
        # Transformer expects (batch, seq_len, d_model)
        x = self.transformer(x)
        
        # Use sequence output for classification
        x = x[:, -1, :]  # Take last token
        return self.classifier(x)


# ============================================
# TRAINING LOOP
# ============================================
def train_model(model, train_loader, val_loader, epochs=50, lr=0.001, device='cuda'):
    """Train the event prediction model"""
    model = model.to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=lr)
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=5, factor=0.5)
    
    history = {'train_loss': [], 'val_loss': [], 'val_acc': []}
    
    for epoch in range(epochs):
        model.train()
        train_loss = 0
        for batch_x, batch_y in train_loader:
            batch_x, batch_y = batch_x.to(device), batch_y.to(device)
            
            optimizer.zero_grad()
            outputs = model(batch_x)
            loss = criterion(outputs, batch_y)
            loss.backward()
            optimizer.step()
            train_loss += loss.item()
        
        # Validation
        model.eval()
        val_loss = 0
        correct = 0
        total = 0
        with torch.no_grad():
            for batch_x, batch_y in val_loader:
                batch_x, batch_y = batch_x.to(device), batch_y.to(device)
                outputs = model(batch_x)
                loss = criterion(outputs, batch_y)
                val_loss += loss.item()
                _, predicted = torch.max(outputs.data, 1)
                total += batch_y.size(0)
                correct += (predicted == batch_y).sum().item()
        
        avg_train_loss = train_loss / len(train_loader)
        avg_val_loss = val_loss / len(val_loader)
        val_acc = 100 * correct / total
        
        history['train_loss'].append(avg_train_loss)
        history['val_loss'].append(avg_val_loss)
        history['val_acc'].append(val_acc)
        
        scheduler.step(avg_val_loss)
        
        if (epoch + 1) % 10 == 0:
            print(f'Epoch {epoch+1}/{epochs} | Train Loss: {avg_train_loss:.4f} | Val Loss: {avg_val_loss:.4f} | Val Acc: {val_acc:.2f}%')
    
    return history


# ============================================
# WORLD MAP VISUALIZATION
# ============================================
def plot_world_map_predictions(predictions, threshold=0.5):
    """
    Plot world map with event predictions
    predictions: dict {country: predicted_event_type}
    """
    import plotly.graph_objects as go
    import plotly.express as px
    
    # Prepare data for choropleth
    df = pd.DataFrame([
        {'country': country, 'event': predictions.get(country, 9), 
         'event_name': EVENT_TYPES[predictions.get(country, 9)]}
        for country in GLOBAL_INDICES.keys()
    ])
    
    # Country name mapping for plotly
    country_mapping = {
        "USA": "United States", "UK": "United Kingdom", "Japan": "Japan",
        "Germany": "Germany", "France": "France", "Hong_Kong": "Hong Kong",
        "India": "India", "China": "China", "South_Korea": "South Korea",
        "Brazil": "Brazil", "Canada": "Canada", "Australia": "Australia"
    }
    df['country_name'] = df['country'].map(country_mapping)
    
    # Color scale based on event type
    fig = go.Figure(data=go.Choropleth(
        locations=df['country_name'],
        locationmode='country names',
        z=df['event'],
        text=df['event_name'],
        colorscale='RdYlGn_r',
        colorbar_title="Predicted Event Type",
        zmin=0,
        zmax=9
    ))
    
    fig.update_layout(
        title='Global Event Forecast Map',
        geo=dict(
            showframe=False,
            showcoastlines=True,
            projection_type='equirectangular'
        ),
        width=1000,
        height=600
    )
    
    fig.show()
    return fig


# ============================================
# MAIN EXECUTION
# ============================================
def main():
    print("="*60)
    print("GLOBAL EVENT FORECASTER - PyTorch Implementation")
    print("="*60)
    
    # Step 1: Collect data
    collector = GlobalEventDataCollector(start_date='2025-01-01')
    market_data = collector.fetch_global_indices()
    macro_data = collector.fetch_macro_indicators()
    event_labels = collector.detect_events()
    
    print(f"\nData collected: {len(market_data)} market records, {len(macro_data)} macro records")
    
    # Step 2: Create dataset
    dataset = EventDataset(market_data, macro_data, event_labels, seq_length=30)
    
    # Split train/val
    train_size = int(0.8 * len(dataset))
    val_size = len(dataset) - train_size
    train_dataset, val_dataset = torch.utils.data.random_split(dataset, [train_size, val_size])
    
    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)
    
    # Step 3: Initialize model
    input_dim = dataset.features.shape[1] if len(dataset.features.shape) > 1 else 128
    model = TimeSeriesEventTransformer(
        input_dim=input_dim,
        d_model=128,
        nhead=8,
        num_layers=3,
        num_classes=10,
        dropout=0.1
    )
    
    print(f"\nModel initialized: {sum(p.numel() for p in model.parameters()):,} parameters")
    
    # Step 4: Train
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"Training on: {device.upper()}")
    
    history = train_model(model, train_loader, val_loader, epochs=50, device=device)
    
    # Step 5: Make predictions for latest data
    model.eval()
    with torch.no_grad():
        latest_features = torch.FloatTensor(dataset.features[-10:]).to(device)
        predictions = model(latest_features)
        _, predicted_events = torch.max(predictions, 1)
    
    # Aggregate by country (simplified - using last prediction)
    print("\n" + "="*60)
    print("LATEST EVENT PREDICTIONS BY COUNTRY")
    print("="*60)
    
    latest_predictions = {}
    for i, country in enumerate(GLOBAL_INDICES.keys()):
        if i < len(predicted_events):
            event_type = predicted_events[i].item()
            latest_predictions[country] = event_type
            print(f"{country:15} -> {EVENT_TYPES[event_type]}")
    
    # Step 6: Visualize on world map
    plot_world_map_predictions(latest_predictions)
    
    return model, history, latest_predictions


if __name__ == "__main__":
    model, history, predictions = main()