factorio_blueprint/factorio-layout/src/misc.rs

use crate::{LayoutInput, LayoutResult};
use factorio_core::prelude::*;
use factorio_pathfinding::{Connection, examples::HashMapMap};
use rand::{
    Rng,
    seq::IndexedRandom,
};

pub fn initally_set_blocks(
    input: &LayoutInput,
    size: Position,
    retries: usize,
    rng: &mut impl Rng,
) -> Option<Vec<Block>> {
    let mut blocks = Vec::new();

    let mut aabbs = Vec::new();

    if place_block(input, size, retries, &mut blocks, &mut aabbs, rng) {
        Some(blocks)
    } else {
        None
    }
}

fn place_block(
    input: &LayoutInput,
    size: Position,
    retries: usize,
    blocks: &mut Vec<Block>,
    aabbs: &mut Vec<AABB>,
    rng: &mut impl Rng,
) -> bool {
    if input.blocks.len() == blocks.len() {
        return true;
    }

    let b = &input.blocks[blocks.len()];
    for _ in 0..retries {
        let dir = rng.random::<Direction>();

        let pos = match dir {
            Direction::Up => {
                if size.x - b.size.x < 0 || size.y - b.size.y < 0 {
                    continue;
                }
                Position::new(
                    rng.random_range(0..=(size.x - b.size.x)),
                    rng.random_range(0..=(size.y - b.size.y)),
                )
            }
            Direction::Right => {
                if size.y - b.size.x < 0 || size.x - b.size.y < 0 {
                    continue;
                }
                Position::new(
                    rng.random_range((b.size.y - 1)..size.x),
                    rng.random_range(0..=(size.y - b.size.x)),
                )
            }
            Direction::Down => {
                if size.x - b.size.x < 0 || size.y - b.size.y < 0 {
                    continue;
                }
                Position::new(
                    rng.random_range((b.size.x - 1)..size.x),
                    rng.random_range((b.size.y - 1)..size.y),
                )
            }
            Direction::Left => {
                if size.y - b.size.x < 0 || size.x - b.size.y < 0 {
                    continue;
                }
                Position::new(
                    rng.random_range(0..=(size.x - b.size.y)),
                    rng.random_range((b.size.x - 1)..size.y),
                )
            }
        };

        let current = Block::new(pos, dir, b.size);
        let current_aabb = current.get_aabb();

        if aabbs.iter().all(|&b| !AABB::collision(b, current_aabb)) {
            blocks.push(current);
            aabbs.push(current_aabb);

            if place_block(input, size, retries, blocks, aabbs, rng) {
                return true;
            }

            blocks.pop();
            aabbs.pop();
        }
    }

    false
}

pub fn path_input_from_blocks_positions(
    input: &LayoutInput,
    size: Position,
    block_positions: &[Block],
) -> (Vec<Connection>, HashMapMap) {
    let mut map = HashMapMap::new(size);

    let mut connections = Vec::new();

    for b in block_positions {
        let aabb = b.get_aabb();

        map.set_blocked_range(aabb.min(), aabb.max(), true);
    }

    for c in &input.connections {
        let start_transform = block_positions[c.startblock].block_to_world();
        let start_pos = input.blocks[c.startblock].output[c.startpoint]
            .offset
            .transform(start_transform);
        let start_dir = input.blocks[c.startblock].output[c.startpoint]
            .dir
            .transform(start_transform);
        let end_transform = block_positions[c.endblock].block_to_world();
        let end_pos = input.blocks[c.endblock].input[c.endpoint]
            .offset
            .transform(end_transform);
        let end_dir = input.blocks[c.endblock].input[c.endpoint]
            .dir
            .transform(end_transform);

        connections.push(Connection {
            start_pos,
            start_dir,
            end_pos,
            end_dir,
            beltspeed: c.beltspeed,
            lanes: c.lanes,
        });
    }

    (connections, map)
}

pub fn score(input: &LayoutInput, output: &LayoutResult) -> usize {
    let _ = input;
    output
        .path_result
        .iter()
        .flatten()
        .map(|p| p.cost())
        .sum::<usize>()
}

pub fn mutate<R: Rng>(
    input: &LayoutInput,
    output: &LayoutResult,
    rng: &mut R,
) -> Option<Vec<Block>> {
    let mut blocks = output.positions.clone();

    #[allow(clippy::type_complexity)]
    let r: &[(
        &dyn Fn(&LayoutInput, &LayoutResult, &mut [Block], &mut R) -> bool,
        _,
    )] = &[
        (&mutate_replace::<R>, 30),
        (&mutate_flip::<R>, 50),
        (&mutate_jiggle::<R>, 160),
    ];

    loop {
        let p = r.choose_weighted(rng, |i| i.1).unwrap().0;

        if p(input, output, &mut blocks, rng) && rng.random_bool(0.5) {
            break;
        }
    }
    Some(blocks)
}

fn mutate_replace<R: Rng>(
    input: &LayoutInput,
    output: &LayoutResult,
    blocks: &mut [Block],
    rng: &mut R,
) -> bool {
    let _ = input;
    let i = rng.random_range(0..blocks.len());
    let block = blocks[i];
    let dir = rng.random::<Direction>();

    let size = output.size;

    let pos = match dir {
        Direction::Up => Position::new(
            rng.random_range(0..=(size.x - block.size().x)),
            rng.random_range(0..=(size.y - block.size().y)),
        ),
        Direction::Right => Position::new(
            rng.random_range((block.size().y - 1)..size.x),
            rng.random_range(0..=(size.y - block.size().x)),
        ),
        Direction::Down => Position::new(
            rng.random_range((block.size().x - 1)..size.x),
            rng.random_range((block.size().y - 1)..size.y),
        ),
        Direction::Left => Position::new(
            rng.random_range(0..=(size.x - block.size().y)),
            rng.random_range((block.size().x - 1)..size.y),
        ),
    };

    let new_block = Block::new(pos, dir, block.size());
    let new_aabb = new_block.get_aabb();

    if blocks
        .iter()
        .enumerate()
        .all(|(j, b)| i == j || !AABB::collision(b.get_aabb(), new_aabb))
    {
        blocks[i] = new_block;
        true
    } else {
        false
    }
}

fn mutate_flip<R: Rng>(
    input: &LayoutInput,
    output: &LayoutResult,
    blocks: &mut [Block],
    rng: &mut R,
) -> bool {
    let _ = output;
    let _ = input;
    let i = rng.random_range(0..blocks.len());
    let block = blocks[i];

    blocks[i] = Block::new(
        match block.dir() {
            Direction::Up => block.pos() + block.size() - Position::new(1, 1),
            Direction::Right => block.pos() + Position::new(1 - block.size().y, block.size().x - 1),
            Direction::Down => block.pos() - block.size() + Position::new(1, 1),
            Direction::Left => block.pos() + Position::new(block.size().y - 1, 1 - block.size().x),
        },
        block.dir().reverse(),
        block.size(),
    );
    true
}

fn mutate_jiggle<R: Rng>(
    input: &LayoutInput,
    output: &LayoutResult,
    blocks: &mut [Block],
    rng: &mut R,
) -> bool {
    let _ = input;
    let i = rng.random_range(0..blocks.len());
    let block = blocks[i];

    let dir = rng.random::<Direction>();
    let step = [(1, 10), (2, 5), (3, 5)]
        .choose_weighted(rng, |i| i.1)
        .unwrap()
        .0;

    let new_block = Block::new(
        block.pos().in_direction(&dir, step),
        block.dir(),
        block.size(),
    );
    let new_aabb = new_block.get_aabb();

    if new_aabb.min().x < 0
        || new_aabb.min().y < 0
        || new_aabb.max().x >= output.size.x
        || new_aabb.max().y >= output.size.y
    {
        return false;
    }

    if blocks
        .iter()
        .enumerate()
        .all(|(j, b)| j == i || !AABB::collision(b.get_aabb(), new_aabb))
    {
        blocks[i] = new_block;
        true
    } else {
        false
    }
}