<?php

require 'vendor/autoload.php';
require 'DeepSeekTokenizer.php';

use OnnxRuntime\InferenceSession;
use OnnxRuntime\ElementType;
use OnnxRuntime\OrtValue;

class ModelGenerator {
    private InferenceSession $session;
    private DeepSeekTokenizer $tokenizer;
    private array $generatedTokens = [];
    private array $pastKeyValues = [];
    private int $maxNewTokens;
    private array $recentTokens = [];
    
    public function __construct(
        string $modelPath, 
        string $tokenizerPath, 
        int $maxNewTokens = 100,
    ) {
        $this->session = new InferenceSession($modelPath);
        $this->tokenizer = new DeepSeekTokenizer($tokenizerPath);
        $this->maxNewTokens = $maxNewTokens;
    }
    
    public function generate(string $prompt): string {
        $this->generatedTokens = $this->tokenizer->encode($prompt);
        $this->pastKeyValues = [];
        $this->recentTokens = [];
        
        $fullText = $prompt;
        
        for ($i = 0; $i < $this->maxNewTokens; $i++) {
            // Prepare input feed
            $inputFeed = $this->prepareInputFeed();
            
            // Run inference
            $outputs = $this->session->run(null, $inputFeed);
            
            // Get next token
            $nextToken = $this->tokenizer->getNextToken($outputs[0]);

            echo str_replace('Ġ', ' ', $nextToken);
            
            // Check for degenerate patterns
            if ($this->isDegenerate($nextToken)) {
                echo "\nStopping due to degenerate pattern\n";
                break;
            }
            
            // Update state
            $tokenId = $this->tokenizer->encode($nextToken)[0];
            $this->generatedTokens[] = $tokenId;
            $this->recentTokens[] = $nextToken;
            if (count($this->recentTokens) > 10) {
                array_shift($this->recentTokens);
            }
            
            // Update past key values
            $this->updatePastKeyValues($outputs);
            
            // Append to text
            $fullText .= $nextToken;
            
            // Check for stopping criteria
            if ($this->shouldStop($nextToken)) {
                break;
            }
        }
        
        return $fullText;
    }
    
    private function prepareInputFeed(): array {
        $inputFeed = [];
        
        // Use only last token if we have context
        $useTokens = empty($this->pastKeyValues) 
            ? $this->generatedTokens 
            : [end($this->generatedTokens)];
        
        // Prepare tensors
        $inputFeed['input_ids'] = OrtValue::fromArray([$useTokens], ElementType::Int64);
        $inputFeed['attention_mask'] = OrtValue::fromArray([array_fill(0, count($useTokens), 1)], ElementType::Int64);
        
        // Position IDs start from current position
        $startPos = count($this->generatedTokens) - count($useTokens);
        $inputFeed['position_ids'] = OrtValue::fromArray([range($startPos, $startPos + count($useTokens) - 1)], ElementType::Int64);
        
        // Add past key values
        if (!empty($this->pastKeyValues)) {
            foreach ($this->pastKeyValues as $layerIdx => $layer) {
                $inputFeed["past_key_values.$layerIdx.key"] = $layer['key'];
                $inputFeed["past_key_values.$layerIdx.value"] = $layer['value'];
            }
        } else {
            // Initialize empty
            for ($i = 0; $i < 28; $i++) {
                $emptyShape = [1, 2, 0, 128];
                $inputFeed["past_key_values.$i.key"] = OrtValue::fromShapeAndType($emptyShape, ElementType::Float);
                $inputFeed["past_key_values.$i.value"] = OrtValue::fromShapeAndType($emptyShape, ElementType::Float);
            }
        }
        
        return $inputFeed;
    }
    
    private function updatePastKeyValues(array $outputs): void {
        $this->pastKeyValues = [];
        for ($i = 0; $i < 28; $i++) {
            $keyIdx = $i * 2 + 1;
            $valueIdx = $i * 2 + 2;
            $this->pastKeyValues[$i] = [
                'key' => $outputs[$keyIdx],
                'value' => $outputs[$valueIdx]
            ];
        }
    }
    
    private function shouldStop(string $token): bool {
        return $token === '<｜end▁of▁sentence｜>';
    }
    
    private function isDegenerate(string $token): bool {
        if (count($this->recentTokens) < 5) {
            return false;
        }
        
        // Check for token repetition
        $lastFive = array_slice($this->recentTokens, -5);
        if (count(array_unique($lastFive)) === 1) {
            return true;
        }
        
        // Check for space repetition
        $recentText = implode('', $this->recentTokens);
        if (substr_count($recentText, '  ') > 2) {
            return true;
        }
        
        // Check for number repetition
        if (preg_match('/(\d)\1{3,}/', $recentText)) {
            return true;
        }
        
        return false;
    }
}