<?php

namespace CleantalkSP\Common\Scanner\HeuristicAnalyser\Vendors\TiktokenPhp\src;

/**
 * @psalm-suppress ReservedWord
 * @psalm-suppress RedundantCondition
 * @psalm-suppress TypeDoesNotContainType
 * @psalm-suppress InvalidArrayOffset
 * @psalm-suppress DuplicateArrayKey
 */
class Encoder
{
    /** @var bool */
    private $initialized = false;

    /** @var array */
    private $bpeCache = array();

    /** @var array */
    private $rawCharacters = array();

    /** @var array */
    private $encoder = array();

    /** @var array */
    private $bpeRanks = array();

    private function initialize()
    {
        if ($this->initialized) {
            return;
        }
        $rawCharacters = file_get_contents(__DIR__ . '/../data/characters.json');
        if (false === $rawCharacters) {
            throw new \RuntimeException('Unable to load characters.json');
        }
        $this->rawCharacters = json_decode($rawCharacters, true, 512);

        $encoder = file_get_contents(__DIR__ . '/../data/encoder.json');
        if (false === $encoder) {
            throw new \RuntimeException('Unable to load encoder.json');
        }
        $this->encoder = json_decode($encoder, true, 512);

        $bpeDictionary = file_get_contents(__DIR__ . '/../data/vocab.bpe');
        if (false === $bpeDictionary) {
            throw new \RuntimeException('Unable to load vocab.bpe');
        }

        $lines = preg_split('#\r\n|\r|\n#', $bpeDictionary);
        if (false === $lines) {
            throw new \RuntimeException('Unable to split vocab.bpe');
        }

        $bpeMerges = array();
        $rawDictionaryLines = array_slice($lines, 1, count($lines));
        foreach ($rawDictionaryLines as $rawDictionaryLine) {
            $splitLine = preg_split('#(\s+)#', (string) $rawDictionaryLine);
            if (false === $splitLine) {
                continue;
            }
            $splitLine = array_filter($splitLine, array(
                $this, 'filterEmpty'
            ));
            if (!empty($splitLine)) {
                $bpeMerges[] = $splitLine;
            }
        }

        $this->bpeRanks = $this->buildBpeRanks($bpeMerges);
        $this->initialized = true;
    }

    /**
     * @return array
     */
    public function encode($text)
    {
        if (empty($text)) {
            return array();
        }

        $this->initialize();

        preg_match_all("#'s|'t|'re|'ve|'m|'ll|'d| ?\\p{L}+| ?\\p{N}+| ?[^\\s\\p{L}\\p{N}]+|\\s+(?!\\S)|\\s+#u", $text, $matches);
        if (!isset($matches[0]) || 0 == (is_array($matches[0]) ? count($matches[0]) : 0)) {
            return array();
        }

        $bpeTokens = array();
        foreach ($matches[0] as $token) {
            $token = mb_convert_encoding((string) $token, "UTF-8", "ISO-8859-1");
            $characters = $this->mb_str_split_compat($token);

            $resultWord = '';
            foreach ($characters as $char) {
                if (!isset($this->rawCharacters[$this->characterToUnicode($char)])) {
                    continue;
                }
                $resultWord .= $this->rawCharacters[$this->characterToUnicode($char)];
            }

            $newTokensBpe = $this->bpe($resultWord);
            $newTokensBpe = explode(' ', $newTokensBpe);
            foreach ($newTokensBpe as $newBpeToken) {
                $encoded = isset($this->encoder[$newBpeToken]) ? $this->encoder[$newBpeToken] : $newBpeToken;
                $bpeTokens[] = $encoded;
            }
        }

        return $bpeTokens;
    }

    private function mb_str_split_compat($string)
    {
        // PHP 7.4+: mb_str_split exists, PHP 5.6: use preg_split
        if (function_exists('mb_str_split')) {
            return mb_str_split($string, 1, 'UTF-8');
        }
        return preg_split('//u', $string, -1, PREG_SPLIT_NO_EMPTY);
    }

    private function filterEmpty($var)
    {
        return null !== $var && false !== $var && '' !== $var;
    }

    private function characterToUnicode($characters)
    {
        $firstCharacterCode = ord($characters[0]);

        if ($firstCharacterCode <= 127) {
            return $firstCharacterCode;
        }

        if ($firstCharacterCode >= 192 && $firstCharacterCode <= 223) {
            return ($firstCharacterCode - 192) * 64 + (ord($characters[1]) - 128);
        }

        if ($firstCharacterCode >= 224 && $firstCharacterCode <= 239) {
            return ($firstCharacterCode - 224) * 4096 + (ord($characters[1]) - 128) * 64 + (ord($characters[2]) - 128);
        }

        if ($firstCharacterCode >= 240 && $firstCharacterCode <= 247) {
            return ($firstCharacterCode - 240) * 262144 + (ord($characters[1]) - 128) * 4096 + (ord($characters[2]) - 128) * 64 + (ord($characters[3]) - 128);
        }

        if ($firstCharacterCode >= 248 && $firstCharacterCode <= 251) {
            return ($firstCharacterCode - 248) * 16777216 + (ord($characters[1]) - 128) * 262144 + (ord($characters[2]) - 128) * 4096 + (ord($characters[3]) - 128) * 64 + (ord($characters[4]) - 128);
        }

        if ($firstCharacterCode >= 252 && $firstCharacterCode <= 253) {
            return ($firstCharacterCode - 252) * 1073741824 + (ord($characters[1]) - 128) * 16777216 + (ord($characters[2]) - 128) * 262144 + (ord($characters[3]) - 128) * 4096 + (ord($characters[4]) - 128) * 64 + (ord($characters[5]) - 128);
        }

        if ($firstCharacterCode >= 254) {
            return 0;
        }

        return 0;
    }

    /**
     * @param array $bpes
     *
     * @return array
     */
    private function buildBpeRanks($bpes)
    {
        $result = array();
        $rank = 0;
        foreach ($bpes as $bpe) {
            if (!isset($bpe[1], $bpe[0])) {
                continue;
            }

            if (!isset($result[$bpe[0]])) {
                $result[$bpe[0]] = array();
            }
            $result[$bpe[0]][$bpe[1]] = $rank;
            ++$rank;
        }

        return $result;
    }

    /**
     * Return set of symbol pairs in a word.
     * Word is represented as tuple of symbols (symbols being variable-length strings).
     *
     * @param array $word
     *
     * @return array
     */
    private function buildSymbolPairs($word)
    {
        $pairs = array();
        $previousPart = null;
        foreach ($word as $i => $part) {
            if ($i > 0) {
                $pairs[] = array($previousPart, $part);
            }

            $previousPart = $part;
        }

        return $pairs;
    }

    private function bpe($token)
    {
        if (isset($this->bpeCache[$token])) {
            return $this->bpeCache[$token];
        }

        $word = $this->mb_str_split_compat($token);
        $initialLength = count($word);
        $pairs = $this->buildSymbolPairs($word);
        if (empty($pairs)) {
            return $token;
        }

        while (true) {
            $minPairs = array();
            foreach ($pairs as $pair) {
                if (isset($this->bpeRanks[$pair[0]][$pair[1]])) {
                    $rank = $this->bpeRanks[$pair[0]][$pair[1]];
                    $minPairs[$rank] = $pair;
                } else {
                    $minPairs[10000000000] = $pair;
                }
            }

            $minPairsKeys = array_keys($minPairs);
            sort($minPairsKeys, SORT_NUMERIC);
            $minimumKey = isset($minPairsKeys[0]) ? $minPairsKeys[0] : null;

            $bigram = $minPairs[$minimumKey];
            if (!isset($this->bpeRanks[$bigram[0]][$bigram[1]])) {
                break;
            }

            $first = $bigram[0];
            $second = $bigram[1];
            $newWord = array();
            $i = 0;
            while ($i < count($word)) {
                $j = $this->indexOf($word, $first, $i);
                if (-1 === $j) {
                    $newWord = array_merge($newWord, array_slice($word, $i));
                    break;
                }

                $slicer = ($i > $j || 0 === $j) ? array() : array_slice($word, $i, $j - $i);

                $newWord = array_merge($newWord, $slicer);
                if (count($newWord) > $initialLength) {
                    break;
                }

                $i = $j;
                if ($word[$i] === $first && $i < count($word) - 1 && $word[$i + 1] === $second) {
                    $newWord[] = $first . $second;
                    $i += 2;
                } else {
                    $newWord[] = $word[$i];
                    ++$i;
                }
            }

            if ($word === $newWord) {
                break;
            }

            $word = $newWord;
            if (count($word) === 1) {
                break;
            }

            $pairs = $this->buildSymbolPairs($word);
        }

        $word = implode(' ', $word);
        $this->bpeCache[$token] = $word;

        return $word;
    }

    /**
     * @param array $array
     */
    private function indexOf($array, $searchElement, $fromIndex)
    {
        $slicedArray = array_slice($array, $fromIndex);
        $indexed = array_search($searchElement, $slicedArray);
        if ($indexed === false) {
            return -1;
        }
        return $indexed + $fromIndex;
    }
}
