use rand::Rng;
use std::fmt;
use std::io;

#[derive(Clone)]
enum TreeStage {
    Sapling,
    Mature,
    Elder,
}

impl fmt::Display for TreeStage {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            TreeStage::Sapling => write!(f, "S"),
            TreeStage::Mature => write!(f, "M"),
            TreeStage::Elder => write!(f, "D"),
        }
    }
}
#[derive(Clone)]
struct Tree {
    age: u16,
    stage: TreeStage,
}

impl Tree {
    fn new(age: u16) -> Tree {
        match age {
            0...11 => Tree {
                age,
                stage: TreeStage::Sapling,
            },
            12...120 => Tree {
                age,
                stage: TreeStage::Mature,
            },
            _ => Tree {
                age,
                stage: TreeStage::Elder,
            },
        }
    }

    fn add_month(&self) -> Tree {
        let new_age = self.age + 1;
        Tree::new(new_age)
    }
}

impl fmt::Display for Tree {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.stage.fmt(f)
    }
}

#[derive(Copy, Clone)]
struct LumberJack {
    lumber_collected: u16,
    // potential to store the amount of times wandered
}

impl fmt::Display for LumberJack {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self)
    }
}
#[derive(Clone)]
enum ForestFeature {
    Tree(Tree),
    LumberJack(LumberJack),
    Bear,
    Empty,
    LumberSap(u16),
    BearTree(Tree),
}

impl fmt::Display for ForestFeature {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            ForestFeature::Bear => write!(f, "B"),
            ForestFeature::Empty => write!(f, "E"),
            ForestFeature::Tree(x) => x.fmt(f),
            ForestFeature::LumberJack(_) => write!(f, "J"),
            ForestFeature::LumberSap(_) => write!(f, "L"),
            ForestFeature::BearTree(_) => write!(f, "R"),
        }
    }
}

struct TreePosition {
    tree: Tree,
    x: usize,
    y: usize,
}

fn main() {
    let dimensions = get_dimensions();

    let forest = generate_forest(dimensions);

    print_forest(&forest);

    simulate(forest);
}

fn simulate(mut forest: Vec<Vec<ForestFeature>>) {
    for year in 0..400 {
        for month in 0..12 {
            forest = age_trees(&forest);
            print_forest(&forest);
        }
    }
}

fn spawn_trees(forest: &[Vec<ForestFeature>]) -> Vec<Vec<ForestFeature>> {
    // get a list of all tree positions
    let trees_positions: Vec<Vec<Option<TreePosition>>> = forest
        .iter()
        .enumerate()
        .map(move |(i, x)| {
            x.iter()
                .enumerate()
                .map(move |(j, y)| match y {
                    ForestFeature::Tree(tre) => Some(TreePosition {
                        tree: tre.clone(),
                        x: i,
                        y: j,
                    }),
                    ForestFeature::BearTree(tre) => Some(TreePosition {
                        tree: tre.clone(),
                        x: i,
                        y: j,
                    }),
                    _ => None,
                })
                .collect()
        })
        .collect();

    let trees_positions: Vec<TreePosition> =
        trees_positions.into_iter().flatten().flatten().collect();

    // set new spawn locations
    let new_spawned_locations: Vec<(usize, usize)> = trees_positions
        .into_iter()
        .map(|x| get_new_spawning_locations(&x, &forest))
        .flatten()
        .collect();

    let newForest = forest.clone();

    //Add new spawned locations
    newForest
        .iter_mut()
        .enumerate()
        .map(|(i, tree_line)| {
            tree_line
                .iter_mut()
                .enumerate()
                .map(|(j, currentTree)| {
                    let tree_info = new_spawned_locations
                        .into_iter()
                        .find(|&(new_y, new_x)| new_y == i && new_x == j);
                    match tree_info {
                        None => currentTree.clone(),
                        Some(loc) => match currentTree {
                            ForestFeature::Empty => ForestFeature::Tree(Tree {
                                age: 0,
                                stage: TreeStage::Sapling,
                            }),
                            ForestFeature::Bear => ForestFeature::BearTree(Tree {
                                age: 0,
                                stage: TreeStage::Sapling,
                            }),
                            _ => currentTree.clone(),
                        },
                    }
                })
                .collect()
        })
        .collect()
}

fn get_new_spawning_locations(
    currentTree: &TreePosition,
    forest: &[Vec<ForestFeature>],
) -> Option<(usize, usize)> {
    let max_y = forest.len();
    let max_x = forest[0].len();

    //generate a list of valid positions
    let adjacent_positions: Vec<Vec<Option<(usize, usize)>>> = (0..2)
        .map(|offset_y| {
            let offset_y = offset_y - 1;
            let new_y = currentTree.y + offset_y;

            if new_y < 0 || new_y > max_y {
                return vec![None, None, None];
            }
            return (0..2)
                .map(|offset_x| {
                    let offset_x = offset_x - 1;
                    let new_x = currentTree.x + offset_x;

                    if new_x < 0 || new_x > max_x {
                        return None;
                    }

                    match &forest[new_y][new_y] {
                        ForestFeature::Tree(tre) | ForestFeature::BearTree(tre) => {
                            return Some((new_y, new_x))
                        }
                        _ => return None,
                    }
                })
                .collect::<_>();
        })
        .collect();

    //Filter out invalid positions
    let potential_positions: Vec<(usize, usize)> =
        adjacent_positions.into_iter().flatten().flatten().collect();

    match potential_positions.len() {
        0 => None,
        _ => {
            let chosen_index = rand::thread_rng().gen_range(0, potential_positions.len());

            let (chosen_y, chosen_x) = potential_positions[chosen_index];

            match forest[chosen_y][chosen_x] {
                ForestFeature::Empty => Some((chosen_y, chosen_x)),
                _ => None,
            }
        }
    }
}

fn age_trees(forest: &[Vec<ForestFeature>]) -> Vec<Vec<ForestFeature>> {
    forest
        .iter()
        .map(|x| {
            x.iter()
                .map(|y| match y {
                    ForestFeature::Tree(tre) => ForestFeature::Tree(tre.add_month()),
                    ForestFeature::LumberSap(age) => ForestFeature::LumberSap(age + 1),
                    ForestFeature::BearTree(tre) => ForestFeature::BearTree(tre.add_month()),
                    ForestFeature::Bear => ForestFeature::Bear,
                    ForestFeature::Empty => ForestFeature::Empty,
                    ForestFeature::LumberJack(l) => ForestFeature::LumberJack(*l),
                })
                .collect()
        })
        .collect::<_>()
}

fn print_forest(forest: &[Vec<ForestFeature>]) {
    for f in forest {
        for t in f {
            print!("{}", t)
        }
        println!()
    }
}

fn get_dimensions() -> isize {
    println!("please enter your dimensions");

    let mut dimensions = String::new();

    io::stdin()
        .read_line(&mut dimensions)
        .expect("Error in input");

    dimensions.trim().parse::<isize>().unwrap()
}

fn generate_forest(dimensions: isize) -> Vec<Vec<ForestFeature>> {
    (0..dimensions)
        .map(|_| {
            (0..dimensions)
                .map(|_| {
                    let num = rand::thread_rng().gen_range(1, 101);

                    match num {
                        1..=50 => ForestFeature::Tree(Tree {
                            age: 12,
                            stage: TreeStage::Mature,
                        }),
                        51..=61 => ForestFeature::LumberJack(LumberJack {
                            lumber_collected: 0,
                        }),
                        62..=64 => ForestFeature::Bear,
                        _ => ForestFeature::Empty,
                    }
                })
                .collect()
        })
        .collect::<_>()
}