#include <iostream>
#include <string>
#include <vector>
#include <fstream>
#include <type_traits>
#include <opencv2/opencv.hpp>

// Read Class Name List
std::vector<std::string> readClassNameList( const std::string list_path )
{
    std::vector<std::string> classes;
    std::ifstream ifs( list_path );
    if( !ifs.is_open() ){
        return classes;
    }
    std::string class_name = "";
    while( std::getline( ifs, class_name ) ){
        classes.push_back( class_name );
    }
    return classes;
}

// Get Class Color Table for Visualize
std::vector<cv::Scalar> getClassColors( const int32_t number_of_colors )
{
    cv::RNG random;
    std::vector<cv::Scalar> colors;
    for( int32_t i = 0; i < number_of_colors; i++ ){
        cv::Scalar color( random.uniform( 0, 255 ), random.uniform( 0, 255 ), random.uniform( 0, 255 ) );
        colors.push_back( color );
    }
    return colors;
}

int main( int argc, char* argv[] )
{
    // Open Video Capture
    cv::VideoCapture capture = cv::VideoCapture( "bicycle.jpg" );
    if( !capture.isOpened() ){
        return -1;
    }

    // Create Model
    const std::string weights = "optimized_graph.pb";
    cv::dnn::SegmentationModel model = cv::dnn::SegmentationModel( weights );

    // Normalize Parameters
    constexpr double scale = 0.007843;
    const cv::Size size = cv::Size( 513, 513 );
    const cv::Scalar mean = cv::Scalar( 127.5, 127.5, 127.5 );
    constexpr bool swap = true;
    constexpr bool crop = false;
    model.setInputParams( scale, size, mean, swap, crop );

    // Read Classes Name and Colors
    const std::string list = "voc.names";
    const std::vector<std::string> classes = readClassNameList( list );
    const std::vector<cv::Scalar> colors = getClassColors( classes.size() );

    while( true ){
        // Capture Frame
        cv::Mat image;
        capture >> image;
        if( image.empty() ){
            cv::waitKey( 0 );
            break;
        }

        // Segmentation
        cv::Mat mask;
        model.segment( image, mask );

        // Generate Color Mask
        const int32_t rows = mask.rows;
        const int32_t cols = mask.cols;
        cv::Mat color_mask = cv::Mat::zeros( rows, cols, CV_8UC3 );
        for( int32_t y = 0; y < rows; y++ ){
            for( int32_t x = 0; x < cols; x++ ){
                const uint8_t class_id = mask.at<uint8_t>( y, x );
                if( class_id == 0 ) continue;
                const cv::Scalar color = colors[class_id];
                color_mask.at<cv::Vec3b>( y, x ) = cv::Vec3b( color[0], color[1], color[2] );
            }
        }
        cv::resize( color_mask, color_mask, image.size() );

        // Alpha Blend
        constexpr double alpha = 0.5;
        constexpr double beta = 1.0 - alpha;
        cv::addWeighted( image, alpha, color_mask, beta, 0.0, image );

        // Show Image
        cv::imshow( "segmentation", image );
        const int32_t key = cv::waitKey( 10 );
        if( key == 'q' ){
            break;
        }
    }

    cv::destroyAllWindows();

    return 0;
}