"""
A collection of functions for performing text analysis.
"""
from typing import List
import re
import random
import numpy as np
#import nltk
# nltk.download('punkt')  # only once
from nltk.tokenize import sent_tokenize
from datetime import datetime
import pandas as pd

def get_sentences_orig(text, minlen=1, language='Chinese', preserve_spaces=False):
    """
    Split input text into sentences, with optional language-specific processing 
    and length filtering, optionally preserving leading/trailing spaces.
    
    This function tokenizes text into sentences. For Chinese text, it uses the 
    `SnowNLP` library; for other languages, it falls back to NLTK's `sent_tokenize`. 
    It can also filter out sentences shorter than a specified minimum length.
    
    ---------------------------------------------------------------------------
    ARGUMENTS:
        text : str
            The input text to be split into sentences.
            
        minlen : int, optional (default=1)
            The minimum length a sentence must have to be included in the output.
            Sentences shorter than this will be discarded. When preserve_spaces=True,
            length is calculated after stripping spaces for filtering purposes.
            
        language : str, optional (default='Chinese')
            The language of the input text. Determines the sentence tokenizer to use.
            'Chinese' uses SnowNLP; all other values use NLTK's `sent_tokenize`.
            
        preserve_spaces : bool, optional (default=False)
            If True, preserves leading and trailing spaces from the original text
            for each sentence. If False, returns sentences as provided by the tokenizer.
    
    ---------------------------------------------------------------------------
    OUTPUT:
        output : list of str
            A list of sentences tokenized from the input text, filtered by 
            minimum length. Spacing depends on preserve_spaces parameter.
    
    ---------------------------------------------------------------------------
    AUTHORS:
        Samuel Davenport
    ---------------------------------------------------------------------------
    """
    if text == None or len(text) == 0:
        output = text
        return 
    
    language_lower = language.lower()
    
    if language_lower == 'chinese':
        from snownlp import SnowNLP
        s = SnowNLP(text)
        sentences = s.sentences
    else:
        sentences = sent_tokenize(text)
    
    # Original functionality (default)
    if not preserve_spaces:
        if minlen > 1:
            output = [sent for sent in sentences if len(sent) >= minlen]
        else:
            output = sentences
        return output

    output = glue_sentences(text, sentences, minlen)
    if minlen > 1:
        output = [sent for sent in output if len(sent) >= minlen]
    
    return output

def glue_sentences(text, sentences, minlen = 1):
    output = []
    search_start = 0
    updated_text = text
    for sent in sentences:
        sent_start = updated_text.find(sent)
        sent_end = sent_start + len(sent)
        
        while sent_end < len(updated_text) and updated_text[sent_end].isspace():
            sent_end += 1
        
        output.append(updated_text[:sent_end])
        updated_text = updated_text[(sent_end):]
        
    noutput = len(output)
    output[noutput-1] += updated_text
        
    return output

def docs2sentences(docs, minlen=1, language='Chinese', preserve_spaces=True, combine = False):
    list_of_sentences = []
    doc_length = []
    for i in range(len(docs)):
        text = docs[i]
        if language == 'English':
            sentences = get_sentences_orig(text, minlen, language=language, preserve_spaces=preserve_spaces)
        else:
            sentences = get_sentences(text, minlen, language=language, preserve_spaces=preserve_spaces)
        doc_length.append(len(sentences))
        if combine:
            list_of_sentences += sentences
        else:
            list_of_sentences.append(sentences)

    return list_of_sentences, doc_length

def parse_list(listofitems, listoflengths):
    out = []
    counter = 0
    for i in range(len(listoflengths)):
        start = counter
        counter += listoflengths[i]
        items = listofitems[start:counter]
        out.append(items)
        
    return out

def get_random_sentence(text, minlen, language = 'chinese'):
    """
    Select a random sentence from text after filtering by minimum length.

    This function extracts sentences from the input text using `get_sentences`,
    filters them based on a minimum length, and returns one randomly selected
    sentence from the filtered list.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    text : str
        The input text from which to extract and select a sentence.
    minlen : int
        Minimum length a sentence must have to be considered for selection.

    ---------------------------------------------------------------------------
    OUTPUT:
    random_sentence : str
        A randomly chosen sentence from the list of filtered sentences.

    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    sentences = get_sentences(text, minlen, language = language)
    random_sentence = random.choice(sentences)
    return random_sentence

def get_chunks(text, jump_by = 1000, language = 'chinese', sentence_splitter = 'new'):
    if sentence_splitter == 'new':
        sentences = get_sentences(text, preserve_spaces=True, language = language)
    else:
        sentences = get_sentences_orig(text, preserve_spaces=True, language = language)

    len_sentences = [len(sentence) for sentence in sentences]
    sentence_cumsum = np.cumsum(len_sentences)
    
    indices = []
    baseline = 0 
    for i, x in enumerate(sentence_cumsum):
        if x >= baseline + jump_by:
            indices.append(i)
            baseline = x

def get_sentences(text: str, minlen: int = 1, language='chinese', preserve_spaces=False, remove_empty: bool = True) -> List[str]:
    if not text or not text.strip():
        return []
    
    # Common abbreviations that shouldn't end sentences (English and Malay)
    ABBREVIATIONS_ENGLISH = [
        'Dr', 'Mr', 'Mrs', 'Ms', 'Sr', 'Jr', 'Prof', 'Rev',  # Titles
        'St', 'Ave', 'Blvd', 'Rd', 'Ln',  # Street types
        'Jan', 'Feb', 'Mar', 'Apr', 'Jun', 'Jul', 'Aug', 'Sep', 'Sept', 'Oct', 'Nov', 'Dec',  # Months
        'Mon', 'Tue', 'Tues', 'Wed', 'Thu', 'Thur', 'Thurs', 'Fri', 'Sat', 'Sun',  # Days
        'Inc', 'Ltd', 'Corp', 'Co', 'LLC',  # Business
        'vs', 'etc', 'al', 'eg', 'ie', 'cf', 'approx', 'apt',  # Latin/common
        'Vol', 'No', 'Fig', 'Ch', 'Sec', 'pp', 'pg',  # References
        'Gen', 'Col', 'Lt', 'Sgt', 'Capt', 'Cmdr', 'Adm',  # Military
        'lb', 'lbs', 'oz', 'ft', 'in', 'cm', 'mm', 'km', 'kg', 'mg',  # Units
    ]
    
    ABBREVIATIONS_MALAY = [
        'W', 'S', 'T', 'A', 'Dr', 'En', 'Cik', 'Pn', 'Tn', 'Hj', 'Hjh', 'Dato', 'Datin', 'Tan', 'Tun', 'Toh',  # Titles
        'Prof', 'Ustaz', 'Ustazah',  # Professional/religious titles
        'Sdn', 'Bhd', 'Plt', 'Sdn Bhd',  # Business
        'Jln', 'Jl', 'Lrg', 'Tmn', 'Kpg', 'Kg', 'Bt',  # Street/place types
        'Ltd', 'Jan', 'Feb', 'Mac', 'Apr', 'Mei', 'Jun', 'Jul', 'Ogos', 'Sep', 'Okt', 'Nov', 'Dis',  # Months
        'Isn', 'Sel', 'Rab', 'Kha', 'Jum', 'Sab', 'Ahd',  # Days
        'No', 'Bil', 'Tel', 'Faks', 'dll', 'dsb', 'etc',  # Common
        'Yg', 'Bhg', 'Kep', 'Kel', 'Neg',  # Administrative
        'cm', 'mm', 'km', 'kg', 'mg', 'm',  # Units
    ]
    
    # Chinese sentence-ending punctuation marks
    # 。 - period, ！ - exclamation, ？ - question mark
    if language.lower() == 'chinese':
        punctuation_pattern = r'([。！？；…])'
        abbreviations = []
    elif language.lower() == 'english':
        # Match sentence-ending punctuation, but not periods between digits like 9.8
        # \.(?!\d) matches a period NOT followed by a digit
        # This allows "Sam 5." to split but not "9.8"
        punctuation_pattern = r'([!?;…]|\.(?!\d))'
        abbreviations = ABBREVIATIONS_ENGLISH
    elif language.lower() == 'malay':
        # Malay uses the same punctuation as English
        punctuation_pattern = r'([!?;…]|\.(?!\d))'
        abbreviations = ABBREVIATIONS_MALAY
    
    # First split by newlines to handle line breaks as sentence boundaries
    lines = text.split('\n')
    sentences = []
    for line in lines:
        line = line.rstrip()
       
        if not line:
            continue
        # Split by punctuation marks while preserving them
        parts = re.split(punctuation_pattern, line)
        current_sentence = ""
        for i, part in enumerate(parts):
            if re.match(r'^[.!?;…。！？；]$', part):
                # This is a punctuation mark
                
                # Check if this period follows an abbreviation
                is_abbreviation = False
                if language.lower() in ['english', 'malay'] and part == '.':
                    # Get the last word before the period
                    words = current_sentence.split()
                    if words and words[-1] in abbreviations:
                        is_abbreviation = True
                
                # Add punctuation to current sentence
                current_sentence += part
                
                if is_abbreviation:
                    # Don't end the sentence - it's an abbreviation
                    # Add a space since we'll continue with more text
                    current_sentence += ' '
                else:
                    # Complete the sentence and add to list
                    cleaned_sentence = current_sentence.strip()
                        
                    if cleaned_sentence and len(cleaned_sentence) >= minlen:
                        sentences.append(cleaned_sentence)
                    current_sentence = ""
            else:
                # This is text content - strip whitespace before adding
                current_sentence += part.strip()
        # Handle any remaining text that doesn't end with punctuation
        if current_sentence.strip():
            cleaned = current_sentence.strip()
            if len(cleaned) >= minlen:
                sentences.append(cleaned)     
    
    if preserve_spaces:
        sentences = glue_sentences(text, sentences)
                
    return sentences

def rep_texts(texts, o, r):
    replaced_texts = [text.replace(o, r) for text in texts]
    return replaced_texts

def extract_first_date(text):
    """
    Extracts the first date from a string and returns it in 'YYYY-MM-DD' format.
    
    Args:
        text (str): The input text containing dates
        
    Returns:
        str: The first date found in 'YYYY-MM-DD' format, or None if no date is found
    """
    # Pattern to match dates like "13 February 2024"
    pattern = r'\b(\d{1,2})\s+([A-Z][a-z]+)\s+(\d{4})\b'
    
    match = re.search(pattern, text)
    
    if match:
        day = match.group(1)
        month = match.group(2)
        year = match.group(3)
        
        # Parse the date and convert to YYYY-MM-DD format
        date_str = f"{day} {month} {year}"
        date_obj = datetime.strptime(date_str, "%d %B %Y")
        
        return date_obj.strftime("%Y-%m-%d")
    
    return None

def contains_chinese(text):
    """
    Check if the text contains any Chinese characters.
    
    Args:
        text (str): The text to check
        
    Returns:
        bool: True if Chinese characters are found, False otherwise
    """
    for char in text:
        # Check if character is in CJK Unified Ideographs ranges
        code_point = ord(char)
        if (0x4E00 <= code_point <= 0x9FFF or      # CJK Unified Ideographs
            0x3400 <= code_point <= 0x4DBF or      # CJK Unified Ideographs Extension A
            0x20000 <= code_point <= 0x2A6DF or    # CJK Unified Ideographs Extension B
            0x2A700 <= code_point <= 0x2B73F or    # CJK Unified Ideographs Extension C
            0x2B740 <= code_point <= 0x2B81F or    # CJK Unified Ideographs Extension D
            0x2B820 <= code_point <= 0x2CEAF or    # CJK Unified Ideographs Extension E
            0x2CEB0 <= code_point <= 0x2EBEF or    # CJK Unified Ideographs Extension F
            0x30000 <= code_point <= 0x3134F):     # CJK Unified Ideographs Extension G
            return True
    return False

def count_chinese(df):
    # Regex range for common Chinese characters
    chinese_regex = r'[\u4e00-\u9fff]'
    
    counts = {}
    for col in df.columns:
        # Convert to string to avoid errors with non-string columns
        counts[col] = df[col].astype(str).str.contains(chinese_regex, regex=True).sum()
    
    return pd.Series(counts, name="rows_with_chinese")

def chunk_texts(texts, max_chars=512):
    all_chunks = []
    num_chunks = []
    
    for text in texts:
        chunks = [text[i:i+max_chars] for i in range(0, len(text), max_chars)]
        all_chunks.extend(chunks)
        num_chunks.append(len(chunks))
    
    return all_chunks, num_chunks

def unchunk(chunks, num_chunks):
    texts = []
    idx = 0
    
    for count in num_chunks:
        combined = ''.join(chunks[idx:idx+count])
        texts.append(combined)
        idx += count
    
    return texts

def split_to_list(s, splitter = ','):
    return [item.strip() for item in s.split(splitter)]

def extract_outputs(text):
    return re.findall(r"\{[^}]*\}", text)

def find_nones(lst):
    return [i for i, x in enumerate(lst) if x is None]

def prefix_strings(string_list, prefix):
    """
    Returns a list of strings that start with the given prefix.
    
    Args:
        string_list: List of strings to search through
        prefix: The prefix string to match
        
    Returns:
        List of strings that start with the prefix
    """
    return [s for s in string_list if s.startswith(prefix)]

def string_line(filepath, search_string):
    """
    Find the line number(s) where a string appears in a Python file.
    
    Args:
        filepath: Path to the .py file
        search_string: String to search for
    
    Returns:
        List of line numbers where the string appears (1-indexed),
        or empty list if not found
    """
    line_numbers = []
    
    try:
        with open(filepath, 'r', encoding='utf-8') as file:
            for line_num, line in enumerate(file, start=1):
                if search_string in line:
                    line_numbers.append(line_num)
    except FileNotFoundError:
        print(f"Error: File '{filepath}' not found")
        return []
    except Exception as e:
        print(f"Error reading file: {e}")
        return []
    
    return line_numbers

def add_spaces(name):
    """
    Takes a camelCase/PascalCase name and returns a spaced out version.
    Also handles underscores and preserves acronyms.
    """
    # Replace underscores with spaces
    name = name.replace('_', ' ')
    
    # Add space before uppercase letters that follow lowercase letters
    # e.g., "MinistryOf" -> "Ministry Of"
    name = re.sub(r'([a-z])([A-Z])', r'\1 \2', name)
    
    # Add space before uppercase letters that are followed by lowercase letters
    # and preceded by uppercase letters (handles acronyms like "NMPATest" -> "NMPA Test")
    name = re.sub(r'([A-Z]+)([A-Z][a-z])', r'\1 \2', name)
    
    # Clean up multiple spaces
    name = re.sub(r' +', ' ', name)
    
    return name.strip()