ck3-regions 0.1.2

mod cli;
mod config;
mod error;
mod model;
mod loading;
mod map;
mod texture;
mod wave;
mod metadata;

use std::{
    collections::HashMap,
    path::Path,
    sync::{Arc, Mutex},
    io::Write
};

use rayon::prelude::*;

use crate::{
    config::ModeSettings,
    map::RegionsMapData,
    model::game::TitleId,
    texture::RegionsMapImageBuffer,
    wave::strategies::{ExpandInitialRegionStrategy, FindNearestNonZeroRegionStrategy}
};

//
// Exports
//

pub use crate::{
    cli::Cli,
    config::Config,
    error::Error,
};

//
// Constants
//

const EXPECT_NO_OTHER_ARC_REFERENCES: &str = "no other Arc references";

//
// Interface
//

pub fn run(config: Config) -> Result<(), Error> {
    let RegionsMapData{
        image_buffer: mut regions_map_image_buffer,
        region_title_indices,
        max_region_index,
        region_initial_pixel_coords,
    } = match &config.mode_settings {
        ModeSettings::TitleBasedRegions { region_title_tier } => {
            map::build_regions_map_from_mod_data(*region_title_tier, &config.mod_paths)?
        },
        ModeSettings::PredefinedRegions { input_regions_map_path } => {
            map::load_predefined_regions_map(&input_regions_map_path)?
        },
    };

    println!("map data loaded successfully");

    if config.must_attempt_early_exit || config.is_dry_run {
        let is_raw_data_current = is_raw_data_current(
            &config.mod_paths.raw_regions_map,
            config.output_metadata_yaml_path.as_ref().map(|p| p.as_path()),
            &regions_map_image_buffer,
            region_title_indices.as_ref()
        );
        println!(
            "existing raw regions map {} title indices are {}",
            if is_raw_data_current { "and" } else { "and/or" },
            if is_raw_data_current { "up to date" } else { "stale" }
        );
        if is_raw_data_current {
            return Ok(());
        } else if config.is_dry_run {
            println!("will not regenerate due to --dry-run");

            return Ok(());
        }
    }

    regions_map_image_buffer.save(config.mod_paths.raw_regions_map.clone())?;
    println!("succesfully wrote {}", config.mod_paths.raw_regions_map.display());

    let map_dimensions = regions_map_image_buffer.dimensions();

    if config.must_expand_regions {
        // Expand regions to fill neutral territory using wave algorithm

        println!("expanding non-zero regions");

        let wave_grid = Arc::new(wave::new_wave_grid(map_dimensions.1 as usize, map_dimensions.0 as usize));
        let expanded_count = Arc::new(Mutex::new(0));
        (1..=max_region_index).into_par_iter().for_each(|initial_region_index| {
            let initial_coords = *region_initial_pixel_coords.get(&initial_region_index)
                .expect("initial pixel coords must be present for every region");

            wave::expand_wave(&wave_grid, initial_coords, &ExpandInitialRegionStrategy::new(initial_region_index, &regions_map_image_buffer));

            {
                let mut count = expanded_count.lock().unwrap();
                *count += 1;
                print!("\rexpanded {}/{} regions", *count, max_region_index);
                std::io::stdout().flush().unwrap();
            }
        });
        let wave_grid = Arc::try_unwrap(wave_grid).expect(EXPECT_NO_OTHER_ARC_REFERENCES);

        println!("\nfilling remaining neutral zones");

        for y in 0..wave_grid.rows() as u32 {
            for x in 0..wave_grid.cols() as u32 {
                let coords = (x, y);

                let initial_cell_region_index = wave_grid.get(y as usize, x as usize).unwrap().read().unwrap().region_index;
                if initial_cell_region_index != 0 {
                    continue;
                }
                //println!("filling neutral pixel at ({}, {})", x, y);

                let nearest_region_index = wave::expand_wave(&wave_grid, coords, &FindNearestNonZeroRegionStrategy::new())
                    .expect("nearest non-zero region must exist");

                wave::set_cell_region_index(&wave_grid, &coords, nearest_region_index);
            }
        }

        for y in 0..wave_grid.rows() as u32 {
            for x in 0..wave_grid.cols() as u32 {
                let region_index = wave_grid.get(y as usize, x as usize).unwrap().read().unwrap().region_index;
                if region_index != 0 {
                    let existing_region_index = texture::decode_index_from_rgb(&regions_map_image_buffer.get_pixel(x, y).0[..3]);
                    //assert!(existing_region_index == 0 || existing_region_index == region_index);
                    if existing_region_index == 0 {
                        regions_map_image_buffer.get_pixel_mut(x, y).0 = texture::encode_index_as_rgba(region_index | (0xFF00 << 16));
                    }
                }
            }
        }
    } else {
        println!("skipped non-zero regions expansion");
    }

    regions_map_image_buffer.save(config.mod_paths.regions_map.clone())?;
    println!("succesfully wrote {}", config.mod_paths.regions_map.display());

    println!("registering adjacencies");

    let mut regions_graph = HashMap::new();

    for y in 0..regions_map_image_buffer.height() {
        for x in 0..regions_map_image_buffer.width() {
            let region_title_index = texture::decode_index_from_rgb(&regions_map_image_buffer.get_pixel(x, y).0[..3]);

            // Register neighbor adjacencies for our top and left neighbours (the ones we already processed in this loop)
            for (neighbour_x, neighbour_y) in [(x.checked_sub(1), Some(y)), (Some(x), y.checked_sub(1))] {
                if neighbour_x.is_none() || neighbour_y.is_none() {
                    continue;
                }

                let (neighbour_x, neighbour_y) = (neighbour_x.unwrap(), neighbour_y.unwrap());

                regions_map_image_buffer.get_pixel_checked(neighbour_x, neighbour_y)
                    .map(|rgba| {
                        register_adjacency(
                            &mut regions_graph,
                            region_title_index,
                            texture::decode_index_from_rgb(&rgba.0[..3]),
                            (neighbour_x, neighbour_y)
                        );
                    });
            }
        }
    }

    println!("building adjacency and proximity maps");

    let adjacency_map_0_image_buffer = RegionsMapImageBuffer::new(map_dimensions.0, map_dimensions.1);
    let proximity_map_image_buffer = RegionsMapImageBuffer::new(map_dimensions.0, map_dimensions.1);
    assert_eq!(adjacency_map_0_image_buffer.dimensions(), map_dimensions);
    assert_eq!(proximity_map_image_buffer.dimensions(), map_dimensions);

    let adjacency_map_0_image_buffer = Arc::new(Mutex::new(adjacency_map_0_image_buffer));
    let proximity_map_image_buffer = Arc::new(Mutex::new(proximity_map_image_buffer));
    let regions_map_image_buffer = Arc::new(regions_map_image_buffer);
    let regions_graph = Arc::new(regions_graph);
    let config = Arc::new(config);
    let processed_rows_count = Arc::new(Mutex::new(0));

    (0..map_dimensions.1).into_par_iter().for_each(|y| {
        for x in 0..map_dimensions.0 {
            let current_region_index = texture::decode_index_from_rgb(
                &regions_map_image_buffer.get_pixel(x, y).0[..3]
            );
            if config.must_expand_regions && current_region_index == 0 {
                continue;
            }

            let adjacent_regions: Vec<_> = regions_graph.get(&current_region_index)
                .expect("every region must be present in the graph")
                .0.iter()
                .map(|(adjacent_region_index, border_pixels_coords)| {
                    let min_sq_distance = border_pixels_coords.iter()
                        .map(|(border_x, border_y)| {
                            (x as f32 - *border_x as f32).powi(2) + (y as f32 - *border_y as f32).powi(2)
                        })
                        .reduce(f32::min)
                        .expect("border pixel set if present cannot be empty");

                    (adjacent_region_index, min_sq_distance)
                })
                .collect();

            let mut adjacent_regions = adjacent_regions;
            adjacent_regions.sort_by(|(_, sq_distance_0), (_, sq_distance_1)| {
                sq_distance_0.partial_cmp(sq_distance_1).expect("all distances must be comparable")
            });

            const DEFAULT_ADJACENCY_PAIR: &(&u32, f32) = &(&0, f32::INFINITY);

            {
                let mut adj_buf = adjacency_map_0_image_buffer.lock().unwrap();
                adj_buf.get_pixel_mut(x, y).0 = texture::encode_index_as_rgba(
                    (adjacent_regions.get(0).unwrap_or(DEFAULT_ADJACENCY_PAIR).0 << 0)
                        | (0xFF00 << 16)
                );
            }
            {
                let mut prox_buf = proximity_map_image_buffer.lock().unwrap();
                prox_buf.get_pixel_mut(x, y).0 = sq_distance_to_proximity_rgba(
                    adjacent_regions.get(0).unwrap_or(DEFAULT_ADJACENCY_PAIR).1,
                    config.proximity_distance_limit
                );
            }
        }

        {
            let mut processed_rows_count = processed_rows_count.lock().unwrap();
            *processed_rows_count += 1;
            print!("\rprocessed {}/{} rows", *processed_rows_count, map_dimensions.1);
            if *processed_rows_count % 128 == 0 {
                std::io::stdout().flush().unwrap();
            }
        }
    });

    println!(); // newline after progress indicator

    let adjacency_map_0_image_buffer = Arc::try_unwrap(adjacency_map_0_image_buffer)
        .expect(EXPECT_NO_OTHER_ARC_REFERENCES)
        .into_inner()
        .unwrap();
    let proximity_map_image_buffer = Arc::try_unwrap(proximity_map_image_buffer)
        .expect(EXPECT_NO_OTHER_ARC_REFERENCES)
        .into_inner()
        .unwrap();

    adjacency_map_0_image_buffer.save(config.mod_paths.adjacency_map_0.clone())?;
    println!("succesfully wrote {}", config.mod_paths.adjacency_map_0.display());
    //adjacency_map_1_image_buffer.save(config.mod_paths.adjacency_map_1.clone())?;
    //println!("succesfully wrote {}", config.mod_paths.adjacency_map_1.display());

    proximity_map_image_buffer.save(config.mod_paths.proximity_map.clone())?;
    println!("succesfully wrote {}", config.mod_paths.proximity_map.display());

    if config.output_metadata_yaml_path.is_some() {
        metadata::write_metadata(
            &config.output_metadata_yaml_path.as_ref().unwrap(),
            &config,
            region_title_indices.as_ref(),
            max_region_index
        ).map_err(Error::from_write_metadata_error)?;
    }

    Ok(())
}

//
// Service types
//

type BorderPixelsSet = Vec<(u32, u32)>;

struct RegionAdjacencies(HashMap<u32, BorderPixelsSet>);

//
// Service
//

fn is_raw_data_current(
    raw_regions_map_path:           &Path,
    region_title_indices_yaml_path: Option<&Path>,
    raw_regions_map_image_buffer:   &RegionsMapImageBuffer,
    region_title_indices:           Option<&HashMap<TitleId, u32>>
) -> bool {
    if !raw_regions_map_path.is_file() {
        println!("existing raw regions map not found");

        return false;
    }

    if !region_title_indices_yaml_path.map_or(false, |path| path.is_file()) {
        println!("no region title indices YAML file available to check");

        return false;
    }

    map::is_regions_data_current(raw_regions_map_path, region_title_indices_yaml_path.unwrap(), raw_regions_map_image_buffer, region_title_indices.unwrap())
}

fn register_adjacency(
    regions_graph:    &mut HashMap<u32, RegionAdjacencies>,
    current_region_index:  u32,
    adjacent_region_index: u32,
    border_pixel_coords:   (u32, u32)
) {
    const BORDER_PIXELS_SET_INITIAL_CAPACITY: usize = 256;

    if current_region_index == adjacent_region_index {
        return;
    }

    // TODO: Optimize so that we don't store the same adjacency info twice?

    regions_graph.entry(current_region_index)
        .or_insert_with(|| RegionAdjacencies(HashMap::new()))
        .0.entry(adjacent_region_index)
        .or_insert_with(|| BorderPixelsSet::with_capacity(BORDER_PIXELS_SET_INITIAL_CAPACITY))
        .push(border_pixel_coords);

    regions_graph.entry(adjacent_region_index)
        .or_insert_with(|| RegionAdjacencies(HashMap::new()))
        .0.entry(current_region_index)
        .or_insert_with(|| BorderPixelsSet::with_capacity(BORDER_PIXELS_SET_INITIAL_CAPACITY))
        .push(border_pixel_coords);
}

fn sq_distance_to_proximity_rgba(sq_distance: f32, proximity_distance_limit: f32) -> [u8; 4] {
    let alpha = (proximity_distance_limit - sq_distance.sqrt()).max(0.0)/proximity_distance_limit;
    let alpha = (alpha*255.0) as u8;

    [0, 0, 0, alpha]
}