water simulation

use std::{
    f32,
    io
};
use minifb::{
    Key,
    WindowOptions,
    Window,
};
use vek::*;

const W: usize = 64;
const H: usize = 64;

// store bottom-left corner + side length
// since its most convenient
#[derive(Copy, Clone)]
struct Square {
    p: Vec2<f32>,
    a: f32
}
impl Square {
    pub fn new(p0: Vec2<f32>, a0: f32) -> Self {
        Square {p: p0, a: a0}
    }

    pub fn shift(&mut self, offset: Vec2<f32>) {
        self.p += offset
    }

    pub fn grow(&mut self, amount: f32) {
        self.a += amount;
        self.p += Vec2::new(-0.5 * amount, -0.5 * amount);
    }

    // copied from
    // https://stackoverflow.com/questions/9324339
    // since it was way better than what i would have come up with
    pub fn intersect_area(&self, other: &Self) -> f32 {
        let x1 = self.p.x.max(other.p.x);
        let y1 = self.p.y.max(other.p.y);
        let x2 = (self.p.x + self.a).min(other.p.x + other.a);
        let y2 = (self.p.y + self.a).min(other.p.y + other.a);
        (x2 - x1).max(0.0) * (y2 - y1).max(0.0)
    }
}

#[derive(Copy, Clone)]
struct Cell {
    pop: f32,
    vel: Vec2<f32>,
    // Watch out when using negative spread; if the group is also
    // moving then it will cause weird behaviour namely if the
    // absolute magnitude of the spread is greater than the largest
    // of the two velocity components then the group will be rooted
    // to the spot.
    spread: f32,
    // In the range 0 to 1 normally, 0 = no flow (if you set it to
    // zero then nothing will ever enter and everything currently
    // there will leave at the end of the tick unless the adjacent
    // cells all have 0 conductivity). 1 = free flow, < 0 is sort of
    // UDB and will make things move the opposite way to the way they
    // normally would (not recommended for use apart from in special
    // circumstances). Everything is relative so you can freely use
    // anything above 1, but 1 should be the baseline for standard
    // cells.
    conductivity: f32, //TODO think of a shorter name for this lol

}
impl Cell {
    pub fn empty() -> Self {
        Cell {
            pop: 0.0,
            vel: Vec2::zero(),
            spread: 0.0,
            conductivity: 1.0,
        }
    }

    pub fn set(&mut self, population: f32, velocity: Vec2<f32>, spread_factor: f32) {
        self.pop = population;
        self.vel = velocity;
        self.spread = spread_factor
    }

    pub fn clean(&mut self) {
        self.pop = 0.0;
    }

    // I switched out the `Option<Vec2<f32>>` since we dont need
    // separate handling for the same cell; now `None` is just
    // replaced by the zero vector. We also need access to delta
    // here.
    #[inline(always)]
    pub fn flow_factor(&self, vec: Vec2<f32>, other: &Self, delta: f32) -> f32 {
        let pop_box = &mut Square::new(Vec2::zero(), 1.0);
        pop_box.grow(self.spread * delta);
        pop_box.shift(self.vel * delta);
        other.conductivity * pop_box.intersect_area(&Square::new(vec, 1.0))
    }

    // I'm guessing we want to inline a function like this
    #[inline(always)]
    fn update_vel(&mut self, new_vel: Vec2<f32>, pop_flow: f32) {
        if self.pop + pop_flow == 0.0 {
            self.vel = Vec2::zero();
            return
        }
        self.vel = (self.vel * self.pop + new_vel * pop_flow) / (self.pop + pop_flow)
    }

    #[inline(always)]
    fn update_spread(&mut self, new_spread: f32, pop_flow: f32) {
        if self.pop + pop_flow == 0.0 {
            self.spread = 0.0;
            return
        }
        self.spread = (self.spread * self.pop + new_spread * pop_flow) / (self.pop + pop_flow)
    }

    pub fn tick(&self, (this, left, right, up, down): (&mut Self, &mut Self, &mut Self, &mut Self, &mut Self), delta: f32) {
        /* save some effort quite often
        if self.pop == 0.0 {
            return
        }*/
        let flow_factors = [
            self.flow_factor(Vec2::zero(), this, delta),
            self.flow_factor(Vec2::left(), left, delta),
            self.flow_factor(Vec2::right(), right, delta),
            self.flow_factor(Vec2::up(), up, delta),
            self.flow_factor(Vec2::down(), down, delta),
        ];
        let flow_sum: f32 = (&flow_factors).iter().sum::<f32>();

        // prevent division by zero; if all flows are zero then
        // all adjacent conductivities must be zero so nothing
        // should move anyway (unless some are negative in which
        // case in theory the net flow should be between the adjacent
        // cells with no flow into/out of here but meh its
        // exceedingly rare anyway)
        if flow_sum == 0.0 {
            return
        }
        let flow_vals = [
            self.pop * flow_factors[1] / flow_sum,
            self.pop * flow_factors[2] / flow_sum,
            self.pop * flow_factors[3] / flow_sum,
            self.pop * flow_factors[4] / flow_sum,
        ];
        let val_sum: f32 = (&flow_vals).iter().sum::<f32>();

        left.update_vel(self.vel, flow_vals[0]);
        right.update_vel(self.vel, flow_vals[1]);
        up.update_vel(self.vel, flow_vals[2]);
        down.update_vel(self.vel, flow_vals[3]);
        left.update_spread(self.spread, flow_vals[0]);
        right.update_spread(self.spread, flow_vals[1]);
        up.update_spread(self.spread, flow_vals[2]);
        down.update_spread(self.spread, flow_vals[3]);
        this.pop -= val_sum;
        left.pop += flow_vals[0];
        right.pop += flow_vals[1];
        up.pop += flow_vals[2];
        down.pop += flow_vals[3];

        this.vel.x += (left.pop - self.pop).powf(3.0) * 0.000000001;
        this.vel.x += (self.pop - right.pop).powf(3.0) * 0.000000001;
        this.vel.y -= (up.pop - self.pop).powf(3.0) * 0.000000001;
        this.vel.y -= (self.pop - down.pop).powf(3.0) * 0.000000001;
    }

    pub fn get_colour(&self) -> u32 {
        ((self.pop as u32).min(255) << 16) + (((1.0 - self.conductivity) * 255.0) as u32).min(255)
    }
}

struct World {
    cells: Box<[[Cell; H]; W]>,
}

impl World {
    pub fn test(option: i32) -> Self {
        let mut this = Self {
            cells: Box::new([[Cell::empty(); H]; W]),
        };

        for i in 0..W {
            for j in 0..H {
                if
                    (20 + i as i32 - W as i32 / 2).wrapping_pow(2) +
                    (j as i32 - H as i32 / 2).wrapping_pow(2) < 100
                {
                    this.cells[i][j].set(250.0, Vec2::new(0.2, 0.1), 0.2);
                }
                if
                    (-20 + i as i32 - W as i32 / 2).wrapping_pow(2) +
                    (j as i32 - H as i32 / 2).wrapping_pow(2) < 100
                {
                    this.cells[i][j].set(250.0, Vec2::new(-0.2, 0.1), 0.2);
                }
            }
        }

        match option {
            2 => for i in 24..W - 24 {
                for j in 24..H - 24 {
                    this.cells[i][j].conductivity = (j as f32) / (H as f32)
                }
            },
            _ => {
                for i in 24..W - 24 {
                    this.cells[i][H / 3].conductivity = 0.0;
                }
                for i in 0..10 {
                    this.cells[20 + i][H / 6].conductivity = 0.0;
                    this.cells[W - 20 - i][H / 6].conductivity = 0.0;
                }
                for i in 0..W {
                    this.cells[i][0].conductivity = 0.0;
                    this.cells[i][0].pop = 400.0;
                    this.cells[i][H - 1].conductivity = 0.0;
                    this.cells[i][H - 1].pop = 400.0;
                }
                for j in 0..H {
                    this.cells[0][j].conductivity = 0.0;
                    this.cells[0][j].pop = 400.0;
                    this.cells[W - 1][j].conductivity = 0.0;
                    this.cells[W - 1][j].pop = 400.0;
                }
            },
        };

        this
    }

    pub fn tick(&mut self, delta: f32) {
        let mut new_cells = self.cells.clone();
        for i in 1..W - 1 {
            for j in 1..H - 1 {
                let mut this = new_cells[i][j];
                let mut left = new_cells[i - 1][j];
                let mut right = new_cells[i + 1][j];
                let mut up = new_cells[i][j - 1];
                let mut down = new_cells[i][j + 1];
                self.cells[i][j].tick((
                    &mut this,
                    &mut left,
                    &mut right,
                    &mut up,
                    &mut down,
                ), delta);

                new_cells[i][j] = this;
                new_cells[i - 1][j] = left;
                new_cells[i + 1][j] = right;
                new_cells[i][j - 1] = up;
                new_cells[i][j + 1] = down;
            }
        }

        self.cells = new_cells;
    }

    pub fn render_to(&self, buf: &mut [u32]) {
        for i in 0..W {
            for j in 0..H {
                buf[j * W + i] = self.cells[i][j].get_colour();
            }
        }
    }
}

fn main() {
    println!("Test mode:");
    let fun = |_| {
        println!("Defaulting to 1.");
        1
    };
    let mut n_str_1 = String::new();
    io::stdin().read_line(&mut n_str_1).expect("Err 1");
    let n: i32 = n_str_1.trim().parse().unwrap_or_else(fun);
    let mut buf = vec![0; W * H];
    let mut opts = WindowOptions::default();
    opts.scale = minifb::Scale::X8;
    let mut win = Window::new("Worldsim", W, H, opts).unwrap();

    let mut world = World::test(n);

    while win.is_open() {
        world.tick(0.01);
        world.render_to(&mut buf);

        win.update_with_buffer(&buf).unwrap();
    }
}