#![feature(array_zip)]
#![feature(generic_const_exprs)]
#![feature(let_chains)]

use clap::Parser;
use rand::prelude::*;
use crate::cli::{drive_system};
use crate::cli::cli::{StickProbabilityCli, InitialCli, BallsCli, PCM, ModelCli, SurfaceProbabilityMeasureCli, KDStickProbabilityCli};
use crate::cli::cli::PCM::Balls2d;
use crate::cli::output::write;
use crate::surface_probability_measure::{LoggerSticker, ReadOnlyVectorStorage};
use crate::system::model::DLASystem;
use crate::system::spaces::continuous_3d::{ContinuousSticker, ContinuousStorage, ContinuousWalker};
use crate::system::spaces::hexagonal::HexPosition;
use crate::system::spaces::kd_grid::{KDSpace};
use crate::system::spaces::square_grid::{Grid2D, Grid3D};
use crate::system::spaces::VectorStorage;
use crate::system::spawner::UniformSpawner;
use crate::system::sticker::{ProbabilisticSticking, SimpleSticking};
use crate::system::walker::{DiagonalRandomWalker, LocalRandomWalker};

mod system;
mod surface_probability_measure;
mod cli;

fn main() {
    let cli = ModelCli::parse();

    println!("Running: {:?}", cli);

    match cli.preconfigured_model {
        PCM::Initial(InitialCli { grid_size }) => {
            let mut sys = DLASystem::<_, Grid2D, _, _, _, _>::new(
                SmallRng::seed_from_u64(cli.seed),
                VectorStorage::new(grid_size, 2),
                LocalRandomWalker,
                UniformSpawner,
                SimpleSticking,
                cli.max_particles,
            );

            drive_system(&mut sys, cli.max_frames, cli.notify_every);
            write(&sys, cli.format, &cli.output);
        }

        PCM::StickProbability(StickProbabilityCli { grid_size, stick_probability }) => {
            let mut sys = DLASystem::<_, Grid2D, _, _, _, _>::new(
                SmallRng::seed_from_u64(cli.seed),
                VectorStorage::new(grid_size, 2),
                LocalRandomWalker,
                UniformSpawner,
                ProbabilisticSticking::new(stick_probability).unwrap(),
                cli.max_particles,
            );

            drive_system(&mut sys, cli.max_frames, cli.notify_every);
            write(&sys, cli.format, &cli.output);
        }

        PCM::Grid3d(KDStickProbabilityCli { stick_probability }) => {
            let mut sys = DLASystem::<_, Grid3D, _, _, _, _>::new(
                SmallRng::seed_from_u64(cli.seed),
                KDSpace::new(),
                LocalRandomWalker,
                UniformSpawner,
                ProbabilisticSticking::new(stick_probability).unwrap(),
                cli.max_particles,
            );

            drive_system(&mut sys, cli.max_frames, cli.notify_every);
            write(&sys, cli.format, &cli.output);
        }

        PCM::Grid3dOffAxis(KDStickProbabilityCli { stick_probability }) => {
            let mut sys = DLASystem::<_, Grid3D, _, _, _, _>::new(
                SmallRng::seed_from_u64(cli.seed),
                KDSpace::new(),
                DiagonalRandomWalker,
                UniformSpawner,
                ProbabilisticSticking::new(stick_probability).unwrap(),
                cli.max_particles,
            );

            drive_system(&mut sys, cli.max_frames, cli.notify_every);
            write(&sys, cli.format, &cli.output);
        }

        PCM::Hex(StickProbabilityCli { grid_size, stick_probability }) => {
            let mut sys = DLASystem::<_, HexPosition, _, _, _, _>::new(
                SmallRng::seed_from_u64(cli.seed),
                VectorStorage::new(grid_size, 2),
                LocalRandomWalker,
                UniformSpawner,
                ProbabilisticSticking::new(stick_probability).unwrap(),
                cli.max_particles,
            );

            drive_system(&mut sys, cli.max_frames, cli.notify_every);
            write(&sys, cli.format, &cli.output);
        }

        PCM::Balls2d(BallsCli { ball_radius, stick_distance, walk_step }) => {
            use system::spaces::continuous_2d;

            let mut sys = DLASystem::new(
                SmallRng::seed_from_u64(cli.seed),
                continuous_2d::ContinuousStorage { inner: kiddo::KdTree::new(), ball_radius_sq: ball_radius * ball_radius },
                continuous_2d::ContinuousWalker { walk_step },
                UniformSpawner,
                continuous_2d::ContinuousSticker { range_sq: stick_distance * stick_distance },
                cli.max_particles,
            );

            drive_system(&mut sys, cli.max_frames, cli.notify_every);
            write(&sys, cli.format, &cli.output);
        }

        PCM::Balls(BallsCli { ball_radius, stick_distance, walk_step }) => {
            let mut sys = DLASystem::new(
                SmallRng::seed_from_u64(cli.seed),
                ContinuousStorage { inner: kiddo::KdTree::new(), ball_radius_sq: ball_radius * ball_radius },
                ContinuousWalker { walk_step },
                UniformSpawner,
                ContinuousSticker { range_sq: stick_distance * stick_distance },
                cli.max_particles,
            );

            drive_system(&mut sys, cli.max_frames, cli.notify_every);
            write(&sys, cli.format, &cli.output);
        }

        PCM::SurfaceProbabilityMeasure(SurfaceProbabilityMeasureCli { grid_size, stick_probability, particles, initial_data }) => {
            let logger_sticker = LoggerSticker::new(stick_probability);

            let mut sys = DLASystem::new(
                SmallRng::seed_from_u64(cli.seed),
                ReadOnlyVectorStorage::new(&initial_data, grid_size),
                LocalRandomWalker,
                UniformSpawner,
                &logger_sticker,
                cli.max_particles,
            );

            let particles = particles as usize;

            while sys.running && logger_sticker.stick_positions.borrow().len() < particles {
                sys.update();
            }

            let mut writer = csv::Writer::from_path(cli.output)
                .expect("Failed to open output path");

            let stick_positions = logger_sticker.stick_positions.borrow();

            stick_positions
                .iter()
                .for_each(|pos|
                    writer.serialize(pos)
                        .expect("Failed to write position")
                );

            writer.flush()
                .unwrap();
        }
    }
}