# -*- coding: utf-8 -*-
from __future__ import annotations

import collections
import os.path
import unicodedata
from functools import reduce
from importlib.resources import files
from typing import Any, Pattern

import numpy
import pandas as pd
import regex as re

from . import xsampa
from .segment import Segment

feature_sets = {
    'spe+': (os.path.join('data', 'ipa_all.csv'),
             os.path.join('data', 'feature_weights.csv'))
}


class SegmentSorter:
    def __init__(self, segments,):
        self._segments = segments
        self._sorted=False

    @property
    def segments(self):
        if not self._sorted:
            self.sort_segments()
        return self._segments

    def sort_segments(self):
        self.segments.sort(key=self.segment_key)

    @staticmethod
    def segment_key(segment_tuple):
        print('segment_tuple=', segment_tuple)
        segment_data=segment_tuple[1]
        return (
            segment_data['syl'], segment_data['son'], segment_data['cons'], segment_data['cont'],
            segment_data['delrel'], segment_data['lat'], segment_data['nas'], segment_data['strid'],
            segment_data['voi'], segment_data['sg'], segment_data['cg'], segment_data['ant'],
            segment_data['cor'], segment_data['distr'], segment_data['lab'], segment_data['hi'],
            segment_data['lo'], segment_data['back'], segment_data['round'], segment_data['velaric'],
            segment_data['tense'], segment_data['long'], segment_data['hitone'], segment_data['hireg']
        )


class FeatureTable(object):
    """The basic PanPhon object for representing the features of sets of segments.

    :param feature_set str: The set of fetures to be used by the FeatureTable object.
    """
    TRIE_LEAF_MARKER = None

    def __init__(self, feature_set: str='spe+'):
        bases_fn, weights_fn = feature_sets[feature_set]
        self.weights = self._read_weights(weights_fn)
        self.segments, self.seg_dict, self.names = self._read_bases(bases_fn, self.weights)
        self.seg_regex = self._build_seg_regex()
        self.seg_trie = self._build_seg_trie()
        self.longest_seg = max([len(x) for x in self.seg_dict.keys()])
        self.xsampa = xsampa.XSampa()

        self.sorted_segments = SegmentSorter(self.segments) # used for quick binary searches

    @staticmethod
    def normalize(data: str) -> str:
        return unicodedata.normalize('NFD', data)

    def _read_bases(self, fn: str, weights) -> tuple[
        list[tuple[str, Segment]],
        dict[str, Segment],
        list[str],
    ]:
        # Define the specification to integer mapping
        spec_to_int = {"+": 1, "0": 0, "-": -1}

        # Read the file name with the phonemes and their feature specification
        with files("panphon").joinpath(fn).open() as f:
            df = pd.read_csv(f)  

        # Normalize the IPA representations
        df["ipa"] = df["ipa"].apply(self.normalize)

        # Compute a list of feature names and convert the corresponding
        # specifications to integers
        feature_names = list(df.columns[1:])
        df[feature_names] = df[feature_names].map(lambda x: spec_to_int[x])
        # Create the segments list
        segments = [
            (row["ipa"], Segment(feature_names, row[1:].to_dict(), weights=weights))
            for (_, row) in df.iterrows()
        ]

        # Convert to dictionary
        seg_dict = dict(segments)

        return segments, seg_dict, feature_names

    def _read_weights(self, weights_fn: str) -> list[float]:
        with files('panphon').joinpath(weights_fn).open() as f:
            df = pd.read_csv(f)
        weights = df.iloc[0].astype(float).tolist()
        return weights

    def _build_seg_regex(self) -> re.Pattern:
        segs = sorted(self.seg_dict.keys(), key=lambda x: len(x), reverse=True)
        return re.compile(r'(?P<all>{})'.format('|'.join(segs)))

    def _build_seg_trie(self) -> dict:
        trie = {}
        for seg in self.seg_dict.keys():
            node = trie
            for char in seg:
                if char not in node:
                    node[char] = {}
                node = node[char]
            node[self.TRIE_LEAF_MARKER] = None
        return trie

    def fts(self, ipa: str, normalize: bool=True) -> dict[str, int]:
        if normalize:
            ipa = FeatureTable.normalize(ipa)
        if ipa in self.seg_dict:
            return self.seg_dict[ipa]
        else:
            return {}

    def longest_one_seg_prefix(self, word: str, normalize: bool=True) -> str:
        """Return longest Unicode IPA prefix of a word

        Args:
            word (unicode): input word as Unicode IPA string
            normalize (bool): whether the word should be pre-normalized

        Returns:
            unicode: longest single-segment prefix of `word` in database
        """
        if normalize:
            word = FeatureTable.normalize(word)
        last_found_length = 0
        node = self.seg_trie
        for pos in range(len(word) + 1):
            if pos == len(word) or word[pos] not in node:
                return word[:last_found_length]
            node = node[word[pos]]
            if self.TRIE_LEAF_MARKER in node:
                last_found_length = pos + 1
        return ''

    def ipa_segs(self, word: str, normalize: bool=True) -> list[str]:
        """Returns a list of segments from a word

        Args:
            word (unicode): input word as Unicode IPA string
            normalize (bool): whether to pre-normalize the word

        Returns:
            list: list of strings corresponding to segments found in `word`
        """
        if normalize:
            word = FeatureTable.normalize(word)
        return self._segs(word, include_invalid=False, normalize=normalize)

    def validate_word(self, word: str, normalize: bool=True):
        """Returns True if `word` consists exhaustively of valid IPA segments

        Args:
            word (unicode): input word as Unicode IPA string
            normalize (bool): whether to pre-normalize the word

        Returns:
            bool: True if `word` can be divided exhaustively into IPA segments
                  that exist in the database

        """
        return not self._segs(word, include_valid=False, include_invalid=True, normalize=normalize)

    def word_fts(self, word: str, normalize: bool=True):
        """Return a list of Segment objects corresponding to the segments in
           word.

        Args:
            word (unicode): word consisting of IPA segments
            normalize (bool): whether to pre-normalize the word

        Returns:
            list: list of Segment objects corresponding to word
        """
        return [self.fts(ipa, False) for ipa in self.ipa_segs(word, normalize)]

    def word_array(self, ft_names: list[str], word: str, normalize: bool=True) -> numpy.ndarray:
        """Return a ndarray of features namd in ft_name for the segments in word

        Args:
            ft_names (list): strings naming subset of features in self.names
            word (unicode): word to be analyzed
            normalize (bool): whether to pre-normalize the word

        Returns:
            ndarray: segments in rows, features in columns as [-1, 0, 1]
        """
        return numpy.array([s.numeric(ft_names) for s in self.word_fts(word, normalize)])

    def bag_of_features(self, word: str, normalize: bool=True) -> numpy.ndarray:
        """Return a vector in which each dimension is the number of times a feature-value pair occurs in the word
        
        Args:
            word (unicode): word consisting of IPA segments
            normalize (bool): whether to pre-normalize the word

        Returns:
            array: array of integers corresponding to a bag of feature-value pair counts
        """
        word_features = self.word_fts(word, normalize)
        features = [v + f for f in self.names for v in ['+', '0', '-']]
        bag = collections.OrderedDict()
        for f in features:
            bag[f] = 0
        vdict = {-1: '-', 0: '0', 1: '+'}
        for w in word_features:
            for (f, v) in w.items():
                bag[vdict[v] + f] += 1
        return numpy.array(list(bag.values()))

    def seg_known(self, segment: str, normalize: bool=True) -> bool:
        """Return True if `segment` is in segment <=> features database

        Args:
            segment (unicode): consonant or vowel
            normalize (bool): whether to pre-normalize the segment

        Returns:
            bool: True, if `segment` is in the database
        """
        if normalize:
            segment = FeatureTable.normalize(segment)
        return segment in self.seg_dict

    def segs_safe(self, word: str, normalize: bool=True):
        """Return a list of segments (as strings) from a word

        Characters that are not valid segments are included in the list as
        individual characters.

        Args:
            word (unicode): word as an IPA string
            normalize (bool): whether to pre-normalize the word

        Returns:
            list: list of Unicode IPA strings corresponding to segments in
                  `word`
        """
        if normalize:
            word = FeatureTable.normalize(word)
        return self._segs(word, include_invalid=True, normalize=normalize)

    def _segs(self, word: str, *, include_valid: bool=True, include_invalid: bool, normalize: bool=True) -> list[str]:
        if normalize:
            word = FeatureTable.normalize(word)
        segs = []
        while word:
            m = self.longest_one_seg_prefix(word, False)
            if m:
                if include_valid:
                    segs.append(m)
                word = word[len(m):]
            else:
                if include_invalid:
                    segs.append(word[0])
                word = word[1:]
        return segs

    def filter_segs(self, segs: list[str], normalize: bool=True) -> list[str]:
        """Given list of strings, return only those which are valid segments

        Args:
            segs (list): list of IPA Unicode strings
            normalize (bool): whether to pre-normalize the segments

        Return:
            list: list of IPA Unicode strings identical to `segs` but with
                  invalid segments filtered out
        """
        return list(filter(lambda seg: self.seg_known(seg, normalize), segs))

    def filter_string(self, word: str, normalize: bool=True) -> str:
        """Return a string like the input but containing only legal IPA segments

        Args:
            word (unicode): input string to be filtered
            normalize (bool): whether to pre-normalize the word (and return a normalized string)

        Returns:
            unicode: string identical to `word` but with invalid IPA segments
                     absent

        """
        return ''.join(self.ipa_segs(word, normalize))

    def fts_intersection(self, segs: list[str], normalize: bool=True) -> Segment:
        """Return a Segment object containing the features shared by all segments

        Args:
            segs (list): IPA segments
            normalize (bool): whether to pre-normalize the segments

        Returns:
            Segment: the features shared by all segments in segs
        """
        return reduce(lambda a, b: a & b,
                      [self.fts(s, normalize) for s in self.filter_segs(segs, normalize)])

    def fts_match_all(self, fts: dict[str, int], inv: list[str], normalize: bool=True) -> bool:
        """Return `True` if all segments in `inv` matches the features in fts

        Args:
            fts (dict): a dictionary of features
            inv (list): a collection of IPA segments represented as Unicode
                        strings
            normalize (bool): whether to pre-normalize the segments

        Returns:
            bool: `True` if all segments in `inv` match the features in `fts`
        """
        return all([self.fts(s, normalize) >= fts for s in inv])

    def fts_match_any(self, fts: dict[str, int], inv: list[str], normalize: bool=True) -> bool:
        """Return `True` if any segments in `inv` matches the features in fts

        Args:
            fts (dict): a dictionary of features
            inv (list): a collection of IPA segments represented as Unicode
                        strings
            normalize (bool): whether to pre-normalize the segments

        Returns:
            bool: `True` if any segments in `inv` matches the features in `fts`
        """
        return any([self.fts(s, normalize) >= fts for s in inv])

    def fts_contrast(self, fs: dict[str, int], ft_name: str, inv: list[str], normalize: bool=True) -> bool:
        """Return `True` if there is a segment in `inv` that contrasts in feature
        `ft_name`.

        Args:
            fs (dict): feature specifications used to filter `inv`.
            ft_name (str): name of the feature where contrast must be present.
            inv (list): collection of segments represented as Unicode strings.
            normalize (bool): whether to pre-normalize the segments

        Returns:
            bool: `True` if two segments in `inv` are identical in features except
                  for feature `ft_name`
        """
        inv_segs = filter(lambda x: x >= fs, map(lambda seg: self.fts(seg, normalize), inv))
        for a in inv_segs:
            for b in inv_segs:
                if a != b:
                    if a.differing_specs(b) == [ft_name]:
                        return True
        return False

    def fts_count(self, fts: dict[str, int], inv: list[str], normalize: bool=True) -> int:
        """Return the count of segments in an inventory matching a given
        feature mask.

        Args:
            fts (dict): feature mask given as a set of (value, feature) tuples
            inv (list): inventory of segments (as Unicode IPA strings)
            normalize (bool): whether to pre-normalize the segments

        Returns:
            int: number of segments in `inv` that match feature mask `fts`
        """
        return len(list(filter(lambda s: self.fts(s, normalize) >= fts, inv)))

    def match_pattern(self, pat: list[str], word: str, normalize: bool=True) -> list[dict[str, int]]:
        """Implements fixed-width pattern matching.

        Matches just in case pattern is the same length (in segments) as the
        word and each of the segments in the pattern is a featural subset of the
        corresponding segment in the word. Matches return the corresponding list
        of feature sets; failed matches return None.

        Args:
           pat (list): pattern consisting of a sequence of feature dicts
           word (unicode): a Unicode IPA string consisting of zero or more
                           segments
           normalize (bool): whether to pre-normalize the word

        Returns:
            list: corresponding list of feature dicts or, if there is no match,
                  None
        """
        segs = self.word_fts(word, normalize)
        if len(pat) != len(segs):
            return None
        else:
            if all([s >= p for (s, p) in zip(segs, pat)]):
                return segs

    def match_pattern_seq(self, pat, const, normalize=True):
        """Implements limited pattern matching. Matches just in case pattern is
        the same length (in segments) as the constituent and each of the
        segments in the pattern is a featural subset of the corresponding
        segment in the word.

        Args:
            pat (list): pattern consisting of a list of feature dicts, e.g.
                        [{'voi': 1}]
            const (list): a sequence of Unicode IPA strings consisting of zero
                          or more segments.
            normalize (bool): whether to pre-normalize the segments

        Returns:
            bool: `True` if `const` matches `pat`
        """
        segs = [self.fts(s, normalize) for s in const]
        if len(pat) != len(segs):
            return False
        else:
            return all([s >= p for (s, p) in zip(segs, pat)])

    def all_segs_matching_fts(self, ft_mask):
        """Return segments matching a feature mask, a dict of features

        Args:
            ft_mask (list): feature mask dict, e.g. {'voi': -1, 'cont': 1}.

        Returns:
            list: segments matching `ft_mask`, sorted in reverse order by length
        """
        matching_segs = [ipa for (ipa, fts) in self.segments if fts >= ft_mask]
        return sorted(matching_segs, key=lambda x: len(x), reverse=True)

    def compile_regex_from_str(self, pat):
        """Given a string describing features masks for a sequence of segments,
        return a compiled regex matching the corresponding strings.

        Args:
            pat (str): feature masks, each enclosed in square brackets, in
            which the features are delimited by any standard delimiter.

        Returns:
           Pattern: regular expression pattern equivalent to `pat`
        """
        s2n = {'-': -1, '0': 0, '+': 1}
        seg_res = []
        for mat in re.findall(r'\[[^]]+\]+', pat):
            ft_mask = {k: s2n[v] for (v, k) in re.findall(r'([+-])(\w+)', mat)}
            segs = self.all_segs_matching_fts(ft_mask)
            seg_res.append('({})'.format('|'.join(segs)))
        regexp = ''.join(seg_res)
        return re.compile(regexp)

    def segment_to_vector(self, seg, normalize=True):
        """Given a Unicode IPA segment, return a list of feature specificiations
        in canonical order.

        Args:
            seg (unicode): IPA consonant or vowel
            normalize: whether to pre-normalize the segment

        Returns:
            list: feature specifications ('+'/'-'/'0') in the order from
            `FeatureTable.names`
        """
        return self.fts(seg, normalize).strings()

    def standardize_tones(self, word, nonstandard_tones=['¹','²','³','⁴','⁵']):
        standard_tones = ['˩', '˨', '˧', '˦', '˥']
        tone_map = dict(zip(nonstandard_tones, standard_tones))
        standardized_word = ''.join(tone_map.get(char, char) for char in word)
        return standardized_word

    def word_to_vector_list(self, word, numeric=False, xsampa=False, nonstandard_tones=['¹','²','³','⁴','⁵'], normalize=True):
        """Return a list of feature vectors, given a Unicode IPA word.

        Args:
            word (unicode): string in IPA (or X-SAMPA, provided `xsampa` is True)
            numeric (bool): if True, return features as numeric values instead
                            of strings
            xsampa (bool): whether the word is in X-SAMPA instead of IPA
            normalize: whether to pre-normalize the word (applies to IPA only)
            nonstandard_tones (list): list of 5 nonstandard tones to be conveted
                            to IPA tone markers.
                            The order and numbering of the tones can be changed to reflect data.
        Returns:
            list: a list of lists of '+'/'-'/'0' or 1/-1/0
        """
        if xsampa:
            word = self.xsampa.convert(word)
        if nonstandard_tones:
            word=self.standardize_tones(word,nonstandard_tones)
        segs = self.word_fts(word, normalize or xsampa)

        if numeric:
            tensor = [x.numeric() for x in segs]
        else:
            tensor = [x.strings() for x in segs]
        return tensor

    def _compare_vectors(self,vector1, vector2):
        """Compare two feature vectors digit by digit.

        Args:
            vector1 (list): First vector to compare.
            vector2 (list): Second vector to compare.

        Returns:
            int: -1 if vector1 < vector2, 1 if vector1 > vector2, 0 if they are equal.
        """
        for v1, v2 in zip(vector1, vector2):
            if v1 < v2:
                return -1
            elif v1 > v2:
                return 1
        return 0  # Vectors are equal

    def _binary_search(self, segment_list, target, fuzzy_search=False):
        """Binary search to find the segment matching the target vector.

        Args:
            segment_list (list): List of segments where each segment is a tuple (IPA, feature vector).
            target (list): Target feature vector to search for.
            fuzzy_search (bool): whether to search for the closest vector match if an exact match is not found.
                If disabled and an exact match is not found, a None value is returned.

        Returns:
            str: The IPA segment matching the target vector, or None if not found.
        """
        low, high = 0, len(segment_list) - 1
        best_match_index = None

        while low <= high:
            mid = (low + high) // 2
            word_vec = self.sorted_segments.segment_key(segment_list[mid])
            comparison = self._compare_vectors(word_vec, target)
            if comparison == 0:
                best_match_index = mid
                break
            elif comparison < 0:
                low = mid + 1
            else:
                high = mid - 1

        if best_match_index is None and fuzzy_search:
            # Used for fuzzy searching
            best_match_index = mid

        if best_match_index is not None:
            # Check neighboring rows within the range of +-5
            best_match = segment_list[best_match_index]
            for offset in range(-9, 5):
                neighbor_index = best_match_index + offset
                if 0 <= neighbor_index < len(segment_list):
                    neighbor_segment = segment_list[neighbor_index]
                    if not self._compare_vectors(self.sorted_segments.segment_key(neighbor_segment),target):
                        # Check if the neighbor segment has a shorter name
                        if len(neighbor_segment[0]) < len(best_match[0]):
                            best_match = neighbor_segment
            return best_match[0]

        return None

    def vector_list_to_word(self, tensor, xsampa=False,fuzzy_search=False):
        """Return a Unicode IPA word, given a list of feature vectors.

        Args:
            tensor (list): a list of lists of '+'/'-'/'0' or 1/-1/0
            xsampa (bool): whether to return the word in X-SAMPA instead of IPA
            fuzzy_search (bool): whether to search for the closest vector match if an exact match is not found.
                If disabled and an exact match is not found, a `ValueError` is raised.
        Returns:
            unicode: string in IPA (or X-SAMPA, provided `xsampa` is True)
        """

        word = ""
        for vector in tensor:
            match = self._binary_search(self.sorted_segments.segments, vector, fuzzy_search)
            if match:
                word += match
            else:
                raise ValueError(f"No matching segment found for vector: {vector}")
        if xsampa:
            word = self.xsampa.convert(word)

        return word
