# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_core.ipynb.

# %% auto 0
__all__ = ['Liwc', 'analyze_df2']

# %% ../nbs/00_core.ipynb 3
from fastcore.utils import patch
import matplotlib.pyplot as plt
import matplotlib.cm as cm
from typing import Union
import pandas as pd
import subprocess
import random
import shutil
import json
import os

# %% ../nbs/00_core.ipynb 5
class Liwc:
    def __init__(self, 
                 liwc_cli_path: str = 'LIWC-22-cli', # LIWC CLI Path.
                 threads: Union[int, None] = None, #Number of threads to use. Defaults to the number of CPU cores minus one.
                 verbose: bool = False): # Display printing and progress bar. Defaults to False.
        """
        Initialize the LIWC Class.
            
        """
        self.liwc_cli_path = liwc_cli_path
        self.threads = str(threads) if threads is not None else str(os.cpu_count() - 1)
        self.verbose = verbose
        
    def _execute_command(self, cmd_to_execute, verbose: Union[bool, None] = None):
        """
        Execute a command and handle verbose output.

        Args:
            cmd_to_execute (list): Command to execute as a list of strings.
            verbose (bool, optional): Override the instance's verbose setting.
        """
        
        if verbose is None:
            verbose = self.verbose
        if self.verbose: # verbose =True
            subprocess.call(cmd_to_execute)
        else: # verbose False
            result = subprocess.run(cmd_to_execute, capture_output=True, text=True)
            # if not result.stdout.strip().endswith('temp_out.csv'): # for df do not print temp csv.
            #     last_line = result.stdout.strip().split('\n')[-1]
            #     print(last_line)


# %% ../nbs/00_core.ipynb 7
@patch
def analyze_csv(self:Liwc, 
                input_file: str, #Path to the input CSV file.
                output_location: str, # Path to save the analysis output.
                row_id_indices: str, # Indices of row IDs in the CSV.
                column_indices: str, # Indices of text columns in the CSV.
                liwc_dict: str = "LIWC22") -> None: # Dictionary to use for analysis. Defaults to "LIWC22".
    """
    Analyze text data from a CSV file using LIWC.

    """
    input_file = rf'{input_file}'
    output_location = rf'{output_location}'
    cmd_to_execute = [
        self.liwc_cli_path, "--mode", "wc", "--input", input_file, 
        "--row-id-indices", row_id_indices, "--column-indices", column_indices, 
        "--output", output_location, "--t", self.threads
    ]
    if liwc_dict != "LIWC22":
        cmd_to_execute.extend(["--dictionary", liwc_dict])
    self._execute_command(cmd_to_execute)

@patch
def analyze_folder(self:Liwc,
                   input_folder: str, # Path to the folder containing text files.
                   output_location: str, # Path to save the analysis output.
                   liwc_dict: str = "LIWC22") -> None: # Dictionary to use for analysis. Defaults to "LIWC22".
    """
    Analyze all text files in a folder using LIWC.

    """
    input_folder = rf'{input_folder}'
    output_location = rf'{output_location}'
    cmd_to_execute = [
        self.liwc_cli_path, "--mode", "wc", "--input", input_folder, 
        "--output", output_location, "--t", self.threads
    ]
    if liwc_dict != "LIWC22":
        cmd_to_execute.extend(["--dictionary", liwc_dict])

    self._execute_command(cmd_to_execute)


@patch
def analyze_df(self:Liwc, 
               text: pd.Series, # Pandas Series containing text data.
               return_input: bool = False, # Whether to return the input text with the output. Defaults to False.
               liwc_dict: str = "LIWC22" # Dictionary to use for analysis. Defaults to "LIWC22".
              ) -> pd.DataFrame: # pd.DataFrame: DataFrame containing the analysis results.
    """
    Analyze text data from a DataFrame using LIWC.

    """
    temp_csv_file_in = f"temp_in_{random.randint(0, 999)}.csv"
    temp_csv_file_out = f"temp_out_{random.randint(0, 999)}.csv"
    text.reset_index().to_csv(temp_csv_file_in, index=False, encoding='utf-8')
    self.analyze_csv(temp_csv_file_in, temp_csv_file_out, "1", "2", liwc_dict=liwc_dict)
    result_df = pd.DataFrame()
    if return_input:
        result_df['text'] = text.values
    result_df = pd.concat([result_df, pd.read_csv(temp_csv_file_out, encoding='utf-8')], axis=1).set_index('Row ID')
    os.remove(temp_csv_file_in)
    os.remove(temp_csv_file_out)
    return result_df
    
def analyze_df2(self:Liwc, 
               df: pd.DataFrame, # name of the dataframe.
               text_col: str, # Name of the column containing text data.
               index_col: str, # Index of the df to be returned
               return_input: bool = False, # Whether to return the input text with the output. Defaults to False.
               liwc_dict: str = "LIWC22" # Dictionary to use for analysis. Defaults to "LIWC22".
              ) -> pd.DataFrame: # pd.DataFrame: DataFrame containing the analysis results.
    """
    Analyze text data from a DataFrame using LIWC.

    """
    temp_csv_file_in = f"temp_in_{random.randint(0, 999)}.csv"
    temp_csv_file_out = f"temp_out_{random.randint(0, 999)}.csv"
    df[text_col].reset_index().to_csv(temp_csv_file_in, index=False, encoding='utf-8')
    self.analyze_csv(temp_csv_file_in, temp_csv_file_out, "1", "2", liwc_dict=liwc_dict)
    result_df = pd.DataFrame()
    if return_input:
        result_df[text_col] = df[text_col].values
        
    if index_col:
        result_df[index_col] = df[index_col].values
        
    result_df = pd.concat([result_df, pd.read_csv(temp_csv_file_out, encoding='utf-8')], axis=1).set_index('Row ID')
    os.remove(temp_csv_file_in)
    os.remove(temp_csv_file_out)
    return result_df
    
@patch
def analyze_string(self:Liwc, 
                   input_string: str, # The string to analyze.
                   output_location: str, # Path to save the analysis output (.csv).
                   liwc_dict: str = "LIWC22") -> None: # Dictionary to use for analysis. Defaults to "LIWC22".
    """
    Analyze a single string using LIWC and save to csv.

    """
    output_location = rf'{output_location}'
    cmd_to_execute = [
        self.liwc_cli_path, "--mode", "wc", "--input", "console", 
        "--console-text", input_string, "--output", output_location
    ]
    if liwc_dict != "LIWC22":
        cmd_to_execute.extend(["--dictionary", liwc_dict])
    subprocess.call(cmd_to_execute)
    
@patch
def analyze_string_to_json(self:Liwc, 
                           input_string: str, # The string to analyze.
                           liwc_dict: str = "LIWC22" # Dictionary to use for analysis. Defaults to "LIWC22".
                          ) -> dict: # Analysis results in JSON format.
    """
    Analyze a single string and return the result as JSON.

    Returns:
        dict: 
    """
    cmd_to_execute = [
        self.liwc_cli_path, "--mode", "wc", "--input", "console", 
        "--console-text", input_string, "--output", "console"
    ]
    
    if liwc_dict != "LIWC22":
        cmd_to_execute.extend(["--dictionary", liwc_dict])
    
    results = subprocess.check_output(cmd_to_execute).strip().splitlines()
    return json.loads(results[1].decode('utf-8'))

# %% ../nbs/00_core.ipynb 15
# local Dictionnaries
import collections

class _LIWCDictionary:
    def __init__(self, filepath):
        """
        Initializes the LIWCDictionary by loading the dictionary file and building the trie.
        
        :param filepath: Path to the LIWC dictionary file.
        """
        self.categories, self.lexicon = self._load_dict_file(filepath)
        self._trie = self._build_char_trie(self.lexicon)

    def search(self, word):
        """
        Search a word in the LIWC dictionary.
        
        :param word: The word to search for.
        :return: A list of the LIWC categories the word belongs to, or an empty list if not found.
        """
        return self._search_trie(self._trie, word)

    def parse(self, tokens):
        """
        Parses a document and extracts raw counts of words that fall into various LIWC categories.
        
        :param tokens: A list of tokens from a tokenized document.
        :return: A counter with the linguistic categories found in the document and the raw count of words in each category.
        """
        cat_counter = collections.Counter()

        for token in tokens:
            cats = self.search(token)
            for cat in cats:
                cat_counter[cat] += 1

        return cat_counter

    def _load_dict_file(self, filepath):
        """
        Loads the LIWC dictionary file and parses the categories and lexicon.
        
        :param filepath: Path to the LIWC dictionary file.
        :return: A tuple containing the categories dictionary and the lexicon dictionary.
        """
        try:
            with open(filepath, 'r') as liwc_file:
                categories = {}
                lexicon = {}
                percent_sign_count = 0

                for line in liwc_file:
                    stp = line.strip()
                    if stp:
                        parts = stp.split('\t')
                        if parts[0] == '%':
                            percent_sign_count += 1
                        else:
                            if percent_sign_count == 1:
                                categories[parts[0]] = parts[1]
                            else:
                                lexicon[parts[0]] = [categories[cat_id] for cat_id in parts[1:]]
        except FileNotFoundError:
            raise Exception(f"The file at {filepath} was not found.")
        except IOError:
            raise Exception(f"An error occurred while reading the file at {filepath}.")

        return categories, lexicon

    @staticmethod
    def _build_char_trie(lexicon):
        """
        Builds a character trie to handle wildcard ('*') matches.
        
        :param lexicon: The lexicon dictionary containing words and their categories.
        :return: A character trie for the lexicon.
        """
        trie = {}
        for pattern, cat_names in lexicon.items():
            cursor = trie
            for char in pattern:
                if char == '*':
                    cursor['*'] = cat_names
                    break
                if char not in cursor:
                    cursor[char] = {}
                cursor = cursor[char]
            cursor['$'] = cat_names
        return trie
    def get_categories(self):
        """
        Returns the dictionary of categories.
    
        :return: A dictionary of categories.
        """
        return self.categories

    
    
    def get_lexicon(self):
        """
        Returns the lexicon dictionary.
    
        :return: A dictionary of words and their associated LIWC categories.
        """
        return self.lexicon


    
    def get_trie(self):
        """
        Returns the character trie.
    
        :return: The character trie built from the lexicon.
        """
        return self._trie
    
    def word_in_categories(self, word, category_list):
        """
        Checks if a word belongs to any of the specified categories.
    
        :param word: The word to check.
        :param category_list: A list of categories to check against.
        :return: True if the word belongs to any of the specified categories, False otherwise.
        """
        word_categories = self.search(word)
        return any(cat in word_categories for cat in category_list)
        
    @staticmethod
    def _search_trie(trie, token, i=0):
        """
        Searches the given character trie for paths that match the token.
        
        :param trie: The character trie.
        :param token: The token to search for.
        :param i: The current index in the token.
        :return: A list of categories if the token matches, otherwise an empty list.
        """
        if '*' in trie:
            return trie['*']
        if '$' in trie and i == len(token):
            return trie['$']
        if i < len(token):
            char = token[i]
            if char in trie:
                return LIWCDictionary._search_trie(trie[char], token, i + 1)
        return []

# Example usage:
# liwc_dict = LIWCDictionary('path_to_liwc_file.dic')
# categories = liwc_dict.search('word')
# counts = liwc_dict.parse(['list', 'of', 'tokens'])


# %% ../nbs/00_core.ipynb 17
#################################################################################
########################### LANGUAGE STYLE MATCHING  ############################
#################################################################################

@patch
def analyze_lsm(self:Liwc, 
                df: pd.DataFrame, 
                calculate_lsm: str = "person-and-group", 
                group_column: str = 'GroupID', 
                person_column: str = 'PersonID', 
                text_column: str = 'Text', 
                output_type: str = "pairwise",
                expanded_output: bool = False, 
                omit_speakers_number_of_turns: int = 0, 
                omit_speakers_word_count: int = 10, 
                segmentation: str = "none"
               ) -> Union[pd.DataFrame, dict]:
    """
    Analyzes Linguistic Style Matching (LSM) based on the provided DataFrame.

    Parameters
    ----------
    df : pd.DataFrame
        Input DataFrame containing the text data to be analyzed.
    calculate_lsm : str, optional
        Sets the type of LSM calculation. Options are:
        - "person": Calculate only person-level LSM.
        - "group": Calculate only group-level LSM.
        - "person-and-group": Calculate both person and group-level LSM.
        Default is "person-and-group".
    group_column : str, optional
        The column name in `df` representing the Group ID. Default is 'GroupID'.
    person_column : str, optional
        The column name in `df` representing the Person ID. Default is 'PersonID'.
    text_column : str, optional
        The column name in `df` representing the text data. Default is 'Text'.
    output_type : str, optional
        Sets the type of output. Options are:
        - "one-to-many": One-to-many comparison.
        - "pairwise": Pairwise comparison.
        Default is "pairwise".
    expanded_output : bool, optional
        Adds an option to get an expanded LSM output. Default is False.
    omit_speakers_word_count : int, optional
        Omit speakers if the word count is less than this value. Default is 10.
    segmentation : str, optional
        Segmentation options for splitting the text. Options are:
        - "none": No segmentation.
        - "not=<number>": Number of turns per segment.
        - "nofst=<number>": Number of segments by speaker turn.
        - "nofwc=<number>": Number of segments by word count.
        - "now=<number>": Number of words per segment.
        - "boc=<characters>": Segmentation based on characters.
        - "regexp=<regex>": Segmentation based on a regular expression.
        Default is "none".

    Returns
    -------
    Union[pd.DataFrame, dict]
        The resulting LSM analysis. The output format depends on the specified `output_format`.
    """
    random_number = random.randint(0, 999)
    input_path = f'temp_input_{random_number}.csv'
    output_dir = 'temp_output_dir'
    
    # Map keywords to command line parameters
    calculate_lsm_map = {
        "person": 1,
        "group": 2,
        "person-and-group": 3
    }
    output_type_map = {
        "one-to-many": 1,
        "pairwise": 2
    }
    
    # Convert keywords to corresponding command line parameters
    calculate_lsm_value = calculate_lsm_map.get(calculate_lsm.lower(), 3)
    output_type_value = output_type_map.get(output_type.lower(), 1)
    
    # Map column names to indices
    group_column_idx = df.columns.get_loc(group_column) + 1
    person_column_idx = df.columns.get_loc(person_column) + 1
    text_column_idx = df.columns.get_loc(text_column) + 1
    
    # Write DataFrame to CSV file
    df.to_csv(input_path, index=False)
    
    # Prepare the command
    cmd_to_execute = [
        self.liwc_cli_path, "--mode", "lsm", "--input", input_path, "--output", output_dir, "--output-format", "csv",
        "--calculate-lsm", str(calculate_lsm_value), "--group-column", str(group_column_idx), "--person-column", str(person_column_idx),
        "--text-column", str(text_column_idx), "--output-type", str(output_type_value)
    ]
    
    if expanded_output:
        cmd_to_execute.append("--expanded-output")
    
    
    if omit_speakers_number_of_turns > 0:
        cmd_to_execute.extend(["--omit-speakers-number-of-turns", str(omit_speakers_number_of_turns)])
    
    if omit_speakers_word_count > 0:
        cmd_to_execute.extend(["--omit-speakers-word-count", str(omit_speakers_word_count)])
    
    if segmentation != "none":
        cmd_to_execute.extend(["--segmentation", segmentation])
    
    # Execute the command
    self._execute_command(cmd_to_execute)
    
    # Determine the appropriate output files based on output_type
    if output_type_value == 1:
        output_files = {
            1: "LSM-Speaker-One-to-many.csv",
            2: "LSM-Group-One-to-many.csv"
        }
    elif output_type_value == 2:
        output_files = {
            1: "LSM-Speaker-Pairwise.csv",
            2: "LSM-Group-Pairwise.csv"
        }
    
    result = {}
    if calculate_lsm_value in [1, 3]:
        speaker_file = os.path.join(output_dir, output_files[1])
        if os.path.exists(speaker_file):
            result['person_level'] = pd.read_csv(speaker_file)
    if calculate_lsm_value in [2, 3]:
        group_file = os.path.join(output_dir, output_files[2])
        if os.path.exists(group_file):
            result['group_level'] = pd.read_csv(group_file)
    
    # Clean up temporary files and directory
    os.remove(input_path)
    shutil.rmtree(output_dir)
    
    # if calculate_lsm == 3:
    #     return result
    # elif calculate_lsm == 1:
    #     return result.get('person_level')
    # elif calculate_lsm == 2:
    #     return result.get('group_level')
    return result


# %% ../nbs/00_core.ipynb 20
#################################################################################
###########################      Narrative arc       ############################
#################################################################################

@patch
def narrative_arc(self:Liwc, df: pd.DataFrame, column_names: Union[list, None] = None, 
                  output_individual_data_points: bool = True, scaling_method: str = '0-100', 
                  segments_number: int = 5, skip_wc: int = 10) -> pd.DataFrame:
    """
    Analyzes the narrative arc of text data based on the provided DataFrame.
    
    Parameters
    ----------
    df : pd.DataFrame
        Input DataFrame containing the text data to be analyzed.
    column_names : Union[list, None], optional
        List of column names in `df` that should be processed. If None, all columns are processed. Default is None.
    output_individual_data_points : bool, optional
        If True, outputs individual data points for each segment. If False, aggregates the data. Default is True.
    scaling_method : str, optional
        Method for scaling the data. Options are:
        - "0-100": Scale values between 0 and 100.
        - "Z-score": Scale values using Z-score normalization.
        Default is "0-100".
    segments_number : int, optional
        Number of segments into which the text is divided for analysis. Default is 5.
    skip_wc : int, optional
        Skip any texts with a word count less than this value. Default is 10.
    
    Returns
    -------
    pd.DataFrame
        The resulting DataFrame with the narrative arc analysis.
    """
    
    random_number = random.randint(0, 999)
    input_path = f'temp_input_{random_number}.csv'
    output_path = f'temp_output_{random_number}.csv'
    
    # Write DataFrame to CSV file
    df.to_csv(input_path, index=False)
    
    # Map column names to indices if column_names is not None
    if column_names is not None:
        column_indices = [str(df.columns.get_loc(col) + 1) for col in column_names]
        column_indices_str = ','.join(column_indices)
    else:
        column_indices_str = None
    
    # Map scaling method names to their respective values
    scaling_method_map = {
        '0-100': 1,
        'z-score': 2
    }
    
    if scaling_method.lower() not in scaling_method_map:
        raise ValueError("Invalid scaling method. Choose '0-100' or 'z-score'.")
    
    scaling_method_value = scaling_method_map[scaling_method.lower()]
    
    # Prepare the command
    cmd_to_execute = [
        self.liwc_cli_path, "--mode", "arc", "--input", input_path, "--output", output_path,
        "--output-format", "csv", 
        "--scaling-method", str(scaling_method_value), "--segments-number", str(segments_number), 
        "--skip-wc", str(skip_wc),
        "--output-individual-data-points", "yes" if output_individual_data_points else "no"
    ]
    
    if column_indices_str is not None:
        cmd_to_execute.extend(["--column-indices", column_indices_str])
    
    # Execute the command
    self._execute_command(cmd_to_execute)
    
    # Read the output into a DataFrame
    result_df = pd.read_csv(output_path)
    
    # Clean up temporary files
    os.remove(input_path)
    os.remove(output_path)
    return result_df
    
@patch
def _arc_get_segment_values(self:Liwc, df, prefix):
    """
    Function to extract segment columns and values for each row
    """
    segment_columns = [col for col in df.columns if col.startswith(prefix)]
    return df[segment_columns]


# Function to plot with an optional legend parameter
@patch
def plot_narrative_arc(self:Liwc, df: pd.DataFrame, legend_labels: list = None) -> plt.Figure:
    """
    Plots the narrative arc for the given DataFrame, showing Staging, Plot Progression, and Cognitive Tension.

    Parameters
    ----------
    df : pd.DataFrame
        Input DataFrame containing the narrative arc data.
        Note: 'output_individual_data_points=True' in narrative_arc to get all required data to plot the narractive arc.
        
    legend_labels : list, optional
        List of labels for the legend, corresponding to each row in the DataFrame.

    Returns
    -------
    plt.Figure
        The resulting plot figure of the narrative arcs.
    """
    
    # Extract values for each segment type
    staging_values = self._arc_get_segment_values(df, 'Staging_')
    plotprog_values = self._arc_get_segment_values(df, 'PlotProg_')
    cogtension_values = self._arc_get_segment_values(df, 'CogTension_')

    # Determine the number of segments
    num_segments = len(staging_values.columns)
    segments = range(1, num_segments + 1)

    # Check if the number of legend labels matches the number of rows
    if legend_labels and len(legend_labels) != len(df):
        raise ValueError("The number of legend labels must match the number of rows in the DataFrame")

    # Generate a colormap for each subplot
    colormaps = [cm.Blues, cm.Greens, cm.Reds]

    # Creating the subplots
    fig, axs = plt.subplots(1, 3, figsize=(15, 5))

    # Plotting Staging
    for i, (index, row) in enumerate(staging_values.iterrows()):
        color = colormaps[0]((i+0.1)*8 / len(df))
        axs[0].plot(segments, row, marker='o', color=color, linewidth=2, 
                    markeredgewidth=2, markerfacecolor='white', markeredgecolor=color, 
                    label=legend_labels[i] if legend_labels else None)
    axs[0].set_title('STAGING')
    axs[0].set_xlabel('Segment')
    axs[0].set_ylabel('')
    axs[0].set_ylim(-5, 110)  # Adding some padding to the bottom
    axs[0].set_xticks(segments)  # Ensuring x-ticks are whole numbers
    if legend_labels:
        axs[0].legend()

    # Plotting Plot Progression
    for i, (index, row) in enumerate(plotprog_values.iterrows()):
        color = colormaps[1]((i+0.1)*8 / len(df))
        axs[1].plot(segments, row, marker='o', color=color, linewidth=2, 
                    markeredgewidth=2, markerfacecolor='white', markeredgecolor=color, 
                    label=legend_labels[i] if legend_labels else None)
    axs[1].set_title('PLOT PROGRESSION')
    axs[1].set_xlabel('Segment')
    axs[1].set_ylabel('')
    axs[1].set_ylim(-5, 110)  # Adding some padding to the bottom
    axs[1].set_xticks(segments)  # Ensuring x-ticks are whole numbers
    if legend_labels:
        axs[1].legend()

    # Plotting Cognitive Tension
    for i, (index, row) in enumerate(cogtension_values.iterrows()):
        color = colormaps[2]((i+0.1)*8 / len(df))
        axs[2].plot(segments, row, marker='o', color=color, linewidth=2, 
                    markeredgewidth=2, markerfacecolor='white', markeredgecolor=color, 
                    label=legend_labels[i] if legend_labels else None)
    axs[2].set_title('COGNITIVE TENSION')
    axs[2].set_xlabel('Segment')
    axs[2].set_ylabel('')
    axs[2].set_ylim(-5, 110)  # Adding some padding to the bottom
    axs[2].set_xticks(segments)  # Ensuring x-ticks are whole numbers
    if legend_labels:
        axs[2].legend()

    # Adjusting layout
    plt.tight_layout()

    # Show plot
    #plt.show()
    plt.close()
    return fig

