Chart Visualization

/**
 * Supported chart types
 */
enum ChartType {
  LINE = 'line',
  SCATTER = 'scatter',
  BAR = 'bar',
  PIE = 'pie',
  BOX_AND_WHISKER = 'box_and_whisker',
  SUPERCHART = 'superchart',
  UNKNOWN = 'unknown'
}

/**
 * Axis scale types for charts
 */
enum ScaleType {
  LINEAR = "linear",
  DATETIME = "datetime",
  CATEGORICAL = "categorical",
  LOG = "log",
  SYMLOG = "symlog",
  LOGIT = "logit",
  FUNCTION = "function",
  FUNCTIONLOG = "functionlog",
  ASINH = "asinh"
}

/**
 * Base chart interface
 */
type Chart = {
  type: ChartType;
  title: string;
  elements: any[];
};

type PointData = {
  label: string;
  points: [number | string, number | string][];
};

type LineChart = Chart & {
  type: ChartType.LINE;
  x_label?: string;
  y_label?: string;
  x_unit?: string;
  y_unit?: string;
  x_ticks: (number | string)[];
  x_scale: ScaleType;
  x_tick_labels: string[];
  y_ticks: (number | string)[];
  y_scale: ScaleType;
  y_tick_labels: string[];
  elements: PointData[];
};

type ScatterChart = Chart & {
  type: ChartType.SCATTER;
  x_label?: string;
  y_label?: string;
  x_unit?: string;
  y_unit?: string;
  x_ticks: (number | string)[];
  x_scale: ScaleType;
  x_tick_labels: string[];
  y_ticks: (number | string)[];
  y_scale: ScaleType;
  y_tick_labels: string[];
  elements: PointData[];
};

type BarData = {
  label: string;
  value: string;
  group: string;
};

type BarChart = Chart & {
  type: ChartType.BAR;
  x_label?: string;
  y_label?: string;
  x_unit?: string;
  y_unit?: string;
  elements: BarData[];
};

type PieData = {
  label: string;
  angle: number;
  radius: number;
};

type PieChart = Chart & {
  type: ChartType.PIE;
  elements: PieData[];
};

type BoxAndWhiskerData = {
  label: string;
  min: number;
  first_quartile: number;
  median: number;
  third_quartile: number;
  max: number;
  outliers: number[];
};

type BoxAndWhiskerChart = Chart & {
  type: ChartType.BOX_AND_WHISKER;
  x_label?: string;
  y_label?: string;
  x_unit?: string;
  y_unit?: string;
  elements: BoxAndWhiskerData[];
};

type SuperChart = Chart & {
  type: ChartType.SUPERCHART;
  elements: Chart[];
};

/**
 * Union of all chart types
 */
type ChartTypes = 
  | LineChart 
  | ScatterChart 
  | BarChart 
  | PieChart 
  | BoxAndWhiskerChart 
  | SuperChart;

/**
 * Deserializes chart data into appropriate chart type
 * @param data - Raw chart data
 * @returns Typed chart object
 */
function deserializeChart(data: any): Chart;

import { Sandbox } from "@e2b/code-interpreter";

const sandbox = await Sandbox.create();

// Generate a line chart with matplotlib
const lineChartResult = await sandbox.runCode(`
import matplotlib.pyplot as plt
import numpy as np

x = np.linspace(0, 10, 50)
y = np.sin(x)

plt.figure(figsize=(10, 6))
plt.plot(x, y, label='sin(x)')
plt.xlabel('X values')
plt.ylabel('Y values')
plt.title('Sine Wave')
plt.legend()
plt.show()
`);

// Access chart data
const result = lineChartResult.results[0];
if (result.chart) {
  console.log('Chart type:', result.chart.type); // 'line'
  console.log('Chart title:', result.chart.title);
  
  if (result.chart.type === 'line') {
    const lineChart = result.chart as LineChart;
    console.log('X label:', lineChart.x_label);
    console.log('Y label:', lineChart.y_label);
    console.log('Data points:', lineChart.elements[0].points.length);
  }
}

// Generate a bar chart
const barChartResult = await sandbox.runCode(`
import matplotlib.pyplot as plt

categories = ['A', 'B', 'C', 'D']
values = [23, 45, 56, 78]

plt.figure(figsize=(8, 6))
plt.bar(categories, values)
plt.xlabel('Categories')
plt.ylabel('Values')
plt.title('Bar Chart Example')
plt.show()
`);

if (barChartResult.results[0].chart?.type === 'bar') {
  const barChart = barChartResult.results[0].chart as BarChart;
  console.log('Bar chart elements:', barChart.elements);
}

// Generate a pie chart
const pieChartResult = await sandbox.runCode(`
import matplotlib.pyplot as plt

sizes = [15, 30, 45, 10]
labels = ['Frogs', 'Hogs', 'Dogs', 'Logs']

plt.figure(figsize=(8, 8))
plt.pie(sizes, labels=labels, autopct='%1.1f%%')
plt.title('Pie Chart Example')
plt.show()
`);

if (pieChartResult.results[0].chart?.type === 'pie') {
  const pieChart = pieChartResult.results[0].chart as PieChart;
  console.log('Pie chart slices:', pieChart.elements.length);
}

// Handle box and whisker plots
const boxPlotResult = await sandbox.runCode(`
import matplotlib.pyplot as plt
import numpy as np

# Generate sample data
data = [np.random.normal(0, std, 100) for std in range(1, 4)]

plt.figure(figsize=(8, 6))
plt.boxplot(data, labels=['Group 1', 'Group 2', 'Group 3'])
plt.ylabel('Values')
plt.title('Box and Whisker Plot')
plt.show()
`);

if (boxPlotResult.results[0].chart?.type === 'box_and_whisker') {
  const boxChart = boxPlotResult.results[0].chart as BoxAndWhiskerChart;
  console.log('Box plot data:', boxChart.elements[0]);
  console.log('Median:', boxChart.elements[0].median);
  console.log('Outliers:', boxChart.elements[0].outliers);
}

// Access chart data programmatically
const chartAnalysisResult = await sandbox.runCode(`
import matplotlib.pyplot as plt
import numpy as np

x = np.arange(10)
y = x ** 2

plt.plot(x, y)
plt.title('Quadratic Function')
plt.show()
`);

// Process chart data
const chartResult = chartAnalysisResult.results[0];
if (chartResult.chart) {
  // Access raw chart data
  console.log('Raw chart data:', chartResult.chart);
  
  // Use deserializeChart if needed (usually handled automatically)
  // const deserializedChart = deserializeChart(rawChartData);
}

// Multiple charts (SuperChart)
const multipleChartsResult = await sandbox.runCode(`
import matplotlib.pyplot as plt
import numpy as np

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))

# First subplot
x = np.linspace(0, 10, 100)
ax1.plot(x, np.sin(x))
ax1.set_title('Sine Wave')

# Second subplot
ax2.bar(['A', 'B', 'C'], [1, 3, 2])
ax2.set_title('Bar Chart')

plt.tight_layout()
plt.show()
`);

if (multipleChartsResult.results[0].chart?.type === 'superchart') {
  const superChart = multipleChartsResult.results[0].chart as SuperChart;
  console.log('Number of sub-charts:', superChart.elements.length);
  superChart.elements.forEach((subChart, index) => {
    console.log(`Sub-chart ${index} type:`, subChart.type);
  });
}

const result = await sandbox.runCode(`
import seaborn as sns
import matplotlib.pyplot as plt

# Generate sample data and plot
tips = sns.load_dataset("tips")
sns.scatterplot(data=tips, x="total_bill", y="tip")
plt.show()
`);

// Chart data is available in result.chart
if (result.results[0].chart) {
  const chart = result.results[0].chart;
  console.log('Chart available with type:', chart.type);
  
  // Also available as part of result.data
  if (result.results[0].data) {
    console.log('Structured data also available');
  }
}

// Access multiple formats
console.log('Available formats:', result.results[0].formats());
// May include: ['png', 'svg', 'text', 'data', 'chart']