<?php

class DeepSeekTokenizer {
    private array $vocab;
    private array $reversedVocab;
    private int $modelMaxLength;

    public function __construct(string $tokenizerJsonPath) {
        $config = json_decode(file_get_contents($tokenizerJsonPath), true);
        $this->vocab = $config['model']['vocab'];
        $this->reversedVocab = array_flip($this->vocab);
        $this->modelMaxLength = $config['model_max_length'] ?? PHP_INT_MAX;
    }

    public function encode(
        string $text,
        ?string $textPair = null,
        bool $addSpecialTokens = true,
        bool|string $padding = false,
        bool|string $truncation = false,
        ?int $maxLength = null
    ): array {
        // First, perform pre-tokenization steps
        $tokens = $this->preTokenize($text);
        
        // Convert to byte-level representation
        $byteLevelTokens = $this->byteLevelEncode($tokens);
        
        // Merge tokens using BPE
        $mergedTokens = $this->mergeBPE($byteLevelTokens);
        
        // Convert tokens to IDs
        return $this->convertTokensToIds($mergedTokens);
    }

    private function preTokenize(string $text): array {
        // Split on numbers (1-3 digits)
        $text = preg_replace('/(\d{1,3})/', ' $1 ', $text);
        
        // Split on CJK characters
        $text = preg_replace('/([\x{4E00}-\x{9FFF}\x{3040}-\x{309F}\x{30A0}-\x{30FF}])/u', ' $1 ', $text);
        
        // Split on other patterns as defined in the tokenizer config
        $pattern = '/[!\"#$%&\'()*+,\-.\/:;<=>?@\[\\\]^_`{|}~][A-Za-z]+|[^\r\n\p{L}\p{P}\p{S}]?[\p{L}\p{M}]+| ?[\p{P}\p{S}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+/u';
        
        preg_match_all($pattern, $text, $matches);
        return $matches[0];
    }

    private function byteLevelEncode(array $tokens): array {
        $byteLevelTokens = [];
        
        foreach ($tokens as $token) {
            // Add space prefix for non-first tokens that don't start with space
            if (!empty($byteLevelTokens) && !preg_match('/^\s/', $token)) {
                $token = ' ' . $token;
            }
            
            // Convert to bytes and add Ġ prefix for spaces
            $bytes = [];
            for ($i = 0; $i < strlen($token); $i++) {
                $char = $token[$i];
                if ($char === ' ') {
                    $bytes[] = 'Ġ';
                } else {
                    $bytes[] = $char;
                }
            }
            
            $byteLevelTokens[] = implode('', $bytes);
        }
        
        return $byteLevelTokens;
    }

    private function mergeBPE(array $tokens): array {
        $mergedTokens = [];
        
        foreach ($tokens as $token) {
            // Split token into characters
            $chars = preg_split('//u', $token, -1, PREG_SPLIT_NO_EMPTY);
            
            // Start with the longest possible tokens that exist in vocab
            $currentTokens = [];
            $currentBuffer = '';
            
            for ($i = 0; $i < count($chars); $i++) {
                $currentBuffer .= $chars[$i];
                
                // Try to find the longest possible token from current position
                $longestMatch = '';
                $tempBuffer = $currentBuffer;
                
                // Look ahead up to 10 characters (adjust this based on your vocabulary)
                for ($j = 1; $j <= 10 && ($i + $j) < count($chars); $j++) {
                    $tempBuffer .= $chars[$i + $j];
                    if (isset($this->vocab[$tempBuffer])) {
                        $longestMatch = $tempBuffer;
                    }
                }
                
                if ($longestMatch !== '') {
                    // Found a match in vocab
                    $currentTokens[] = $longestMatch;
                    $i += strlen($longestMatch) - strlen($currentBuffer);
                    $currentBuffer = '';
                } elseif (isset($this->vocab[$currentBuffer])) {
                    // Current buffer is a valid token
                    $currentTokens[] = $currentBuffer;
                    $currentBuffer = '';
                }
            }
            
            // Add any remaining buffer if it's a valid token
            if ($currentBuffer !== '' && isset($this->vocab[$currentBuffer])) {
                $currentTokens[] = $currentBuffer;
            }
            
            $mergedTokens = array_merge($mergedTokens, $currentTokens);
        }
        
        return $mergedTokens;
    }

    private function convertTokensToIds(array $tokens): array {
        $ids = [];
        
        foreach ($tokens as $token) {
            if (isset($this->vocab[$token])) {
                $ids[] = $this->vocab[$token];
            }
        }
        
        return $ids;
    }

    public function decode(array $ids, bool $skipSpecialTokens = false): string {
        $tokens = [];
        
        foreach ($ids as $id) {
            if (isset($this->reversedVocab[$id])) {
                $token = $this->reversedVocab[$id];
                if (!$skipSpecialTokens || !$this->isSpecialToken($token)) {
                    $tokens[] = $token;
                }
            }
        }
        
        $text = implode('', $tokens);
        
        // Clean up the text
        return $this->cleanUp($text);
    }

    private function isSpecialToken(string $token): bool {
        return strpos($token, '<') === 0 && strpos($token, '>') === strlen($token) - 1;
    }

    /**
     * Get next token prediction from model output
     */
    public function getNextToken(array $modelOutput, ?int $lastUsedLogitSet = null): string {
        $bestProb = 0;
        $bestToken = '';
        $bestSet = -1;
        $bestLogits = [];
        
        // Look at each set of logits
        foreach ($modelOutput[0] as $setIndex => $logits) {
            // Skip if not an array of logits
            if (!is_array($logits) || empty($logits)) {
                continue;
            }
            
            // Filter logits to only include tokens in our vocabulary
            $validLogits = [];
            foreach ($logits as $idx => $logit) {
                if (isset($this->reversedVocab[$idx])) {
                    $validLogits[$idx] = $logit;
                }
            }
            
            // Convert to probabilities using softmax
            $probs = $this->softmax($validLogits);
            
            // Get the most likely token
            $maxProb = max($probs);
            $tokenId = array_search($maxProb, $probs);
            
            // If this is the most confident prediction so far, save it
            if (($maxProb > $bestProb) && $setIndex !== $lastUsedLogitSet) {
                $bestProb = $maxProb;
                $bestToken = $this->reversedVocab[$tokenId];
                $bestSet = $setIndex;
                $bestLogits[$tokenId] = $maxProb;
            }
        }

        // die(var_dump($bestLogits));
        
        // Debug output for the best set
        //  "Using logit set: " . $bestSet . "\n";
        // echo "Best probability: " . $bestProb . "\n\n";
        
        // Get top 5 from the best set
        $validLogits = [];
        foreach ($bestLogits as $idx => $logit) {
            if (isset($this->reversedVocab[$idx])) {
                $validLogits[$idx] = $logit;
            }
        }
        $probs = $this->softmax($validLogits);
        arsort($probs);
        $topProbs = array_slice($probs, 0, 5, true);

        $temperature = 0.5;
        $scaled = [];
        foreach ($topProbs as $id => $prob) {
            $scaled[$id] = pow($prob, 1/$temperature);
        }

        // die(var_dump($scaled));
        
        // Normalize scaled probabilities
        $sum = array_sum($scaled);
        $scaled = array_map(function($v) use ($sum) { return $v/$sum; }, $scaled);
        
        // Sample token
        $r = mt_rand() / mt_getrandmax();
        $cumsum = 0;
        foreach ($scaled as $id => $prob) {
            $cumsum += $prob;
            if ($r <= $cumsum) {
                $selectedId = $id;
                break;
            }
        }
        
        // Fallback to highest probability
        if (!isset($selectedId)) {
            $selectedId = array_key_first($topProbs);
        }

        // die(var_dump($topProbs));

        $selectedId = array_key_first($topProbs);

        // die(var_dump($selectedId));
        
        // echo "Top 5 tokens from best set:\n";
        // foreach ($top5 as $tid => $prob) {
        //     $token = $this->reversedVocab[$tid];
        //     $originalLogit = $bestLogits[$tid];
        //     echo sprintf(
        //         "ID: %d Token: '%s' Logit: %.4f Prob: %.6f\n",
        //         $tid,
        //         str_replace('Ġ', ' ', $token),
        //         $originalLogit,
        //         $prob
        //     );
        // }
        
        // return [$this->cleanUp($bestToken), $bestSet];
        // return $bestToken;
        return $this->reversedVocab[$selectedId];
    }

    /**
     * Convert logits to probabilities using softmax
     */
    private function softmax(array $x): array {
        if (!is_array($x) || !count($x)) {
            throw new Exception('The softmax function requires an array of real numbers.');
        }
        
        // Find max for numerical stability
        $max = max($x);
        $exp = [];
        $sum = 0;
        
        foreach ($x as $idx => $value) {
            $e = exp($value - $max);
            $exp[$idx] = $e;
            $sum += $e;
        }
        
        return array_map(function ($value) use ($sum) {
            return $value / $sum;
        }, $exp);
    }

    private function cleanUp(string $text): string {
        // Replace Ġ with space and normalize spaces
        // return trim(str_replace('Ġ', ' ', $text));
        return str_replace('Ġ', ' ', $text);
    }
}