#include <iostream>
#include <stdexcept>
#include <vector>
#include <map>
#include <algorithm>

#define cimg_display 0
#define cimg_use_png
#define cimg_use_zlib
#include <CImg.h>
using namespace cimg_library;

int main(int argc, char *argv[]) {
    if (argc != 4) {
        std::cerr << "Usage: " << argv[0] << " infile outfile rgba" << std::endl;
        return 1;
    }

    const CImg<unsigned char> input(argv[1]);
    std::string mapping = argv[3];

    std::map<char, CImg<unsigned char>> in_channels{
        {'r', input.get_channel(0)},
        {'g', input.get_channel(1)},
        {'b', input.get_channel(2)},
        {'a', input.get_channel(3)},
        {'0', CImg<unsigned char>(input.width(), input.height(), 1, 1, 0)},
        {'1', CImg<unsigned char>(input.width(), input.height(), 1, 1, 1)},
    };

    std::vector<CImg<unsigned char>> channels;
    for (char c : mapping) {
        CImg<unsigned char> subc;
        if (!in_channels.count(c)) {
            std::cerr << "Unknown channel " << c << std::endl;
            return 1;
        }
        channels.push_back(in_channels[c]);
    }

    CImg<unsigned char> output(input.width(), input.height(), 1, channels.size());

    for (int i = 0; i < channels.size(); i++) {
        std::copy(channels[i].begin(), channels[i].end(), &output.atX(0, 0, 0, i));
    }

    output.save(argv[2]);

    return 0;
}