visualizer.py

import plotly.graph_objects as go
import plotly.express as px
from plotly.subplots import make_subplots
import numpy as np
import pandas as pd

class BrainwaveVisualizer:
    """
    Visualization tools for brainwave data and analysis results.
    """
    
    def __init__(self):
        self.colors = {
            'raw': '#1f77b4',
            'processed': '#ff7f0e',
            'tokens': '#2ca02c',
            'predictions': '#d62728',
            'confidence': '#9467bd'
        }
    
    def plot_raw_eeg(self, data, sampling_rate=256, title="Raw EEG Data"):
        """
        Plot raw EEG data for multiple channels.
        
        Args:
            data: EEG data (samples x channels)
            sampling_rate: Sampling rate in Hz
            title: Plot title
            
        Returns:
            Plotly figure
        """
        num_samples, num_channels = data.shape
        time_vector = np.linspace(0, num_samples / sampling_rate, num_samples)
        
        fig = make_subplots(
            rows=num_channels,
            cols=1,
            shared_xaxes=True,
            subplot_titles=[f'Channel {i+1}' for i in range(num_channels)],
            vertical_spacing=0.02
        )
        
        for ch in range(num_channels):
            fig.add_trace(
                go.Scatter(
                    x=time_vector,
                    y=data[:, ch],
                    mode='lines',
                    name=f'Channel {ch+1}',
                    line=dict(color=self.colors['raw'], width=1),
                    showlegend=ch == 0
                ),
                row=ch+1,
                col=1
            )
        
        fig.update_layout(
            title=title,
            xaxis_title='Time (seconds)',
            height=150 * num_channels,
            showlegend=True
        )
        
        # Update y-axis labels
        for ch in range(num_channels):
            fig.update_yaxes(title_text='Amplitude (μV)', row=ch+1, col=1)
        
        return fig
    
    def plot_processing_comparison(self, raw_data, processed_data, sampling_rate=256):
        """
        Compare raw and processed EEG data.
        """
        time_vector = np.linspace(0, len(raw_data) / sampling_rate, len(raw_data))
        
        fig = make_subplots(
            rows=2,
            cols=1,
            subplot_titles=['Raw Data', 'Processed Data'],
            shared_xaxes=True
        )
        
        # Raw data
        fig.add_trace(
            go.Scatter(
                x=time_vector,
                y=raw_data,
                mode='lines',
                name='Raw',
                line=dict(color=self.colors['raw'], width=1)
            ),
            row=1,
            col=1
        )
        
        # Processed data
        fig.add_trace(
            go.Scatter(
                x=time_vector,
                y=processed_data,
                mode='lines',
                name='Processed',
                line=dict(color=self.colors['processed'], width=1)
            ),
            row=2,
            col=1
        )
        
        fig.update_layout(
            title='Signal Processing Comparison',
            xaxis_title='Time (seconds)',
            height=500
        )
        
        fig.update_yaxes(title_text='Amplitude (μV)', row=1, col=1)
        fig.update_yaxes(title_text='Amplitude (μV)', row=2, col=1)
        
        return fig
    
    def plot_frequency_spectrum(self, data, sampling_rate=256, title="Frequency Spectrum"):
        """
        Plot frequency spectrum of EEG data.
        """
        from scipy import signal
        
        # Compute power spectral density
        freqs, psd = signal.welch(data, sampling_rate, nperseg=min(256, len(data)//4))
        
        fig = go.Figure()
        
        fig.add_trace(
            go.Scatter(
                x=freqs,
                y=10 * np.log10(psd),  # Convert to dB
                mode='lines',
                name='PSD',
                line=dict(color=self.colors['raw'], width=2)
            )
        )
        
        # Mark frequency bands
        band_colors = ['rgba(255,0,0,0.2)', 'rgba(0,255,0,0.2)', 'rgba(0,0,255,0.2)', 
                      'rgba(255,255,0,0.2)', 'rgba(255,0,255,0.2)']
        bands = [
            ('Delta', 0.5, 4),
            ('Theta', 4, 8),
            ('Alpha', 8, 13),
            ('Beta', 13, 30),
            ('Gamma', 30, 50)
        ]
        
        for i, (name, low, high) in enumerate(bands):
            if high <= freqs.max():
                fig.add_vrect(
                    x0=low, x1=high,
                    fillcolor=band_colors[i],
                    opacity=0.3,
                    line_width=0,
                    annotation_text=name,
                    annotation_position="top left"
                )
        
        fig.update_layout(
            title=title,
            xaxis_title='Frequency (Hz)',
            yaxis_title='Power Spectral Density (dB)',
            showlegend=True
        )
        
        return fig
    
    def plot_token_distribution(self, tokens, title="Token Distribution"):
        """
        Plot distribution of tokens.
        """
        # Calculate token frequencies
        unique_tokens, counts = np.unique(tokens, return_counts=True)
        
        fig = go.Figure()
        
        fig.add_trace(
            go.Bar(
                x=unique_tokens,
                y=counts,
                name='Token Frequency',
                marker_color=self.colors['tokens']
            )
        )
        
        fig.update_layout(
            title=title,
            xaxis_title='Token Value',
            yaxis_title='Frequency',
            showlegend=False
        )
        
        return fig
    
    def plot_token_sequence(self, tokens, title="Token Sequence"):
        """
        Plot token sequence over time.
        """
        fig = go.Figure()
        
        fig.add_trace(
            go.Scatter(
                x=list(range(len(tokens))),
                y=tokens,
                mode='lines+markers',
                name='Tokens',
                line=dict(color=self.colors['tokens'], width=2),
                marker=dict(size=4)
            )
        )
        
        fig.update_layout(
            title=title,
            xaxis_title='Token Position',
            yaxis_title='Token Value',
            showlegend=False
        )
        
        return fig
    
    def plot_brain_connectivity(self, connectivity_matrix, channel_names=None):
        """
        Plot brain connectivity matrix.
        """
        if channel_names is None:
            channel_names = [f'Ch{i+1}' for i in range(connectivity_matrix.shape[0])]
        
        fig = go.Figure(data=go.Heatmap(
            z=connectivity_matrix,
            x=channel_names,
            y=channel_names,
            colorscale='RdBu',
            zmid=0,
            text=np.round(connectivity_matrix, 3),
            texttemplate="%{text}",
            textfont={"size": 10},
            hoverongaps=False
        ))
        
        fig.update_layout(
            title='Brain Connectivity Matrix',
            xaxis_title='Channels',
            yaxis_title='Channels'
        )
        
        return fig
    
    def plot_training_progress(self, training_history, title="Model Training Progress"):
        """
        Plot model training progress.
        """
        epochs = list(range(1, len(training_history) + 1))
        
        fig = go.Figure()
        
        fig.add_trace(
            go.Scatter(
                x=epochs,
                y=training_history,
                mode='lines+markers',
                name='Training Loss',
                line=dict(color=self.colors['predictions'], width=2),
                marker=dict(size=6)
            )
        )
        
        fig.update_layout(
            title=title,
            xaxis_title='Epoch',
            yaxis_title='Loss',
            showlegend=False
        )
        
        return fig
    
    def plot_prediction_confidence(self, predictions, confidences, title="Prediction Confidence"):
        """
        Plot predictions with confidence intervals.
        """
        x_values = list(range(len(predictions)))
        
        fig = go.Figure()
        
        # Predictions
        fig.add_trace(
            go.Scatter(
                x=x_values,
                y=predictions,
                mode='lines+markers',
                name='Predictions',
                line=dict(color=self.colors['predictions'], width=2),
                marker=dict(size=6)
            )
        )
        
        # Confidence bands
        upper_bound = np.array(predictions) + np.array(confidences)
        lower_bound = np.array(predictions) - np.array(confidences)
        
        fig.add_trace(
            go.Scatter(
                x=x_values + x_values[::-1],
                y=np.concatenate([upper_bound, lower_bound[::-1]]),
                fill='toself',
                fillcolor='rgba(0,100,80,0.2)',
                line=dict(color='rgba(255,255,255,0)'),
                hoverinfo="skip",
                showlegend=False,
                name='Confidence'
            )
        )
        
        fig.update_layout(
            title=title,
            xaxis_title='Prediction Step',
            yaxis_title='Predicted Value',
            showlegend=True
        )
        
        return fig
    
    def plot_brain_states_timeline(self, timestamps, brain_states, title="Brain States Over Time"):
        """
        Plot brain states over time.
        """
        # Convert brain states to numerical values for plotting
        state_map = {
            'relaxed': 1,
            'focused': 2,
            'excited': 3,
            'drowsy': 0,
            'meditative': 1.5
        }
        
        state_values = [state_map.get(state, 1) for state in brain_states]
        
        fig = go.Figure()
        
        fig.add_trace(
            go.Scatter(
                x=timestamps,
                y=state_values,
                mode='lines+markers',
                name='Brain State',
                line=dict(color=self.colors['raw'], width=2),
                marker=dict(size=6),
                text=brain_states,
                hovertemplate='Time: %{x}<br>State: %{text}<extra></extra>'
            )
        )
        
        # Add state labels
        fig.update_layout(
            title=title,
            xaxis_title='Time',
            yaxis_title='Brain State',
            yaxis=dict(
                tickmode='array',
                tickvals=list(state_map.values()),
                ticktext=list(state_map.keys())
            ),
            showlegend=False
        )
        
        return fig
    
    def create_dashboard_layout(self, raw_data, processed_data, tokens, predictions=None):
        """
        Create a comprehensive dashboard layout.
        """
        from plotly.subplots import make_subplots
        
        fig = make_subplots(
            rows=2,
            cols=2,
            subplot_titles=[
                'Raw EEG Signal',
                'Processed Signal',
                'Token Distribution',
                'Token Sequence'
            ],
            specs=[[{"secondary_y": False}, {"secondary_y": False}],
                   [{"secondary_y": False}, {"secondary_y": False}]]
        )
        
        # Raw signal (first channel only)
        time_vector = np.linspace(0, len(raw_data) / 256, len(raw_data))
        fig.add_trace(
            go.Scatter(
                x=time_vector,
                y=raw_data[:, 0] if raw_data.ndim > 1 else raw_data,
                mode='lines',
                name='Raw',
                line=dict(color=self.colors['raw'], width=1)
            ),
            row=1, col=1
        )
        
        # Processed signal
        fig.add_trace(
            go.Scatter(
                x=time_vector,
                y=processed_data[:, 0] if processed_data.ndim > 1 else processed_data,
                mode='lines',
                name='Processed',
                line=dict(color=self.colors['processed'], width=1)
            ),
            row=1, col=2
        )
        
        # Token distribution
        unique_tokens, counts = np.unique(tokens, return_counts=True)
        fig.add_trace(
            go.Bar(
                x=unique_tokens,
                y=counts,
                name='Token Freq',
                marker_color=self.colors['tokens']
            ),
            row=2, col=1
        )
        
        # Token sequence
        fig.add_trace(
            go.Scatter(
                x=list(range(len(tokens))),
                y=tokens,
                mode='lines+markers',
                name='Token Seq',
                line=dict(color=self.colors['tokens'], width=2),
                marker=dict(size=3)
            ),
            row=2, col=2
        )
        
        fig.update_layout(
            title='Brainwave Analysis Dashboard',
            height=800,
            showlegend=False
        )
        
        return fig