use aoc2019::*;
use rpds::list;
use rpds::list::List;
use rpds::vector::Vector;
use std::io;
use std::io::prelude::*;
use std::iter;
use std::process;
use structopt::StructOpt;

#[derive(Debug, StructOpt)]
#[structopt(name = "Day 7: Amplification Circuit")]
/// Executes an Intcode program on 5 amplifiers, and finds the input that gives the max output
///
/// See https://adventofcode.com/2019/day/7 for details.
struct Opt {
    #[structopt(short = "f", long = "feedback-loop")]
    feedback_loop_mode: bool,
}

fn main() {
    let stdin = io::stdin();
    let opt = Opt::from_args();

    let program: IntcodeProgram = stdin
        .lock()
        .split(b',')
        .map(|x| exit_on_failed_assertion(x, "Error reading input"))
        .map(|x| exit_on_failed_assertion(String::from_utf8(x), "Input was not valid UTF-8"))
        .map(|x| exit_on_failed_assertion(x.trim().parse::<Intcode>(), "Invalid number"))
        .collect::<IntcodeProgram>();

    let result = exit_on_failed_assertion(
        find_max_power(&program, opt.feedback_loop_mode),
        "Program errored",
    );
    println!("{}", result);
}

fn exit_on_failed_assertion<A, E: std::error::Error>(data: Result<A, E>, message: &str) -> A {
    match data {
        Ok(data) => data,
        Err(e) => {
            eprintln!("{}: {}", message, e);
            process::exit(1);
        }
    }
}

fn find_max_power(
    program: &IntcodeProgram,
    feedback_loop_mode: bool,
) -> Result<Intcode, IntcodeProgramError> {
    PhaseSetting::all(feedback_loop_mode)
        .map(|phase| AmplifierArray::new(program, &phase).execute())
        .collect::<Result<Vec<Intcode>, _>>()
        .map(|powers| powers.into_iter().max().unwrap_or(Intcode::from(0)))
}

#[derive(Debug, Clone)]
struct PhaseSetting([Intcode; 5]);

impl PhaseSetting {
    fn all(feedback_loop_mode: bool) -> impl Iterator<Item = PhaseSetting> {
        if feedback_loop_mode {
            PhaseSetting::permute(5, 10)
        } else {
            PhaseSetting::permute(0, 5)
        }
    }

    fn permute(min: i32, max: i32) -> impl Iterator<Item = PhaseSetting> {
        // This is an absolutely atrocious way to do the permutation,
        // but luckily it's only 5 elements long.
        (min..max)
            .flat_map(move |a| {
                (min..max).flat_map(move |b| {
                    (min..max).flat_map(move |c| {
                        (min..max).flat_map(move |d| {
                            (min..max).map(move |e| {
                                PhaseSetting([a.into(), b.into(), c.into(), d.into(), e.into()])
                            })
                        })
                    })
                })
            })
            .filter(move |phase| (min..max).all(|x| phase.0.contains(&x.into())))
    }
}

#[derive(Debug, Clone)]
struct AmplifierArray {
    amplifiers: Vector<IntcodeProgram>,
}

impl AmplifierArray {
    fn new(program: &IntcodeProgram, phase: &PhaseSetting) -> AmplifierArray {
        AmplifierArray {
            amplifiers: (0..5)
                .map(|n| AmplifierArray::new_amp(program, phase, n))
                .collect(),
        }
    }

    fn new_amp(program: &IntcodeProgram, phase: &PhaseSetting, n: usize) -> IntcodeProgram {
        if n == 0 {
            program.with_input(list![phase.0[n].clone(), Intcode::from(0)])
        } else {
            program.with_input(list![phase.0[n].clone()])
        }
    }

    fn execute(&self) -> Result<Intcode, IntcodeProgramError> {
        self.run_to_termination().output_into_result()
    }

    fn run_to_termination(&self) -> AmplifierArray {
        iter::successors(Some(self.clone()), |p| Some(p.next()))
            .find(|p| p.is_terminated())
            .unwrap() // successors doesn't terminate, so this will never be none.
    }

    fn output_into_result(&self) -> Result<Intcode, IntcodeProgramError> {
        self.amplifiers
            .first()
            .and_then(|amp| amp.input.first().cloned())
            .ok_or(IntcodeProgramError::Unknown)
    }

    fn is_terminated(&self) -> bool {
        self.amplifiers.last().map(|amp| amp.halted).unwrap_or(true)
    }

    fn next(&self) -> AmplifierArray {
        self.run_amplifiers().update_inputs()
    }

    fn run_amplifiers(&self) -> AmplifierArray {
        AmplifierArray {
            amplifiers: self
                .amplifiers
                .iter()
                .map(|amp| amp.run_to_termination_or_input())
                .collect(),
        }
    }

    fn update_inputs(&self) -> AmplifierArray {
        AmplifierArray {
            amplifiers: self
                .amplifiers
                .iter()
                .fold(
                    (
                        Vector::new(),
                        self.amplifiers
                            .last()
                            .map(|a| a.output.iter().cloned().collect::<List<Intcode>>())
                            .unwrap_or(List::new()),
                    ),
                    |(amps, piped_input), next_amp| {
                        (
                            amps.push_back(
                                next_amp
                                    .with_additional_input(piped_input)
                                    .with_cleared_output(),
                            ),
                            next_amp.output.iter().cloned().collect::<List<Intcode>>(),
                        )
                    },
                )
                .0,
        }
    }
}