#!/usr/bin/env perl

use strict;
use warnings;
use FindBin qw($Bin);
use lib "$Bin/../lib";
use Getopt::Long qw(GetOptions);
use File::Basename qw(basename);
use IO::Handle ();
use JSON::PP ();
use Time::HiRes qw(time);
use LaptopNN::Network;
use LaptopNN::TextVectorizer qw(read_labeled_text_csv);
use LaptopNN::Data qw(
    build_label_encoder encode_labels stratified_split_indices
    take_rows class_weights csv_quote
);

binmode STDOUT, ':encoding(UTF-8)';
binmode STDERR, ':encoding(UTF-8)';

my %opt = (
    dimensions => 256, ngram_max => 2, profile => 'perl', hidden => '48',
    activation => 'relu', dropout => 0.10, epochs => 160, batch_size => 24,
    learning_rate => 0.003, validation_split => 0.20, patience => 25,
    l2 => 0.0002, clip_norm => 5, label_smoothing => 0.02,
    class_weight => 'balanced', seed => 42, workers => 1, quiet => 0,
    events_json => 0,
);
GetOptions(
    'data=s' => \$opt{data}, 'model=s' => \$opt{model}, 'history=s' => \$opt{history},
    'dimensions=i' => \$opt{dimensions}, 'ngram-max=i' => \$opt{ngram_max},
    'profile=s' => \$opt{profile}, 'hidden=s' => \$opt{hidden},
    'activation=s' => \$opt{activation}, 'dropout=f' => \$opt{dropout},
    'epochs=i' => \$opt{epochs}, 'batch-size=i' => \$opt{batch_size},
    'learning-rate=f' => \$opt{learning_rate},
    'validation-split=f' => \$opt{validation_split}, 'patience=i' => \$opt{patience},
    'l2=f' => \$opt{l2}, 'clip-norm=f' => \$opt{clip_norm},
    'label-smoothing=f' => \$opt{label_smoothing},
    'class-weight=s' => \$opt{class_weight}, 'seed=i' => \$opt{seed},
    'workers=i' => \$opt{workers}, 'quiet!' => \$opt{quiet},
    'events-json!' => \$opt{events_json}, 'cancel-file=s' => \$opt{cancel_file},
    'help|h' => \$opt{help},
) or _usage(2);
_usage(0) if $opt{help};
_usage(2, '--data is required') unless defined $opt{data};
_usage(2, '--model is required') unless defined $opt{model};

my @hidden = lc($opt{hidden}) eq 'none' || !length($opt{hidden})
    ? () : split(/,/, $opt{hidden});
die "--hidden must be comma-separated positive integers or none\n"
    if grep { $_ !~ /^\d+$/ || $_ < 1 } @hidden;
die "--workers must be between 1 and 16\n"
    unless $opt{workers} >= 1 && $opt{workers} <= 16;
STDERR->autoflush(1);

srand($opt{seed});
my $dataset = read_labeled_text_csv(path => $opt{data});
my ($labels, $label_index) = build_label_encoder($dataset->{labels});
my $encoded = encode_labels($dataset->{labels}, $label_index);
my ($train_indices, $validation_indices) = stratified_split_indices(
    $encoded, $opt{validation_split}
);
my $train_text = take_rows($dataset->{texts}, $train_indices);
my $train_y = take_rows($encoded, $train_indices);
my $validation_text = take_rows($dataset->{texts}, $validation_indices);
my $validation_y = take_rows($encoded, $validation_indices);
my $vectorizer = LaptopNN::TextVectorizer->new(
    dimensions => $opt{dimensions}, ngram_max => $opt{ngram_max}, profile => $opt{profile},
)->fit($train_text);
my $train_x = $vectorizer->transform_many($train_text);
my $validation_x = $vectorizer->transform_many($validation_text);
my $weights = class_weights($train_y, scalar(@$labels), $opt{class_weight});
my @sizes = ($opt{dimensions}, @hidden, scalar(@$labels));
my $network = LaptopNN::Network->new(
    sizes => \@sizes, activation => $opt{activation}, dropout => $opt{dropout}, seed => $opt{seed},
);
my $started = time;
my $summary = $network->train(
    $train_x, $train_y,
    validation_x => $validation_x, validation_y => $validation_y,
    epochs => $opt{epochs}, batch_size => $opt{batch_size},
    learning_rate => $opt{learning_rate}, patience => $opt{patience},
    lr_patience => 6, lr_factor => 0.5, min_delta => 0.0001,
    monitor => @$validation_x ? 'val_loss' : 'train_loss',
    l2 => $opt{l2}, clip_norm => $opt{clip_norm},
    label_smoothing => $opt{label_smoothing}, class_weights => $weights,
    workers => $opt{workers}, worker_program => "$Bin/plnn-gradient-worker",
    parallel_seed => $opt{seed},
    cancellation_requested => sub {
        return defined($opt{cancel_file}) && -e $opt{cancel_file};
    },
    callback => sub {
        my ($row) = @_;
        if ($opt{events_json}) {
            my %event = (
                type => 'epoch', epoch => 0 + $row->{epoch},
                learning_rate => 0 + $row->{learning_rate},
                train_loss => 0 + $row->{train_loss},
                train_accuracy => 0 + $row->{train_accuracy},
            );
            $event{val_loss} = 0 + $row->{val_loss} if exists $row->{val_loss};
            $event{val_accuracy} = 0 + $row->{val_accuracy} if exists $row->{val_accuracy};
            print STDERR 'PLNN_EVENT ', JSON::PP->new->canonical(1)->encode(\%event), "\n";
        } elsif (!$opt{quiet}) {
            printf STDERR "epoch %d train %.4f validation %s\n", $row->{epoch},
                $row->{train_accuracy}, exists($row->{val_accuracy})
                    ? sprintf('%.4f', $row->{val_accuracy}) : '-';
        }
    },
);
$summary->{wall_seconds} = 0 + (time - $started);
my $dimensions = $vectorizer->dimensions;
my $model = {
    format => 'perl-laptop-nn', format_version => 1,
    software_version => $LaptopNN::Network::VERSION,
    model_type => 'text_intent_classifier', network => $network->to_hash,
    labels => [ @$labels ], text_vectorizer => $vectorizer->to_hash,
    preprocessing => {
        feature_names => [ map { sprintf('text_hash_%04d', $_) } 0 .. $dimensions - 1 ],
        missing_policy => 'error', impute_values => [ (0) x $dimensions ],
        centers => [ (0) x $dimensions ], scales => [ (1) x $dimensions ],
        standardize => JSON::PP::false,
    },
    training => {
        source_name => basename($opt{data}), total_rows => 0 + $dataset->{rows},
        training_rows => 0 + @$train_x, validation_rows => 0 + @$validation_x,
        options => {
            dimensions => 0 + $dimensions, ngram_max => 0 + $opt{ngram_max},
            profile => $opt{profile}, hidden => [ map { 0 + $_ } @hidden ],
            activation => $opt{activation}, dropout => 0 + $opt{dropout},
            epochs_requested => 0 + $opt{epochs}, batch_size => 0 + $opt{batch_size},
            learning_rate => 0 + $opt{learning_rate},
            validation_split => 0 + $opt{validation_split},
            patience => 0 + $opt{patience}, seed => 0 + $opt{seed},
            workers => 0 + $opt{workers}, l2 => 0 + $opt{l2},
            clip_norm => 0 + $opt{clip_norm}, label_smoothing => 0 + $opt{label_smoothing},
            class_weight => $opt{class_weight},
        },
        summary => $summary,
    },
};
LaptopNN::Network->save_model($opt{model}, $model);
_write_history($opt{history}, $summary->{history}) if defined $opt{history};
print "$opt{model}\n" unless $opt{quiet};

sub _write_history {
    my ($path, $history) = @_;
    # The history format is ASCII and package evidence is canonical LF on every OS.
    open my $fh, '>:raw', $path or die "cannot write history '$path': $!\n";
    print {$fh} "epoch,learning_rate,train_loss,train_accuracy,val_loss,val_accuracy\n";
    for my $row (@$history) {
        print {$fh} join(',', map { csv_quote($_) } (
            $row->{epoch}, $row->{learning_rate}, $row->{train_loss},
            $row->{train_accuracy}, $row->{val_loss}, $row->{val_accuracy}
        )), "\n";
    }
    close $fh or die "cannot close history '$path': $!\n";
}

sub _usage {
    my ($status, $message) = @_;
    print STDERR "$message\n\n" if defined $message;
    print STDERR <<'USAGE';
Usage: plnn-text-train --data training.csv --model model.json [options]
Pure-Perl backpropagation trainer for a reviewed text-intent candidate.

Workflow integration options:
  --events-json       emit one PLNN_EVENT JSON line per completed epoch
  --cancel-file PATH  stop cooperatively when PATH appears
USAGE
    exit $status;
}
