use std::fmt;
use std::io;
use std::io::prelude::*;
use std::process;
use structopt::StructOpt;

#[derive(Debug, StructOpt)]
#[structopt(name = "Day 8: Space Image Format")]
/// Executes an Intcode program on 5 amplifiers, and finds the input that gives the max output
///
/// See https://adventofcode.com/2019/day/8 for details.
struct Opt {
    /// Rather than rendering the image, calculate and print its checksum
    #[structopt(short = "c", long = "checksum")]
    checksum_mode: bool,
    #[structopt(short = "w", long = "width")]
    width: u32,
    #[structopt(short = "h", long = "height")]
    height: u32,
}

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

    let image: Image = {
        let mut buffer = String::new();
        exit_on_failed_assertion(
            io::stdin().read_to_string(&mut buffer),
            "Error reading input",
        );

        Image::from_str(&buffer.trim(), opt.width, opt.height)
    };

    if opt.checksum_mode {
        println!("{}", image.checksum());
    } else {
        println!("{}", image);
    }
}

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);
        }
    }
}

#[derive(Debug)]
struct Image {
    width: u32,
    height: u32,
    layers: Vec<ImageLayer>,
}

impl fmt::Display for Image {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.flatten()
            .pixels
            .chunks(self.width as usize)
            .map(|line| {
                line.iter()
                    .map(|c| write!(f, "{}", if *c == 0 { ' ' } else { '1' }))
                    .collect::<fmt::Result>()
                    .and_then(|_| writeln!(f))
            })
            .collect()
    }
}

impl Image {
    fn from_str(s: &str, width: u32, height: u32) -> Image {
        Image {
            width,
            height,
            layers: s
                .as_bytes()
                .chunks((width * height) as usize)
                .map(|chunk| ImageLayer::new(chunk))
                .collect(),
        }
    }

    fn checksum(&self) -> usize {
        self.layers
            .iter()
            .min_by_key(|layer| layer.pixel_count(0))
            .map(|layer| layer.pixel_count(1) * layer.pixel_count(2))
            .unwrap_or(0)
    }

    fn flatten(&self) -> ImageLayer {
        self.layers
            .iter()
            .fold(ImageLayer::empty(self.width, self.height), |acc, next| {
                ImageLayer {
                    pixels: acc
                        .pixels
                        .iter()
                        .zip(next.pixels.iter())
                        .map(|(a, b)| if *a == 2 { *b } else { *a })
                        .collect(),
                }
            })
    }
}

#[derive(Debug)]
struct ImageLayer {
    pixels: Vec<u8>,
}

impl ImageLayer {
    fn new(char_pixels: &[u8]) -> ImageLayer {
        ImageLayer {
            pixels: char_pixels
                .iter()
                .map(|c| c.overflowing_sub(b'0').0)
                .collect(),
        }
    }

    fn empty(width: u32, height: u32) -> ImageLayer {
        ImageLayer {
            pixels: vec![2; (width * height) as usize],
        }
    }

    fn pixel_count(&self, value: u8) -> usize {
        self.pixels.iter().filter(|p| **p == value).count()
    }
}