use rayon::prelude::*;
use std::ops::{Add, Index, IndexMut, Sub};

pub enum ReadError {
    BadSymbol,
}

pub enum Error {
    Read(ReadError),
    BadPopSize,
    BadMut,
    BadCrossRatio,
    BadMove(Pos),
}

type Result<T> = std::result::Result<T, Error>;

#[derive(PartialEq, Copy, Clone)]
enum Direction {
    U,
    D,
    L,
    R,
}

#[derive(PartialEq, Copy, Clone)]
enum Cell {
    Wall,
    Box,
    Goal(bool),
    None,
}

#[derive(PartialEq, Clone)]
enum Solution {
    Ok,
    Win,
    None,
}

#[derive(PartialEq, Default, Clone, Copy)]
pub struct Pos(isize, isize);

impl Add for Pos {
    type Output = Self;

    fn add(self, other: Self) -> Self {
        Self(self.0 + other.0, self.1 + other.1)
    }
}

impl Sub for Pos {
    type Output = Self;

    fn sub(self, other: Self) -> Self {
        Self(self.0 - other.0, self.1 - other.1)
    }
}

impl std::fmt::Display for Pos {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{} {}", self.0, self.1)
    }
}

impl std::fmt::Display for ReadError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ReadError::BadSymbol => write!(f, "bad symbol"),
        }
    }
}

impl std::fmt::Display for Error {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Error::Read(e) => write!(f, "{e}"),
            Error::BadPopSize => write!(f, "population size cannot be 0"),
            Error::BadMut => write!(f, "mutation length too large"),
            Error::BadCrossRatio => write!(f, "crossover ratio cannot be 0"),
            Error::BadMove(pos) => write!(f, "invalid move ({pos})"),
        }
    }
}

#[derive(Clone, Default)]
struct Level {
    size: Pos,
    player: Pos,
    map: Vec<Cell>,
}

impl Index<Pos> for Level {
    type Output = Cell;

    fn index(&self, pos: Pos) -> &Self::Output {
        &self.map[((pos.1 * self.size.0) + pos.0) as usize]
    }
}

impl IndexMut<Pos> for Level {
    fn index_mut(&mut self, pos: Pos) -> &mut Self::Output {
        &mut self.map[((pos.1 * self.size.0) + pos.0) as usize]
    }
}

impl std::fmt::Display for Level {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        for i in 0..self.size.1 {
            for j in 0..self.size.0 {
                if Pos(j, i) == self.player {
                    write!(f, "@")?;
                } else {
                    write!(
                        f,
                        "{}",
                        match self[Pos(j, i)] {
                            Cell::Wall => 'X',
                            Cell::Box => '*',
                            Cell::Goal(true) => '!',
                            Cell::Goal(false) => '.',
                            Cell::None => ' ',
                        }
                    )?;
                }
            }
            writeln!(f)?;
        }
        Ok(())
    }
}

impl Level {
    fn read<R: std::io::BufRead>(file: R) -> std::result::Result<Self, ReadError> {
        let (mut player, mut size) = (Pos::default(), Pos::default());

        let map = file
            .lines()
            .flat_map(|l| {
                let l = l.unwrap();

                if l.len() > size.0 as usize {
                    size.0 = l.len() as isize;
                }

                size.1 += 1;
                l.char_indices()
                    .map(|(i, c)| match c {
                        '@' => {
                            player.0 = i as isize;
                            player.1 = size.1 - 1;
                            Ok(Cell::None)
                        }
                        'X' => Ok(Cell::Wall),
                        '*' => Ok(Cell::Box),
                        'g' => Ok(Cell::Goal(false)),
                        '.' => Ok(Cell::None),
                        _ => Err(ReadError::BadSymbol),
                    })
                    .collect::<Vec<std::result::Result<_, _>>>()
            })
            .collect::<std::result::Result<Vec<_>, _>>()?;
        Ok(Self { size, player, map })
    }

    fn read_rle(map: String) -> std::result::Result<Self, ReadError> {
        let mut n: usize = 1;

        Self::read(
            map.chars()
                .filter_map(|c| match c {
                    '0'..='9' => {
                        n = c.to_digit(10).unwrap() as usize;
                        None
                    }
                    '|' => Some("\n".into()),
                    _ => Some(c.to_string().repeat(n)),
                })
                .collect::<Vec<String>>()
                .join("")
                .as_bytes(),
        )
    }
}

#[derive(Clone, Default)]
struct State {
    lvl: Level,
}

impl State {
    fn new(lvl: Level) -> Self {
        State { lvl }
    }

    fn toggle_box(&mut self, pos: Pos) {
        self.lvl[pos] = if self.lvl[pos] == Cell::Goal(true) {
            Cell::Goal(false)
        } else {
            Cell::None
        };
    }

    fn move_pos(&mut self, pos: Pos) -> Result<(bool, bool)> {
        let dpos = self.lvl.player + pos;

        if ((dpos.1 * self.lvl.size.0) + dpos.0) as usize >= self.lvl.map.len() {
            return Err(Error::BadMove(dpos));
        }

        match self.lvl[dpos] {
            Cell::None | Cell::Goal(false) => self.lvl.player = dpos,
            Cell::Box | Cell::Goal(true) => {
                let bpos = dpos + pos;

                match self.lvl[bpos] {
                    Cell::None => {
                        self.toggle_box(dpos);
                        self.lvl[bpos] = Cell::Box;
                        self.lvl.player = dpos;
                        return Ok((false, false));
                    }
                    Cell::Goal(false) => {
                        self.toggle_box(dpos);
                        self.lvl[bpos] = Cell::Goal(true);
                        self.lvl.player = dpos;
                        return Ok((true, false));
                    }
                    _ => return Ok((false, true)),
                }
            }
            _ => return Ok((false, true)),
        }

        Ok((false, false))
    }

    fn move_dir(&mut self, dir: Direction) -> Result<(bool, bool)> {
        match dir {
            Direction::U => self.move_pos(Pos(0, -1)),
            Direction::D => self.move_pos(Pos(0, 1)),
            Direction::L => self.move_pos(Pos(-1, 0)),
            Direction::R => self.move_pos(Pos(1, 0)),
        }
    }

    fn test_box(&self, pos: Pos, bt: bool) -> bool {
        let dir: &[Pos] = &[Pos(0, -1), Pos(0, 1), Pos(-1, 0), Pos(1, 0)];
        let mut mbc = 0;

        for i in dir {
            match self.lvl[pos + *i] {
                Cell::Goal(false) | Cell::None => {
                    let dpos = pos - *i;

                    match self.lvl[dpos] {
                        Cell::Goal(false) | Cell::None => mbc += 1,
                        Cell::Box => {
                            if bt {
                                mbc += self.test_box(dpos, false) as usize
                            }
                        }
                        _ => (),
                    }
                }
                _ => (),
            }
        }

        mbc != 0
    }

    fn dist_box(&self, pos: Pos) -> usize {
        let mut d = 0;

        for i in 0..self.lvl.size.1 {
            for j in 0..self.lvl.size.0 {
                if self.lvl[Pos(j, i)] == Cell::Goal(false) {
                    d += ((j as i32 - pos.0 as i32).abs() + (i as i32 - pos.1 as i32).abs())
                        as usize;
                }
            }
        }

        d
    }

    fn dist_boxes(&self) -> usize {
        let (mut max, mut rb) = (0, 0);

        for i in 0..self.lvl.size.1 {
            for j in 0..self.lvl.size.0 {
                if self.lvl[Pos(j, i)] == Cell::Box {
                    max += self.dist_box(Pos(j, i));
                    rb += 1;
                }
            }
        }

        max * rb
    }

    fn evaluate(&self) -> Solution {
        let mut cond = Solution::Win;

        for i in 0..self.lvl.size.1 {
            for j in 0..self.lvl.size.0 {
                match self.lvl[Pos(j, i)] {
                    Cell::Goal(false) => cond = Solution::Ok,
                    Cell::Box => {
                        if !self.test_box(Pos(j, i), true) {
                            return Solution::None;
                        }
                    }
                    _ => (),
                }
            }
        }

        cond
    }
}

#[derive(Clone)]
struct Agent {
    state: State,
    dna: Vec<Direction>,
    score: f32,
    dir: usize,
    goal: Solution,
}

impl Agent {
    const DIR_TABLE: &'static [Direction] =
        &[Direction::U, Direction::D, Direction::L, Direction::R];

    fn chromo() -> Direction {
        Self::DIR_TABLE[rand::random::<usize>() % Self::DIR_TABLE.len()]
    }

    fn with_dna(state: State, dna: Vec<Direction>) -> Self {
        Self {
            state,
            dna,
            score: 0.0,
            dir: 0,
            goal: Solution::Ok,
        }
    }

    fn new(state: State, size: usize) -> Self {
        let dna: Vec<Direction> = (0..size).map(|_| Self::chromo()).collect();

        Self::with_dna(state, dna)
    }

    fn cross(&mut self, seed: &Self, cr: usize) {
        let ratio = (rand::random::<usize>() % cr) + 1;
        let len = self.dna.len();

        self.dna[..len / ratio].copy_from_slice(&seed.dna[..len / ratio]);
    }

    fn mutate(&mut self, max: usize) {
        let len = self.dna.len();
        let ni = rand::random::<usize>() % max;
        let nd = rand::random::<usize>() % len;

        for x in 0..ni {
            self.dna[(nd + x) % len] = Agent::chromo();
        }
    }

    fn move_dna(&mut self) -> Result<Solution> {
        if self.dir == self.dna.len() || self.goal == Solution::None {
            Ok(Solution::None)
        } else {
            let (mb, hit) = self.state.move_dir(self.dna[self.dir])?;

            if mb {
                self.score += 0.00021;
            }

            if hit {
                self.score -= 0.00080;
            }

            self.dir += 1;
            Ok(self.state.evaluate())
        }
    }

    fn evaluate(&mut self) -> Result<()> {
        loop {
            let player = self.state.lvl.player;

            match self.move_dna()? {
                Solution::Ok => (),
                Solution::Win => {
                    self.score += 100000.0;
                    self.goal = Solution::Win;
                    break;
                }
                Solution::None => {
                    self.score -= 0.0032;
                    self.goal = Solution::None;
                    break;
                }
            }

            if player == self.state.lvl.player {
                self.score -= 0.0010;
            }
        }

        self.score -= 1024.0 + self.state.dist_boxes() as f32;
        Ok(())
    }

    fn solution(&self) -> String {
        (0..self.dir)
            .map(|i| match self.dna[i] {
                Direction::U => "U",
                Direction::D => "D",
                Direction::L => "L",
                Direction::R => "R",
            })
            .collect()
    }
}

struct Population {
    gen: usize,
    epop_max: usize,
    mut_max: usize,
    cross_ratio: usize,
    map: Level,
    life: Vec<Agent>,
}

pub struct Config {
    pub num_threads: usize,
    pub dna_size: usize,
    pub pop_size: usize,
    pub epop_max: usize,
    pub mut_max: usize,
    pub gen_max: usize,
    pub gen_depth: usize,
    pub cross_ratio: usize,
    pub map_output: bool,
}

impl Population {
    fn new(config: &Config, map: Level) -> Self {
        let life: Vec<Agent> = (0..config.pop_size)
            .map(|_| Agent::new(State::new(map.clone()), config.dna_size))
            .collect();

        Self {
            gen: 0,
            epop_max: config.epop_max,
            mut_max: config.mut_max,
            cross_ratio: config.cross_ratio,
            map,
            life,
        }
    }

    fn step(&mut self) -> Result<()> {
        self.life
            .par_iter_mut()
            .map(|a| a.evaluate())
            .collect::<Result<Vec<()>>>()?;
        self.life.par_sort_by(|a, b| b.score.total_cmp(&a.score));

        if self.life[0].goal != Solution::Win {
            let epop: Vec<Agent> = (0..self.epop_max)
                .map(|i| {
                    let mut a = self.life[i].clone();

                    a.cross(
                        if i % 2 == 0 {
                            &self.life[(i + 1) % self.epop_max]
                        } else {
                            &self.life[(i - 1) % self.epop_max]
                        },
                        self.cross_ratio,
                    );
                    a
                })
                .collect();

            for i in epop.len()..self.life.len() {
                let mut next = Agent::with_dna(
                    State::new(self.map.clone()),
                    if i % 2 == 0 {
                        epop[i % self.epop_max].dna.clone()
                    } else {
                        epop[(i + 1) % self.epop_max].dna.clone()
                    },
                );

                next.mutate(self.mut_max);
                self.life[i] = next;
            }
        }

        self.gen += 1;
        Ok(())
    }
}

pub fn run(config: Config, rle: Option<String>) -> Result<()> {
    rayon::ThreadPoolBuilder::new()
        .num_threads(config.num_threads)
        .build_global()
        .unwrap();

    if config.pop_size == 0 {
        return Err(Error::BadPopSize);
    }

    if config.mut_max > config.dna_size {
        return Err(Error::BadMut);
    }

    if config.cross_ratio == 0 {
        return Err(Error::BadCrossRatio);
    }

    let map = if let Some(rle) = rle {
        Level::read_rle(rle)
    } else {
        Level::read(std::io::stdin().lock())
    }
    .map_err(Error::Read)?;

    let mut agents: Vec<Agent> = Vec::new();

    for _ in 0..config.gen_depth {
        let mut pop = Population::new(&config, map.clone());

        if let Some(agent) = || -> Result<Option<Agent>> {
            while pop.gen < config.gen_max {
                pop.step()?;

                if config.map_output {
                    print!("{}", pop.life[0].state.lvl);
                }

                println!("gen {} score {}", pop.gen, pop.life[0].score);
                if pop.life[0].goal == Solution::Win {
                    return Ok(Some(pop.life[0].clone()));
                }
            }

            Ok(None)
        }()? {
            agents.push(agent);
        }
    }

    agents.par_sort_by_key(|a| a.dir);

    if !agents.is_empty() {
        println!("{}", agents[0].solution());
    }

    Ok(())
}