import numpy as np
import ncnn
import cv2
import re
import jaconv
from transformers import AutoTokenizer
import os
from pathlib import Path
from typing import Optional, Union, List


import re
import numpy as np


# we we're kinda forced to implement the argmax by hand because of two factors:
# 1- the argmax from this model always takes the last axis, and this is not possible in numcpp , it only goes right or left.
# 2- numpy's implementation of argmax was too obscure , so we couldn't replicate it in c++ .
# well, now we can rewrite in c++, and also we dont need to give numpy any credits for this.
def argmax_last_axis(mat: List[List[int]]) -> List[int]:
    """Return first max index for each row in matrices of ANY size/shape."""    
    return [
        max(
            ((val, idx) for idx, val in enumerate(row)),
            default=(-1, -1)
        )[1]
        for row in mat
    ]

def get_most_recent_file(
    directory: Union[str, Path],
    pattern: str = "*",
    recursive: bool = False,
    return_path: bool = True
) -> Optional[Union[str, Path]]:
   
    dir_path = Path(directory)
    
    if not dir_path.exists():
        raise FileNotFoundError(f"Directory not found: {directory}")
    if not dir_path.is_dir():
        raise ValueError(f"Path is not a directory: {directory}")

    # Get all matching files
    if recursive:
        files = list(dir_path.rglob(pattern))
    else:
        files = list(dir_path.glob(pattern))
    
    # Filter out directories
    files = [f for f in files if f.is_file()]
    
    if not files:
        return None
    
    # Find most recent file by modification time
    latest = max(files, key=lambda f: f.stat().st_mtime)
    
    return str(latest) if return_path else latest


class TextProcessor:
    def __init__(self, vocab_path: str):
        self.vocab = self._load_vocab(vocab_path)
        self.unk_token = "[UNK]"
        self.special_token_pattern = re.compile(r'\[.*?\]')  # Pattern to match [ANYTHING]
        
    def _load_vocab(self, path: str) -> list:
        """Load vocab with 1-based indexing"""
        try:
            with open(path, 'r', encoding='utf-8') as f:
                return [line.strip() for line in f]
        except FileNotFoundError:
            raise RuntimeError(f"Vocab file not found: {path}")

    def _remove_specials(self, text: str) -> str:
        """Remove special tokens and consecutive duplicates"""
        # Filter out special tokens
        cleaned = self.special_token_pattern.sub('', text)
        
        # Remove consecutive duplicates
        if not cleaned:
            return ''
            
        result = [cleaned[0]]
        for char in cleaned[1:]:
            if char != result[-1]:
                result.append(char)
        return ''.join(result)


    def decode(self, indices) -> str:
        """
        Convert indices to cleaned text
        """
        if isinstance(indices, (int, np.integer)):
            indices = [indices]
            
        raw_text = ''.join(
            self.vocab[idx] if 1 <= idx <= len(self.vocab) else self.unk_token
            for idx in indices
        )
        almost_done = self._remove_specials(raw_text)
        final_text = post_process(almost_done)
        return final_text





def preprocess_image(image_path):
    img = cv2.imread(image_path)
    mat_in = ncnn.Mat.from_pixels_resize(
        img, ncnn.Mat.PixelType.PIXEL_RGB, img.shape[1], img.shape[0], 224, 224
    )
    mean_vals = [0.5, 0.5, 0.5] * 2
    norm_vals = [1/255, 1/255, 1/255] * 2
    mat_in.substract_mean_normalize(mean_vals, norm_vals)
    return mat_in


def run_inference(mat_in, model_path="model"):
    net = ncnn.Net()
    net.opt.num_threads = 32
    net.opt.use_vulkan_compute = False
    net.load_param(f"{model_path}/ncnn_model.param")
    net.load_model(f"{model_path}/ncnn_model.bin")

    ex = net.create_extractor()
    ex.input("in0", mat_in)
    ex.input("in1", ncnn.Mat(np.full((1, 300), 0, dtype=np.int64)))

    try:
        ret, out0 = ex.extract("out0")
        print("Output shape:", out0.shape)
        predictions = np.array(out0)
      #  ids = np.argmax(predictions, axis=-1)
        ids = argmax_last_axis(predictions) 
        print("argmax'ed: ", ids)
        return ids
    finally:
        del ex, net


def decode_tokens(logits, tokenizer_path="model"):
    tkz = AutoTokenizer.from_pretrained(tokenizer_path)
    return tkz.decode(token_ids, skip_special_tokens=True)


def post_process(text):
    text = "".join(text.split()).replace("…", "...")
    text = re.sub("[・.]{2,}", lambda x: (x.end() - x.start()) * ".", text)
    return jaconv.h2z(text, ascii=True, digit=True)


def main(image_path):
    mat_in = preprocess_image(image_path)
    logits = run_inference(mat_in)
    processor = TextProcessor("model/vocab.txt")
    result = processor.decode(logits)
    print(f"Decoded text: {result}")
#    raw_text = decode_tokens(logits)
#    return post_process(raw_text)
   

if __name__ == "__main__":
    recent_img = get_most_recent_file("/home/hashirama/screenshots/", pattern="*.png", recursive=True)
    print("PROCESSED OUT:", main(recent_img))