<?php

declare(strict_types=1);

namespace Classifier\Lib\Model;

use Classifier\Lib\Enum\MessageType;
use Rubix\ML\Pipeline;
use Rubix\ML\PersistentModel;
use Rubix\ML\Persisters\Filesystem;
use Rubix\ML\Datasets\Labeled;
use Rubix\ML\Datasets\Unlabeled;
use Rubix\ML\Classifiers\KNearestNeighbors;
use Rubix\ML\Kernels\Distance\Cosine;
use Rubix\ML\Transformers\TextNormalizer;
use Rubix\ML\Transformers\WordCountVectorizer;
use Rubix\ML\Transformers\TfIdfTransformer;

/**
 * Gère la classification des messages via un modèle de machine learning.
 * Utilise Rubix ML pour l'apprentissage et la prédiction.
 */
class MessageQualifier
{

    private const MODEL_CONFIDENCE_THRESHOLD = 0.60;

    /** @var PersistentModel */
    private $estimator;

    /** @var string */
    private $modelPath;

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

        if (file_exists($this->modelPath)) {
            $this->estimator = PersistentModel::load($persister);
        } else {
            $this->estimator = new PersistentModel(
                new Pipeline(
                    [
                        // * Normalisation de texte pour réduire la variance (ex: "Error" vs "error")
                        new TextNormalizer(),
                        // * Unigrammes + Bigrammes pour capturer les mots composés
                        new WordCountVectorizer(10000, 1, 2),
                        // * Pondération TF-IDF pour mieux différencier les mots fréquents mais peu informatifs
                        new TfIdfTransformer(),
                    ],
                    // * Classifieur KNN simple (peut être remplacé par un SVM ou un Random Forest pour de meilleures performances)
                    new KNearestNeighbors(5, false, new Cosine())
                ),
                $persister
            );
        }
    }

    /**
     * Entraîne le modèle et sauvegarde le fichier rbx
     */
    public function train(array $samples, array $labels): void
    {
        if (file_exists($this->modelPath)) {
            unlink($this->modelPath);
        }

        $dataset = new Labeled($samples, $labels);
        $this->estimator->train($dataset);
        $this->estimator->save();
    }

    /**
     * Prédiction simplifiée (Retourne juste le label)
     */
    public function predict(string $message): MessageType
    {
        $result = $this->predictWithConfidence($message);

        return $result->decision;
    }

    public function predictWithConfidence(string $message): Prediction
    {
        // NOTE: Rubix attend un tableau de tableaux pour les échantillons
        $dataset = new Unlabeled([[$message]]);

        // * Récupération des probabilités (ex: ['ALERTE' => 0.8, 'INFO' => 0.2])
        // var_dump($this->estimator->proba($dataset));
        // die;
        $probabilities = $this->estimator->proba($dataset)[0];

        // * Trie pour obtenir le meilleur score en premier
        arsort($probabilities);
        $bestLabel = (string) key($probabilities);
        $confidence = (float) current($probabilities);

        // var_dump(
        //     $bestLabel,
        //     $confidence,
        //     $probabilities
        // );
        // die;

        $label      = $this->fromLabelToMessageType($bestLabel);
        $finalLabel = $this->applyDecisionThreshold($label, $confidence);

        return new Prediction(
            $finalLabel,
            $label,
            'machine_learning',
            $confidence,
            $probabilities
        );
    }

    /**
     * Logique métier pour décider si le score est suffisant
     */
    private function applyDecisionThreshold(MessageType $label, float $confidence): MessageType
    {
        if ($confidence < self::MODEL_CONFIDENCE_THRESHOLD) {
            return MessageType::DOUBT();
        }

        return $label;
    }

    /**
     * Vérifie si le modèle a déjà été entraîné
     */
    public function isTrained(): bool
    {
        return file_exists($this->modelPath);
    }

    private function fromLabelToMessageType(string $label): MessageType
    {
        switch ($label) {
            case 'ALERTE':
                return MessageType::ALERTE();
            case 'INFO':
                return MessageType::INFO();
            default:
                return MessageType::DOUBT();
        }
    }
}