Probability and Statistics

class SampleSpace(Rectangle):
    def __init__(
        self,
        width: float = 3,
        height: float = 3,
        fill_color: ManimColor = GREY_D,
        fill_opacity: float = 1,
        stroke_width: float = 0.5,
        stroke_color: ManimColor = GREY_B,
        default_label_scale_val: float = 1,
        **kwargs
    ):
        """
        Create a visual sample space for probability demonstrations.
        
        Parameters:
        - width: Width of the sample space rectangle
        - height: Height of the sample space rectangle
        - fill_color: Base color for the sample space
        - fill_opacity: Opacity of the filled region
        - stroke_width: Border thickness
        - stroke_color: Border color
        - default_label_scale_val: Scale factor for labels
        """
    
    def add_title(self, title: str = "Sample space", buff: float = MED_SMALL_BUFF) -> None:
        """Add a title above the sample space."""
    
    def add_label(self, label: str) -> None:
        """Add a label to the sample space."""
    
    def complete_p_list(self, p_list: list[float]) -> list[float]:
        """
        Ensure probability list sums to 1 by adding remainder.
        
        Parameters:
        - p_list: List of probability values
        
        Returns:
        Completed probability list that sums to 1
        """
    
    def get_horizontal_division(
        self, 
        p_list, 
        colors=[GREEN_E, BLUE_E], 
        vect=DOWN
    ) -> VGroup:
        """
        Create horizontal subdivision of sample space.
        
        Parameters:
        - p_list: List of probability values for division
        - colors: Colors for each subdivision
        - vect: Direction vector for subdivision
        
        Returns:
        VGroup of divided regions
        """
    
    def get_vertical_division(
        self, 
        p_list, 
        colors=[MAROON_B, YELLOW], 
        vect=RIGHT
    ) -> VGroup:
        """
        Create vertical subdivision of sample space.
        
        Parameters:
        - p_list: List of probability values for division
        - colors: Colors for each subdivision  
        - vect: Direction vector for subdivision
        
        Returns:
        VGroup of divided regions
        """
    
    def divide_horizontally(self, *args, **kwargs) -> None:
        """Divide the sample space horizontally and add to scene."""
    
    def divide_vertically(self, *args, **kwargs) -> None:
        """Divide the sample space vertically and add to scene."""
    
    def get_subdivision_braces_and_labels(
        self, 
        parts, 
        labels, 
        direction, 
        buff=SMALL_BUFF
    ) -> VGroup:
        """
        Create braces and labels for subdivisions.
        
        Parameters:
        - parts: List of subdivision parts
        - labels: List of label strings
        - direction: Direction for brace placement
        - buff: Buffer distance from subdivision
        
        Returns:
        VGroup containing braces and labels
        """
    
    def get_side_braces_and_labels(
        self, 
        labels, 
        direction=LEFT, 
        **kwargs
    ) -> VGroup:
        """Create side braces and labels for sample space."""
    
    def get_top_braces_and_labels(self, labels, **kwargs) -> VGroup:
        """Create top braces and labels for sample space."""
    
    def get_bottom_braces_and_labels(self, labels, **kwargs) -> VGroup:
        """Create bottom braces and labels for sample space."""
    
    def __getitem__(self, index: int | slice) -> VGroup:
        """Access subdivisions by index."""

class BarChart(VGroup):
    def __init__(
        self,
        values: Iterable[float],
        height: float = 4,
        width: float = 6,
        n_ticks: int = 4,
        include_x_ticks: bool = False,
        tick_width: float = 0.2,
        tick_height: float = 0.15,
        label_y_axis: bool = True,
        y_axis_label_height: float = 0.25,
        max_value: float = 1,
        bar_colors: list[ManimColor] = [BLUE, YELLOW],
        bar_fill_opacity: float = 0.8,
        bar_stroke_width: float = 3,
        bar_names: list[str] = [],
        bar_label_scale_val: float = 0.75,
        **kwargs
    ):
        """
        Create a statistical bar chart with axes and labels.
        
        Parameters:
        - values: Iterable of numeric values for bar heights
        - height: Total height of the chart
        - width: Total width of the chart
        - n_ticks: Number of tick marks on y-axis
        - include_x_ticks: Whether to include x-axis tick marks
        - tick_width: Width of tick marks
        - tick_height: Height of tick marks
        - label_y_axis: Whether to label the y-axis
        - y_axis_label_height: Height of y-axis labels
        - max_value: Maximum value for scaling bars
        - bar_colors: Colors for bars (gradient applied)
        - bar_fill_opacity: Opacity of bar fills
        - bar_stroke_width: Width of bar borders
        - bar_names: Labels for individual bars
        - bar_label_scale_val: Scale factor for bar labels
        """
    
    def add_axes(self) -> None:
        """Create and add x and y axes with tick marks."""
    
    def add_bars(self, values: Iterable[float]) -> None:
        """
        Create bars proportional to input values.
        
        Parameters:
        - values: Numeric values determining bar heights
        """
    
    def change_bar_values(self, values: Iterable[float]) -> None:
        """
        Update bar heights dynamically while maintaining positions.
        
        Parameters:
        - values: New values for bar heights
        """

from manimlib import *

class BasicProbability(Scene):
    def construct(self):
        # Create a sample space
        sample_space = SampleSpace(width=4, height=3)
        sample_space.add_title("Rolling a Die")
        sample_space.set_fill(BLUE_E, opacity=0.3)
        
        self.play(ShowCreation(sample_space))
        self.wait()
        
        # Show equally likely outcomes
        probabilities = [1/6] * 6
        horizontal_div = sample_space.get_horizontal_division(
            probabilities, 
            colors=[RED, ORANGE, YELLOW, GREEN, BLUE, PURPLE]
        )
        
        self.play(ShowCreation(horizontal_div))
        
        # Add labels for outcomes
        outcome_labels = ["1", "2", "3", "4", "5", "6"]
        braces_and_labels = sample_space.get_bottom_braces_and_labels(
            outcome_labels
        )
        
        self.play(Write(braces_and_labels))
        self.wait()

class ConditionalProbability(Scene):
    def construct(self):
        # Create sample space for conditional probability
        sample_space = SampleSpace(width=5, height=4)
        sample_space.add_title("P(B|A) - Conditional Probability")
        
        self.add(sample_space)
        
        # First division: P(A) and P(A')
        sample_space.divide_horizontally([0.4, 0.6])
        
        # Add horizontal labels
        h_labels = sample_space.get_side_braces_and_labels(["A", "A'"])
        self.play(Write(h_labels))
        self.wait()
        
        # Second division: P(B|A) and P(B'|A) within each region
        sample_space.divide_vertically([0.7, 0.3])
        
        # Add vertical labels
        v_labels = sample_space.get_top_braces_and_labels(["B", "B'"])
        self.play(Write(v_labels))
        self.wait()
        
        # Highlight the intersection P(A ∩ B)
        intersection = sample_space[0][0]  # Top-left region
        intersection.set_fill(YELLOW, opacity=0.8)
        
        # Add calculation text
        calc_text = TexText(
            "P(A ∩ B) = P(A) × P(B|A) = 0.4 × 0.7 = 0.28",
            font_size=36
        ).to_edge(DOWN)
        
        self.play(
            intersection.animate.set_fill(YELLOW, opacity=0.8),
            Write(calc_text)
        )
        self.wait()

class ProbabilityDistribution(Scene):
    def construct(self):
        # Binomial distribution example
        n, p = 4, 0.3
        
        # Calculate binomial probabilities
        from math import comb
        values = []
        labels = []
        
        for k in range(n + 1):
            prob = comb(n, k) * (p ** k) * ((1 - p) ** (n - k))
            values.append(prob)
            labels.append(f"X = {k}")
        
        # Create bar chart
        chart = BarChart(
            values,
            height=4,
            width=8,
            max_value=0.5,
            bar_names=labels,
            bar_colors=[BLUE, RED],
            include_x_ticks=True
        )
        
        # Add title and axis labels
        title = Text("Binomial Distribution: n=4, p=0.3", font_size=32)
        title.to_edge(UP)
        
        y_label = Text("P(X = k)", font_size=24)
        y_label.rotate(PI/2)
        y_label.next_to(chart, LEFT)
        
        x_label = Text("Number of Successes (k)", font_size=24)
        x_label.next_to(chart, DOWN)
        
        self.play(ShowCreation(chart))
        self.play(Write(title), Write(x_label), Write(y_label))
        
        # Animate probability calculation
        for i, (val, label) in enumerate(zip(values, labels)):
            prob_text = Text(f"P({label.split()[-1]}) = {val:.3f}", font_size=20)
            prob_text.next_to(chart[1][i], UP, buff=0.1)
            
            self.play(Write(prob_text), run_time=0.5)
        
        self.wait()

from manimlib.mobject.interactive import LinearNumberSlider

class InteractiveStats(Scene):
    def setup(self):
        # Create sliders for parameters
        self.n_slider = LinearNumberSlider(
            value=10, min_value=5, max_value=50, step=1
        )
        self.p_slider = LinearNumberSlider(
            value=0.5, min_value=0.1, max_value=0.9, step=0.05
        )
        
        self.n_slider.to_edge(DOWN).shift(UP * 0.5)
        self.p_slider.to_edge(DOWN)
        
        self.add(self.n_slider, self.p_slider)
    
    def construct(self):
        # Create responsive chart
        def get_binomial_chart():
            n = int(self.n_slider.get_value())
            p = self.p_slider.get_value()
            
            # Calculate probabilities
            from math import comb
            values = []
            for k in range(min(n + 1, 20)):  # Limit for display
                prob = comb(n, k) * (p ** k) * ((1 - p) ** (n - k))
                values.append(prob)
            
            return BarChart(
                values,
                height=3,
                width=6,
                max_value=0.4,
                bar_colors=[BLUE, YELLOW]
            )
        
        chart = get_binomial_chart()
        chart.to_edge(UP)
        
        # Add updater
        def update_chart(c):
            new_chart = get_binomial_chart()
            new_chart.to_edge(UP)
            c.become(new_chart)
        
        chart.add_updater(update_chart)
        
        # Add parameter labels
        n_label = Text("n (trials)", font_size=20).next_to(self.n_slider, LEFT)
        p_label = Text("p (success prob)", font_size=20).next_to(self.p_slider, LEFT)
        
        # Dynamic title
        title = Text("Interactive Binomial Distribution", font_size=28)
        title.to_edge(UP).shift(DOWN * 0.5)
        
        def update_title(t):
            n = int(self.n_slider.get_value())
            p = self.p_slider.get_value()
            new_title = Text(
                f"Binomial Distribution: n={n}, p={p:.2f}", 
                font_size=28
            )
            new_title.to_edge(UP).shift(DOWN * 0.5)
            t.become(new_title)
        
        title.add_updater(update_title)
        
        self.add(chart, n_label, p_label, title)
        self.wait(15)  # Interactive exploration time

class MonteCarloDemo(Scene):
    def construct(self):
        # Set up sample space for coin flips
        sample_space = SampleSpace(width=3, height=2)
        sample_space.add_title("Coin Flip Outcomes")
        sample_space.to_edge(LEFT)
        
        # Divide for heads/tails
        sample_space.divide_horizontally([0.5, 0.5])
        sample_space[0].set_fill(GREEN, opacity=0.7)
        sample_space[1].set_fill(RED, opacity=0.7)
        
        labels = sample_space.get_side_braces_and_labels(["H", "T"])
        
        # Results chart
        chart = BarChart(
            [0, 0],  # Start with no results
            height=3,
            width=4,
            max_value=1,
            bar_names=["Heads", "Tails"],
            bar_colors=[GREEN, RED]
        )
        chart.to_edge(RIGHT)
        
        self.add(sample_space, labels, chart)
        
        # Simulate coin flips
        import random
        heads_count = 0
        total_flips = 0
        
        counter_text = Text("Flips: 0", font_size=24).to_edge(UP)
        self.add(counter_text)
        
        for flip in range(100):
            total_flips += 1
            result = random.choice([0, 1])  # 0 = heads, 1 = tails
            
            if result == 0:
                heads_count += 1
            
            # Update proportions
            heads_prop = heads_count / total_flips
            tails_prop = 1 - heads_prop
            
            # Update chart
            chart.change_bar_values([heads_prop, tails_prop])
            
            # Update counter
            counter_text.become(
                Text(f"Flips: {total_flips}", font_size=24).to_edge(UP)
            )
            
            if flip % 10 == 0:  # Update every 10 flips for animation
                self.wait(0.1)
        
        # Final result
        final_text = Text(
            f"Final: {heads_count}/{total_flips} = {heads_prop:.3f}",
            font_size=24
        ).to_edge(DOWN)
        
        self.play(Write(final_text))
        self.wait()

class HypothesisTest(Scene):
    def construct(self):
        # Set up hypothesis testing scenario
        title = Text("Hypothesis Testing", font_size=36).to_edge(UP)
        
        # Null and alternative hypotheses
        h0 = TexText("H₀: p = 0.5 (fair coin)", font_size=24)
        h1 = TexText("H₁: p ≠ 0.5 (biased coin)", font_size=24)
        
        hypotheses = VGroup(h0, h1).arrange(DOWN, buff=0.3)
        hypotheses.to_edge(LEFT).shift(UP)
        
        # Sample space showing expected vs observed
        expected_space = SampleSpace(width=2.5, height=1.5)
        expected_space.add_title("Expected (H₀)")
        expected_space.divide_horizontally([0.5, 0.5])
        expected_space[0].set_fill(BLUE, opacity=0.6)
        expected_space[1].set_fill(RED, opacity=0.6)
        
        observed_space = SampleSpace(width=2.5, height=1.5)
        observed_space.add_title("Observed")
        observed_space.divide_horizontally([0.3, 0.7])  # Biased result
        observed_space[0].set_fill(BLUE, opacity=0.6)
        observed_space[1].set_fill(RED, opacity=0.6)
        
        spaces = VGroup(expected_space, observed_space)
        spaces.arrange(RIGHT, buff=1).to_edge(DOWN).shift(UP * 0.5)
        
        # Statistical analysis
        analysis = VGroup(
            Text("Sample size: n = 100", font_size=20),
            Text("Observed heads: 30", font_size=20),
            Text("Expected heads: 50", font_size=20),
            Text("Test statistic: z = -4.0", font_size=20),
            Text("p-value < 0.001", font_size=20),
        ).arrange(DOWN, aligned_edge=LEFT, buff=0.2)
        analysis.to_edge(RIGHT)
        
        # Animate the analysis
        self.play(Write(title))
        self.play(Write(hypotheses))
        self.play(ShowCreation(spaces))
        
        for line in analysis:
            self.play(Write(line), run_time=0.8)
        
        # Conclusion
        conclusion = Text("Reject H₀: Strong evidence of bias", 
                         font_size=24, color=YELLOW)
        conclusion.to_edge(DOWN)
        
        self.play(Write(conclusion))
        self.wait()

# Create custom sample space divisions
def create_custom_distribution(probabilities, colors):
    space = SampleSpace()
    division = space.get_horizontal_division(probabilities, colors)
    return space, division

# Multi-level conditioning
def create_tree_diagram(space, levels):
    for level in levels:
        # Recursive subdivision logic
        pass

# Smooth transitions between data sets
def animate_data_change(chart, old_values, new_values, run_time=2):
    # Custom animation for smooth bar transitions
    pass

# Real-time data streaming
def add_streaming_updater(chart, data_source):
    def updater(c):
        new_data = data_source.get_latest()
        c.change_bar_values(new_data)
    chart.add_updater(updater)