"""
Core Woseco functionality for extracting concept associations from text.
"""
import spacy
import pandas as pd
from typing import List, Tuple, Optional
from tqdm import tqdm


class WosecoExtractor:
    """
    Main class for extracting WOrd SEntence COoccurrence associations.

    Woseco extracts noun-based concepts from consecutive sentences and identifies
    semantic relationships between concepts using similarity matching.
    """

    def __init__(self, language: str = "en", model_size: str = "lg"):
        """
        Initialize the Woseco extractor with a spacy language model.

        Args:
            language: Language code ('en' for English, 'de' for German)
            model_size: Spacy model size ('sm', 'md', 'lg')
        """
        self.language = language
        self.model_size = model_size
        self.nlp = self._load_model()

    def _load_model(self):
        """Load the appropriate spacy model based on language and size."""
        model_name = f"{self.language}_core_{'web' if self.language == 'en' else 'news'}_{self.model_size}"
        try:
            return spacy.load(model_name)
        except OSError:
            print(f"Model {model_name} not found. Downloading...")
            spacy.cli.download(model_name)
            return spacy.load(model_name)

    def get_concept_matches(
        self,
        concept_list: List[List[str]],
        similarity_threshold: Optional[float] = None
    ) -> List[str]:
        """
        Match concepts between consecutive sentences using semantic similarity.

        This function finds the most semantically similar noun pairs between consecutive
        sentences to create concept chains representing narrative flow.

        Args:
            concept_list: List of concept lists, one per sentence
            similarity_threshold: Minimum similarity score (0-1). If None, uses max similarity.

        Returns:
            List of linked concepts forming a chain through the narrative
        """
        if not concept_list or len(concept_list) < 2:
            return []

        chain = []

        # Add first concept from first sentence
        if concept_list[0]:
            chain.append(concept_list[0][0])

        # Match concepts between consecutive sentences
        for i in range(len(concept_list) - 1):
            current_concepts = concept_list[i]
            next_concepts = concept_list[i + 1]

            if not current_concepts or not next_concepts:
                continue

            max_similarity = -1
            best_match = None

            # Find pair with maximum similarity
            for current_concept in current_concepts:
                current_token = self.nlp(current_concept)

                for next_concept in next_concepts:
                    next_token = self.nlp(next_concept)

                    # Calculate semantic similarity using word embeddings
                    similarity = current_token.similarity(next_token)

                    if similarity > max_similarity:
                        max_similarity = similarity
                        best_match = next_concept

            # Add best matching concept if it meets threshold
            if best_match:
                if similarity_threshold is None or max_similarity >= similarity_threshold:
                    chain.append(best_match)

        # Add last concept from last sentence
        if concept_list[-1] and (not chain or chain[-1] != concept_list[-1][-1]):
            chain.append(concept_list[-1][-1])

        return chain

    def extract_nouns(self, text: str) -> List[str]:
        """
        Extract and lemmatize nouns from text.

        Args:
            text: Input text

        Returns:
            List of lemmatized nouns
        """
        doc = self.nlp(text)
        nouns = [token.lemma_ for token in doc if token.pos_ == "NOUN"]
        return nouns

    def get_woseco(
        self,
        df: pd.DataFrame,
        text_column: str = "Story",
        similarity_threshold: Optional[float] = None,
        show_progress: bool = True
    ) -> List[List[str]]:
        """
        Extract Woseco associations from a dataframe of stories.

        Args:
            df: DataFrame containing stories
            text_column: Name of column containing text
            similarity_threshold: Minimum similarity for concept matching
            show_progress: Whether to show progress bar

        Returns:
            List of concept chains, one per story
        """
        woseco_lists = []

        iterator = tqdm(df[text_column], disable=not show_progress, desc="Extracting Woseco")

        for story in iterator:
            # Split story into sentences by periods
            sentences = [s.strip() for s in story.split('.') if s.strip()]

            # Extract nouns from each sentence
            concept_list = []
            for sentence in sentences:
                nouns = self.extract_nouns(sentence)
                if nouns:
                    concept_list.append(nouns)

            # Match concepts between sentences
            concept_chain = self.get_concept_matches(concept_list, similarity_threshold)
            woseco_lists.append(concept_chain)

        return woseco_lists

    def process_csv(
        self,
        input_path: str,
        text_column: str = "Story",
        output_path: Optional[str] = None,
        similarity_threshold: Optional[float] = None
    ) -> pd.DataFrame:
        """
        Process a CSV file and extract Woseco associations.

        Args:
            input_path: Path to input CSV file
            text_column: Name of column containing text
            output_path: Optional path to save output CSV
            similarity_threshold: Minimum similarity for concept matching

        Returns:
            DataFrame with added 'woseco_concepts' column
        """
        df = pd.read_csv(input_path)

        woseco_lists = self.get_woseco(
            df,
            text_column=text_column,
            similarity_threshold=similarity_threshold
        )

        df['woseco_concepts'] = woseco_lists

        if output_path:
            df.to_csv(output_path, index=False)
            print(f"Saved results to {output_path}")

        return df
