<?php

declare(strict_types=1);

namespace Classifier\Lib;

use Rubix\ML\PersistentModel;
use Rubix\ML\Pipeline;
use Rubix\ML\Transformers\TextNormalizer;
use Rubix\ML\Transformers\WordCountVectorizer;
use Rubix\ML\Transformers\TfIdfTransformer;
use Rubix\ML\Classifiers\NaiveBayes;
use Rubix\ML\Persisters\Filesystem;
use Rubix\ML\Datasets\Unlabeled;

class MessageQualifierShort
{
    private $estimator;
    private $modelPath;

    public function __construct(string $modelPath = __DIR__ . '/model.rbx')
    {
        $this->modelPath = $modelPath;

        if (file_exists($this->modelPath)) {
            $this->estimator = PersistentModel::load(new Filesystem($this->modelPath));
        } else {
            // Pipeline optimisée pour le texte
            $this->estimator = new PersistentModel(
                new Pipeline([
                    new TextNormalizer(), // Supprime ponctuation, minuscules
                    new WordCountVectorizer(5000, 1, 2), // Unigrammes et Bigrammes (ex: "alerte" et "alerte incendie")
                    new TfIdfTransformer(), // Donne du poids aux mots rares et importants
                ], new NaiveBayes()),
                new Filesystem($this->modelPath)
            );
        }
    }

    public function predict(string $message): array
    {
        $dataset = new Unlabeled([[$message]]);
        $probabilities = $this->estimator->proba($dataset)[0];

        arsort($probabilities);
        $label = key($probabilities);
        $confidence = current($probabilities);

        return [
            'label' => $label,
            'confidence' => $confidence,
            'all_scores' => $probabilities
        ];
    }
}