visualizers.py

import numpy as np
import pandas as pd
import plotly.graph_objects as go
import plotly.express as px
from plotly.subplots import make_subplots
from pathlib import Path
from collections import defaultdict


class PowerballVisualizer:
    """Creates visualizations for Powerball data and predictions"""
    
    def __init__(self, data_handler):
        self.data_handler = data_handler
        
    def create_visualizations(self, save_path="visualizations"):
        """Create comprehensive visualizations"""
        viz_dir = Path(save_path)
        viz_dir.mkdir(exist_ok=True)
        
        print("Creating visualizations...")
        
        # Create different types of visualizations
        self._create_frequency_heatmap(viz_dir)
        self._create_temporal_plots(viz_dir)
        self._create_correlation_matrix(viz_dir)
        self._create_distribution_plots(viz_dir)
        
        print(f"Visualizations saved to: {viz_dir}")
        return viz_dir
    
    def _create_frequency_heatmap(self, viz_dir):
        """Create frequency heatmap for number positions"""
        # Prepare frequency matrix
        freq_matrix = np.zeros((6, 10))  # 6 rows (5 whites + PB), 10 columns for binned numbers
        
        for _, row in self.data_handler.data.iterrows():
            # White balls (binned into groups of 7: 1-7, 8-14, ..., 64-69)
            for i, col in enumerate(self.data_handler.white_cols):
                white_num = row[col]
                bin_idx = min(9, (white_num - 1) // 7)  # 0-9 bins
                freq_matrix[i, bin_idx] += 1
            
            # Powerball (binned into groups of 3: 1-3, 4-6, ..., 25-26)
            pb = row['Powerball']
            pb_bin = min(9, (pb - 1) // 3)
            freq_matrix[5, pb_bin] += 1
        
        # Create heatmap
        fig = go.Figure(data=go.Heatmap(
            z=freq_matrix,
            x=[f"{i*7+1}-{min(69, (i+1)*7)}" for i in range(10)],
            y=["White 1", "White 2", "White 3", "White 4", "White 5", "Powerball"],
            colorscale='Viridis',
            title="Number Frequency Heatmap by Position"
        ))
        
        fig.update_layout(
            title="Powerball Number Frequency Analysis",
            xaxis_title="Number Ranges",
            yaxis_title="Position"
        )
        
        fig.write_html(viz_dir / "frequency_heatmap.html")
    
    def _create_temporal_plots(self, viz_dir):
        """Create temporal analysis plots"""
        # Prepare temporal data
        df_temporal = self.data_handler.data.copy()
        
        # Create subplots
        fig = make_subplots(
            rows=3, cols=2,
            subplot_titles=("Number Frequency Over Time", "Sum Trends", 
                          "High/Low Ratio", "Powerball Trend", "Draw Frequency"),
            specs=[[{"secondary_y": False}, {"secondary_y": False}],
                   [{"secondary_y": False}, {"secondary_y": False}],
                   [{"secondary_y": False}, {"secondary_y": False}]]
        )
        
        # 1. Rolling averages for each position
        for i, col in enumerate(self.data_handler.white_cols):
            if col in df_temporal.columns:
                rolling_avg = df_temporal[col].rolling(window=20).mean()
                fig.add_trace(go.Scatter(
                    x=df_temporal['Date'], y=rolling_avg,
                    name=f"White {i+1}", mode='lines'
                ), row=1, col=1)
        
        # 2. Sum trends
        if all(col in df_temporal.columns for col in self.data_handler.white_cols):
            df_temporal['Sum'] = df_temporal[self.data_handler.white_cols].sum(axis=1)
            rolling_sum = df_temporal['Sum'].rolling(window=20).mean()
            fig.add_trace(go.Scatter(
                x=df_temporal['Date'], y=rolling_sum,
                name="Sum Trend", line=dict(color='orange')
            ), row=1, col=2)
        
        # 3. High/Low ratio over time
        if all(col in df_temporal.columns for col in self.data_handler.white_cols):
            high_count = df_temporal[self.data_handler.white_cols].apply(
                lambda row: sum(1 for x in row if x > 34), axis=1
            )
            high_ratio = high_count.rolling(window=20).mean()
            fig.add_trace(go.Scatter(
                x=df_temporal['Date'], y=high_ratio,
                name="High Numbers Ratio", line=dict(color='red')
            ), row=2, col=1)
        
        # 4. Powerball trends
        if 'Powerball' in df_temporal.columns:
            pb_rolling = df_temporal['Powerball'].rolling(window=20).mean()
            fig.add_trace(go.Scatter(
                x=df_temporal['Date'], y=pb_rolling,
                name="Powerball", line=dict(color='purple')
            ), row=2, col=2)
        
        # 5. Draw frequency histogram
        if 'Date' in df_temporal.columns:
            df_temporal['Month'] = df_temporal['Date'].dt.month
            month_counts = df_temporal['Month'].value_counts().sort_index()
            fig.add_trace(go.Bar(
                x=month_counts.index, y=month_counts.values,
                name="Draws per Month"
            ), row=3, col=1)
        
        fig.update_layout(height=800, title_text="Temporal Analysis Dashboard")
        fig.write_html(viz_dir / "temporal_analysis.html")
    
    def _create_correlation_matrix(self, viz_dir):
        """Create correlation matrix for white ball positions"""
        if all(col in self.data_handler.data.columns for col in self.data_handler.white_cols):
            corr_matrix = self.data_handler.data[self.data_handler.white_cols].corr()
            
            fig = go.Figure(data=go.Heatmap(
                z=corr_matrix.values,
                x=corr_matrix.columns,
                y=corr_matrix.columns,
                colorscale='RdBu',
                zmid=0,
                text=np.round(corr_matrix.values, 3),
                texttemplate="%{text}",
                textfont={"size": 10}
            ))
            
            fig.update_layout(
                title="White Ball Position Correlations",
                xaxis_title="Position",
                yaxis_title="Position"
            )
            
            fig.write_html(viz_dir / "correlation_matrix.html")
    
    def _create_distribution_plots(self, viz_dir):
        """Create distribution plots for numbers"""
        fig = make_subplots(
            rows=2, cols=3,
            subplot_titles=("White Ball Distributions", "Powerball Distribution",
                          "Sum Distribution", "Gap Analysis", "Hot/Cold Analysis", "Pattern Analysis")
        )
        
        # 1. White ball distributions
        for i, col in enumerate(self.data_handler.white_cols):
            if col in self.data_handler.data.columns:
                fig.add_trace(go.Histogram(
                    x=self.data_handler.data[col],
                    name=f"White {i+1}",
                    opacity=0.7,
                    nbinsx=20
                ), row=1, col=1)
        
        # 2. Powerball distribution
        if 'Powerball' in self.data_handler.data.columns:
            fig.add_trace(go.Histogram(
                x=self.data_handler.data['Powerball'],
                name="Powerball",
                marker_color='red',
                nbinsx=26
            ), row=1, col=2)
        
        # 3. Sum distribution
        if all(col in self.data_handler.data.columns for col in self.data_handler.white_cols):
            sums = self.data_handler.data[self.data_handler.white_cols].sum(axis=1)
            fig.add_trace(go.Histogram(
                x=sums,
                name="Sum Distribution",
                marker_color='green'
            ), row=1, col=3)
        
        # 4. Gap analysis (simplified)
        recent_data = self.data_handler.data.head(50)
        if all(col in recent_data.columns for col in self.data_handler.white_cols):
            gaps = []
            for num in range(1, 70):
                for col in self.data_handler.white_cols:
                    last_appearance = None
                    for idx, row in recent_data.iterrows():
                        if row[col] == num:
                            last_appearance = idx
                            break
                    gap = last_appearance if last_appearance is not None else len(recent_data)
                    gaps.append(gap)
            
            fig.add_trace(go.Histogram(
                x=gaps,
                name="Number Gaps",
                marker_color='orange'
            ), row=2, col=1)
        
        # 5. Hot/Cold analysis
        hot_window = min(30, len(self.data_handler.data))
        if hot_window > 0 and all(col in self.data_handler.data.columns for col in self.data_handler.white_cols):
            hot_data = self.data_handler.data.head(hot_window)
            hot_counts = []
            for col in self.data_handler.white_cols:
                hot_counts.extend(hot_data[col].tolist())
            
            fig.add_trace(go.Histogram(
                x=hot_counts,
                name="Hot Numbers",
                marker_color='red',
                opacity=0.7
            ), row=2, col=2)
        
        # 6. Pattern analysis (odd/even distribution)
        if all(col in self.data_handler.data.columns for col in self.data_handler.white_cols):
            odd_counts = []
            for _, row in self.data_handler.data.iterrows():
                odd_count = sum(1 for col in self.data_handler.white_cols if row[col] % 2 == 1)
                odd_counts.append(odd_count)
            
            fig.add_trace(go.Histogram(
                x=odd_counts,
                name="Odd Numbers per Draw",
                marker_color='purple'
            ), row=2, col=3)
        
        fig.update_layout(height=800, title_text="Statistical Distribution Analysis")
        fig.write_html(viz_dir / "distributions.html")
    
    def create_prediction_comparison(self, predictions, historical_matches=None, save_path="visualizations"):
        """Create visualization comparing predictions with historical data"""
        viz_dir = Path(save_path)
        viz_dir.mkdir(exist_ok=True)
        
        if not predictions:
            print("No predictions provided for visualization")
            return
        
        # Create comparison plots
        fig = make_subplots(
            rows=2, cols=2,
            subplot_titles=("Predicted vs Historical Frequencies", "Sum Comparison",
                          "Odd/Even Distribution", "High/Low Distribution")
        )
        
        # 1. Frequency comparison
        pred_frequencies = defaultdict(int)
        hist_frequencies = defaultdict(int)
        
        # Count prediction frequencies
        for pred in predictions:
            for num in pred[:5]:  # White balls only
                pred_frequencies[num] += 1
        
        # Count historical frequencies (recent 100 draws)
        recent_data = self.data_handler.data.head(100)
        for _, row in recent_data.iterrows():
            for col in self.data_handler.white_cols:
                if col in row:
                    hist_frequencies[row[col]] += 1
        
        # Plot frequency comparison
        numbers = sorted(set(list(pred_frequencies.keys()) + list(hist_frequencies.keys())))
        pred_counts = [pred_frequencies.get(num, 0) for num in numbers]
        hist_counts = [hist_frequencies.get(num, 0) for num in numbers]
        
        fig.add_trace(go.Bar(
            x=numbers, y=pred_counts, name="Predictions", opacity=0.7
        ), row=1, col=1)
        fig.add_trace(go.Bar(
            x=numbers, y=hist_counts, name="Historical", opacity=0.7
        ), row=1, col=1)
        
        # 2. Sum comparison
        pred_sums = [sum(pred[:5]) for pred in predictions]
        hist_sums = recent_data[self.data_handler.white_cols].sum(axis=1).tolist()
        
        fig.add_trace(go.Histogram(
            x=pred_sums, name="Predicted Sums", opacity=0.7, nbinsx=20
        ), row=1, col=2)
        fig.add_trace(go.Histogram(
            x=hist_sums, name="Historical Sums", opacity=0.7, nbinsx=20
        ), row=1, col=2)
        
        # 3. Odd/Even comparison
        pred_odd_counts = [sum(1 for num in pred[:5] if num % 2 == 1) for pred in predictions]
        hist_odd_counts = []
        for _, row in recent_data.iterrows():
            odd_count = sum(1 for col in self.data_handler.white_cols if row[col] % 2 == 1)
            hist_odd_counts.append(odd_count)
        
        fig.add_trace(go.Histogram(
            x=pred_odd_counts, name="Predicted Odd Count", opacity=0.7
        ), row=2, col=1)
        fig.add_trace(go.Histogram(
            x=hist_odd_counts, name="Historical Odd Count", opacity=0.7
        ), row=2, col=1)
        
        # 4. High/Low comparison (>34 is high)
        pred_high_counts = [sum(1 for num in pred[:5] if num > 34) for pred in predictions]
        hist_high_counts = []
        for _, row in recent_data.iterrows():
            high_count = sum(1 for col in self.data_handler.white_cols if row[col] > 34)
            hist_high_counts.append(high_count)
        
        fig.add_trace(go.Histogram(
            x=pred_high_counts, name="Predicted High Count", opacity=0.7
        ), row=2, col=2)
        fig.add_trace(go.Histogram(
            x=hist_high_counts, name="Historical High Count", opacity=0.7
        ), row=2, col=2)
        
        fig.update_layout(height=600, title_text="Prediction vs Historical Analysis")
        fig.write_html(viz_dir / "prediction_comparison.html")
        
        print(f"Prediction comparison visualization saved to: {viz_dir / 'prediction_comparison.html'}")
        return viz_dir / "prediction_comparison.html"