ComputerAssignments/CA18_memory_systems/ai_visualization.py

#!/usr/bin/env python3

import os
import json
import asyncio
from typing import Dict, List, Tuple, Any, Optional
from dataclasses import dataclass
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import plotly.graph_objects as go
import plotly.express as px
from plotly.subplots import make_subplots
import plotly.figure_factory as ff
from datetime import datetime, timedelta
import logging


import dash
from dash import dcc, html, Input, Output, State, callback_context
import dash_bootstrap_components as dbc
import dash_cytoscape as cyto
from dash import dash_table


from ai_benchmark import AIBenchmarkResult
from llm_memory_systems import LLMMemoryItem

@dataclass
class VisualizationConfig:
    

    output_dir: str = "visualizations"
    figure_size: tuple = (15, 12)
    dpi: int = 300
    style: str = "seaborn-v0_8"
    color_palette: str = "husl"
    save_format: str = "png"
    show_plots: bool = False
    

    enable_dash_app: bool = True
    dash_port: int = 8050
    dash_host: str = "127.0.0.1"
    

    enable_real_time: bool = True
    update_interval: int = 1000
    

    dashboard_title: str = "AI Memory Systems Dashboard"
    dashboard_theme: str = "BOOTSTRAP"

class AIVisualizationSuite:
    
    def __init__(self, config: VisualizationConfig = None):
        self.config = config or VisualizationConfig()
        self.logger = logging.getLogger(self.__class__.__name__)
        self._ensure_output_dir()
        

        plt.style.use(self.config.style)
        sns.set_palette(self.config.color_palette)
        

        self.dash_app = None
        if self.config.enable_dash_app:
            self._initialize_dash_app()
    
    def _ensure_output_dir(self):
        if not os.path.exists(self.config.output_dir):
            os.makedirs(self.config.output_dir)
    
    def _initialize_dash_app(self):
        self.dash_app = dash.Dash(
            __name__,
            external_stylesheets=[dbc.themes.BOOTSTRAP],
            title=self.config.dashboard_title
        )
        

        self._create_dash_layout()
        

        self._create_dash_callbacks()
    
    def _create_dash_layout(self):
        
        self.dash_app.layout = dbc.Container([
            dbc.Row([
                dbc.Col([
                    html.H1("🧠 AI Memory Systems Dashboard", className="text-center mb-4"),
                    html.P("Real-time monitoring and analysis of AI-enhanced memory systems", 
                           className="text-center text-muted")
                ])
            ]),
            

            dbc.Row([
                dbc.Col([
                    dbc.Card([
                        dbc.CardHeader("📊 Control Panel"),
                        dbc.CardBody([
                            dbc.Row([
                                dbc.Col([
                                    html.Label("System Type:"),
                                    dcc.Dropdown(
                                        id="system-type-dropdown",
                                        options=[
                                            {"label": "Semantic Memory", "value": "semantic"},
                                            {"label": "Conversational Memory", "value": "conversational"},
                                            {"label": "Adaptive Memory", "value": "adaptive"}
                                        ],
                                        value="semantic"
                                    )
                                ], width=6),
                                dbc.Col([
                                    html.Label("Time Range:"),
                                    dcc.Dropdown(
                                        id="time-range-dropdown",
                                        options=[
                                            {"label": "Last Hour", "value": "1h"},
                                            {"label": "Last Day", "value": "1d"},
                                            {"label": "Last Week", "value": "1w"},
                                            {"label": "All Time", "value": "all"}
                                        ],
                                        value="1d"
                                    )
                                ], width=6)
                            ]),
                            dbc.Row([
                                dbc.Col([
                                    dbc.Button("🔄 Refresh Data", id="refresh-button", 
                                             color="primary", className="mt-3")
                                ], width=12)
                            ])
                        ])
                    ])
                ], width=3),
                

                dbc.Col([

                    dbc.Row([
                        dbc.Col([
                            dbc.Card([
                                dbc.CardHeader("⚡ Performance Metrics"),
                                dbc.CardBody([
                                    dcc.Graph(id="performance-metrics-graph")
                                ])
                            ])
                        ], width=6),
                        dbc.Col([
                            dbc.Card([
                                dbc.CardHeader("🎯 AI Metrics"),
                                dbc.CardBody([
                                    dcc.Graph(id="ai-metrics-graph")
                                ])
                            ])
                        ], width=6)
                    ], className="mb-4"),
                    

                    dbc.Row([
                        dbc.Col([
                            dbc.Card([
                                dbc.CardHeader("📈 Detailed Analysis"),
                                dbc.CardBody([
                                    dcc.Graph(id="detailed-analysis-graph")
                                ])
                            ])
                        ], width=12)
                    ], className="mb-4"),
                    

                    dbc.Row([
                        dbc.Col([
                            dbc.Card([
                                dbc.CardHeader("🔄 Real-time Monitoring"),
                                dbc.CardBody([
                                    dcc.Graph(id="realtime-graph"),
                                    dcc.Interval(
                                        id="interval-component",
                                        interval=self.config.update_interval,
                                        n_intervals=0
                                    )
                                ])
                            ])
                        ], width=12)
                    ])
                ], width=9)
            ])
        ], fluid=True)
    
    def _create_dash_callbacks(self):
        
        @self.dash_app.callback(
            [Output("performance-metrics-graph", "figure"),
             Output("ai-metrics-graph", "figure"),
             Output("detailed-analysis-graph", "figure"),
             Output("realtime-graph", "figure")],
            [Input("refresh-button", "n_clicks"),
             Input("system-type-dropdown", "value"),
             Input("time-range-dropdown", "value"),
             Input("interval-component", "n_intervals")]
        )
        def update_dashboard(n_clicks, system_type, time_range, n_intervals):
            

            performance_fig = self._create_performance_metrics_chart(system_type, time_range)
            ai_fig = self._create_ai_metrics_chart(system_type, time_range)
            detailed_fig = self._create_detailed_analysis_chart(system_type, time_range)
            realtime_fig = self._create_realtime_chart(system_type, n_intervals)
            
            return performance_fig, ai_fig, detailed_fig, realtime_fig
    
    def _create_performance_metrics_chart(self, system_type: str, time_range: str) -> go.Figure:
        

        metrics = ["Execution Time", "Memory Usage", "CPU Usage", "Error Rate"]
        values = np.random.uniform(0.1, 1.0, len(metrics))
        
        fig = go.Figure(data=[
            go.Bar(
                x=metrics,
                y=values,
                marker_color=['#FF6B6B', '#4ECDC4', '#45B7D1', '#96CEB4'],
                text=[f"{v:.3f}" for v in values],
                textposition='auto'
            )
        ])
        
        fig.update_layout(
            title=f"Performance Metrics - {system_type.title()} System",
            xaxis_title="Metrics",
            yaxis_title="Normalized Values",
            template="plotly_white",
            height=400
        )
        
        return fig
    
    def _create_ai_metrics_chart(self, system_type: str, time_range: str) -> go.Figure:
        

        ai_metrics = ["Semantic Accuracy", "LLM Response Time", "Embedding Quality", "Workflow Efficiency"]
        values = np.random.uniform(0.5, 1.0, len(ai_metrics))
        
        fig = go.Figure(data=[
            go.Scatter(
                x=ai_metrics,
                y=values,
                mode='markers+lines',
                marker=dict(size=15, color='#FF6B6B'),
                line=dict(color='#4ECDC4', width=3)
            )
        ])
        
        fig.update_layout(
            title=f"AI Metrics - {system_type.title()} System",
            xaxis_title="AI Metrics",
            yaxis_title="Score (0-1)",
            template="plotly_white",
            height=400
        )
        
        return fig
    
    def _create_detailed_analysis_chart(self, system_type: str, time_range: str) -> go.Figure:
        

        dates = pd.date_range(start='2025-01-01', periods=30, freq='D')
        values = np.cumsum(np.random.randn(30)) + 100
        
        fig = go.Figure(data=[
            go.Scatter(
                x=dates,
                y=values,
                mode='lines+markers',
                name='Memory Usage',
                line=dict(color='#45B7D1', width=2)
            )
        ])
        
        fig.update_layout(
            title=f"Memory Usage Over Time - {system_type.title()} System",
            xaxis_title="Date",
            yaxis_title="Memory Usage (MB)",
            template="plotly_white",
            height=400
        )
        
        return fig
    
    def _create_realtime_chart(self, system_type: str, n_intervals: int) -> go.Figure:
        

        time_points = list(range(max(0, n_intervals-20), n_intervals))
        values = np.random.uniform(0.1, 1.0, len(time_points))
        
        fig = go.Figure(data=[
            go.Scatter(
                x=time_points,
                y=values,
                mode='lines+markers',
                name='Real-time Performance',
                line=dict(color='#96CEB4', width=2)
            )
        ])
        
        fig.update_layout(
            title=f"Real-time Performance - {system_type.title()} System",
            xaxis_title="Time Steps",
            yaxis_title="Performance Score",
            template="plotly_white",
            height=300
        )
        
        return fig
    
    def plot_ai_benchmark_results(self, results: List[AIBenchmarkResult], save: bool = True) -> str:
        
        if not results:
            self.logger.warning("No benchmark results to visualize")
            return ""
        

        df = pd.DataFrame([
            {
                "system": r.system_name,
                "test_type": r.test_type,
                "execution_time": r.execution_time,
                "memory_usage": r.memory_usage,
                "semantic_accuracy": r.semantic_accuracy,
                "llm_response_time": r.llm_response_time,
                "embedding_quality": r.embedding_quality,
                "workflow_efficiency": r.workflow_efficiency,
                "error_rate": r.error_rate,
                "timestamp": r.timestamp
            }
            for r in results
        ])
        

        fig = make_subplots(
            rows=2, cols=2,
            subplot_titles=("Performance Metrics", "AI Quality Metrics", 
                          "Error Analysis", "Time Series Analysis"),
            specs=[[{"secondary_y": False}, {"secondary_y": False}],
                   [{"secondary_y": False}, {"secondary_y": False}]]
        )
        

        performance_cols = ["execution_time", "memory_usage"]
        for i, col in enumerate(performance_cols):
            fig.add_trace(
                go.Bar(
                    x=df["system"].unique(),
                    y=df.groupby("system")[col].mean(),
                    name=col.replace("_", " ").title(),
                    marker_color=['#FF6B6B', '#4ECDC4'][i]
                ),
                row=1, col=1
            )
        

        ai_cols = ["semantic_accuracy", "embedding_quality", "workflow_efficiency"]
        for i, col in enumerate(ai_cols):
            fig.add_trace(
                go.Scatter(
                    x=df["system"].unique(),
                    y=df.groupby("system")[col].mean(),
                    mode='markers+lines',
                    name=col.replace("_", " ").title(),
                    marker=dict(size=10),
                    line=dict(width=2)
                ),
                row=1, col=2
            )
        

        error_data = df.groupby("test_type")["error_rate"].mean()
        fig.add_trace(
            go.Pie(
                labels=error_data.index,
                values=error_data.values,
                name="Error Rate by Test Type"
            ),
            row=2, col=1
        )
        

        df["timestamp"] = pd.to_datetime(df["timestamp"])
        time_series = df.groupby(df["timestamp"].dt.date)["execution_time"].mean()
        
        fig.add_trace(
            go.Scatter(
                x=time_series.index,
                y=time_series.values,
                mode='lines+markers',
                name="Execution Time Trend",
                line=dict(color='#45B7D1', width=2)
            ),
            row=2, col=2
        )
        

        fig.update_layout(
            title="AI Memory Systems Benchmark Analysis",
            height=800,
            showlegend=True,
            template="plotly_white"
        )
        

        if save:
            filename = f"ai_benchmark_analysis_{datetime.now().strftime('%Y%m%d_%H%M%S')}.html"
            filepath = os.path.join(self.config.output_dir, filename)
            fig.write_html(filepath)
            self.logger.info(f"AI benchmark visualization saved to {filepath}")
            return filepath
        
        return ""
    
    def plot_memory_usage_heatmap(self, memory_items: Dict[str, LLMMemoryItem], save: bool = True) -> str:
        
        if not memory_items:
            self.logger.warning("No memory items to visualize")
            return ""
        

        data = []
        for key, item in memory_items.items():
            data.append({
                "key": key,
                "access_count": item.access_count,
                "importance_score": item.importance_score,
                "days_since_created": (datetime.now() - item.timestamp).days,
                "days_since_accessed": (datetime.now() - item.last_accessed).days
            })
        
        df = pd.DataFrame(data)
        

        correlation_matrix = df[["access_count", "importance_score", "days_since_created", "days_since_accessed"]].corr()
        
        fig = go.Figure(data=go.Heatmap(
            z=correlation_matrix.values,
            x=correlation_matrix.columns,
            y=correlation_matrix.columns,
            colorscale='RdYlBu_r',
            text=correlation_matrix.round(3).values,
            texttemplate="%{text}",
            textfont={"size": 12}
        ))
        
        fig.update_layout(
            title="Memory Item Correlation Heatmap",
            xaxis_title="Metrics",
            yaxis_title="Metrics",
            template="plotly_white",
            height=500
        )
        

        if save:
            filename = f"memory_heatmap_{datetime.now().strftime('%Y%m%d_%H%M%S')}.html"
            filepath = os.path.join(self.config.output_dir, filename)
            fig.write_html(filepath)
            self.logger.info(f"Memory heatmap saved to {filepath}")
            return filepath
        
        return ""
    
    def plot_semantic_similarity_matrix(self, similarity_matrix: np.ndarray, labels: List[str], save: bool = True) -> str:
        
        fig = go.Figure(data=go.Heatmap(
            z=similarity_matrix,
            x=labels,
            y=labels,
            colorscale='Viridis',
            text=similarity_matrix.round(3),
            texttemplate="%{text}",
            textfont={"size": 10}
        ))
        
        fig.update_layout(
            title="Semantic Similarity Matrix",
            xaxis_title="Memory Items",
            yaxis_title="Memory Items",
            template="plotly_white",
            height=600,
            width=800
        )
        

        if save:
            filename = f"semantic_similarity_matrix_{datetime.now().strftime('%Y%m%d_%H%M%S')}.html"
            filepath = os.path.join(self.config.output_dir, filename)
            fig.write_html(filepath)
            self.logger.info(f"Semantic similarity matrix saved to {filepath}")
            return filepath
        
        return ""
    
    def create_interactive_dashboard(self, results: List[AIBenchmarkResult] = None) -> str:
        
        if not self.config.enable_dash_app:
            self.logger.warning("Dash app is disabled in configuration")
            return ""
        
        if not self.dash_app:
            self._initialize_dash_app()
        

        dashboard_url = f"http://{self.config.dash_host}:{self.config.dash_port}"
        
        self.logger.info(f"Starting interactive dashboard at {dashboard_url}")
        


        return dashboard_url
    
    def generate_comprehensive_report(self, results: List[AIBenchmarkResult], 
                                   memory_items: Dict[str, LLMMemoryItem] = None) -> str:
        
        report_files = []
        

        if results:
            benchmark_file = self.plot_ai_benchmark_results(results, save=True)
            if benchmark_file:
                report_files.append(benchmark_file)
        

        if memory_items:
            heatmap_file = self.plot_memory_usage_heatmap(memory_items, save=True)
            if heatmap_file:
                report_files.append(heatmap_file)
        

        if memory_items:

            embeddings = []
            labels = []
            for key, item in memory_items.items():
                if item.embedding:
                    embeddings.append(item.embedding)
                    labels.append(key)
            
            if embeddings and len(embeddings) > 1:
                embeddings_array = np.array(embeddings)
                similarity_matrix = np.dot(embeddings_array, embeddings_array.T)
                
                similarity_file = self.plot_semantic_similarity_matrix(
                    similarity_matrix, labels, save=True
                )
                if similarity_file:
                    report_files.append(similarity_file)
        

        report_summary = {
            "generated_at": datetime.now().isoformat(),
            "total_visualizations": len(report_files),
            "files": report_files,
            "dashboard_url": self.create_interactive_dashboard(results) if self.config.enable_dash_app else None
        }
        

        summary_file = os.path.join(self.config.output_dir, f"visualization_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json")
        with open(summary_file, "w") as f:
            json.dump(report_summary, f, indent=2)
        
        self.logger.info(f"Comprehensive visualization report generated: {summary_file}")
        return summary_file


def test_ai_visualization():
    
    print("🎨 Testing AI Visualization System")
    print("=" * 50)
    

    config = VisualizationConfig(
        output_dir="visualizations",
        enable_dash_app=False,
        enable_real_time=False
    )
    
    visualizer = AIVisualizationSuite(config)
    

    from ai_benchmark import AIBenchmarkResult
    
    sample_results = [
        AIBenchmarkResult(
            system_name="semantic_memory",
            test_type="semantic_accuracy",
            timestamp=datetime.now(),
            execution_time=1.5,
            memory_usage=1024,
            cpu_usage=25.0,
            semantic_accuracy=0.85,
            llm_response_time=2.3,
            embedding_quality=0.78,
            workflow_efficiency=0.92,
            error_rate=0.05
        ),
        AIBenchmarkResult(
            system_name="conversational_memory",
            test_type="workflow_efficiency",
            timestamp=datetime.now(),
            execution_time=2.1,
            memory_usage=1536,
            cpu_usage=30.0,
            semantic_accuracy=0.78,
            llm_response_time=3.1,
            embedding_quality=0.82,
            workflow_efficiency=0.88,
            error_rate=0.08
        )
    ]
    

    from llm_memory_systems import LLMMemoryItem
    
    sample_memory_items = {
        "user_prefs": LLMMemoryItem(
            key="user_prefs",
            content="User prefers dark mode",
            access_count=15,
            importance_score=0.8,
            semantic_tags=["preferences", "ui"]
        ),
        "project_info": LLMMemoryItem(
            key="project_info",
            content="AI memory systems project",
            access_count=8,
            importance_score=0.9,
            semantic_tags=["project", "ai"]
        )
    }
    

    print("\n📊 Generating AI benchmark visualization...")
    benchmark_file = visualizer.plot_ai_benchmark_results(sample_results, save=True)
    print(f"✓ Benchmark visualization: {benchmark_file}")
    
    print("\n🔥 Generating memory usage heatmap...")
    heatmap_file = visualizer.plot_memory_usage_heatmap(sample_memory_items, save=True)
    print(f"✓ Memory heatmap: {heatmap_file}")
    
    print("\n📈 Generating comprehensive report...")
    report_file = visualizer.generate_comprehensive_report(sample_results, sample_memory_items)
    print(f"✓ Comprehensive report: {report_file}")
    
    print("\n✅ AI visualization testing completed!")
    
    return visualizer, report_file

if __name__ == "__main__":

    visualizer, report_file = test_ai_visualization()
    
    print(f"\n📁 All visualizations saved to: {visualizer.config.output_dir}")
    print(f"📊 Report file: {report_file}")