<?php

/**
 * Naive Bayes untuk klasifikasi kelayakan PIP.
 * Fitur: pekerjaan_ortu, penghasilan_ortu (kategorik), jumlah_saudara (numerik), kps, kip (kategorik).
 */
class NaiveBayes
{
    private array $training = [];
    private array $classes = ['Layak', 'Tidak Layak'];
    private array $classCounts = [];
    private int $total = 0;

    /** @var array<string, array<string, int>> */
    private array $catCounts = [];

    /** @var array<string, array{mean: float, var: float}> */
    private array $numStats = [];

    private const NUMERIC = [];
    private const CATEGORIC = ['pekerjaan_ortu', 'penghasilan_ortu', 'jumlah_saudara', 'kepemilikan_kps', 'kepemilikan_kip'];

    public function __construct(array $trainingRows)
    {
        $this->training = $trainingRows;
        $this->train();
    }

    private function train(): void
    {
        $this->total = count($this->training);
        if ($this->total === 0) {
            return;
        }

        foreach ($this->classes as $c) {
            $this->classCounts[$c] = 0;
        }

        foreach ($this->training as $row) {
            $kelas = $row['kelas'];
            $this->classCounts[$kelas] = ($this->classCounts[$kelas] ?? 0) + 1;

            foreach (self::CATEGORIC as $attr) {
                $val = (string) $row[$attr];
                $this->catCounts[$kelas][$attr][$val] =
                    ($this->catCounts[$kelas][$attr][$val] ?? 0) + 1;
            }

            foreach (self::NUMERIC as $attr) {
                $this->numStats[$kelas][$attr][] = (float) $row[$attr];
            }
        }

        foreach ($this->classes as $c) {
            foreach (self::NUMERIC as $attr) {
                $values = $this->numStats[$c][$attr] ?? [0.0];
                $mean = array_sum($values) / count($values);
                $variance = 0.0;
                if (count($values) > 1) {
                    foreach ($values as $v) {
                        $variance += ($v - $mean) ** 2;
                    }
                    $variance /= count($values);
                }
                $this->numStats[$c][$attr] = [
                    'mean' => $mean,
                    'var' => max($variance, 1e-6),
                ];
            }
        }
    }

    private function gaussian(float $x, float $mean, float $var): float
    {
        $exponent = -(($x - $mean) ** 2) / (2 * $var);
        return (1 / sqrt(2 * M_PI * $var)) * exp($exponent);
    }

    private function laplace(int $count, int $classTotal, int $vocabSize): float
    {
        return ($count + 1) / ($classTotal + $vocabSize);
    }

    public function predict(array $input): array
    {
        if ($this->total === 0) {
            return [
                'kelas' => 'Tidak Layak',
                'prob_layak' => 0,
                'prob_tidak_layak' => 100,
            ];
        }

        $logScores = [];

        foreach ($this->classes as $c) {
            $classCount = $this->classCounts[$c] ?? 0;
            if ($classCount === 0) {
                $logScores[$c] = -INF;
                continue;
            }

            $logP = log($classCount / $this->total);

            foreach (self::CATEGORIC as $attr) {
                $val = (string) $input[$attr];
                $attrCounts = $this->catCounts[$c][$attr] ?? [];
                $count = $attrCounts[$val] ?? 0;
                $vocab = count($attrCounts) + 1;
                $logP += log($this->laplace($count, $classCount, $vocab));
            }

            foreach (self::NUMERIC as $attr) {
                $x = (float) $input[$attr];
                $mean = $this->numStats[$c][$attr]['mean'];
                $var = $this->numStats[$c][$attr]['var'];
                $g = $this->gaussian($x, $mean, $var);
                $logP += log(max($g, 1e-300));
            }

            $logScores[$c] = $logP;
        }

        $maxLog = max($logScores);
        $expScores = [];
        $sum = 0.0;
        foreach ($logScores as $c => $logS) {
            $expScores[$c] = exp($logS - $maxLog);
            $sum += $expScores[$c];
        }

        $probLayak = round(($expScores['Layak'] / $sum) * 100, 2);
        $probTidak = round(($expScores['Tidak Layak'] / $sum) * 100, 2);

        $predicted = $expScores['Layak'] >= $expScores['Tidak Layak'] ? 'Layak' : 'Tidak Layak';

        return [
            'kelas' => $predicted,
            'prob_layak' => $probLayak,
            'prob_tidak_layak' => $probTidak,
        ];
    }
}
