use strict;
use warnings;
use utf8;
use constant IN_IPERL => !! $ENV{PERL_IPERL_RUNNING};
no if IN_IPERL, warnings => 'redefine'; # fewer messages when re-running cells

use feature qw(say state);
use Syntax::Construct qw(each-array);

use lib::projectroot qw(lib);

BEGIN {
    if( IN_IPERL ) {
        $ENV{TF_CPP_MIN_LOG_LEVEL} = 3;
    }
    require AI::TensorFlow::Libtensorflow;
}

use URI ();
use HTTP::Tiny ();
use Path::Tiny qw(path);

use File::Which ();

use List::Util ();

use Data::Printer ( output => 'stderr', return_value => 'void', filters => ['PDL'] );
use Data::Printer::Filter::PDL ();
use Text::Table::Tiny qw(generate_table);

use Imager;

my $s = AI::TensorFlow::Libtensorflow::Status->New;
sub AssertOK {
    die "Status $_[0]: " . $_[0]->Message
        unless $_[0]->GetCode == AI::TensorFlow::Libtensorflow::Status::OK;
    return;
}
AssertOK($s);

use PDL;
use AI::TensorFlow::Libtensorflow::DataType qw(FLOAT);

use FFI::Platypus::Memory qw(memcpy);
use FFI::Platypus::Buffer qw(scalar_to_pointer);

sub FloatPDLTOTFTensor {
    my ($p) = @_;
    return AI::TensorFlow::Libtensorflow::Tensor->New(
        FLOAT, [ reverse $p->dims ], $p->get_dataref, sub { undef $p }
    );
}

sub FloatTFTensorToPDL {
    my ($t) = @_;

    my $pdl = zeros(float,reverse( map $t->Dim($_), 0..$t->NumDims-1 ) );

    memcpy scalar_to_pointer( ${$pdl->get_dataref} ),
        scalar_to_pointer( ${$t->Data} ),
        $t->ByteSize;
    $pdl->upd_data;

    $pdl;
}

use HTML::Tiny;

sub my_table {
    my ($data, $cb) = @_;
    my $h = HTML::Tiny->new;
    $h->table( { style => 'width: 100%' },
        [
            $h->tr(
                map {
                    [
                        $h->td( $cb->($_, $h) )
                    ]
                } @$data
            )
        ]
    )
}

sub show_in_gnuplot {
    my ($p) = @_;
    require PDL::Graphics::Gnuplot;
    PDL::Graphics::Gnuplot::image( square => 1, $p );
}

# image_size => [width, height] (but usually square images)
my %model_name_to_params = (
    mobilenet_v2_100_224 => {
        handle     => 'https://tfhub.dev/google/imagenet/mobilenet_v2_100_224/classification/5',
        image_size => [ 224, 224 ],
    },
    mobilenet_v2_140_224 => {
        handle => "https://tfhub.dev/google/imagenet/mobilenet_v2_140_224/classification/5",
        image_size => [ 224, 224 ],
    },
);

my $model_name = 'mobilenet_v2_100_224';

say "Selected model: $model_name : $model_name_to_params{$model_name}{handle}";

my $model_uri = URI->new( $model_name_to_params{$model_name}{handle} );
$model_uri->query_form( 'tf-hub-format' => 'compressed' );
my $model_base = substr( $model_uri->path, 1 ) =~ s,/,_,gr;
my $model_archive_path = "${model_base}.tar.gz";

use constant IMAGENET_LABEL_COUNT_WITH_BG => 1001;
my $labels_uri = URI->new('https://storage.googleapis.com/download.tensorflow.org/data/ImageNetLabels.txt');
my $labels_path = ($labels_uri->path_segments)[-1];

my $http = HTTP::Tiny->new;

for my $download ( [ $model_uri  => $model_archive_path ],
                   [ $labels_uri => $labels_path        ]) {
    my ($uri, $path) = @$download;
    say "Downloading $uri to $path";
    next if -e $path;
    $http->mirror( $uri, $path );
}

use Archive::Extract;
my $ae = Archive::Extract->new( archive => $model_archive_path );
die "Could not extract archive" unless $ae->extract( to => $model_base );

my $saved_model = path($model_base)->child('saved_model.pb');
say "Saved model is in $saved_model" if -f $saved_model;

my @labels = path($labels_path)->lines( { chomp => 1 });
die "Labels should have @{[ IMAGENET_LABEL_COUNT_WITH_BG ]} items"
    unless @labels == IMAGENET_LABEL_COUNT_WITH_BG;
say "Got labels: ", join( ", ", List::Util::head(5, @labels) ), ", etc.";

my @tags = ( 'serve' );

if( File::Which::which('saved_model_cli')) {
    local $ENV{TF_CPP_MIN_LOG_LEVEL} = 3; # quiet the TensorFlow logger for the following command
    system(qw(saved_model_cli show),
        qw(--dir)           => $model_base,
        qw(--tag_set)       => join(',', @tags),
        qw(--signature_def) => 'serving_default'
    ) == 0 or die "Could not run saved_model_cli";
} else {
    say "Install the tensorflow Python package to get the `saved_model_cli` command.";
}

my $opt = AI::TensorFlow::Libtensorflow::SessionOptions->New;

my $graph = AI::TensorFlow::Libtensorflow::Graph->New;
my $session = AI::TensorFlow::Libtensorflow::Session->LoadFromSavedModel(
    $opt, undef, $model_base, \@tags, $graph, undef, $s
);
AssertOK($s);

my %ops = (
    in  => $graph->OperationByName('serving_default_inputs'),
    out => $graph->OperationByName('StatefulPartitionedCall'),
);

die "Could not get all operations" unless List::Util::all(sub { defined }, values %ops);

my %outputs = map { $_ => [ AI::TensorFlow::Libtensorflow::Output->New( { oper => $ops{$_}, index => 0 } ) ] }
    keys %ops;

p %outputs;

say "Input: " , $outputs{in}[0];
say "Output: ", $outputs{out}[0];

my %images_for_test_to_uri = (
    "tiger" => "https://upload.wikimedia.org/wikipedia/commons/b/b0/Bengal_tiger_%28Panthera_tigris_tigris%29_female_3_crop.jpg",
    #by Charles James Sharp, CC BY-SA 4.0 <https://creativecommons.org/licenses/by-sa/4.0>, via Wikimedia Commons
    "bus" => "https://upload.wikimedia.org/wikipedia/commons/6/63/LT_471_%28LTZ_1471%29_Arriva_London_New_Routemaster_%2819522859218%29.jpg",
    #by Martin49 from London, England, CC BY 2.0 <https://creativecommons.org/licenses/by/2.0>, via Wikimedia Commons
    "car" => "https://upload.wikimedia.org/wikipedia/commons/4/49/2013-2016_Toyota_Corolla_%28ZRE172R%29_SX_sedan_%282018-09-17%29_01.jpg",
    #by EurovisionNim, CC BY-SA 4.0 <https://creativecommons.org/licenses/by-sa/4.0>, via Wikimedia Commons
    "cat" => "https://upload.wikimedia.org/wikipedia/commons/4/4d/Cat_November_2010-1a.jpg",
    #by Alvesgaspar, CC BY-SA 3.0 <https://creativecommons.org/licenses/by-sa/3.0>, via Wikimedia Commons
    "dog" => "https://upload.wikimedia.org/wikipedia/commons/archive/a/a9/20090914031557%21Saluki_dog_breed.jpg",
    #by Craig Pemberton, CC BY-SA 3.0 <https://creativecommons.org/licenses/by-sa/3.0>, via Wikimedia Commons
    "apple" => "https://upload.wikimedia.org/wikipedia/commons/1/15/Red_Apple.jpg",
    #by Abhijit Tembhekar from Mumbai, India, CC BY 2.0 <https://creativecommons.org/licenses/by/2.0>, via Wikimedia Commons
    "banana" => "https://upload.wikimedia.org/wikipedia/commons/1/1c/Bananas_white_background.jpg",
    #by fir0002  flagstaffotos [at] gmail.com		Canon 20D + Tamron 28-75mm f/2.8, GFDL 1.2 <http://www.gnu.org/licenses/old-licenses/fdl-1.2.html>, via Wikimedia Commons
    "turtle" => "https://upload.wikimedia.org/wikipedia/commons/8/80/Turtle_golfina_escobilla_oaxaca_mexico_claudio_giovenzana_2010.jpg",
    #by Claudio Giovenzana, CC BY-SA 3.0 <https://creativecommons.org/licenses/by-sa/3.0>, via Wikimedia Commons
    "flamingo" => "https://upload.wikimedia.org/wikipedia/commons/b/b8/James_Flamingos_MC.jpg",
    #by Christian Mehlführer, User:Chmehl, CC BY 3.0 <https://creativecommons.org/licenses/by/3.0>, via Wikimedia Commons
    "piano" => "https://upload.wikimedia.org/wikipedia/commons/d/da/Steinway_%26_Sons_upright_piano%2C_model_K-132%2C_manufactured_at_Steinway%27s_factory_in_Hamburg%2C_Germany.png",
    #by "Photo: © Copyright Steinway & Sons", CC BY-SA 3.0 <https://creativecommons.org/licenses/by-sa/3.0>, via Wikimedia Commons
    "honeycomb" => "https://upload.wikimedia.org/wikipedia/commons/f/f7/Honey_comb.jpg",
    #by Merdal, CC BY-SA 3.0 <http://creativecommons.org/licenses/by-sa/3.0/>, via Wikimedia Commons
    "teapot" => "https://upload.wikimedia.org/wikipedia/commons/4/44/Black_tea_pot_cropped.jpg",
    #by Mendhak, CC BY-SA 2.0 <https://creativecommons.org/licenses/by-sa/2.0>, via Wikimedia Commons
);

my @image_names = sort keys %images_for_test_to_uri;


if( IN_IPERL ) {
    IPerl->html(
        my_table( \@image_names, sub {
            my ($image_name, $h) = @_;
            (
                $h->tt($image_name),
                $h->a( { href => $images_for_test_to_uri{$image_name} },
                    $h->img({
                        src => $images_for_test_to_uri{$image_name},
                        alt => $image_name,
                        width => '50%',
                    })
                ),
            )
        })
    );
}

sub imager_paste_center_pad {
    my ($inner, $padded_sz, @rest) = @_;

    my $outer = Imager->new( List::Util::mesh( [qw(xsize ysize)], $padded_sz ),
        @rest
    );

    $outer->paste(
        left => int( ($outer->getwidth  - $inner->getwidth ) / 2 ),
        top  => int( ($outer->getheight - $inner->getheight) / 2 ),
        src  => $inner,
    );

    $outer;
}

sub imager_scale_to {
    my ($img, $image_size) = @_;
    my $rescaled = $img->scale(
        List::Util::mesh( [qw(xpixels ypixels)], $image_size ),
        type => 'min',
        qtype => 'mixing', # 'mixing' seems to work better than 'normal'
    );
}

sub load_image_to_pdl {
    my ($uri, $image_size) = @_;

    my $http = HTTP::Tiny->new;
    my $response = $http->get( $uri );
    die "Could not fetch image from $uri" unless $response->{success};
    say "Downloaded $uri";

    my $img = Imager->new;
    $img->read( data => $response->{content} );

    my $rescaled = imager_scale_to($img, $image_size);

    say sprintf "Rescaled image from [ %d x %d ] to [ %d x %d ]",
        $img->getwidth, $img->getheight,
        $rescaled->getwidth, $rescaled->getheight;

    my $padded = imager_paste_center_pad($rescaled, $image_size,
        # ARGB fits in 32-bits (uint32_t)
        channels => 4
    );

    say sprintf "Padded to [ %d x %d ]", $padded->getwidth, $padded->getheight;

    # Create PDL ndarray from Imager data in-memory.
    my $data;
    $padded->write( data => \$data, type => 'raw' )
        or die "could not write ". $padded->errstr;

    # $data is packed as PDL->dims == [w,h] with ARGB pixels
    #   $ PDL::howbig(ulong) # 4
    my $pdl_raw = zeros(ulong, $padded->getwidth, $padded->getheight);
    ${ $pdl_raw->get_dataref } = $data;
    $pdl_raw->upd_data;

    # Split uint32_t pixels into first dimension with 3 channels (R,G,B) with values 0-255.
    my @shifts = map 8*$_, 0..2;
    my $pdl_channels = $pdl_raw->dummy(0)
        ->and2(ulong(map 0xFF << $_, @shifts)->slice(':,*,*') )
        ->shiftright( ulong(@shifts)->slice(':,*,*') )
        ->byte;

    my $pdl_scaled = (
            # Scale to [ 0, 1 ].
            ( $pdl_channels / float(255) )
        );

    ## flip vertically to see image right way up
    #show_in_gnuplot( $pdl_channels->slice(':,:,-1:0')         ); #DEBUG
    #show_in_gnuplot(   $pdl_scaled->slice(':,:,-1:0') * 255.0 ); #DEBUG

    $pdl_scaled;
}

my @pdl_images = map {
    load_image_to_pdl(
        $images_for_test_to_uri{$_},
        $model_name_to_params{$model_name}{image_size}
    );
} @image_names;

my $pdl_image_batched = cat(@pdl_images);
my $t = FloatPDLTOTFTensor($pdl_image_batched);

p $pdl_image_batched;
p $t;

my $RunSession = sub {
    my ($session, $t) = @_;
    my @outputs_t;

    $session->Run(
        undef,
        $outputs{in}, [$t],
        $outputs{out}, \@outputs_t,
        undef,
        undef,
        $s
    );
    AssertOK($s);

    return $outputs_t[0];
};

say "Warming up the model";
use PDL::GSL::RNG;
my $rng = PDL::GSL::RNG->new('default');
my $image_size = $model_name_to_params{$model_name}{image_size};
my $warmup_input = zeros(float, 3, @$image_size, 1 );
$rng->get_uniform($warmup_input);

p $RunSession->($session, FloatPDLTOTFTensor($warmup_input));

my $output_pdl_batched = FloatTFTensorToPDL($RunSession->($session, $t));
my $softmax = sub { ( map $_/sumover($_)->dummy(0), exp($_[0]) )[0] };
my $probabilities_batched = $softmax->($output_pdl_batched);
p $probabilities_batched;

my $N = 5; # number to select

my $top_batched = $probabilities_batched->qsorti->slice([-1, -$N]);

my @top_lists   = dog($top_batched);

my $includes_background_class = $probabilities_batched->dim(0) == IMAGENET_LABEL_COUNT_WITH_BG;

if( IN_IPERL ) {
    my $html = IPerl->html(
        my_table( [0..$#image_names], sub {
            my ($batch_idx, $h) = @_;
            my $image_name = $image_names[$batch_idx];
            my @top_for_image = $top_lists[$batch_idx]->list;
            (
                    $h->tt($image_name),
                    $h->a( { href => $images_for_test_to_uri{$image_name} },
                        $h->img({
                            src => $images_for_test_to_uri{$image_name},
                            alt => $image_name,
                            width => '50%',
                        })
                    ),
                    do {
                        my @tr;
                        push @tr, [ $h->th('Rank', 'Label No', 'Label', 'Prob') ];
                        while( my ($i, $label_index) = each @top_for_image ) {
                            my $class_index = $includes_background_class ? $label_index : $label_index + 1;
                            push @tr, [ $h->td(
                                    $i + 1,
                                    $class_index,
                                    $labels[$class_index],
                                    $probabilities_batched->at($label_index,$batch_idx),
                            ) ];

                        }
                        $h->table([$h->tr(@tr)])
                    },
                )
        })
    );
    IPerl->display($html);
} else {
    for my $batch_idx (0..$#image_names) {
        my $image_name = $image_names[$batch_idx];
        my @top_for_image = $top_lists[$batch_idx]->list;
        my @td;
        say "Image name: `$image_name`";
        my $header = [ ('Rank', 'Label No', 'Label', 'Prob') ];
        my @rows;
        while( my ($i, $label_index) = each @top_for_image ) {
            my $class_index = $includes_background_class ? $label_index : $label_index + 1;
            push @rows, [ (
                    $i + 1,
                    $class_index,
                    $labels[$class_index],
                    $probabilities_batched->at($label_index,$batch_idx),
            ) ];
        }
        say generate_table( rows => [ $header, @rows ], header_row => 1 );
        print "\n";
    }
}

my $p_approx_batched = $probabilities_batched->sumover->approx(1, 1e-5);
p $p_approx_batched;
say "All probabilities sum up to approximately 1" if $p_approx_batched->all->sclr;

use Filesys::DiskUsage qw/du/;

my $total = du( { 'human-readable' => 1, dereference => 1 },
    $model_archive_path, $model_base, $labels_path );

say "Disk space usage: $total"; undef;

my @solid_channel_uris = (
    'https://upload.wikimedia.org/wikipedia/commons/thumb/6/62/Solid_red.svg/480px-Solid_red.svg.png',
    'https://upload.wikimedia.org/wikipedia/commons/thumb/1/1d/Green_00FF00_9x9.svg/480px-Green_00FF00_9x9.svg.png',
    'https://upload.wikimedia.org/wikipedia/commons/thumb/f/ff/Solid_blue.svg/480px-Solid_blue.svg.png',
);
undef;