# Der Quellcode zum Video 'Bedingte Wahrscheinlichkeiten'
#
# wurde entwickelt von
#
# Oğulcan Aşık <ogasi100@hhu.de>
# Chavoush Kalhor <chavoush.kalhor@hhu.de>
# Christian Müller <c.mueller@hhu.de>
#
# an der Heinrich-Heine-Universität Düsseldorf.
#
# Dieses Werk ist lizenziert unter einer Creative Commons Namensnennung-Weitergabe unter gleichen Bedingungen 4.0 International Lizenz.
# Um eine Kopie der Lizenz zu erhalten, besuchen Sie https://creativecommons.org/licenses/by-sa/4.0/.
#
# SPDX-License-Identifier: CC-BY-SA-4.0

from manim import *
import math

ORCAblue = "#002b44"
ORCAred = "#e61300"
ORCAgreen = "#4cb011"
ORCAgrey = "#f5f5f5"
ORCAwhite = "#ffffff"

config.background_color = "#c4d1d7"
# TODO change background color to #c9ffff?
config.max_files_cached = -1

# video sections and learning outcomes

sections = Tex(
    r"\textsf{\textbf{Bedingte Wahrscheinlichkeiten}}\\",
    r"\textsf{\textbf{Kapitel 1: \\ Ein Zahlenbeispiel}}\\",
    r"\textsf{\textbf{Kapitel 2: \\ Vertauschung der Sichtweisen}}\\",
    r"\textsf{\textbf{Kapitel 3: \\ Vergleich beider Sichtweisen}}\\",
    r"\textsf{\textbf{Zusammenfassung}}",
    color=ORCAblue,
)

outcomes = Tex(
    r"\textsf{\textbf{In diesem Video}}",
    r"\begin{itemize} \item \textsf{Verwendung bedingter Wahrscheinlichkeiten motivieren} \end{itemize}",
    r"\begin{itemize} \item \textsf{Formale Definition herleiten} \end{itemize}",
    r"\begin{itemize} \item \textsf{Rechnungen durchführen} \end{itemize}",
    color=ORCAblue,
)

# animate video sections and learning outcomes


class CondProbSection0(Scene):
    # scene 0: introduction and learning outcomes
    def construct(self):
        sections[0].set(font_size=56).move_to(ORIGIN)
        outcomes[0].to_edge(UL)
        outcomes[1:4].to_edge(LEFT)

        logo_license = SVGMobject(file_name="..\Anderes\Logo_CC_BY_SA.svg", width=1.5)
        self.add(logo_license.to_edge(DR))

        self.wait()
        self.play(Write(sections[0]))
        self.wait()
        self.play(FadeOut(sections[0]))
        self.wait()
        self.play(FadeIn(outcomes[0]))
        self.wait()
        self.play(FadeIn(outcomes[1]))
        self.wait()
        self.play(FadeIn(outcomes[2]))
        self.wait()
        self.play(FadeIn(outcomes[3]))
        self.wait()

        self.play(*[FadeOut(mobject) for mobject in self.mobjects])


class CondProbSection1(Scene):
    def construct(self):
        sections[1].move_to(ORIGIN)

        self.wait()
        self.play(Write(sections[1]))
        self.wait()
        self.play(FadeOut(sections[1]))
        self.wait()


class CondProbSection2(Scene):
    def construct(self):
        sections[2].move_to(ORIGIN)

        self.wait()
        self.play(Write(sections[2]))
        self.wait()
        self.play(FadeOut(sections[2]))
        self.wait()


class CondProbSection3(Scene):
    def construct(self):
        sections[3].move_to(ORIGIN)

        self.wait()
        self.play(Write(sections[3]))
        self.wait()
        self.play(FadeOut(sections[3]))
        self.wait()


class CondProbSection4(Scene):
    def construct(self):
        sections[4].move_to(ORIGIN)

        self.wait()
        self.play(Write(sections[4]))
        self.wait()
        self.play(FadeOut(sections[4]))
        self.wait()


# begin with regular animations

# display vertices x and y of graph as fraction x / y
def fraction(x, y, graph):
    table = MobjectTable(
        [
            [graph.vertices[x].copy()],
            [graph.vertices[y].copy()],
        ],
        v_buff=0.25,
        h_buff=0.25,
    )
    table.get_horizontal_lines().set_stroke(color=ORCAblue, width=2)
    return table


# stuff for scene CondProb1 (number example)

# define colors
color_mat = BLUE
color_phy = ORANGE
color_inf = PURPLE
color_pass = ORCAgreen
color_fail = ORCAred
color_total = LIGHT_GREY

# define numbers
n_mat_pass = 38
n_pass_mat = 38
n_mat_fail = 7
n_fail_mat = 7
n_phy_pass = 10
n_pass_phy = 10
n_phy_fail = 3
n_fail_phy = 3
n_inf_pass = 70
n_pass_inf = 70
n_inf_fail = 16
n_fail_inf = 16
n_mat = n_mat_pass + n_mat_fail
n_phy = n_phy_pass + n_phy_fail
n_inf = n_inf_pass + n_inf_fail
n_pass = n_mat_pass + n_phy_pass + n_inf_pass
n_fail = n_mat_fail + n_phy_fail + n_inf_fail
n_total = n_pass + n_fail

# summarize numbers in a table
table = (
    IntegerTable(
        [
            [n_mat_pass, n_mat_fail, n_mat],
            [n_phy_pass, n_phy_fail, n_phy],
            [n_inf_pass, n_inf_fail, n_inf],
            [n_pass, n_fail, n_total],
        ],
        col_labels=[Tex("be"), Tex("nb"), MathTex("\\Sigma")],
        row_labels=[Tex("mat"), Tex("phy"), Tex("inf"), MathTex("\\Sigma")],
    )
    .set_color(ORCAblue)
    .scale(0.75)
)
table.remove(*table.get_vertical_lines()[1])
table.remove(*table.get_horizontal_lines()[1:3])

# generate questions
questions = Tex(
    r"\begin{itemize} \item \textsf{Mit welcher Wahrscheinlichkeit hat eine zufällig ausgewählte Person aus dem Studiengang Mathematik die Klausur bestanden?} \end{itemize}",
    r"\begin{itemize} \item \textsf{Wie groß ist die Wahrscheinlichkeit, dass eine zufällig ausgewählte bestandene Klausur von einer Person aus dem Studiengang Mathematik stammt?} \end{itemize}",
    color=ORCAblue,
    font_size=40,
)


class CondProb1(Scene):
    # scene 1: number example
    def construct(self):
        # create table
        self.wait()
        self.play(Write(table), run_time=5)
        self.wait(1.5)

        # indicate events
        self.play(
            Indicate(table.get_entries((2, 1)), color=color_mat, scale_factor=2),
            Indicate(table.get_entries((3, 1)), color=color_phy, scale_factor=2),
            Indicate(table.get_entries((4, 1)), color=color_inf, scale_factor=2),
        )
        self.play(
            Indicate(table.get_entries((1, 2)), color=color_pass, scale_factor=2),
            Indicate(table.get_entries((1, 3)), color=color_fail, scale_factor=2),
        )
        self.wait(6)

        # indicate numbers
        self.play(
            Indicate(table.get_entries((5, 4)), color=color_total, scale_factor=2)
        )
        self.wait(2)
        self.play(Indicate(table.get_entries((2, 2)), color=ORCAblue, scale_factor=2))
        self.wait(0.5)
        self.play(Indicate(table.get_entries((2, 3)), color=ORCAblue, scale_factor=2))
        self.wait(4.5)
        self.play(Indicate(table.get_entries((3, 2)), color=ORCAblue, scale_factor=2))
        self.play(Indicate(table.get_entries((3, 3)), color=ORCAblue, scale_factor=2))
        self.wait(3.5)
        self.play(Indicate(table.get_entries((4, 2)), color=ORCAblue, scale_factor=2))
        self.play(Indicate(table.get_entries((4, 3)), color=ORCAblue, scale_factor=2))
        self.wait(3)
        self.play(Indicate(table.get_entries((2, 4)), color=color_mat, scale_factor=2))
        self.wait()
        self.play(Indicate(table.get_entries((3, 4)), color=color_phy, scale_factor=2))
        self.wait()
        self.play(Indicate(table.get_entries((4, 4)), color=color_inf, scale_factor=2))
        self.wait(3)
        self.play(Indicate(table.get_entries((5, 2)), color=color_pass, scale_factor=2))
        self.wait(2)
        self.play(Indicate(table.get_entries((5, 3)), color=color_fail, scale_factor=2))
        self.wait(2)
        self.play(table.animate.to_edge(DOWN))
        self.wait(2)

        # display questions
        questions.to_edge(UP)
        self.play(FadeIn(questions[0]))
        self.wait(6)
        self.play(FadeIn(questions[1]))
        self.wait(16)

        # indicate events
        self.play(Indicate(table.get_rows()[1], color=color_mat))
        self.wait(4.5)
        self.play(Indicate(table.get_columns()[1], color=color_pass))
        self.wait(9)


# stuff for scene CondProb2 (conditioning on exam result)

# tree diagram
vertices1 = [
    n_total,
    n_fail,
    n_pass,
    n_inf_fail,
    n_phy_fail,
    n_mat_fail,
    n_inf_pass,
    n_phy_pass,
    n_mat_pass,
]
edges1 = [
    (n_total, n_fail),
    (n_total, n_pass),
    (n_fail, n_inf_fail),
    (n_fail, n_phy_fail),
    (n_fail, n_mat_fail),
    (n_pass, n_inf_pass),
    (n_pass, n_phy_pass),
    (n_pass, n_mat_pass),
]
vertices1_config = {
    n_total: {"fill_color": color_total},
    n_fail: {"fill_color": color_fail},
    n_pass: {"fill_color": color_pass},
    n_inf_fail: {"fill_color": color_inf},
    n_phy_fail: {"fill_color": color_phy},
    n_mat_fail: {"fill_color": color_mat},
    n_inf_pass: {"fill_color": color_inf},
    n_phy_pass: {"fill_color": color_phy},
    n_mat_pass: {"fill_color": color_mat},
}
edges1_config = {
    (n_total, n_fail): {"stroke_color": color_fail},
    (n_total, n_pass): {"stroke_color": color_pass},
    (n_fail, n_inf_fail): {"stroke_color": color_inf},
    (n_fail, n_phy_fail): {"stroke_color": color_phy},
    (n_fail, n_mat_fail): {"stroke_color": color_mat},
    (n_pass, n_inf_pass): {"stroke_color": color_inf},
    (n_pass, n_phy_pass): {"stroke_color": color_phy},
    (n_pass, n_mat_pass): {"stroke_color": color_mat},
}
graph1 = Graph(
    vertices1,
    edges1,
    vertex_config=vertices1_config,
    edge_config=edges1_config,
    layout="tree",
    root_vertex=n_total,
    layout_scale=(3, 2),
    labels=True,
)

# edge labels
label1_fail = (
    MathTex(r"P(\text{nb})", color=ORCAblue)
    .scale(0.4)
    .rotate(-PI / 4)
    .next_to(graph1.edges[(n_total, n_fail)].get_center(), RIGHT, buff=0.1)
)
label1_pass = (
    MathTex(r"P(\text{be})", color=ORCAblue)
    .scale(0.4)
    .rotate(PI / 4)
    .next_to(graph1.edges[(n_total, n_pass)].get_center(), LEFT, buff=0.1)
)
label1_inf_fail = (
    MathTex(r"P(\text{inf}|\text{nb})", color=ORCAblue)
    .scale(0.4)
    .rotate(-PI / 3)
    .next_to(graph1.edges[(n_fail, n_inf_fail)].get_center(), RIGHT, buff=0.1)
)
label1_phy_fail = (
    MathTex(r"P(\text{phy}|\text{nb})", color=ORCAblue)
    .scale(0.4)
    .rotate(PI / 2)
    .next_to(graph1.edges[(n_fail, n_phy_fail)].get_center(), LEFT, buff=0.1)
)
label1_mat_fail = (
    MathTex(r"P(\text{mat}|\text{nb})", color=ORCAblue)
    .scale(0.4)
    .rotate(PI / 3)
    .next_to(graph1.edges[(n_fail, n_mat_fail)].get_center(), LEFT, buff=0.1)
)
label1_inf_pass = (
    MathTex(r"P(\text{inf}|\text{be})", color=ORCAblue)
    .scale(0.4)
    .rotate(-PI / 3)
    .next_to(graph1.edges[(n_pass, n_inf_pass)].get_center(), RIGHT, buff=0.1)
)
label1_phy_pass = (
    MathTex(r"P(\text{phy}|\text{be})", color=ORCAblue)
    .scale(0.4)
    .rotate(PI / 2)
    .next_to(graph1.edges[(n_pass, n_phy_pass)].get_center(), LEFT, buff=0.1)
)
label1_mat_pass = (
    MathTex(r"P(\text{mat}|\text{be})", color=ORCAblue)
    .scale(0.4)
    .rotate(PI / 3)
    .next_to(graph1.edges[(n_pass, n_mat_pass)].get_center(), LEFT, buff=0.1)
)
label1 = VGroup(
    label1_pass,
    label1_fail,
    label1_mat_pass,
    label1_phy_pass,
    label1_inf_pass,
    label1_mat_fail,
    label1_phy_fail,
    label1_inf_fail,
)

# calculation with frequentist probability
calculation1A = MathTex("P(", "\\text{mat}", "|", "\\text{be}", ")", "=")
calculation1A.set_color(ORCAblue).shift(LEFT * 5)

calculation1B = fraction(n_mat_pass, n_pass, graph1).next_to(calculation1A, RIGHT)
calculation1B_value = MathTex(r"\approx0.322", color=ORCAblue).next_to(
    calculation1B, RIGHT
)

calculation1C = MobjectTable(
    [[fraction(n_mat_pass, n_total, graph1)], [fraction(n_pass, n_total, graph1)]],
    v_buff=0.25,
    h_buff=1,
).next_to(calculation1A, RIGHT)
calculation1C.get_horizontal_lines().set_stroke(color=ORCAblue, width=2)

calculation1D = MathTex(r"\frac{P(\text{mat}\cap\text{be})}{P(\text{be})}")
calculation1D.set_color(ORCAblue).next_to(calculation1A, RIGHT)

# general formula
formula1A = MathTex(r"P(A|\text{be})", "=", r"\frac{P(A\cap\text{be})}{P(\text{be})}")
formula1B = MathTex(r"P(A|B)", "=", r"\frac{P(A\cap B)}{P(B)}")
formula1C = MathTex(r"P(A\cap B)", "=", r"P(A|B)\cdot P(B)")
formula1D = MathTex(r"P(B)", "=", r"\frac{P(A\cap B)}{P(A|B)}")

formula1A.set_color(ORCAblue)
formula1B.set_color(ORCAblue)
formula1C.set_color(ORCAblue)
formula1D.set_color(ORCAblue)


class CondProb2(Scene):
    # scene 2: conditioning on exam result
    def construct(self):
        # set up scene
        table.to_edge(DOWN)
        questions.to_edge(UP)
        self.add(table, questions)
        self.wait(2)
        self.play(FadeOut(questions[0]))
        self.wait()
        self.play(
            table.animate.move_to(ORIGIN).to_edge(LEFT),
            questions[1].animate.to_edge(UP),
        )
        self.wait()

        # animate creation of tree diagramm
        graph1.to_edge(RIGHT)
        table.add_highlighted_cell((5, 4), color=color_total)
        self.play(
            ReplacementTransform(
                table.get_cell((5, 4)).copy(), graph1.vertices[n_total]
            )
        )
        # first stage
        table.add_highlighted_cell((5, 2), color=color_pass)
        self.play(
            ReplacementTransform(
                table.get_cell((5, 2)).copy(), graph1.vertices[n_pass]
            ),
            Create(graph1.edges[(n_total, n_pass)]),
        )
        table.add_highlighted_cell((5, 3), color=color_fail)
        self.play(
            ReplacementTransform(
                table.get_cell((5, 3)).copy(), graph1.vertices[n_fail]
            ),
            Create(graph1.edges[(n_total, n_fail)]),
        )
        self.wait()
        # second stage mat
        table.add_highlighted_cell((2, 4), color=color_mat)
        self.play(
            ReplacementTransform(
                table.get_cell((2, 2)).copy(), graph1.vertices[n_mat_pass]
            ),
            Create(graph1.edges[(n_pass, n_mat_pass)]),
        )
        self.play(
            ReplacementTransform(
                table.get_cell((2, 3)).copy(), graph1.vertices[n_mat_fail]
            ),
            Create(graph1.edges[(n_fail, n_mat_fail)]),
        )
        self.wait()
        # second stage phy
        table.add_highlighted_cell((3, 4), color=color_phy)
        self.play(
            ReplacementTransform(
                table.get_cell((3, 2)).copy(), graph1.vertices[n_phy_pass]
            ),
            Create(graph1.edges[(n_pass, n_phy_pass)]),
        )
        self.play(
            ReplacementTransform(
                table.get_cell((3, 3)).copy(), graph1.vertices[n_phy_fail]
            ),
            Create(graph1.edges[(n_fail, n_phy_fail)]),
        )
        self.wait()
        # second stage inf
        table.add_highlighted_cell((4, 4), color=color_inf)
        self.play(
            ReplacementTransform(
                table.get_cell((4, 2)).copy(), graph1.vertices[n_inf_pass]
            ),
            Create(graph1.edges[(n_pass, n_inf_pass)]),
        )
        self.play(
            ReplacementTransform(
                table.get_cell((4, 3)).copy(), graph1.vertices[n_inf_fail]
            ),
            Create(graph1.edges[(n_fail, n_inf_fail)]),
        )
        self.wait(6)

        # animate calculation with frequentist probability
        self.play(FadeOut(table))
        self.wait(3)
        self.play(
            ReplacementTransform(
                graph1.vertices[n_mat_pass].copy(), calculation1B[0][0]
            )
        )
        self.wait(2)
        self.play(
            Create(calculation1B[1]),
            ReplacementTransform(graph1.vertices[n_pass].copy(), calculation1B[0][1]),
        )
        self.wait(4)
        # display numeric value
        self.play(FadeIn(calculation1B_value))
        self.wait(2)
        self.play(FadeOut(calculation1B_value))
        self.wait(2)
        self.play(FadeIn(calculation1A))
        self.wait(16)
        # highlight notation
        self.play(Indicate(calculation1A[2], color=ORCAblue, scale_factor=2))
        self.wait(2)
        self.play(Indicate(calculation1A[1], color=color_mat, scale_factor=2))
        self.wait(2)
        self.play(Indicate(calculation1A[3], color=color_pass, scale_factor=2))
        self.wait(8)
        self.play(ReplacementTransform(calculation1B, calculation1C))
        self.wait()
        self.play(Indicate(calculation1C[0][0], color=None))
        self.wait()
        self.play(Indicate(calculation1C[0][1], color=None))
        self.wait(9)
        self.play(ReplacementTransform(calculation1C, calculation1D))
        self.wait(19)

        # animate general formula
        self.play(
            FadeOut(questions[1], graph1),
            VGroup(calculation1A, calculation1D).animate.move_to(ORIGIN),
        )
        self.wait(10)
        self.play(
            TransformMatchingShapes(calculation1A, formula1A[0:2]),
            TransformMatchingShapes(calculation1D, formula1A[2]),
            run_time=2,
        )
        self.wait()
        self.play(TransformMatchingShapes(formula1A, formula1B))
        self.wait(28)
        # draw box
        self.play(Create(SurroundingRectangle(formula1B, color=ORCAblue)))
        self.wait()


# stuff for scene CondProb3 (conditioning on subject of study)

# tree diagram
vertices2 = [
    n_total,
    n_inf,
    n_phy,
    n_mat,
    n_fail_inf,
    n_pass_inf,
    n_fail_phy,
    n_pass_phy,
    n_fail_mat,
    n_pass_mat,
]
edges2 = [
    (n_total, n_inf),
    (n_total, n_phy),
    (n_total, n_mat),
    (n_inf, n_fail_inf),
    (n_inf, n_pass_inf),
    (n_phy, n_fail_phy),
    (n_phy, n_pass_phy),
    (n_mat, n_fail_mat),
    (n_mat, n_pass_mat),
]
vertices2_config = {
    n_total: {"fill_color": color_total},
    n_inf: {"fill_color": color_inf},
    n_phy: {"fill_color": color_phy},
    n_mat: {"fill_color": color_mat},
    n_fail_inf: {"fill_color": color_fail},
    n_pass_inf: {"fill_color": color_pass},
    n_fail_phy: {"fill_color": color_fail},
    n_pass_phy: {"fill_color": color_pass},
    n_fail_mat: {"fill_color": color_fail},
    n_pass_mat: {"fill_color": color_pass},
}
edges2_config = {
    (n_total, n_inf): {"stroke_color": color_inf},
    (n_total, n_phy): {"stroke_color": color_phy},
    (n_total, n_mat): {"stroke_color": color_mat},
    (n_inf, n_fail_inf): {"stroke_color": color_fail},
    (n_inf, n_pass_inf): {"stroke_color": color_pass},
    (n_phy, n_fail_phy): {"stroke_color": color_fail},
    (n_phy, n_pass_phy): {"stroke_color": color_pass},
    (n_mat, n_fail_mat): {"stroke_color": color_fail},
    (n_mat, n_pass_mat): {"stroke_color": color_pass},
}
graph2 = Graph(
    vertices2,
    edges2,
    vertex_config=vertices2_config,
    edge_config=edges2_config,
    layout="tree",
    root_vertex=n_total,
    layout_scale=(3, 2),
    labels=True,
)

# edge labels
label2_inf = (
    MathTex(r"P(\text{inf})", color=ORCAblue)
    .scale(0.4)
    .rotate(-math.radians(40))
    .next_to(graph2.edges[(n_total, n_inf)].get_center(), RIGHT, buff=0.1)
)
label2_phy = (
    MathTex(r"P(\text{phy})", color=ORCAblue)
    .scale(0.4)
    .rotate(PI / 2)
    .next_to(graph2.edges[(n_total, n_phy)].get_center(), LEFT, buff=0.1)
)
label2_mat = (
    MathTex(r"P(\text{mat})", color=ORCAblue)
    .scale(0.4)
    .rotate(math.radians(40))
    .next_to(graph2.edges[(n_total, n_mat)].get_center(), LEFT, buff=0.1)
)
label2_fail_inf = (
    MathTex(r"P(\text{nb}|\text{inf})", color=ORCAblue)
    .scale(0.4)
    .rotate(-math.radians(73))
    .next_to(graph2.edges[(n_inf, n_fail_inf)].get_center(), RIGHT, buff=0.1)
)
label2_pass_inf = (
    MathTex(r"P(\text{be}|\text{inf})", color=ORCAblue)
    .scale(0.4)
    .rotate(math.radians(73))
    .next_to(graph2.edges[(n_inf, n_pass_inf)].get_center(), LEFT, buff=0.1)
)
label2_fail_phy = (
    MathTex(r"P(\text{nb}|\text{phy})", color=ORCAblue)
    .scale(0.4)
    .rotate(-math.radians(73))
    .next_to(graph2.edges[(n_phy, n_fail_phy)].get_center(), RIGHT, buff=0.1)
)
label2_pass_phy = (
    MathTex(r"P(\text{be}|\text{phy})", color=ORCAblue)
    .scale(0.4)
    .rotate(math.radians(73))
    .next_to(graph2.edges[(n_phy, n_pass_phy)].get_center(), LEFT, buff=0.1)
)
label2_fail_mat = (
    MathTex(r"P(\text{nb}|\text{mat})", color=ORCAblue)
    .scale(0.4)
    .rotate(-math.radians(73))
    .next_to(graph2.edges[(n_mat, n_fail_mat)].get_center(), RIGHT, buff=0.1)
)
label2_pass_mat = (
    MathTex(r"P(\text{be}|\text{mat})", color=ORCAblue)
    .scale(0.4)
    .rotate(math.radians(73))
    .next_to(graph2.edges[(n_mat, n_pass_mat)].get_center(), LEFT, buff=0.1)
)
label2 = VGroup(
    label2_mat,
    label2_phy,
    label2_inf,
    label2_pass_mat,
    label2_fail_mat,
    label2_pass_phy,
    label2_fail_phy,
    label2_pass_inf,
    label2_fail_inf,
)

# calculation with frequentist probability
calculation2A = MathTex(r"P(\text{be}|\text{mat})", "=")
calculation2A.set_color(ORCAblue).shift(LEFT * 5)

calculation2B = fraction(n_pass_mat, n_mat, graph2).next_to(calculation2A, RIGHT)
calculation2B_value = MathTex(r"\approx0.844", color=ORCAblue).next_to(
    calculation2B, RIGHT
)

calculation2C = MobjectTable(
    [[fraction(n_pass_mat, n_total, graph2)], [fraction(n_mat, n_total, graph2)]],
    v_buff=0.25,
    h_buff=1,
).next_to(calculation2A, RIGHT)
calculation2C.get_horizontal_lines().set_stroke(color=ORCAblue, width=2)

calculation2D = MathTex(r"\frac{P(\text{be}\cap\text{mat})}{P(\text{mat})}")
calculation2D.set_color(ORCAblue).next_to(calculation2A, RIGHT)

# general formula
formula2A = MathTex(
    r"P(B|\text{mat})", "=", r"\frac{P(B\cap\text{mat})}{P(\text{mat})}"
)
formula2B = MathTex(r"P(B|A)", "=", r"\frac{P(B\cap A)}{P(A)}")
formula2C = MathTex(r"P(B\cap A)", "=", r"P(B|A)\cdot P(A)")
formula2D = MathTex(r"P(A)", "=", r"\frac{P(B\cap A)}{P(B|A)}")

formula2A.set_color(ORCAblue)
formula2B.set_color(ORCAblue)
formula2C.set_color(ORCAblue)
formula2D.set_color(ORCAblue)


class CondProb3(Scene):
    # scene 3: conditioning on subject of study
    def construct(self):
        # set up scene
        questions[0].to_edge(UP)
        self.wait()
        self.play(FadeIn(table))
        self.wait(3)
        self.play(FadeIn(questions[0]))
        self.wait(10)
        self.play(table.animate.to_edge(LEFT))
        self.wait(2)

        # animate creation of tree diagramm
        graph2.to_edge(RIGHT)
        table.add_highlighted_cell((5, 4), color=color_total)
        self.play(
            ReplacementTransform(
                table.get_cell((5, 4)).copy(), graph2.vertices[n_total]
            )
        )
        # first stage
        table.add_highlighted_cell((2, 4), color=color_mat)
        self.play(
            ReplacementTransform(table.get_cell((2, 4)).copy(), graph2.vertices[n_mat]),
            Create(graph2.edges[(n_total, n_mat)]),
        )
        table.add_highlighted_cell((3, 4), color=color_phy)
        self.play(
            ReplacementTransform(table.get_cell((3, 4)).copy(), graph2.vertices[n_phy]),
            Create(graph2.edges[(n_total, n_phy)]),
        )
        table.add_highlighted_cell((4, 4), color=color_inf)
        self.play(
            ReplacementTransform(table.get_cell((4, 4)).copy(), graph2.vertices[n_inf]),
            Create(graph2.edges[(n_total, n_inf)]),
        )
        self.wait()
        # second stage pass
        table.add_highlighted_cell((5, 2), color=color_pass)
        self.play(
            ReplacementTransform(
                table.get_cell((2, 2)).copy(), graph2.vertices[n_pass_mat]
            ),
            Create(graph2.edges[(n_mat, n_pass_mat)]),
        )
        self.play(
            ReplacementTransform(
                table.get_cell((3, 2)).copy(), graph2.vertices[n_pass_phy]
            ),
            Create(graph2.edges[(n_phy, n_pass_phy)]),
        )
        self.play(
            ReplacementTransform(
                table.get_cell((4, 2)).copy(), graph2.vertices[n_pass_inf]
            ),
            Create(graph2.edges[(n_inf, n_pass_inf)]),
        )
        self.wait()
        # second stage fail
        table.add_highlighted_cell((5, 3), color=color_fail)
        self.play(
            ReplacementTransform(
                table.get_cell((2, 3)).copy(), graph2.vertices[n_fail_mat]
            ),
            Create(graph2.edges[(n_mat, n_fail_mat)]),
        )
        self.play(
            ReplacementTransform(
                table.get_cell((3, 3)).copy(), graph2.vertices[n_fail_phy]
            ),
            Create(graph2.edges[(n_phy, n_fail_phy)]),
        )
        self.play(
            ReplacementTransform(
                table.get_cell((4, 3)).copy(), graph2.vertices[n_fail_inf]
            ),
            Create(graph2.edges[(n_inf, n_fail_inf)]),
        )
        self.wait(3)

        # animate calculation with frequentist probability
        self.play(FadeOut(table))
        self.wait()
        self.add(calculation2A)
        self.wait(8)
        self.play(
            ReplacementTransform(
                graph2.vertices[n_pass_mat].copy(), calculation2B[0][0]
            ),
            Create(calculation2B[1]),
            ReplacementTransform(graph2.vertices[n_mat].copy(), calculation2B[0][1]),
        )
        self.wait()
        # display numeric value
        self.play(FadeIn(calculation2B_value))
        self.wait(2)
        self.play(FadeOut(calculation2B_value))
        self.wait()
        self.play(ReplacementTransform(calculation2B, calculation2C))
        self.wait(5)
        self.play(ReplacementTransform(calculation2C, calculation2D))
        self.wait(4)

        # animate general formula
        self.play(
            FadeOut(questions[0], graph2),
            VGroup(calculation2A, calculation2D).animate.move_to(ORIGIN),
        )
        self.wait(5)
        self.play(
            TransformMatchingShapes(calculation2A[0:2], formula2A[0:2]),
            TransformMatchingShapes(calculation2D, formula2A[2]),
        )
        self.wait()
        self.play(TransformMatchingShapes(formula2A, formula2B))
        self.wait(4)
        # draw box
        self.play(Create(SurroundingRectangle(formula2B, color=ORCAblue)))
        self.wait()


class CondProb4(Scene):
    # scene 4: comparison between both points of view
    def construct(self):
        # compare graph1 and graph2
        VGroup(graph1, label1).to_edge(LEFT, buff=0.25)
        VGroup(graph2, label2).to_edge(RIGHT, buff=0.25)
        self.wait()
        self.play(FadeIn(graph1))
        self.wait()
        self.play(FadeIn(graph2))
        self.wait(2)
        self.play(*[FadeIn(label) for label in label1[:2]])
        self.wait(2)
        self.play(*[FadeIn(label) for label in label2[:3]])
        self.wait(14)
        self.play(*[FadeIn(label) for label in label1[2:]])
        self.wait(2)
        self.play(*[FadeIn(label) for label in label2[3:]])
        self.wait(2)

        # display formulas
        self.play(VGroup(graph1, label1, graph2, label2).animate.shift(DOWN))
        self.wait()
        self.play(
            FadeIn(
                calculation1A[:5].shift(2.5 * UP),
                calculation2A[0].shift(2.5 * UP + 7 * RIGHT),
                formula1B[1].next_to(calculation1A[:5], RIGHT),
                formula2B[1].next_to(calculation2A[0], RIGHT),
                calculation1D.next_to(formula1B[1], RIGHT),
                calculation2D.next_to(formula2B[1], RIGHT),
            )
        )
        self.wait(2)
        self.play(
            TransformMatchingShapes(
                calculation1A[:5],
                formula1B[0].next_to(formula1B[1], LEFT),
            ),
            TransformMatchingShapes(
                calculation1D,
                formula1B[2].next_to(formula1B[1], RIGHT),
            ),
        )
        self.wait()
        self.play(
            TransformMatchingShapes(
                calculation2A[0],
                formula2B[0].next_to(formula2B[1], LEFT),
            ),
            TransformMatchingShapes(
                calculation2D,
                formula2B[2].next_to(formula2B[1], RIGHT),
            ),
        )
        self.wait(2)

        # center both formulas
        self.play(FadeOut(graph1, label1, graph2, label2))
        self.wait()
        self.play(
            formula1B.animate.move_to(ORIGIN + UP),
            formula2B.animate.move_to(ORIGIN + DOWN),
        )
        self.wait(35)

        # solve for different terms
        formula1C.shift(UP)
        formula1D.shift(UP)
        formula2C.shift(DOWN)
        formula2D.shift(DOWN)
        self.play(TransformMatchingShapes(formula1B, formula1C))
        self.wait()
        self.play(TransformMatchingShapes(formula2B, formula2C))
        self.wait()
        self.play(TransformMatchingShapes(formula1C, formula1D))
        self.wait()
        self.play(TransformMatchingShapes(formula2C, formula2D))
        self.wait()