ck3-regions 0.0.0

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

use std::{
    collections::HashMap,
};

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

//
// Exports
//

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

//
// Constants
//

// TODO?

//
// 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)?
        },
    };

    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 mut wave_grid = WaveGrid::new(map_dimensions.1 as usize, map_dimensions.0 as usize);
        for initial_region_index in 1..=max_region_index {
            println!("expanding region {}/{}", initial_region_index, max_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(&mut wave_grid, initial_coords, &ExpandInitialRegionStrategy::new(initial_region_index, &regions_map_image_buffer));
        }

        println!("filling 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::get_cell_region_index(&wave_grid, &coords);
                if initial_cell_region_index != 0 {
                    continue;
                }
                println!("filling neutral pixel at ({}, {})", x, y);

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

                wave::set_cell_region_index(&mut 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().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 mut adjacency_map_0_image_buffer = RegionsMapImageBuffer::new(map_dimensions.0, map_dimensions.1);
    //let mut adjacency_map_1_image_buffer = RegionsMapImageBuffer::new(map_dimensions.0, map_dimensions.1);
    let mut proximity_map_image_buffer = RegionsMapImageBuffer::new(map_dimensions.0, map_dimensions.1);
    assert_eq!(adjacency_map_0_image_buffer.dimensions(), map_dimensions);
    //assert_eq!(adjacency_map_1_image_buffer.dimensions(), map_dimensions);
    assert_eq!(proximity_map_image_buffer.dimensions(), map_dimensions);

    for y in 0..map_dimensions.1 {
        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 mut 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)
                            //(x as i32 - *border_x as i32).abs() as u32 + (y as i32 - *border_y as i32).abs() as u32
                        })
                        .reduce(f32::min)
                        //.min()
                        .expect("border pixel set if present cannot be empty");

                    (adjacent_region_index, min_sq_distance)
                })
                .collect();
            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);

            adjacency_map_0_image_buffer.get_pixel_mut(x, y).0 = texture::encode_index_as_rgba(
                (adjacent_regions.get(0).unwrap_or(DEFAULT_ADJACENCY_PAIR).0 << 0)
                    | (0xFF00 << 16)
            );
            // adjacency_map_1_image_buffer.get_pixel_mut(x, y).0 = texture::encode_index_as_rgba(
            //     (adjacent_regions.get(2).unwrap_or(DEFAULT_ADJACENCY_PAIR).0 << 0)
            //         | (adjacent_regions.get(3).unwrap_or(DEFAULT_ADJACENCY_PAIR).0 << 16)
            // );

            proximity_map_image_buffer.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
            );

            if x == 0 {
                println!("pixel row {}/{}", y, map_dimensions.1);
            }
        }
    }

    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 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]
}