Update visualizer.py

visualizer.py

@@ -1,14 +1,47 @@
+import streamlit as st
 import plotly.express as px
+import plotly.graph_objects as go
 
 class Visualizer:
-    def create_plot(self, df, x_axis, y_axis, chart_type):
-        if chart_type == "Scatter":
-            fig = px.scatter(df, x=x_axis, y=y_axis)
-        elif chart_type == "Line":
-            fig = px.line(df, x=x_axis, y=y_axis)
-        elif chart_type == "Bar":
-            fig = px.bar(df, x=x_axis, y=y_axis)
-        else:
-            raise ValueError("Invalid chart type")
-        
-        return fig
+    def create_visualizations(self, df):
+        chart_type = st.selectbox("Select chart type", 
+                                  ["Scatter", "Line", "Bar", "Histogram", "Box Plot", "Heatmap", "Pie Chart", "3D Scatter"])
+        
+        if chart_type in ["Scatter", "Line", "Bar"]:
+            x_axis = st.selectbox("Select X-axis", df.columns)
+            y_axis = st.selectbox("Select Y-axis", df.columns)
+            color = st.selectbox("Select color (optional)", ["None"] + df.columns.tolist())
+            
+            if chart_type == "Scatter":
+                fig = px.scatter(df, x=x_axis, y=y_axis, color=None if color == "None" else color)
+            elif chart_type == "Line":
+                fig = px.line(df, x=x_axis, y=y_axis, color=None if color == "None" else color)
+            else:  # Bar
+                fig = px.bar(df, x=x_axis, y=y_axis, color=None if color == "None" else color)
+        
+        elif chart_type == "Histogram":
+            column = st.selectbox("Select column", df.columns)
+            fig = px.histogram(df, x=column)
+        
+        elif chart_type == "Box Plot":
+            y_axis = st.selectbox("Select Y-axis", df.columns)
+            x_axis = st.selectbox("Select X-axis (optional)", ["None"] + df.columns.tolist())
+            fig = px.box(df, y=y_axis, x=None if x_axis == "None" else x_axis)
+        
+        elif chart_type == "Heatmap":
+            corr_matrix = df.corr()
+            fig = px.imshow(corr_matrix, color_continuous_scale='RdBu_r')
+        
+        elif chart_type == "Pie Chart":
+            values = st.selectbox("Select values", df.columns)
+            names = st.selectbox("Select names", df.columns)
+            fig = px.pie(df, values=values, names=names)
+        
+        elif chart_type == "3D Scatter":
+            x_axis = st.selectbox("Select X-axis", df.columns)
+            y_axis = st.selectbox("Select Y-axis", df.columns)
+            z_axis = st.selectbox("Select Z-axis", df.columns)
+            color = st.selectbox("Select color (optional)", ["None"] + df.columns.tolist())
+            fig = px.scatter_3d(df, x=x_axis, y=y_axis, z=z_axis, color=None if color == "None" else color)
+
+        st.plotly_chart(fig)