"""The `LIWC` module provides a Python interface to the *Linguistic Inquiry and Word Count* (LIWC) tool, allowing users to perform text analysis."""

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

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

# %% ../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
import re

import tempfile
from pathlib import Path
import platform

# %% ../nbs/00_core.ipynb 5
def to_windows_path(path: str) -> str:
    """Convert WSL path to Windows path using wslpath -w. Raises if impossible."""
    wslpath_bin = shutil.which("wslpath")
    if not wslpath_bin:
        raise FileNotFoundError("wslpath not found")
    res = subprocess.run([wslpath_bin, "-w", path], capture_output=True, text=True, check=True)
    return res.stdout.strip()

def _is_wsl() -> bool:
    """Return True if running inside WSL."""
    if platform.system() != "Linux":
        return False
    # check env var or /proc/version content
    if os.environ.get("WSL_DISTRO_NAME"):
        return True
    try:
        with open("/proc/version", "r", encoding="utf-8") as f:
            v = f.read()
        return "microsoft" in v.lower() or "wsl" in v.lower()
    except Exception:
        return False
    

# helper: detect Windows absolute paths (C:\ or C:/ or UNC \\)
def _looks_like_windows_path(p: str) -> bool:
    return bool(re.match(r'^[A-Za-z]:[\\/]', p)) or p.startswith('\\\\')

# %% ../nbs/00_core.ipynb 6
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 8
@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.

    """
    # preserve raw input strings
    input_file = rf'{input_file}'
    output_location = rf'{output_location}'

    # use absolute Python-side paths for reliable reading/writing
    abs_in = os.path.abspath(input_file)
    abs_out = os.path.abspath(output_location)

    # paths passed to the CLI. default to Python absolute paths.
    in_cli, out_cli = abs_in, abs_out

    # If running in WSL, convert to Windows paths (wslpath -w) when available.
    if _is_wsl():
        try:
            in_cli = to_windows_path(abs_in)
            out_cli = to_windows_path(abs_out)
        except Exception:
            # fallback: keep Linux absolute paths (may fail if LIWC is a Windows exe)
            in_cli, out_cli = abs_in, abs_out

    cmd_to_execute = [
        self.liwc_cli_path, "--mode", "wc", "--input", in_cli, 
        "--row-id-indices", row_id_indices, "--column-indices", column_indices, 
        "--output", out_cli, "--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.
    """

    # preserve raw input strings
    input_folder = rf'{input_folder}'
    output_location = rf'{output_location}'


    # use absolute Python-side paths for reliable reading/writing
    abs_in = input_folder if _looks_like_windows_path(input_folder) else os.path.abspath(input_folder)
    abs_out = output_location if _looks_like_windows_path(output_location) else os.path.abspath(output_location)

    # paths passed to the CLI. default to Python absolute paths.
    in_cli, out_cli = abs_in, abs_out

    # If running in WSL, convert to Windows paths (wslpath -w) when available.
    # Do not reconvert if already Windows-style paths.
    if _is_wsl():
        try:
            if not _looks_like_windows_path(abs_in):
                in_cli = to_windows_path(abs_in)
            if not _looks_like_windows_path(abs_out):
                out_cli = to_windows_path(abs_out)
        except Exception:
            # fallback: keep Linux absolute paths (may fail if LIWC is a Windows exe)
            in_cli, out_cli = abs_in, abs_out

    cmd_to_execute = [
        self.liwc_cli_path, "--mode", "wc", "--input", in_cli,
        "--output", out_cli, "--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.

    """
    with tempfile.TemporaryDirectory() as tmpdir:
        in_path = os.path.join(tmpdir, "temp_in.csv")
        out_path = os.path.join(tmpdir, "temp_out.csv")

        # écriture input (index + texte)
        text.reset_index().to_csv(in_path, index=False, encoding='utf-8')

        # chemins absolus côté Python (pour lecture après exécution)
        abs_in = os.path.abspath(in_path)
        abs_out = os.path.abspath(out_path)

        # chemins passés au CLI (par défaut les mêmes que Python)
        in_cli, out_cli = abs_in, abs_out

        # si demandé et si WSL, convertir en chemins Windows (wslpath -w)
        if _is_wsl():
            try:
                in_cli = to_windows_path(abs_in)
                out_cli = to_windows_path(abs_out)
            except Exception:
                # fallback : garder les chemins Linux/absolus
                in_cli, out_cli = abs_in, abs_out

        # appel du CLI (utilise in_cli/out_cli)
        self.analyze_csv(in_cli, out_cli, "1", "2", liwc_dict=liwc_dict)

        # lire les résultats via le chemin Python (abs_out)
        result_df = pd.read_csv(abs_out, encoding='utf-8').set_index('Row ID')

        if return_input:
            result_df.insert(0, 'text', text.values)

        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 16
# 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) # type: ignore
        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 18
#################################################################################
########################### 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",
                wsl_mode: bool = True,
               ) -> 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".
    wsl_mode : bool, optional
        Whether to convert paths for WSL. Defaults to True.

    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)
    
    # --- ONLY PATH HANDLING ADDED FOR WSL ---
    # compute Python-side absolute paths (used for reading/cleanup)
    abs_in = os.path.abspath(input_path)
    abs_out_dir = os.path.abspath(output_dir)

    # default CLI-visible paths are the Python abs paths
    in_cli, out_cli = abs_in, abs_out_dir

    # If running in WSL and wsl_mode requested, convert paths to Windows form for the CLI
    if _is_wsl():
        try:
            in_cli = to_windows_path(abs_in)
            out_cli = to_windows_path(abs_out_dir)
        except Exception:
            # fallback: keep Linux absolute paths (CLI may fail if it's a Windows exe)
            in_cli, out_cli = abs_in, abs_out_dir
    # --- END PATH HANDLING ---

    # Prepare the command
    cmd_to_execute = [
        self.liwc_cli_path, "--mode", "lsm", "--input", in_cli, "--output", out_cli, "--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 = {}
    # Read results from the Python-side output directory (abs_out_dir)
    if calculate_lsm_value in [1, 3]:
        speaker_file = os.path.join(abs_out_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(abs_out_dir, output_files[2])
        if os.path.exists(group_file):
            result['group_level'] = pd.read_csv(group_file)
    
    # Clean up temporary files and directory (use Python-side paths)
    try:
        os.remove(abs_in)
    except Exception:
        pass
    try:
        shutil.rmtree(abs_out_dir)
    except Exception:
        pass
    
    # 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 22
#################################################################################
###########################      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

