import ast
import json

import pandas as pd

from ..core import helpers


def _deserialize_if_needed(cells, cell_data_type="list"):
    """
    Ensure cell data is a list of dictionaries or a dictionary.

    Args:
        cells (str, list, or dict): The cell data to deserialize.
        cell_data_type (str): The expected type of the cell data. One of ["list","dict"]. Defaults to "list".

    Returns:
        Out: "list" or "dict". The deserialized cell data.
    """
    if cell_data_type == "list":
        if isinstance(cells, str):
            try:
                # Try JSON first
                cells = json.loads(cells)
            except json.JSONDecodeError:
                try:
                    # Fallback to Python literal eval
                    cells = ast.literal_eval(cells)
                except (ValueError, SyntaxError):
                    return []  # Return empty list on failure

        return cells if isinstance(cells, list) else []

    elif cell_data_type == "dict":
        if isinstance(cells, str):
            try:
                # Try JSON first
                cells = json.loads(cells)
            except json.JSONDecodeError:
                try:
                    # Fallback to Python literal eval
                    cells = ast.literal_eval(cells)
                except (ValueError, SyntaxError):
                    return {}  # Return empty dict on failure

        return cells if isinstance(cells, dict) else {}

    else:
        raise ValueError(f"Unknown cell_data_type '{cell_data_type}'")


def _normalize_angle(angle):
    """
    Earlier versions of the task could potentially use both 0 or 180 as the angles for a completely horizontal line, but did not recognize them as equivalent when scoring.
    This function normalizes the angle by returning 0 if the angle value is 180, or the inputted value if any other angle. This is needed prior to re-scoring older data.
    Not needed if data comes from task version 1.2.0 or above.

    Args:
        angle (int/float): The angle value

    Returns:
        int/float: Normalized angle (180 -> 0)
    """
    return 0 if angle == 180 else angle


def score_accuracy(row):
    """
    Score trial correctness using stimuli_array and response_array with 180/0 bug fix.
    This fixes the bug where 180 and 0 are treated as different when they should be equivalent.
    Making sure that order doesn't matter, only that the contents match.

    Args:
        row (pd.Series): A single trial's data.

    Returns:
        bool: True if response matches stimulus (accounting for 180/0 bug), False otherwise.
    """
    try:
        # Get stimuli and response arrays
        stimuli_array = _deserialize_if_needed(
            row["stimuli_array"], cell_data_type="list"
        )
        response_array = _deserialize_if_needed(
            row["response_array"], cell_data_type="list"
        )

        if not stimuli_array or not response_array:
            return None

        # Check and fix both arrays (180 -> 0) and convert to sets for comparison
        normalized_stimuli = set(_normalize_angle(x) for x in stimuli_array)
        normalized_response = set(_normalize_angle(x) for x in response_array)

        # Check if sets match (order doesn't matter)
        return normalized_stimuli == normalized_response

    except Exception as e:
        print(f"Error in score_response_correct_fixed: {e}")
        return None


def score_response_array_matches(row):
    """
    Count how many responses in response_array match any stimulus in stimuli_array.
    Accounts for the 180/0 equivalency bug.

    Args:
        row (pd.Series): A single trial's data.

    Returns:
        int: Number of responses that match any stimulus. 0 if neither response matches stimuli, 1 if at least one response matches, 2 if all responses match (trial correct).
    """
    try:
        # Get stimuli and response arrays
        stimuli_array = _deserialize_if_needed(
            row["stimuli_array"], cell_data_type="list"
        )
        response_array = _deserialize_if_needed(
            row["response_array"], cell_data_type="list"
        )

        if not stimuli_array or not response_array:
            return None

        # Check and fix arrays (180 -> 0)
        normalized_stimuli = set(_normalize_angle(x) for x in stimuli_array)
        normalized_response = [_normalize_angle(x) for x in response_array]

        # Count matches
        matches = sum(
            1 for response in normalized_response if response in normalized_stimuli
        )
        return matches

    except Exception as e:
        print(f"Error in score_response_array_matches: {e}")
        return None


def score_first_tap_rt(row):
    """
    Extract the response time to first tap from actions_array.

    Args:
        row (pd.Series): A single trial's data.

    Returns:
        int/float: elapsed_time_ms to first tap/response, or None if no actions.
    """
    try:
        # Get actions array
        actions_array = _deserialize_if_needed(
            row["actions_array"], cell_data_type="list"
        )

        if not actions_array:
            return None

        # Return elapsed_time_ms from first action
        first_action = actions_array[0]
        return first_action.get("elapsed_time_ms")

    except Exception as e:
        print(f"Error in score_first_tap_response_time: {e}")
        return None


def summarize(x, trials_expected=3, rt_outlier_low=100, rt_outlier_high=10000):
    """
    Summarizes Jolo task performance.

    Args:
        x (pd.DataFrame): Trial-level scored dataset.
        trials_expected (int): Number of trials expected. Defaults to 3.
        rt_outlier_low (int): Threshold for Low RT Outliers in milliseconds. Defaults to 100.
        rt_outlier_high (int): Threshold for High RT Outliers in milliseconds. Defaults to 10000.

    Returns:
        pd.Series: Summary statistics.
    """
    # ABSTRACTION TO APPEAR IN EACH SCORING SCRIPT
    d = helpers.summarize_common_metadata(x, trials_expected)

    # Trial accuracy
    d["n_trials_correct"] = sum(x["metric_accuracy"] == True)
    d["n_trials_incorrect"] = sum(x["metric_accuracy"] == False)

    # Response-level accuracy (sum across all trials)
    d["n_responses_total"] = x.apply(
        lambda row: len(
            _deserialize_if_needed(row["response_array"], cell_data_type="list")
        ),
        axis=1,
    ).sum()
    d["n_responses_correct"] = x["metric_response_array_matches"].sum()
    d["n_responses_incorrect"] = d["n_responses_total"] - d["n_responses_correct"]

    # Response Times - Handle null/invalid response times first
    valid_rt_mask = pd.notna(x["response_time_ms"]) & (x["response_time_ms"] > 0)
    valid_rt_data = x.loc[valid_rt_mask, "response_time_ms"]

    # Number of trials that had null/invalid response times
    d["n_trials_rt_invalid"] = x["response_time_ms"].shape[0] - valid_rt_data.shape[0]

    # Overall Response Times (all trials)
    d["median_response_time_all"] = valid_rt_data.median()
    d["sd_response_time_all"] = valid_rt_data.std()

    # Response time filtered overall
    rt_filtered_overall = x.loc[
        valid_rt_mask
        & (x["response_time_ms"] >= rt_outlier_low)
        & (x["response_time_ms"] <= rt_outlier_high),
        "response_time_ms",
    ]
    d["median_response_time_all_filtered"] = rt_filtered_overall.median()
    d["std_response_time_all_filtered"] = rt_filtered_overall.std()

    # First tap response time
    valid_first_tap_rt_mask = pd.notna(x["metric_first_tap_rt"]) & (
        x["metric_first_tap_rt"] > 0
    )
    valid_first_tap_rt_data = x.loc[valid_first_tap_rt_mask, "metric_first_tap_rt"]

    d["median_first_tap_rt"] = valid_first_tap_rt_data.median()
    d["std_first_tap_rt"] = valid_first_tap_rt_data.std()

    # First tap rt filtered
    first_tap_rt_filtered = x.loc[
        valid_first_tap_rt_mask
        & (x["metric_first_tap_rt"] >= rt_outlier_low)
        & (x["metric_first_tap_rt"] <= rt_outlier_high),
        "metric_first_tap_rt",
    ]
    d["median_first_tap_rt_filtered"] = first_tap_rt_filtered.median()
    d["std_first_tap_rt_filtered"] = first_tap_rt_filtered.std()

    # return as series
    return pd.Series(d)
