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 pyinotify
import fnmatch
import glob
import subprocess
from collections import Counter
import matplotlib.pyplot as plt
import re
import numpy as np

WATCH_DIR = "/tmp/mote-ocr-screenshots/"



def contrast_stretch(gray):
    p2, p98 = np.percentile(gray, (2, 98))
    return cv2.normalize(gray, None, p2, p98, cv2.NORM_MINMAX)


def copy_to_clipboard(text):
    subprocess.run(["xsel", "--clipboard", "--input"], input=text.encode(), check=True)


# 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
    ]


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
        )
        text = self._remove_specials(raw_text)
        
        return text


def bicubic(img):
    h, w = img.shape[:2]
    M = np.array([[1, 0, 0],
                  [0, 1, 0]], dtype=np.float32)
    return cv2.warpAffine(img, M, (w, h), flags=cv2.INTER_CUBIC)


def adjust_gamma(image, gamma=1.0, alpha=1.0, beta=0):
    """
    Adjust the gamma and contrast of a grayscale image.

    Parameters:
        image (numpy.ndarray): Input grayscale image.
        gamma (float): Gamma correction factor. Values < 1 darken the image, > 1 brighten it.
        alpha (float): Contrast control. 1.0 means no change, < 1.0 reduces contrast, > 1.0 increases contrast.
        beta (int): Brightness control. Positive values brighten the image, negative values darken it.

    Returns:
        numpy.ndarray: Adjusted image.
    """
    # Apply gamma correction
    invGamma = 1.0 / gamma
    table = np.array([((i / 255.0) ** invGamma) * 255 for i in np.arange(0, 256)]).astype("uint8")
    gamma_corrected = cv2.LUT(image, table)

    # Adjust contrast and brightness
    adjusted = cv2.convertScaleAbs(gamma_corrected, alpha=alpha, beta=beta)

    return adjusted



def preprocess_image(image_path):
    # Load image in grayscale (1-channel)
    gray = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    resampled = bicubic(gray) 
    # Convert grayscale to 3-channel RGB by duplicating the gray channel
    img1 = contrast_stretch(adjust_gamma(cv2.cvtColor(resampled, cv2.COLOR_GRAY2RGB),gamma=0.3,alpha=1.6,beta=-5)) # perfect tunning
    
    img = cv2.fastNlMeansDenoising(img1, None, 30, 10, 20)
    
    plt.imshow(img)
    plt.axis("off")
    plt.savefig("/tmp/output.png", bbox_inches="tight", pad_inches=0)

    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] * 3
    norm_vals = [1/255] * 3
    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.squeeze(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 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}")
    copy_to_clipboard(result)
#    raw_text = decode_tokens(logits)
#    return post_process(raw_text)
   

def get_most_recent_file(directory, pattern="*.png", recursive=True):
    search_path = os.path.join(directory, "**", pattern) if recursive else os.path.join(directory, pattern)
    files = glob.glob(search_path, recursive=recursive)
    return max(files, key=os.path.getctime) if files else None


class EventHandler(pyinotify.ProcessEvent):
    def process_IN_CLOSE_WRITE(self, event):
        if fnmatch.fnmatch(event.pathname, "*.png"):
            recent_img = get_most_recent_file(WATCH_DIR, pattern="*.png", recursive=True)
            print("PROCESSED OUT:", main(recent_img))

if __name__ == "__main__":
    wm = pyinotify.WatchManager()
    handler = EventHandler()
    notifier = pyinotify.Notifier(wm, handler)
    mask = pyinotify.IN_CLOSE_WRITE  # triggers when a file is written and closed

    wm.add_watch(WATCH_DIR, mask, rec=True, auto_add=True)
    print(f"Watching {WATCH_DIR} for new PNG files...")

    notifier.loop()