content: """
  This module implements visualization operators for checks in uptrain.
  
  The module provides the `BarChart`, `LineChart`, `ScatterPlot`, `Table`, `Histogram`, `MultiPlot`, `CustomPlotlyChart` and `BubbleChart`  classes, 
  which allow generating various types of visualizations. The charts are created based on 
  the input DataFrame.
  
  """
  
  from __future__ import annotations
  import typing as t
  
  from pydantic import Field
  import polars as pl
  
  if t.TYPE_CHECKING:
      from uptrain.framework import Settings
  from uptrain.operators.base import (
      OpBaseModel,
      register_op,
      TYPE_TABLE_OUTPUT,
  )
  from uptrain.utilities import lazy_load_dep, polars_to_pandas
  
  px = lazy_load_dep("plotly.express", "plotly>=5.0.0")
  ps = lazy_load_dep("plotly.subplots", "plotly>=5.0.0")
  
  __all__ = [
      "BarChart",
      "LineChart",
      "ScatterPlot",
      "Histogram",
      "MultiPlot",
      "CustomPlotlyChart",
      "BubbleChart",
  ]
  
  
  # Not to be used as an operator, only as a base class
  class Chart(OpBaseModel):
      props: dict = Field(default_factory=dict)
      title: str = ""
      x: str = ""
      y: str = ""
      color: str = ""
      description: str = ""
  
      def setup(self, settings: Settings = None):
          self.props = self.props | {
              k: v
              for k, v in [
                  ("x", self.x),
                  ("y", self.y),
                  ("color", self.color),
                  ("title", self.title),
                  ("description", self.description),
              ]
              if v
          }
          return self
  
      def run(self, data: pl.DataFrame) -> TYPE_TABLE_OUTPUT:
          chart = getattr(px, self.kind)(polars_to_pandas(data), **self.props)
  
          # Add annotation based on the description
          if self.description:
              chart.update_layout(
                  annotations=[
                      dict(
                          text=self.description,
                          x=0.5,
                          y=-0.25,
                          xref="paper",
                          yref="paper",
                          showarrow=False,
                          align="center",
                      )
                  ]
              )
  
          return {"output": None, "extra": {"chart": chart}}
  
  
  @register_op
  class CustomPlotlyChart(Chart):
      """
      Operator to generate a custom plotly chart, other than the ones provided by uptrain.
  
      Attributes:
          props (dict): Additional properties to pass to the PlotlyChart constructor.
          title (str): The title of the chart.
          x (str): The name of the column to use for the x-axis.
          y (str): The name of the column to use for the y-axis.
          color (str): The name of the column to use for the color.
          kind (str): The type of chart to generate.
          description (str): Add a description of the chart being created.
  
      Returns:
          dict: A dictionary containing the chart object.
  
      Example:
          ```
          import polars as pl
          from uptrain.operators import CustomPlotlyChart
  
          # Create a DataFrame
          df = pl.DataFrame({
              "x": [1, 2, 3, 4, 5],
              "y": [10, 20, 15, 25, 30]
          })
  
          # Create a funnel chart using the CustomPlotlyChart class
          funnel_chart = CustomPlotlyChart(x="x", y="y", title="Funnel Chart")
  
          # Generate the funnel chart
          chart = funnel_chart.run(df)["extra"]["chart"]
  
          # Show the chart
          chart.show()
          ```
  
      """
  
      props: dict = Field(default_factory=dict)
      title: str = ""
      x: str = ""
      y: str = ""
      description: str = ""
      color: str = ""
      kind: str = Field(default_factory=str)
  
  
  @register_op
  class BarChart(Chart):
      """
      Operator to generate a bar chart.
  
      Attributes:
          props (dict): Additional properties to pass to the BarChart constructor.
          x (str): The name of the column to use for the x-axis.
          y (str): The name of the column to use for the y-axis.
          color (str): The name of the column to use for the color.
          barmode (str): The type of bar chart to generate.
          title (str): The title of the chart.
          description (str): Add a description of the chart being created.
  
      Returns:
          dict: A dictionary containing the chart object.
  
      Example:
          ```
          import polars as pl
          from uptrain.operators import BarChart
  
          # Create a DataFrame
          df = pl.DataFrame({
              "x": [1, 2, 3, 4, 5],
              "y": [10, 20, 15, 25, 30]
          })
  
          # Create a bar chart using the BarChart class
          bar_chart = BarChart(x="x", y="y", title="Bar Chart")
  
          # Generate the bar chart
          chart = bar_chart.run(df)["extra"]["chart"]
  
          # Show the chart
          chart.show()
          ```
  
      """
  
      props: dict = Field(default_factory=dict)
      title: str = ""
      x: str = ""
      y: str = ""
      description: str = ""
      color: str = ""
  
      barmode: str = "group"
  
      kind: str = "bar"
  
  
  @register_op
  class LineChart(Chart):
      """
      Operator to generate a line chart.
  
      Attributes:
          props (dict): Additional properties to pass to the LineChart constructor.
          x (str): The name of the column to use for the x-axis.
          y (str): The name of the column to use for the y-axis.
          color (str): The name of the column to use for the color.
          title (str): The title of the chart.
          description (str): Add a description of the chart being created.
  
      Returns:
          dict: A dictionary containing the chart object.
  
      Example:
          ```
          import polars as pl
          from uptrain.operators import LineChart
  
          # Create a DataFrame
          df = pl.DataFrame({
              "x": [1, 2, 3, 4, 5],
              "y": [10, 20, 15, 25, 30]
          })
  
          # Create a line chart using the LineChart class
          line_chart = LineChart(x="x", y="y", title="Line Chart")
  
          # Generate the line chart
          chart = line_chart.run(df)["extra"]["chart"]
  
          # Show the chart
          chart.show()
          ```
  
      """
  
      props: dict = Field(default_factory=dict)
      title: str = ""
      x: str = ""
      y: str = ""
      description: str = ""
      color: str = ""
  
      kind: str = "line"
  
  
  @register_op
  class ScatterPlot(Chart):
      """
      Operator to generate a scatter chart.
  
      Attributes:
          props (dict): Additional properties to pass to the ScatterPlot constructor.
          title (str): The title of the chart.
          x (str): The name of the column to use for the x-axis.
          y (str): The name of the column to use for the y-axis.
          color (str): The name of the column to use for the color.
          description (str): Add a description of the chart being created.
  
      Returns:
          dict: A dictionary containing the chart object.
  
      Example:
          ```
          import polars as pl
          from uptrain.operators import ScatterPlot
  
          # Create a DataFrame
          df = pl.DataFrame({
              "x": [1, 2, 3, 4, 5],
              "y": [10, 20, 15, 25, 30]
          })
  
          # Create a scatter chart using the ScatterPlot class
          scatter_plot = ScatterPlot(x="x", y="y", title="Scatter Plot")
  
          # Generate the scatter plot
          chart = scatter_plot.run(df)["extra"]["chart"]
  
          # Show the chart
          chart.show()
          ```
  
      """
  
      props: dict = Field(default_factory=dict)
      title: str = ""
      x: str = ""
      y: str = ""
      description: str = ""
      color: str = ""
      symbol: str = "circle"
  
      kind: str = "scatter"
  
  
  @register_op
  class Scatter3DPlot(Chart):
      """
      Operator to generate a 3D scatter chart.
  
      Attributes:
          props (dict): Additional properties to pass to the ScatterPlot constructor.
          title (str): The title of the chart.
          x (str): The name of the column to use for the x-axis.
          y (str): The name of the column to use for the y-axis.
          z (str): The name of the column to use for the z-axis.
          color (str): The name of the column to use for the color.
          description (str): Add a description of the chart being created.
  
      Returns:
          dict: A dictionary containing the chart object.
  
      Example:
          ```
          import polars as pl
          from uptrain.operators import ScatterPlot
  
          # Create a DataFrame
          df = pl.DataFrame({
              "x": [1, 2, 3, 4, 5],
              "y": [10, 20, 15, 25, 30],
              "z": [30, 40, 55, 65, 70],
          })
  
          # Create a scatter chart using the ScatterPlot class
          scatter_plot = Scatter3DPlot(x="x", y="y", z="z", title="Scatter 3D Plot")
  
          # Generate the scatter plot
          chart = scatter_plot.run(df)["extra"]["chart"]
  
          # Show the chart
          chart.show()
          ```
  
      """
  
      props: dict = Field(default_factory=dict)
      title: str = ""
      x: str = ""
      y: str = ""
      z: str = ""
      description: str = ""
      color: str = ""
      symbol: str = ""
  
      kind: str = "scatter_3d"
  
      def setup(self, settings: Settings = None):
          super(Scatter3DPlot, self).setup()
  
          self.props = self.props | {
              k: v
              for k, v in [
                  ("z", self.z),
                  ("symbol", self.symbol),
              ]
              if v
          }
          return self
  
  
  @register_op
  class Histogram(Chart):
      """
      Operator to generate a histogram.
  
      Attributes:
          props (dict): Additional properties to pass to the Histogram chart constructor.
          title (str): The title of the chart.
          x (str): The name of the column to use for the x-axis.
          y (str): The name of the column to use for the y-axis.
          color (str): The name of the column to use for the color.
          nbins (int): The maximum number of bins to use for the histogram.
          description (str): Add a description of the chart being created.
  
      Returns:
          dict: A dictionary containing the chart object.
  
      Example:
          ```
          import polars as pl
          from uptrain.operators import Histogram
  
          # Create a DataFrame
          df = pl.DataFrame({
              "x": [1, 2, 3, 4, 5],
              "y": [10, 20, 15, 25, 30]
          })
  
          # Create a histogram chart using the Histogram class
          histogram_chart = Histogram(x="x", y="y", title="Histogram Chart")
  
          # Generate the histogram chart
          chart = histogram_chart.run(df)["extra"]["chart"]
  
          # Show the chart
          chart.show()
          ```
  
      """
  
      props: dict = Field(default_factory=dict)
      title: str = ""
      x: str = ""
      y: str = ""
      description: str = ""
      color: str = ""
      nbins: int = 20
  
      kind: str = "histogram"
  
  
  @register_op
  class MultiPlot(Chart):
      """
      Operator to generate a subplot that can display multiple charts side-by-side.
  
      Attributes:
          props (dict): Additional properties to pass to the MultiPlot constructor.
          title (str): The title of the chart.
          charts (list): A list of charts to display in the subplot.
          description (str): Add a description of the chart being created. (supports upto 70 characters without overlapping on graph)
  
      Returns:
          dict: A dictionary containing the chart object.
  
      Example:
          ```
          import polars as pl
          from uptrain.operators import MultiPlot, LineChart, ScatterPlot, BarChart, Histogram
  
          # Create a DataFrame
          df = pl.DataFrame({
              "x": [1, 2, 3, 4, 5],
              "y": [10, 20, 15, 25, 30]
          })
  
          # Create a multiplot using the MultiPlot class
          multiplot = MultiPlot(
              props={},
              charts=[
                  LineChart(
                      x="x",
                      y="y",
                      title="Line Chart"
                  ),
                  ScatterPlot(
                      x="x",
                      y="y",
                      title="Scatter Plot"
                  ),
                  BarChart(
                      x="x",
                      y="y",
                      title="Bar Chart"
                  ),
                  Histogram(
                      x="x",
                      y="y",
                      title="Histogram"
                  )
              ],
              title="MultiPlot",
          )
  
          # Generate the multiplot
          chart = multiplot.run(df)["extra"]["chart"]
  
          # Show the chart
          chart.show()
          ```
  
      """
  
      props: dict = Field(default_factory=dict)
      title: str = ""
      description: str = ""
      charts: list
  
      kind: str = "multiplot"
  
      def run(self, data: pl.DataFrame) -> TYPE_TABLE_OUTPUT:
          if isinstance(self.charts[0], dict):
              self.charts = [Chart(**chart).setup() for chart in self.charts]
  
          fig = ps.make_subplots(
              rows=1,
              cols=len(self.charts),
              shared_xaxes=True,
              shared_yaxes=True,
              horizontal_spacing=0.05,
              vertical_spacing=0.2,  # Adjust this value for spacing between graph and annotation
              subplot_titles=[chart.title for chart in self.charts],
          )
  
          annotation_single_height = (
              -0.1
          )  # Adjust this value for single-line annotation position
          annotation_multi_height = (
              -0.3
          )  # Adjust this value for multiline annotation position
  
          for idx, chart in enumerate(self.charts):
              plot = getattr(px, chart.kind)(polars_to_pandas(data), **chart.props)
  
              # Use only the first trace from the plot
              trace = plot.data[0]
              fig.add_trace(trace, row=1, col=idx + 1)
  
              # Break down description into segments of four words each for long descriptions
              words = chart.description.split()
              if len(words) > 12:
                  description_lines = [words[i : i + 4] for i in range(0, len(words), 4)]
                  multiline_description = "<br>".join(
                      " ".join(line) for line in description_lines
                  )
                  annotation_text = multiline_description
                  annotation_height = annotation_multi_height
              else:
                  annotation_text = chart.description
                  annotation_height = annotation_single_height
  
              # Add annotation for the description
              annotation = dict(
                  text=annotation_text,
                  align="center",
                  showarrow=False,
                  xref=f"x{idx + 1}",
                  yref="paper",
                  x=0.5,
                  y=annotation_height,
                  font=dict(size=10),
              )
              fig.add_annotation(annotation)
  
          fig.update_layout(
              title_text=self.title,
              showlegend=False,
              width=800,
              height=400,
              margin=dict(t=100, b=50),  # Adjust the bottom margin to center the title
          )
  
          return {"output": None, "extra": {"chart": fig}}
  
  @register_op
  class BubbleChart(Chart):
      """
      Operator to generate a bubble chart.
  
      Attributes:
          props (dict): Additional properties to pass to the BubbleChart constructor.
          title (str): The title of the chart.
          x (str): The name of the column to use for the x-axis.
          y (str): The name of the column to use for the y-axis.
          color (str): The name of the column to use for the bubble color.
          description (str): Description of the chart being created.
  
      Returns:
          dict: A dictionary containing the chart object.
  
      Example:
          ```
          import polars as pl
          from uptrain.operators import BubbleChart
  
          # Create a DataFrame
          df = pl.DataFrame({
              "x": [1, 2, 3, 4, 5],
              "y": [10, 20, 15, 25, 30],
              "color": [0.2, 0.5, 0.8, 0.3, 0.6]  
          })
  
          # Create a bubble chart using the BubbleChart class
          bubble_chart = BubbleChart(x="x", y="y", color="color", title="Bubble Chart")
  
          # Generate the bubble chart
          chart = bubble_chart.run(df)["extra"]["chart"]
  
          # Show the chart
          chart.show()
          ```
  
      """
  
      props: dict = Field(default_factory=dict)
      title: str = ""
      x: str = ""
      y: str = ""
      color: str = ""
      description: str = ""
  
      kind: str = "bubble"