"""The one idea: exchangeability makes a rank uniform, and a uniform
rank is a guarantee.

The scores on screen are real conformity scores from this repository's
bike-sharing fit. Nothing else is needed for the theorem: no Gaussian, no
large-n limit, no assumption the model is any good.

Render:
  PATH="$HOME/Library/TinyTeX/bin/universal-darwin:/usr/local/bin:$PATH" \\
  ../ember/.venv/bin/manim -qh assets/manim_ranks.py Ranks
"""

import json
from pathlib import Path

import numpy as np
from manim import *

GROUND = "#0B1412"
INK = "#E9F0EA"
DIM = "#7FA093"
GREEN = "#58C29A"
RED = "#E05545"
GOLD = "#D9B54A"
BLUE = "#6FAED9"

# twelve real calibration scores from the bike fit: the 12 evenly spaced
# order statistics of the 4,395 scores (|y - yhat|, log-count scale); n is
# kept small on screen so every slot is visible
SCORES = [0.02, 0.08, 0.16, 0.25, 0.34, 0.44, 0.52, 0.62, 0.72, 0.84,
          1.04, 2.91]


class Ranks(Scene):
    def construct(self):
        self.camera.background_color = GROUND

        title = Text("Why conformal prediction works", font_size=42,
                     color=INK)
        sub = Text("one counting argument, no distributions anywhere",
                   font_size=22, color=DIM).next_to(title, DOWN, buff=0.3)
        self.play(FadeIn(title, shift=UP * 0.3), FadeIn(sub), run_time=1.1)
        self.wait(1.2)
        self.play(FadeOut(title), FadeOut(sub), run_time=0.5)

        # the sorted calibration scores as dots on a number line
        line = NumberLine(x_range=[0, 3.0, 1.0], length=11,
                          color=DIM, include_numbers=False,
                          stroke_width=2).shift(UP * 1.4)
        lab = Text("calibration scores |y - yhat|, sorted", font_size=19,
                   color=DIM).next_to(line, UP, buff=0.4)
        self.play(Create(line), FadeIn(lab), run_time=0.9)

        dots = VGroup(*[Dot(line.n2p(s), radius=0.09, color=GREEN)
                        for s in SCORES])
        nums = VGroup(*[Text(str(i + 1), font_size=15, color=DIM)
                        .next_to(d, DOWN, buff=0.18)
                        for i, d in enumerate(dots)])
        self.play(LaggedStart(*[FadeIn(d, scale=0.4) for d in dots],
                              lag_ratio=0.06), FadeIn(nums), run_time=1.6)
        n_lab = Text("n = 12 held-out points the model has already been "
                     "wrong about", font_size=20, color=INK)
        n_lab.next_to(line, DOWN, buff=0.9)
        self.play(FadeIn(n_lab), run_time=0.7)
        self.wait(1.2)
        self.play(FadeOut(n_lab), run_time=0.4)

        # the new point arrives
        q = Dot(line.n2p(0.65), radius=0.11, color=GOLD)
        qlab = Text("tomorrow's point", font_size=19, color=GOLD)
        qlab.next_to(q, UP, buff=0.3)
        self.play(FadeIn(q, scale=0.3), FadeIn(qlab), run_time=0.9)

        claim = Text("if it is exchangeable with the twelve, its rank "
                     "among all 13 is uniform:",
                     font_size=21, color=INK).shift(DOWN * 0.4)
        claim2 = Text("slot 1, slot 2, ... slot 13, each with probability "
                      "1/13. That is the whole theorem.",
                      font_size=21, color=INK).next_to(claim, DOWN,
                                                       buff=0.25)
        self.play(FadeIn(claim), run_time=0.8)

        # slide the new point through several slots to show the symmetry
        for x in (0.13, 0.48, 0.95, 2.2, 0.65):
            self.play(q.animate.move_to(line.n2p(x)),
                      qlab.animate.next_to(line.n2p(x), UP, buff=0.42),
                      run_time=0.55)
        self.play(FadeIn(claim2), run_time=0.8)
        self.wait(1.6)
        self.play(FadeOut(claim), FadeOut(claim2), FadeOut(qlab),
                  run_time=0.5)

        # the quantile: paint the first k dots
        k = 12  # ceil(13 * 0.9) = 12
        kmath = MathTex(r"q=\text{score at rank } \lceil (n+1)(1-\alpha)"
                        r"\rceil = \lceil 13 \times 0.9\rceil = 12",
                        font_size=34, color=INK).shift(DOWN * 0.5)
        self.play(Write(kmath), run_time=1.2)
        ring = SurroundingRectangle(dots[k - 1], color=GOLD, buff=0.16,
                                    corner_radius=0.1)
        self.play(Create(ring), run_time=0.6)
        prob = MathTex(r"P(\text{new score} \le q) \;\ge\; \frac{12}{13}"
                       r" = 0.923 \;\ge\; 1-\alpha",
                       font_size=38, color=GREEN).shift(DOWN * 1.6)
        self.play(Write(prob), run_time=1.3)
        self.wait(2.0)
        self.play(*[FadeOut(m) for m in
                    [line, lab, dots, nums, q, kmath, ring, prob]],
                  run_time=0.6)

        # what breaks it
        b1 = Text("The guarantee never asked the model to be right.",
                  font_size=26, color=INK)
        b2 = Text("It asked tomorrow to be exchangeable with yesterday.",
                  font_size=26, color=INK)
        b3 = Text("Drift breaks exactly that, and only that:", font_size=26,
                  color=RED)
        b4 = Text("measured on real 2012 data: 90% promised, 78% delivered",
                  font_size=22, color=DIM)
        grp = VGroup(b1, b2, b3, b4).arrange(DOWN, buff=0.4)
        self.play(LaggedStart(*[FadeIn(m, shift=UP * 0.15) for m in grp],
                              lag_ratio=0.4), run_time=2.4)
        self.wait(2.2)
        self.play(FadeOut(grp), run_time=0.5)

        end = Text("idaniloju", font_size=52, color=GREEN)
        end2 = Text("assurance, its exact law, and its honest repair",
                    font_size=21, color=DIM).next_to(end, DOWN, buff=0.3)
        self.play(FadeIn(end, shift=UP * 0.2), FadeIn(end2), run_time=1.0)
        self.wait(1.6)