use rayon::prelude::*;
use std::f32::consts::PI;
use std::io::Write;
use std::process::{Command, Stdio};

const WIDTH: usize = 2000;
const HEIGHT: usize = 2000;

const FPS: usize = 1;
const DURATION: usize = 25;
const TOTAL_FRAMES: usize = FPS * DURATION;

const SUBSTEPS_PER_FRAME: usize = 100;

const DT: f32 = 0.002;

const EPS: f32 = 1e-6;

const OUTPUT: &str = "double_pendulum_lyapunov.mp4";

const G: f32 = 9.81;

#[derive(Clone, Copy)]
struct State {
    th1: f32,
    th2: f32,
    w1: f32,
    w2: f32,
}

#[derive(Clone, Copy)]
struct Cell {
    a: State,
    b: State,
    lyap: f32,
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let mut cells = initialize();

    let mut frame = vec![0u8; WIDTH * HEIGHT * 3];

    let mut ffmpeg = Command::new("ffmpeg")
        .args([
            "-y",
            "-f",
            "rawvideo",
            "-pix_fmt",
            "rgb24",
            "-s",
            &format!("{}x{}", WIDTH, HEIGHT),
            "-r",
            &FPS.to_string(),
            "-i",
            "-",
            "-c:v",
            "libx264",
            "-preset",
            "slow",
            "-crf",
            "18",
            "-pix_fmt",
            "yuv420p",
            OUTPUT,
        ])
        .stdin(Stdio::piped())
        .spawn()?;

    let stdin = ffmpeg.stdin.as_mut().unwrap();

    for frame_i in 0..TOTAL_FRAMES {
        for _ in 0..SUBSTEPS_PER_FRAME {
            cells.par_iter_mut().for_each(|c| {
                c.a = rk4(c.a);
                c.b = rk4(c.b);

                let dth1 = wrap(c.a.th1 - c.b.th1);
                let dth2 = wrap(c.a.th2 - c.b.th2);

                let dist = (dth1 * dth1 + dth2 * dth2).sqrt();

                if dist > 0.0 {
                    c.lyap += (dist / EPS).ln();

                    let scale = EPS / dist;

                    c.b.th1 = c.a.th1 + dth1 * scale;
                    c.b.th2 = c.a.th2 + dth2 * scale;
                    c.b.w1 = c.a.w1;
                    c.b.w2 = c.a.w2;
                }
            });
        }

        render(&cells, &mut frame);

        stdin.write_all(&frame)?;

        println!("frame {}/{}", frame_i + 1, TOTAL_FRAMES);
    }

    drop(ffmpeg.stdin.take());

    ffmpeg.wait()?;

    Ok(())
}

fn initialize() -> Vec<Cell> {
    let mut v = vec![
        Cell {
            a: State {
                th1: 0.0,
                th2: 0.0,
                w1: 0.0,
                w2: 0.0
            },
            b: State {
                th1: 0.0,
                th2: 0.0,
                w1: 0.0,
                w2: 0.0
            },
            lyap: 0.0
        };
        WIDTH * HEIGHT
    ];

    v.par_chunks_mut(WIDTH).enumerate().for_each(|(y, row)| {
        let fy = y as f32 / (HEIGHT - 1) as f32;

        for (x, c) in row.iter_mut().enumerate() {
            let fx = x as f32 / (WIDTH - 1) as f32;

            let phi1 = -PI + fx * 2.0 * PI;
            let phi2 = -PI + fy * 2.0 * PI;

            let a = State {
                th1: phi1,
                th2: phi2,
                w1: 0.0,
                w2: 0.0,
            };

            let b = State {
                th1: phi1 + EPS,
                th2: phi2 + EPS,
                w1: 0.0,
                w2: 0.0,
            };

            *c = Cell { a, b, lyap: 0.0 };
        }
    });

    v
}

fn render(cells: &[Cell], frame: &mut [u8]) {
    frame
        .par_chunks_mut(WIDTH * 3)
        .enumerate()
        .for_each(|(y, row)| {
            let src = &cells[y * WIDTH..(y + 1) * WIDTH];

            for (x, c) in src.iter().enumerate() {
                let i = x * 3;

                let val = (c.lyap * 0.02).tanh();

                let v = (val * 255.0) as u8;

                row[i] = v;
                row[i + 1] = v;
                row[i + 2] = v;
            }
        });
}

fn rk4(s: State) -> State {
    let k1 = deriv(s);

    let s2 = add(s, k1, 0.5 * DT);
    let k2 = deriv(s2);

    let s3 = add(s, k2, 0.5 * DT);
    let k3 = deriv(s3);

    let s4 = add(s, k3, DT);
    let k4 = deriv(s4);

    State {
        th1: s.th1 + DT * (k1.th1 + 2.0 * k2.th1 + 2.0 * k3.th1 + k4.th1) / 6.0,
        th2: s.th2 + DT * (k1.th2 + 2.0 * k2.th2 + 2.0 * k3.th2 + k4.th2) / 6.0,
        w1: s.w1 + DT * (k1.w1 + 2.0 * k2.w1 + 2.0 * k3.w1 + k4.w1) / 6.0,
        w2: s.w2 + DT * (k1.w2 + 2.0 * k2.w2 + 2.0 * k3.w2 + k4.w2) / 6.0,
    }
}

#[derive(Copy, Clone)]
struct Deriv {
    th1: f32,
    th2: f32,
    w1: f32,
    w2: f32,
}

fn deriv(s: State) -> Deriv {
    let d = s.th1 - s.th2;

    let sin = d.sin();
    let cos = d.cos();

    let denom = 3.0 - (2.0 * d).cos();

    let w1dot = (-9.81 * (3.0 * s.th1.sin() + (s.th1 - 2.0 * s.th2).sin())
        - 2.0 * sin * (s.w2 * s.w2 + s.w1 * s.w1 * cos))
        / denom;

    let w2dot =
        (2.0 * sin * (2.0 * s.w1 * s.w1 + 2.0 * 9.81 * s.th1.cos() + s.w2 * s.w2 * cos)) / denom;

    Deriv {
        th1: s.w1,
        th2: s.w2,
        w1: w1dot,
        w2: w2dot,
    }
}

fn add(s: State, k: Deriv, h: f32) -> State {
    State {
        th1: s.th1 + k.th1 * h,
        th2: s.th2 + k.th2 * h,
        w1: s.w1 + k.w1 * h,
        w2: s.w2 + k.w2 * h,
    }
}

fn wrap(a: f32) -> f32 {
    let tau = 2.0 * PI;

    (a + PI).rem_euclid(tau) - PI
}