use std::{
    collections::{HashMap, HashSet, VecDeque},
    hash::Hash,
};

use itertools::Itertools;
use regex::Regex;

use crate::util;

#[derive(Debug, Clone)]
pub struct Valve {
    pub flow_rate: usize,
    pub tunnels: Vec<usize>,
}

pub fn parse_valves(input: &String) -> Vec<Valve> {
    let indices: HashMap<String, usize> = input
        .lines()
        .map(|l| l.split(' ').nth(1).unwrap())
        .sorted()
        .enumerate()
        .map(|(i, n)| (n.to_owned(), i))
        .collect();

    let mut valves = HashMap::new();

    let re =
        Regex::new(r"^Valve (\w+) has flow rate=(\d+); tunnels? leads? to valves? (.*)$").unwrap();
    for line in input.lines() {
        let captures = re
            .captures(line)
            .expect(&format!("No captures for {}", line));
        let name = captures.get(1).unwrap().as_str().to_owned();
        let flow_rate = captures.get(2).unwrap().as_str().parse().unwrap();
        let tunnels = captures
            .get(3)
            .unwrap()
            .as_str()
            .split(", ")
            .map(|s| indices[s])
            .collect();
        let valve = Valve { flow_rate, tunnels };
        valves.insert(name, valve);
    }

    valves
        .iter()
        .map(|(name, v)| (indices[name], v))
        .sorted_by_key(|(i, _)| *i)
        .map(|(_, v)| v.clone())
        .collect()
}

#[derive(Debug, Clone)]
pub struct Path {
    pub dist: usize,
    pub path: Vec<usize>,
}

pub fn calc_paths(valves: &Vec<Valve>) -> Vec<Vec<Path>> {
    let n = valves.len();

    // result[from][to] = (dist, path)
    let mut result = vec![
        vec![
            Path {
                dist: usize::MAX,
                path: Vec::new()
            };
            n
        ];
        n
    ];

    for i in 0..n {
        result[i][i].dist = 0;
        let mut next = VecDeque::new();
        next.push_back(i);
        while let Some(j) = next.pop_front() {
            let d_next = result[i][j].dist + 1;
            for tunnel in &valves[j].tunnels {
                let d_tunnel = &result[i][*tunnel];
                if d_next < d_tunnel.dist {
                    let mut path = result[i][j].path.clone();
                    path.push(*tunnel);
                    result[i][*tunnel] = Path { dist: d_next, path };
                    next.push_back(*tunnel);
                }
            }
        }
    }

    result
}

#[derive(Debug, Clone, Hash, PartialEq)]
pub enum Goal {
    Idle,
    NothingReachable,
    Opening(usize),
}

#[derive(Debug, Clone)]
pub struct State {
    pub opened: u64,
    pub locations: Vec<usize>,
    pub goals: Vec<Goal>,
    pub flow_rate: usize,
    pub time_remaining: usize,
    pub pressure_released: usize,
    pub history: Vec<Vec<usize>>,

    interesting: HashSet<usize>,
}

impl State {
    pub fn new(valves: &Vec<Valve>, initial_time: usize, agents: usize) -> State {
        let interesting = valves
            .iter()
            .enumerate()
            .filter_map(|(i, v)| if v.flow_rate > 0 { Some(i) } else { None })
            .collect();
        State {
            opened: 0,
            locations: vec![0; agents],
            goals: vec![Goal::Idle; agents],
            flow_rate: 0,
            time_remaining: initial_time,
            pressure_released: 0,
            interesting,
            history: vec![vec![0]; agents],
        }
    }

    fn find_goals(&self, (_, paths): &ExtraArgs) -> Vec<State> {
        let target_dists: HashMap<usize, usize> = self
            .goals
            .iter()
            .enumerate()
            .filter_map(|(i, g)| {
                if let Goal::Opening(t) = g {
                    Some((*t, paths[self.locations[i]][*t].dist))
                } else {
                    None
                }
            })
            .collect();

        let new_goals: HashMap<usize, Vec<usize>> = self
            .goals
            .iter()
            .positions(|g| matches!(g, Goal::Idle))
            .map(|i| {
                let loc = self.locations[i];
                (
                    i,
                    self.interesting
                        .iter()
                        .filter_map(|target| {
                            let agent_dist = paths[loc][*target].dist;
                            if agent_dist + 2 >= self.time_remaining {
                                // Can't reach + open the valve as well as
                                // letting anything flow through
                                return None;
                            }

                            if let Some(dist) = target_dists.get(target) {
                                if *dist > agent_dist {
                                    // Someone else is targeting this, but they're slower
                                    Some(*target)
                                } else {
                                    // Someone else is targeting this, and they're faster
                                    None
                                }
                            } else {
                                // Noone else is targeting this
                                Some(*target)
                            }
                        })
                        .collect_vec(),
                )
            })
            .collect();

        if new_goals.len() == 1 {
            let (i, targets) = new_goals.iter().next().unwrap();
            if targets.is_empty() {
                let mut state = self.clone();
                state.goals[*i] = Goal::NothingReachable;
                return vec![state];
            }

            targets
                .iter()
                .map(|t| {
                    let mut state = self.clone();
                    state.goals[*i] = Goal::Opening(*t);
                    state
                })
                .collect()
        } else if new_goals.len() == 2 {
            let mut it = new_goals.iter();
            let (i1, targets1) = it.next().unwrap();
            let (i2, targets2) = it.next().unwrap();
            if targets1.len() == 0 {
                if targets2.len() == 0 {
                    let mut state = self.clone();
                    state.goals[*i1] = Goal::NothingReachable;
                    state.goals[*i2] = Goal::NothingReachable;
                    return vec![state];
                }

                return targets2
                    .iter()
                    .map(|t| {
                        let mut state = self.clone();
                        state.goals[*i1] = Goal::NothingReachable;
                        state.goals[*i2] = Goal::Opening(*t);
                        state
                    })
                    .collect();
            } else if targets2.len() == 0 {
                return targets1
                    .iter()
                    .map(|t| {
                        let mut state = self.clone();
                        state.goals[*i1] = Goal::Opening(*t);
                        state.goals[*i2] = Goal::NothingReachable;
                        state
                    })
                    .collect();
            } else if targets1.len() == 1 && targets1 == targets2 {
                // Both agents can only reach one and the same goal
                let mut state = self.clone();
                let target = targets1[0];
                if paths[self.locations[*i1]][target].dist < paths[self.locations[*i2]][target].dist
                {
                    state.goals[*i1] = Goal::Opening(target);
                    state.goals[*i2] = Goal::NothingReachable;
                } else {
                    state.goals[*i1] = Goal::NothingReachable;
                    state.goals[*i2] = Goal::Opening(target);
                }
                return vec![state];
            }

            let mut result = Vec::new();

            for target1 in targets1 {
                for target2 in targets2 {
                    if target1 == target2 {
                        continue;
                    }

                    let mut state = self.clone();
                    state.goals[*i1] = Goal::Opening(*target1);
                    state.goals[*i2] = Goal::Opening(*target2);
                    result.push(state);
                }
            }

            result
        } else {
            unimplemented!("Can't find goals for >2 agents at the same time");
        }
    }
}

type ExtraArgs<'a> = (&'a Vec<Valve>, &'a Vec<Vec<Path>>);
impl util::BnBState<ExtraArgs<'_>> for State {
    fn possible_actions(&self, (valves, paths): &ExtraArgs) -> Vec<State> {
        assert!(!self.finished());

        let goalless = self
            .goals
            .iter()
            .enumerate()
            .filter(|(_, g)| matches!(g, Goal::Idle))
            .collect_vec();
        if goalless.len() > 0 {
            return self.find_goals(&(valves, paths));
        }

        let min_goal_dist = self
            .goals
            .iter()
            .enumerate()
            .filter_map(|(i, g)| {
                if let Goal::Opening(target) = g {
                    Some(paths[self.locations[i]][*target].dist)
                } else {
                    None
                }
            })
            .min()
            .unwrap();

        let time_spent = min_goal_dist + 1;
        let mut completed = self.clone();
        completed.pressure_released += completed.flow_rate * time_spent;
        completed.time_remaining -= time_spent;
        for (i, goal) in self.goals.iter().enumerate() {
            if let Goal::Opening(goal) = goal {
                let path = &paths[self.locations[i]][*goal];
                if path.dist == min_goal_dist {
                    completed.locations[i] = *goal;
                    completed.opened |= 1 << goal;
                    completed.goals[i] = Goal::Idle;
                    if completed.interesting.remove(goal) {
                        // Only increase flow rate if no other agent already
                        // opened this valve
                        completed.flow_rate += valves[*goal].flow_rate;
                    }
                    completed.history[i].push(*goal);
                } else {
                    completed.locations[i] = path.path[time_spent - 1];
                }
            }
        }
        return vec![completed];
    }

    fn finished(&self) -> bool {
        self.time_remaining == 0
            || self
                .goals
                .iter()
                .all(|g| matches!(g, Goal::NothingReachable)) //self.interesting.is_empty()
    }

    fn lower_bound(&self, _: &ExtraArgs) -> usize {
        self.pressure_released + self.flow_rate * self.time_remaining
    }

    fn upper_bound(&self, (valves, paths): &ExtraArgs) -> usize {
        let mut additional_flow = 0;
        for i in &self.interesting {
            let i = i.clone();
            if (self.opened & (1 << i)) != 0 {
                continue;
            }
            let dist = self
                .locations
                .iter()
                .map(|loc| paths[*loc][i].dist)
                .min()
                .unwrap();
            if dist < self.time_remaining {
                let flow = valves[i].flow_rate;
                additional_flow += flow * (self.time_remaining - dist - 1);
            }
        }

        self.lower_bound(&(valves, paths)) + additional_flow
    }
}

impl Hash for State {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.opened.hash(state);
        self.locations.hash(state);
        self.goals.hash(state);
        self.time_remaining.hash(state);
        self.pressure_released.hash(state);
    }
}

impl PartialEq for State {
    fn eq(&self, other: &Self) -> bool {
        self.opened == other.opened
            && self.locations == other.locations
            && self.goals == other.goals
            && self.time_remaining == other.time_remaining
            && self.pressure_released == other.pressure_released
    }
}

impl Eq for State {}