import json
import time
from typing import Tuple, Optional, Union, Callable, List
import numpy as np
import pandas as pd
from tqdm import tqdm
from scipy.signal import lfilter


def _validate_columns(df: pd.DataFrame, columns: List[str]):
    """检查列是否存在"""
    missing = [c for c in columns if c not in df.columns]
    if missing:
        raise ValueError(f"Missing required columns: {missing}")

def prepare_data(
        data: pd.DataFrame,
        positions: Tuple[str, ...],
        source: str = 'source',
        target: str = 'target',
        interaction: str = 'interaction',
        weight: Optional[str] = None,
):
    """
    Prepare the input data for analysis by validating, sorting, and renaming columns.

    Args:
        data (pd.DataFrame): The input dataframe containing interaction data.
        positions (Tuple[str, ...]): Column names used for positional indexing.
        source (str, optional): Name of the source column. Defaults to 'source'.
        target (str, optional): Name of the target column. Defaults to 'target'.
        interaction (str, optional): Name of the interaction column. Defaults to 'interaction'.
        weight (Optional[str], optional): Name of the weight column, if any. Defaults to None.

    Returns:
        pd.DataFrame: The processed dataframe ready for analysis.

    Raises:
        ValueError: If required columns are missing or position columns are not numeric.
    """
    # make a copy, keep original data safe
    df = data.copy().reset_index(drop=True)
    # check required columns
    required_columns = list(positions) + [source, target]
    if weight: required_columns.append(weight)
    _validate_columns(df, required_columns)
    # examine data
    _check_numeric_position(df, positions)

    # rename columns
    rename_map = {source: 'source', target: 'target'}
    if interaction: rename_map[interaction] = 'interaction'
    if weight: rename_map[weight] = 'weight'
    df = df.rename(columns=rename_map)

    # Process 'source', 'target' columns to ensure lists
    df['source'] = _process_column(df['source'])
    df['target'] = _process_column(df['target'])

    if weight:
        df['weight'] = df['weight'].astype(float)
    df = _flatten_interactions(df)

    df = df.sort_values(by=list(positions)).reset_index(drop=True)

    return df

def segment(data,
            start=None,
            end=None,
            positions=None):
    """
    Extract a subset of interactions based on a specified interval.

    Args:
        data (pd.DataFrame): The dataframe containing interaction data.
        start (tuple, optional): The starting position of the interval. Defaults to None.
        end (tuple, optional): The ending position of the interval. Defaults to None.
        positions (tuple): Column names used for positional indexing. Required.

    Returns:
        pd.DataFrame: A dataframe containing interactions within the specified interval.

    Raises:
        ValueError: If position columns are not provided.
    """
    if positions is None:
        raise ValueError("Must provide 'positions' column names for segmentation.")

    _validate_columns(data, list(positions))
    df_temp = data.copy().set_index(list(positions)).sort_index()
    idx = pd.IndexSlice
    try:
        return df_temp.loc[idx[start:end], :].reset_index()
    except KeyError:
        raise KeyError(f"Interval {start}-{end} not found. Ensure data is sorted.")


def calculate_weights(data,
                      weight_func=lambda x: len(x)):
    """
    Calculate weights for each interaction based on a provided weight function.

    Args:
        data (pd.DataFrame): The dataframe containing interaction data.
        weight_func (Callable, optional): A function that takes an interaction entry and returns a numerical weight. Defaults to lambda x: len(x).

    Returns:
        pd.DataFrame: The dataframe with an added 'weight' column and flattened weights.
    """
    _validate_columns(data, ['interaction'])
    df = data.copy()
    # work on interactions
    unique_interactions = df['interaction'].unique()
    weight_map = {text: weight_func(text) for text in unique_interactions}
    df['weight'] = df['interaction'].map(weight_map)
    return _flatten_interactions(df)

def _flatten_interactions(df):
    """
    Expand the dataframe so that each source-target pair has its own row.

    This is useful when the 'source' and 'target' columns contain lists,
    ensuring that each interaction is represented as a single row.

    Args:
        df (pd.DataFrame): The dataframe with potential list-like 'source' and 'target' columns.

    Returns:
        pd.DataFrame: The exploded dataframe with individual source-target pairs.
    """
    return df.explode('source').explode('target')


def agg_weights(data,
                positions,
                agg_func: Union[str, Callable] = 'sum'):
    """
    Aggregate weights by grouping interactions based on positional columns and source-target pairs.

    Args:
        data (pd.DataFrame): The dataframe containing interaction data with weights.
        positions (Tuple[str, ...]): Column names used for positional indexing (excluding the line identifier).
        agg_func (Union[str, Callable], optional): The aggregation function to apply to the weights.
            Can be a string like 'sum', 'mean', 'max', 'min', or a callable. Defaults to 'sum'.

    Returns:
        pd.DataFrame: The aggregated dataframe with summed weights for each group.

    Raises:
        ValueError: If the 'weight' column is missing from the dataframe.
    """
    cols = list(positions) + ['source', 'target', 'weight']
    _validate_columns(data, cols)
    grouped = data.groupby(list(positions) + ['source', 'target'])['weight'].agg(agg_func)
    return grouped.reset_index()

def _process_column(series):
    """
    Process a dataframe column to ensure each cell contains a list of unique, stripped string items.

    Args:
        series (pd.Series): The column to process.

    Returns:
        pd.Series: The processed column with each cell as a list of unique, stripped strings.
    """
    def clean_cell(cell):
        if isinstance(cell, list):
            items = cell
        elif isinstance(cell, str):
            items = [item.strip() for item in cell.strip('[]').split(',')]
        else:
            items = [str(cell).strip()]
        # return np.unique([str(item).strip() for item in items])
        return np.unique([str(item).strip() for item in items if str(item).strip()])
    return series.apply(clean_cell)

# get accumulate data
def accumulate_weights(
        data: pd.DataFrame,
        decay: float = 0.0,
        decay_mode: str = 'exponential',  # 'exponential' | 'power_law' | 'linear'
        accum_mode: str = 'auto',  # 'auto' | 'recursive' | 'convolution'
        ltm_rate: float = 1.0,
        time: Optional[str] = None,
        positions: Optional[Tuple[str, ...]] = None
) -> pd.DataFrame:
    """
    Accumulate weights over time with optional decay and long-term memory.

    This function supports two calculation kernels (accum_mode):
    1. Recursive Mode ('recursive'): Calculates S_t = S_{t-1} * (1-decay) + X_t
       - Characteristics: Maintains "state memory", high computational efficiency O(N).
       - Applicable: Exponential decay.

    2. Convolution Mode ('convolution'): Calculates S_t = Sum(X_{t-k} * Kernel(k))
       - Characteristics: Superimposes independent decay of events, higher complexity O(N*Window).
       - Applicable: Power Law, Linear, or other non-Markovian decays, or scenarios requiring history modification.

    Args:
        data (pd.DataFrame): Input DataFrame containing interaction data.
        decay (float, optional): Decay rate between 0 and 1. Defaults to 0.0.
        decay_mode (str, optional): Shape of the decay function ('exponential', 'power_law', 'linear'). Defaults to 'exponential'.
        accum_mode (str, optional): Accumulation mode ('recursive', 'convolution', 'auto').
                                    'auto' selects the optimal algorithm based on decay_mode. Defaults to 'auto'.
        ltm_rate (float, optional): Long Term Memory retention rate, contribution from cumulative sum. Defaults to 1.0.
        time (Optional[str], optional): Name of the time column (Time Mode). Defaults to None.
        positions (Optional[Tuple[str, ...]], optional): List of column names defining position/stages (Position Mode). Defaults to None.

    Returns:
        pd.DataFrame: DataFrame with accumulated weights.

    Raises:
        ValueError: If required columns are missing or parameters are invalid.
    """
    if positions is None:
        raise ValueError("Must provide 'positions' columns.")

    pos_cols = list(positions)
    required = pos_cols + ['source', 'target', 'weight']
    if time:
        if time not in data.columns:
            raise ValueError(f"Time column '{time}' not found in data.")
        required.append(time)

    missing = [c for c in required if c not in data.columns]
    if missing:
        raise ValueError(f"Missing columns: {missing}")

    if decay_mode == 'exponential' and not (0 <= decay < 1):
        raise ValueError("Exponential decay must be in [0, 1).")
    if ltm_rate < 0:
        raise ValueError("ltm_rate must be >= 0.")

    if accum_mode == 'auto':
        accum_mode = 'recursive' if decay_mode == 'exponential' else 'convolution'
    
    valid_accum_modes = ['recursive', 'convolution', 'auto']
    if accum_mode not in valid_accum_modes:
        raise ValueError(f"Invalid accum_mode '{accum_mode}'. Must be one of {valid_accum_modes}.")
    
    if accum_mode == 'recursive' and decay_mode != 'exponential':
        raise ValueError(f"Conflict: 'recursive' supports only 'exponential'. Got '{decay_mode}'.")
        
    valid_decay_modes = ['exponential', 'power_law', 'linear']
    if decay_mode not in valid_decay_modes:
        raise ValueError(f"Invalid decay_mode '{decay_mode}'. Must be one of {valid_decay_modes}.")

    df = data.copy().reset_index(drop=True)
    temp_axis_col = '_calc_axis_'

    # Save original time/position mapping
    axis_metadata_map = None

    if time:
        if not pd.api.types.is_integer_dtype(df[time]):
            df[temp_axis_col] = df[time].astype(int)
        else:
            df[temp_axis_col] = df[time]
        axis_metadata_map = df[[temp_axis_col] + pos_cols].drop_duplicates(subset=[temp_axis_col])
    else:
        unique_pos = df[pos_cols].drop_duplicates().sort_values(by=pos_cols)
        unique_pos['_rank'] = np.arange(len(unique_pos))

        df = df.merge(unique_pos, on=pos_cols, how='left')
        df[temp_axis_col] = df['_rank']
        axis_metadata_map = unique_pos.rename(columns={'_rank': temp_axis_col})

    group_keys = ['source', 'target', temp_axis_col]
    df_agg = df.groupby(group_keys)['weight'].sum().reset_index()

    min_t, max_t = df_agg[temp_axis_col].min(), df_agg[temp_axis_col].max()
    full_axis_index = np.arange(min_t, max_t + 1)

    def process_pair(group):
        s_series = group.set_index(temp_axis_col)['weight']
        s_expanded = s_series.reindex(full_axis_index, fill_value=0.0)
        vals = s_expanded.values

        stm = np.zeros_like(vals, dtype=float)

        if decay == 0:
            stm = np.cumsum(vals)
        elif accum_mode == 'recursive':
            # Conflict check moved to start of function
            stm = lfilter([1.0], [1.0, -(1.0 - decay)], vals)
        elif accum_mode == 'convolution':
            lags = np.arange(len(vals))
            kernel = np.zeros_like(vals, dtype=float)
            if decay_mode == 'exponential':
                kernel = (1.0 - decay) ** lags
            elif decay_mode == 'power_law':
                kernel = (lags + 1.0) ** (-decay)
            elif decay_mode == 'linear':
                kernel = np.maximum(0.0, 1.0 - decay * lags)
            stm = np.convolve(vals, kernel, mode='full')[:len(vals)]

        final_weights = stm
        if ltm_rate > 0:
            final_weights += (np.cumsum(vals) * ltm_rate)

        # Fill calculation results back into Series
        s_expanded[:] = final_weights
        return s_expanded

    # --- 6. Parallel Application ---
    desc_str = f"Accumulating [Mode: {decay_mode} | Algo: {accum_mode}]"
    tqdm.pandas(desc=desc_str)

    # result_obj can be Series (Long) or DataFrame (Wide)
    result_obj = df_agg.groupby(['source', 'target'])[group_keys + ['weight']].progress_apply(process_pair)

    # --- 7. Restructure and Merge (Critical Fix!) ---

    # Check: If Pandas automatically converts it to a wide table (DataFrame), we need to stack it back.
    if isinstance(result_obj, pd.DataFrame):
        # Columns are 0, 1, 2... (time steps), Index is (source, target)
        # stack() converts to MultiIndex Series: (source, target, time)
        result_obj = result_obj.stack()

    # Now result_obj is guaranteed to be a Series.
    result_df = result_obj.reset_index()

    # Force column renaming to ensure safety.
    # At this point, column order must be: [source, target, time_axis, value]
    # Sometimes index has level_n, but it's usually clean after stack+reset
    if len(result_df.columns) == 4:
        result_df.columns = ['source', 'target', temp_axis_col, 'weight']
    else:
        raise KeyError(f"Unexpected columns structure after stack/reset: {result_df.columns}")

    # --- 8. Restore Metadata ---
    if axis_metadata_map is not None:
        df_final = result_df.merge(axis_metadata_map, on=temp_axis_col, how='left')
    else:
        df_final = result_df

    df_final = df_final.drop(columns=[temp_axis_col, '_rank'], errors='ignore')

    final_cols = ['source', 'target', 'weight'] + pos_cols
    if time and time in df_final.columns and time not in final_cols:
        final_cols.append(time)

    valid_cols = [c for c in final_cols if c in df_final.columns]

    return df_final[valid_cols].sort_values(by=pos_cols).reset_index(drop=True)

def _check_numeric_position(data, position):
    """
    Validate that all position columns contain numeric data types.

    Args:
        data (pd.DataFrame): The dataframe containing position columns.
        position (Tuple[str, ...]): Column names used for positional indexing.

    Raises:
        ValueError: If any position column does not have a numeric data type.
    """
    # check required columns
    for col in position:
        if not pd.api.types.is_numeric_dtype(data[col]):
            raise ValueError(f"Position column '{col}' must be numeric.")

# Infer listener LLM

def generate_prompt(scene_data):
    """
    Generate a formatted prompt for the language model to identify listeners in dialogue lines.

    Args:
        scene_data (Dict[str, Any]): A dictionary containing scene descriptions and dialogue lines.

    Returns:
        str: A formatted prompt string to be sent to the language model.
    """
    prompt = f"""
    Scene Description: {scene_data['Scene_Description']}

    Dialogue Lines:
    """
    for line in scene_data['Lines']:
        prompt += f"Line {line['Line_ID']} - {line['Speaker']}: \"{line['Line']}\"\n"

    prompt += """
    For each line, identify the listeners from the characters present. Provide the results in JSON format only, following this exact structure:

    {
      "listeners": {
        "1": ["Listener1", "Listener2"],
        "2": ["Listener3"],
        ...
      }
    }
    """

    # print(">> prompt:\n", prompt)

    return prompt


def infer_scene_listeners(client, model_name, prompt, *, max_tokens=200):
    """
    Use a language model to infer listeners for each dialogue line in a scene.

    Args:
        client: The API client for interacting with the language model (OpenAI-compatible).
        model_name (str): The name of the language model to use.
        prompt (str): The prompt string generated for the language model.
        max_tokens (int, optional): The maximum number of tokens for the response. Defaults to 200.

    Returns:
        Dict[str, List[str]]: A dictionary mapping line IDs to lists of listeners.

    Notes:
        - Ensures the response is valid JSON and adheres to the expected structure.
        - Implements error handling for JSON decoding and other exceptions.
    """
    listeners_json = ""
    listeners_data = {}
    try:
        response = client.chat.completions.create(
            # model="gpt-4o-mini",
            model=model_name,
            messages=[
                {"role": "system",
                    "content": "You are an assistant that identifies listeners for each line in a conversation. Respond only with JSON following the specified format."},
                {"role": "user", "content": prompt}
            ],
            max_tokens=max_tokens,
            n=1,
            stop=None,
            temperature=0.3
        )
        # print(">> response:\n", response)
        # print(">> response.choices[0].message.content:\n", response.choices[0].message.content)
        listeners_json = response.choices[0].message.content

        # Validate that the response is proper JSON and matches the expected structure
        listeners_data = json.loads(listeners_json)

        if "listeners" in listeners_data and isinstance(listeners_data["listeners"], dict):
            return listeners_data["listeners"]
        else:
            print("JSON response does not match the expected format.")
            return {}
    except json.JSONDecodeError:
        print("Failed to decode JSON. Response was:", listeners_json)
        return {}
    except Exception as e:
        print(f"Error inferring listeners for scene: {e}")
        return {}

def infer_listeners(data: pd.DataFrame,
                    position: Tuple[str, ...],
                    speaker: str = 'speaker',
                    dialogue: str = 'dialogue',
                    action: Optional[str] = None,
                    scene_description: Optional[str] = None,
                    client: Optional = None,
                    model: Optional[str] = None,
                    max_tokens: int = 200,
                    gap: float = 1.0,
                    base_url: Optional[str] = None,
                    api_key: Optional[str] = None
                    ) -> pd.DataFrame:
    """
    Infer listeners for each dialogue line across different scenes using a language model.

    Args:
        data (pd.DataFrame): The dataframe containing interaction data.
        position (Tuple[str, ...]): Column names used for positional indexing.
        speaker (str, optional): Name of the speaker column. Defaults to 'speaker'.
        dialogue (str, optional): Name of the dialogue column. Defaults to 'dialogue'.
        action (Optional[str], optional): Name of the action notes column. Defaults to None.
        scene_description (Optional[str], optional): Name of the scene description column. Defaults to None.
        client (Optional, optional): The API client for interacting with the language model. Defaults to None.
        model (Optional[str], optional): The name of the language model to use. Defaults to None.
        max_tokens (int, optional): The maximum number of tokens for the response. Defaults to 200.
        gap (float, optional): Delay between API requests to prevent rate limiting. Defaults to 1.0.
        base_url (str, optional): The base URL for the OpenAI API. Defaults to None.
        api_key (str, optional): The API key for the OpenAI API. Defaults to None.

    Returns:
        pd.DataFrame: The dataframe with an added 'listener' column containing inferred listeners.

    Raises:
        ValueError: If required columns are missing or position columns are not numeric.
    """
    if client is None:
        import openai
        client = openai.OpenAI(
            base_url=base_url,
            api_key=api_key
        )

    required_columns = list(position) + [speaker, dialogue]
    if not all(col in data.columns for col in required_columns):
        raise ValueError(f"""Missing required columns.\nExpected: {required_columns}\nFound: {data.columns.tolist()}""")

    _check_numeric_position(data, position)
    df = data.sort_values(by=list(position)).reset_index(drop=True)

    df = df.rename(columns={speaker: 'speaker'})
    if dialogue:
        df = df.rename(columns={dialogue: 'dialogue'})
    else:
        df['dialogue'] = None
    if action:
        df = df.rename(columns={action: 'action'})
    else:
        df['action'] = None
    if scene_description:
        df = df.rename(columns={scene_description: 'scene_description'})
    else:
        df['scene_description'] = None

    # Process 'source', 'target', 'observer' columns to ensure lists
    df['speaker'] = _process_column(df['speaker'])

    # print(df)
    # check data type of speaker column elements
    # print(df.iloc[0]['speaker'])
    # print(type(df.iloc[0]['speaker']))

    # return

    grouped = df.groupby(list(position[:-1]))
    scene_json = []
    df['listener'] = None
    # use tqdm for progress bar
    # for scene_number, group in grouped
    for scene_number, group in tqdm(grouped, desc='Processing scenes', unit='scene'):
        # print(">> scene_number:", scene_number)
        # print(">> group:\n", group)
        scene_data = {
            "Scene_Description": group['scene_description'].iloc[0],
            "Lines": [
                {
                    "Line_ID": int(row[position[-1]]),
                    "Speaker": list(row['speaker']),
                    "Line": row['dialogue'],
                    "Action_Notes": row['action']
                }
                for _, row in group.iterrows()
            ]
        }
        # print(">> scene_data before inference:\n", json.dumps(scene_data, indent=4))

        # Infer listeners for the entire scene
        prompt = generate_prompt(scene_data)
        listeners = infer_scene_listeners(client, model, prompt, max_tokens=max_tokens)

        # Assign listeners to each line
        for line in scene_data["Lines"]:
            line_id = line["Line_ID"]
            line["Listeners"] = listeners.get(str(line_id), [])
            # for row with column position[-1] == line_id, assign listeners
            df.loc[(df[list(position[:-1])].eq(scene_number)).all(axis=1) &
                   (df[position[-1]] == line_id), 'listener'] = str(line["Listeners"])

        # print(">> scene_data after inference:\n", json.dumps(scene_data, indent=4))
        scene_json.append(scene_data)
        time.sleep(gap)  # Adjust delay as needed
    return df

