use rand::{SeedableRng, seq::SliceRandom};
use ranim::{
    anims::morph::MorphAnim,
    color::palettes::manim,
    glam::{DVec3, dvec2, dvec3},
    items::vitem::geometry::Rectangle,
    prelude::*,
    utils::rate_functions::linear,
};

fn bubble_sort(r: &mut RanimScene, num: usize) {
    let _r_cam = r.insert(CameraFrame::default());

    let frame_size = dvec2(8.0 * 16.0 / 9.0, 8.0);
    let padded_frame_size = frame_size * 0.9;

    let anim_step_duration = 15.0 / num.pow(2) as f64;

    let width_unit = padded_frame_size.x / num as f64;
    let height_unit = padded_frame_size.y / num as f64;

    let mut rng = rand_chacha::ChaChaRng::seed_from_u64(114514);
    let mut heights = (1..=num)
        .map(|x| x as f64 * height_unit)
        .collect::<Vec<f64>>();
    heights.shuffle(&mut rng);

    let padded_frame_bl = dvec2(padded_frame_size.x / -2.0, padded_frame_size.y / -2.0);
    let mut r_rects = heights
        .iter()
        .enumerate()
        .map(|(i, &height)| {
            let target_bc_coord =
                padded_frame_bl.extend(0.0) + DVec3::X * (width_unit * i as f64 + width_unit / 2.0);
            let rect = Rectangle::new(width_unit, height).with(|rect| {
                rect.stroke_width = 0.0;
                rect.set_fill_color(manim::WHITE.with_alpha(0.5))
                    .scale(DVec3::splat(0.8))
                    .move_anchor_to(AabbPoint(dvec3(0.0, -1.0, 0.0)), target_bc_coord);
            });
            (r.insert(rect.clone()), rect)
        })
        .collect::<Vec<_>>();

    let anim_highlight = |rect: &mut Rectangle| {
        rect.morph(|data| {
            data.set_fill_color(manim::BLUE_C.with_alpha(0.5));
        })
        .with_duration(anim_step_duration)
        .with_rate_func(linear)
    };
    let anim_unhighlight = |rect: &mut Rectangle| {
        rect.morph(|data| {
            data.set_fill_color(manim::WHITE.with_alpha(0.5));
        })
        .with_duration(anim_step_duration)
        .with_rate_func(linear)
    };
    let shift_right = DVec3::X * width_unit;
    let swap_shift = [shift_right, -shift_right];
    let anim_swap = |r: &mut RanimScene, r_rectab: [&mut (TimelineId, Rectangle); 2]| {
        let timelines = r.timeline_mut(r_rectab).unwrap();
        timelines
            .into_iter()
            .zip(swap_shift.iter())
            .for_each(|((timeline, rect), shift)| {
                timeline.play(
                    rect.morph(|data| {
                        data.shift(*shift);
                    })
                    .with_duration(anim_step_duration)
                    .with_rate_func(linear),
                );
            });
    };

    for i in (1..num).rev() {
        for j in 0..i {
            r.timeline_mut(r_rects.get_disjoint_mut([j, j + 1]).unwrap())
                .unwrap()
                .into_iter()
                .for_each(|(timeline, rect)| {
                    timeline.play(anim_highlight(rect));
                });
            if heights[j] > heights[j + 1] {
                anim_swap(r, r_rects.get_disjoint_mut([j, j + 1]).unwrap());
                r.timelines_mut().sync();
                heights.swap(j, j + 1);
                r_rects.swap(j, j + 1);
            }
            r.timeline_mut(r_rects.get_disjoint_mut([j, j + 1]).unwrap())
                .unwrap()
                .into_iter()
                .for_each(|(timeline, rect)| {
                    timeline.play(anim_unhighlight(rect));
                });
            r.timelines_mut().sync();
        }
    }

    r.insert_time_mark(
        r.timelines().max_total_secs(),
        TimeMark::Capture(format!("preview-{num}.png")),
    );
}

#[scene]
#[output(dir = "./output/bubble_sort")]
fn bubble_sort_10(r: &mut RanimScene) {
    bubble_sort(r, 10);
}

#[scene(name = "bubble_sort")]
#[output(dir = "./output/bubble_sort")]
fn bubble_sort_100(r: &mut RanimScene) {
    bubble_sort(r, 100);
}
