use std::{
    sync::{
        atomic::{AtomicU64, AtomicUsize, Ordering::Relaxed},
        Arc, RwLock,
    },
    thread,
};

struct Point2d {
    x: f64,
    y: f64,
}

#[derive(Debug, Default, Clone, Copy)]
struct PVal {
    point: usize,
    val: f64,
}

impl PartialEq for PVal {
    fn eq(&self, other: &Self) -> bool {
        self.val == other.val
    }
}

impl Eq for PVal {}

impl PartialOrd for PVal {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for PVal {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.val.partial_cmp(&other.val).unwrap()
    }
}

struct Spindex2d {
    xpoints: Arc<Vec<AtomicUsize>>,
    ypoints: Arc<Vec<AtomicUsize>>,
    xs: Arc<RwLock<Vec<PVal>>>,
    ys: Arc<RwLock<Vec<PVal>>>,
}

impl Spindex2d {
    pub fn new(points: &[Point2d]) -> Self {
        let mut xpoints = Vec::with_capacity(points.len());
        let mut ypoints = Vec::with_capacity(points.len());
        let mut xs = Vec::with_capacity(points.len());
        let mut ys = Vec::with_capacity(points.len());
        for (i, point) in points.iter().enumerate() {
            let x = PVal {
                point: i,
                val: point.x,
            };
            xs.push(x);
            let y = PVal {
                point: i,
                val: point.y,
            };
            ys.push(y);

            xpoints.push(AtomicUsize::new(i));
            ypoints.push(AtomicUsize::new(i));
        }

        let index = Self {
            xpoints: xpoints.into(),
            ypoints: ypoints.into(),
            xs: Arc::new(RwLock::new(xs)),
            ys: Arc::new(RwLock::new(ys)),
        };

        index.sort();

        index
    }

    /// This method assumes that the points have the same cardinality and order as the points used
    /// when constructing this index. May panic if the updated points argument is longer than the
    /// original points.
    pub fn update(&self, points: &[Point2d]) {
        let mut xs = self.xs.write().unwrap();
        let mut ys = self.ys.write().unwrap();
        for (i, point) in points.iter().enumerate() {
            let xi = self.xpoints[i].load(Relaxed);
            let yi = self.ypoints[i].load(Relaxed);
            let x = PVal {
                point: i,
                val: point.x,
            };
            xs[xi] = x;

            let y = PVal {
                point: i,
                val: point.y,
            };
            ys[yi] = y;
        }
        self.sort();
    }

    fn sort(&self) {
        let xs = self.xs.clone();
        let points = self.xpoints.clone();
        let xhandle = thread::spawn(move || {
            let mut xs = xs.write().unwrap();
            xs.sort_unstable();
            for (i, x) in xs.iter().enumerate() {
                let p = x.point;
                points[p].store(i, Relaxed);
            }
        });

        let ys = self.ys.clone();
        let points = self.ypoints.clone();
        let yhandle = thread::spawn(move || {
            let mut ys = ys.write().unwrap();
            ys.sort_unstable();
            for (i, y) in ys.iter().enumerate() {
                let p = y.point;
                points[p].store(i, Relaxed);
            }
        });

        xhandle.join().unwrap();
        yhandle.join().unwrap();
    }
}

use rand::{Rng, SeedableRng};
use rand_hc::Hc128Rng;

fn create_random_points(num_points: usize, seed: &[u8; 32]) -> Vec<[f64; 2]> {
    let mut rng = Hc128Rng::from_seed(*seed);
    (0..num_points).map(|_| rng.random()).collect()
}

use criterion::Criterion;
use criterion::{criterion_group, criterion_main};

const SEED_1: &[u8; 32] = b"Gv0aHMtHkBGsUXNspGU9fLRuCWkZWHZx";
// const SEED_2: &[u8; 32] = b"km7DO4GeaFZfTcDXVpnO7ZJlgUY7hZiS";

const DEFAULT_NUM_POINTS: usize = 1_000_000;

fn new_index(c: &mut Criterion) {
    c.bench_function("new index", move |b| {
        let points: Vec<_> = create_random_points(DEFAULT_NUM_POINTS, SEED_1)
            .into_iter()
            .map(|[x, y]| Point2d { x, y })
            .collect();

        b.iter(|| {
            Spindex2d::new(&points);
        });
    });
}

fn update_positions(c: &mut Criterion) {
    c.bench_function("update positions", move |b| {
        let points: Vec<_> = create_random_points(DEFAULT_NUM_POINTS, SEED_1)
            .into_iter()
            .map(|[x, y]| Point2d { x, y })
            .collect();

        let mut rng = Hc128Rng::from_seed(*SEED_1);

        let jitters: Vec<_> = points
            .iter()
            .map(|p| {
                let x = p.x + rng.random_range(-5.0..=5.0);
                let y = p.y + rng.random_range(-5.0..=5.0);
                Point2d { x, y }
            })
            .collect();

        let index = Spindex2d::new(&points);
        let mut round = 0;
        b.iter(|| {
            if round % 2 == 0 {
                index.update(&jitters);
            } else {
                index.update(&points);
            }
            round += 1;
        });
    });
}

criterion_group!(benches, new_index);
criterion_main!(benches);