"""
Line chart curriculum generator
"""

import os
import json
import copy
from typing import List, Dict
from metachart import MetaGenerator


class PieChartGenerator(MetaGenerator):
    def __init__(self, args, chart_id):
        super().__init__(args)
        self.chart_type = args.chart_type
        self.chart_id = chart_id
        self.all_qa_data_list = []
        self.round_num = 2
        self.qa_idx = 0
        self.detailed_reasoning = args.detailed_reasoning
        self.force_ground = args.force_ground
    
    
    ############################################################
    #   One-step Operator: h(list[Bar] | list[v]) → list[v]
    ############################################################
    
    def _one_step_statistics(self, chart_metadata: Dict):
        """
        Statistics: sum, mean, median, count
        For pie charts, we need to consider that the chart shows percentages,
        but we can ask about total values when given one slice's value.
        """
        pie_data_category_singular = chart_metadata['pie_data_category']['singular']
        pie_data_category_plural = chart_metadata['pie_data_category']['plural']
        pie_label_category_singular = chart_metadata['pie_label_category']['singular']
        pie_label_category_plural = chart_metadata['pie_label_category']['plural']
        
        pie_data_indices = [slice_idx for slice_idx in range(len(chart_metadata["pie_data"]))]
        target_pie_data = copy.deepcopy(chart_metadata["pie_data"])
        pie_labels = copy.deepcopy(chart_metadata["pie_labels"])

        # Read percentages (what's shown in the chart)
        total_data = sum(target_pie_data)
        pie_percentages = [(data/total_data)*100 for data in target_pie_data]
        read_percent_reason = '\n'.join([f"* {pie_labels[slice_idx]}: {pie_percentages[slice_idx]:.2f}%" for slice_idx in range(len(target_pie_data))]).strip()

        # Sum - needs to be calculated from given information
        sum_answer = self._compute_data_sum(target_pie_data)
        
        # Mean
        mean_answer = sum_answer / len(pie_data_indices)
        
        # Median
        median_value, median_indices = self._compute_data_median(target_pie_data)
        
        # Count
        count_answer = len(pie_data_indices)
        
        # Chart QA Pool
        easy_qa_pool = {}
        
        # For each slice, create a sum question that uses the given value to calculate total
        for k in range(len(target_pie_data)):
            kth_data = target_pie_data[k]
            kth_label = pie_labels[k]
            kth_percent = pie_percentages[k]
            
            easy_qa_pool[f"one_step__statistics__sum__given_slice_{k+1}"] = {
                "question": [
                    f"If given that the {pie_data_category_singular} of the {pie_label_category_singular} {kth_label} is {kth_data}, what is the total {pie_data_category_plural} of all {pie_label_category_plural} in this chart? Please round to two decimal places.",
                    f"Given that {kth_label} has {kth_data} {pie_data_category_singular}, what is the sum of {pie_data_category_plural} for all {pie_label_category_plural}? Please round to two decimal places.",
                    f"If {kth_label} represents {kth_data} {pie_data_category_singular}, calculate the total {pie_data_category_plural} for all {pie_label_category_plural} in this pie chart. Please round to two decimal places.",
                ],
                "reasoning": [
                    {
                        "step_1": f"First, I need to find the percentage of {kth_label} in the pie chart" + (f": {kth_percent:.2f}%" if self.detailed_reasoning else "."),
                        "step_2": f"Second, I need to calculate the total using the given value" + (f": if {kth_percent:.2f}% = {kth_data}, then 100% = {kth_data} \u00F7 {kth_percent:.2f}% \u00D7 100% = {sum_answer:.2f}" if self.detailed_reasoning else "."),
                    },
                ],
                "constraint": "none",
                "answer": sum_answer,
                "mask": {
                    "step_1": [k],
                    "step_2": [k] if self.force_ground else [],
                    "answer": pie_data_indices,
                },
            }

            easy_qa_pool[f"one_step__statistics__mean__given_slice_{k+1}"] = {
                "question": [
                    f"If given that the {pie_data_category_singular} of the {pie_label_category_singular} '{kth_label}' is {kth_data}, what is the mean {pie_data_category_singular} of all {pie_label_category_plural} in this chart? Please round to two decimal places.",
                    f"Given that {kth_label} has {kth_data} {pie_data_category_singular}, what is the average {pie_data_category_singular} per {pie_label_category_singular}? Please round to two decimal places.",
                ],
                "reasoning": [
                    {
                        "step_1": f"First, I need to find the percentage of {kth_label} in the pie chart" + (f": {kth_percent:.2f}%" if self.detailed_reasoning else "."),
                        "step_2": f"Second, I need to calculate the total {pie_data_category_singular}" + (f": if {kth_percent:.2f}% = {kth_data}, then 100% = {sum_answer}" if self.detailed_reasoning else "."),
                        "step_3": f"Third, I need to calculate the mean {pie_data_category_singular} of all {len(pie_data_indices)} {pie_data_category_plural}" + (f": {sum_answer} \u00F7 {len(pie_data_indices)} = {mean_answer:.2f}" if self.detailed_reasoning else "."),
                    },
                ],
                "constraint": "none",
                "answer": mean_answer,
                "mask": {
                    "step_1": [k],
                    "step_2": [k] if self.force_ground else [],
                    "step_3": pie_data_indices if self.force_ground else [],
                    "answer": pie_data_indices,
                },
            }

        # Regular percentage-based questions (no calculation of exact values needed)
        easy_qa_pool["one_step__statistics__count"] = {
            "question": [
                f"How many {pie_label_category_plural} are included in this pie chart?",
                f"What is the number of {pie_label_category_plural} shown in this chart?",
                f"Please help count the total number of {pie_label_category_plural} displayed in this pie chart.",
            ],
            "reasoning": [
                {
                    "step_1": f"First, I need to count the number of slices in this pie chart.",
                },
            ],
            "constraint": "none",
            "answer": count_answer,
            "mask": {
                "step_1": pie_data_indices,
                "answer": pie_data_indices,
            },
        }

        return easy_qa_pool
    
    
    ############################################################
    #                     Two-step Operator
    ############################################################

    def _two_step_statistics(self, chart_metadata: Dict, target_slice_indices: List, constraint: str):
        """
        Statistics: sum, mean, median, count for constrained slices
        """
        target_slice_indices.sort()
        pie_data_category_singular = chart_metadata['pie_data_category']['singular']
        pie_data_category_plural = chart_metadata['pie_data_category']['plural']
        pie_label_category_singular = chart_metadata['pie_label_category']['singular']
        pie_label_category_plural = chart_metadata['pie_label_category']['plural']
        
        target_pie_data = [chart_metadata["pie_data"][slice_idx] for slice_idx in target_slice_indices]
        target_pie_labels = [chart_metadata["pie_labels"][slice_idx] for slice_idx in target_slice_indices]

        # Calculate percentages for target slices
        total_data = sum(chart_metadata["pie_data"])
        target_percentages = [(chart_metadata["pie_data"][slice_idx]/total_data)*100 for slice_idx in target_slice_indices]
        
        # Read
        read_percent_reason = '\n'.join([f"* {target_pie_labels[i]}: {target_percentages[i]:.2f}%" for i in range(len(target_pie_labels))]).strip()
        read_slice_reason = ', '.join(target_pie_labels).strip()
        
        # Sum of percentages
        sum_percentage = sum(target_percentages)
        sum_percentage_reason = " + ".join([f"{p:.2f}%" for p in target_percentages]) + f" = {sum_percentage:.2f}%"

        # Count
        count_answer = len(target_slice_indices)

        # Chart QA Pool
        medium_qa_pool = {
            "two_step__statistics__sum_percentage": {
                "question": [
                    f"What is the total percentage for {pie_label_category_plural} {constraint}? Please round to two decimal places.",
                    f"For the {pie_label_category_plural} {constraint}, what is the sum of their percentages? Please round to two decimal places.",
                    f"Can you help calculate the total percentage for {pie_label_category_plural} {constraint}? Please round to two decimal places.",
                    f"Please compute the sum of percentages for {pie_label_category_plural} {constraint}. Please round to two decimal places.",
                ],
                "reasoning": [
                    {
                        "step_1": f"First, I need to identify the {pie_label_category_plural} {constraint}" + (f":\n{read_percent_reason}" if self.detailed_reasoning else "."),
                        "step_2": f"Second, I need to calculate their total percentage" + (f":\n{sum_percentage_reason}" if self.detailed_reasoning else "."),
                    },
                ],
                "constraint": constraint,
                "answer": f"{sum_percentage:.2f}%",
                "mask": {
                    "step_1": target_slice_indices,
                    "step_2": target_slice_indices if self.force_ground else [],
                    "answer": target_slice_indices,
                },
            },
            "two_step__statistics__count": {
                "question": [
                    f"How many {pie_label_category_plural} {constraint} are shown in this pie chart?",
                    f"What is the number of {pie_label_category_plural} {constraint}?",
                    f"Please help count the number of {pie_label_category_plural} {constraint}.",
                ],
                "reasoning": [
                    {
                        "step_1": f"First, I need to find all the {pie_label_category_plural} {constraint}" + (f":\n{read_slice_reason}" if self.detailed_reasoning else "."),
                        "step_2": f"Second, I need to count the total number of these {pie_label_category_plural}.",
                    },
                ],
                "constraint": constraint,
                "answer": count_answer,
                "mask": {
                    "step_1": target_slice_indices,
                    "step_2": target_slice_indices if self.force_ground else [],
                    "answer": target_slice_indices,
                },
            },
        }

        # For specific slice: ask about calculating total values
        if len(target_slice_indices) == 1:
            k = target_slice_indices[0]
            kth_data = chart_metadata["pie_data"][k]
            kth_label = chart_metadata["pie_labels"][k]
            kth_percent = target_percentages[0]
            total_sum = sum(chart_metadata["pie_data"])
            
            medium_qa_pool["two_step__statistics__calculate_total_from_constraint"] = {
                "question": [
                    f"Among {pie_label_category_plural} {constraint}, if the {pie_data_category_singular} of {kth_label} is {kth_data}, what is the total {pie_data_category_plural} of all {pie_label_category_plural} in this chart? Please round to two decimal places.",
                    f"Given that {kth_label} ({constraint}) has {kth_data} {pie_data_category_singular}, calculate the total {pie_data_category_plural} for all {pie_label_category_plural}. Please round to two decimal places.",
                ],
                "reasoning": [
                    {
                        "step_1": f"First, I need to identify the {pie_label_category_plural} {constraint}" + (f": {kth_label}" if self.detailed_reasoning else "."),
                        "step_2": f"Second, I need to read the percentage of the {pie_label_category_singular} '{kth_label}'" + (f": {kth_percent:.2f}%" if self.detailed_reasoning else "."),
                        "step_3": f"Third, I need to calculate the total" + (f": if {kth_percent:.2f}% = {kth_data}, then 100% = {total_sum:.2f}" if self.detailed_reasoning else "."),
                    },
                ],
                "constraint": constraint,
                "answer": f"{total_sum:.2f}",
                "mask": {
                    "step_1": target_slice_indices,
                    "step_2": target_slice_indices,
                    "step_3": list(range(len(chart_metadata["pie_data"]))) if self.force_ground else [],
                    "answer": list(range(len(chart_metadata["pie_data"]))),
                },
            }

        return medium_qa_pool


    ############################################################
    #                     Multi-step Operator
    ############################################################

    def _multi_step_threshold(self, chart_metadata: Dict):
        """
        Threshold: above / below mean percentage, differences
        """
        pie_data_category_singular = chart_metadata['pie_data_category']['singular']
        pie_data_category_plural = chart_metadata['pie_data_category']['plural']
        pie_label_category_singular = chart_metadata['pie_label_category']['singular']
        pie_label_category_plural = chart_metadata['pie_label_category']['plural']
        
        target_pie_data = copy.deepcopy(chart_metadata["pie_data"])
        target_slice_labels = copy.deepcopy(chart_metadata["pie_labels"])
        target_slice_indices = [iii for iii in range(len(target_slice_labels))]

        # Calculate percentages and mean
        total_data = sum(target_pie_data)
        target_percentages = [(data/total_data)*100 for data in target_pie_data]
        mean_percentage = 100.0 / len(target_pie_data)  # Equal distribution percentage

        # Above & below grouping
        above_mean_num, below_mean_num, slice_idx = 0, 0, 0
        above_mean_indices, below_mean_indices = [], []
        above_mean_percentages, below_mean_percentages = [], []
        for percentage in target_percentages:
            if percentage > mean_percentage:
                above_mean_num += 1
                above_mean_indices.append(slice_idx)
                above_mean_percentages.append(percentage)
            elif percentage < mean_percentage:
                below_mean_num += 1
                below_mean_indices.append(slice_idx)
                below_mean_percentages.append(percentage)
            slice_idx += 1

        # Sum sublist percentages
        above_sum_percentage = sum(above_mean_percentages)
        below_sum_percentage = sum(below_mean_percentages)
        above_sum_str = " + ".join([f"{p:.2f}%" for p in above_mean_percentages]) + f" = {above_sum_percentage:.2f}%"
        below_sum_str = " + ".join([f"{p:.2f}%" for p in below_mean_percentages]) + f" = {below_sum_percentage:.2f}%"

        # Mean sublist percentages
        above_mean_avg = above_sum_percentage / above_mean_num if above_mean_num > 0 else 0
        below_mean_avg = below_sum_percentage / below_mean_num if below_mean_num > 0 else 0
        above_mean_str = f"({' + '.join([f'{p:.2f}%' for p in above_mean_percentages])}) ÷ {above_mean_num} = {above_sum_percentage:.2f}% ÷ {above_mean_num} = {above_mean_avg:.2f}%"
        below_mean_str = f"({' + '.join([f'{p:.2f}%' for p in below_mean_percentages])}) ÷ {below_mean_num} = {below_sum_percentage:.2f}% ÷ {below_mean_num} = {below_mean_avg:.2f}%"

        # Max/min among above-mean sublist
        if above_mean_percentages:
            max_above_mean_percentage = max(above_mean_percentages)
            min_above_mean_percentage = min(above_mean_percentages)
            max_above_percentage_indices = self._find_indices_in_list(target_percentages, max_above_mean_percentage)
            min_above_percentage_indices = self._find_indices_in_list(target_percentages, min_above_mean_percentage)
            max_above_mean_label = self._convert_answer_idx_to_str(target_slice_labels, max_above_percentage_indices)
            min_above_mean_label = self._convert_answer_idx_to_str(target_slice_labels, min_above_percentage_indices)

        # Max/min among below-mean sublist
        if below_mean_percentages:
            max_below_mean_percentage = max(below_mean_percentages)
            min_below_mean_percentage = min(below_mean_percentages)
            max_below_percentage_indices = self._find_indices_in_list(target_percentages, max_below_mean_percentage)
            min_below_percentage_indices = self._find_indices_in_list(target_percentages, min_below_mean_percentage)
            max_below_mean_label = self._convert_answer_idx_to_str(target_slice_labels, max_below_percentage_indices)
            min_below_mean_label = self._convert_answer_idx_to_str(target_slice_labels, min_below_percentage_indices)

        # Differences
        sum_diff = abs(above_sum_percentage - below_sum_percentage)
        mean_diff = abs(above_mean_avg - below_mean_avg)

        # Reason
        read_percent_reason = '\n'.join([f"* {target_slice_labels[slice_idx]}: {target_percentages[slice_idx]:.2f}%" for slice_idx in range(len(target_pie_data))]).strip()
        reason_mean_percentage = f"100% ÷ {len(target_pie_data)} = {mean_percentage:.2f}%"

        # Chart QA Pool
        hard_qa_pool = {}

        if above_mean_percentages:
            hard_qa_pool["multi_step__threshold__above_mean__max__value"] = {
                "question": [
                    f"What is the highest percentage among {pie_label_category_plural} that have their percentages above the average percentage? Please round to two decimal places.",
                    f"Among {pie_label_category_plural} whose percentages are above the average percentage, what is the highest percentage? Please round to two decimal places.",
                    f"What is the largest percentage among {pie_label_category_plural} that have their percentages above the mean percentage? Please round to two decimal places.",
                    f"Among {pie_label_category_plural} whose percentages are above the mean percentage, what is the largest percentage? Please round to two decimal places.",
                ],
                "reasoning": [
                    {
                        "step_1": f"First, I need to read the percentage of each {pie_label_category_singular} in this chart" + (f":\n{read_percent_reason}" if self.detailed_reasoning else "."),
                        "step_2": f"Second, I need to compute the average percentage" + (f":\n{reason_mean_percentage}" if self.detailed_reasoning else "."),
                        "step_3": f"Third, I need to find all the {pie_label_category_plural} whose percentage is higher than {mean_percentage:.2f}%" + (f":\n{', '.join(target_slice_labels[iii] for iii in above_mean_indices)}" if self.detailed_reasoning else "."),
                        "step_4": f"Fourth, I need to identify the {pie_label_category_singular} with the highest percentage among these {pie_label_category_plural}" + (f":\n{max_above_mean_label}: {max_above_mean_percentage:.2f}%" if self.detailed_reasoning else "."),
                    },
                ],
                "constraint": f"max percentage among {pie_label_category_plural} above mean",
                "answer": f"{max_above_mean_percentage:.2f}%",
                "mask": {
                    "step_1": target_slice_indices,
                    "step_2": target_slice_indices if self.force_ground else [],
                    "step_3": above_mean_indices,
                    "step_4": max_above_percentage_indices,
                    "answer": max_above_percentage_indices,
                },
            }

            hard_qa_pool["multi_step__threshold__above_mean__max__label"] = {
                "question": [
                    f"What is the label of the {pie_label_category_singular} that has the highest percentage among {pie_label_category_plural} that have their percentages above the average percentage?",
                    f"Among {pie_label_category_plural} whose percentages are above the average percentage, what is the label of the {pie_label_category_singular} that has the highest percentage?",
                    f"What is the label of the {pie_label_category_singular} that has the largest percentage among {pie_label_category_plural} that have their percentages above the mean percentage?",
                    f"Among {pie_label_category_plural} whose percentages are above the mean percentage, what is the label of the {pie_label_category_singular} that has the largest percentage?",
                ],
                "reasoning": [
                    {
                        "step_1": f"First, I need to read the percentage of each {pie_label_category_singular} in this chart" + (f":\n{read_percent_reason}" if self.detailed_reasoning else "."),
                        "step_2": f"Second, I need to compute the average percentage" + (f":\n{reason_mean_percentage}" if self.detailed_reasoning else "."),
                        "step_3": f"Third, I need to find all the {pie_label_category_plural} whose percentage is higher than {mean_percentage:.2f}%" + (f":\n{', '.join(target_slice_labels[iii] for iii in above_mean_indices)}" if self.detailed_reasoning else "."),
                        "step_4": f"Fourth, I need to identify the {pie_label_category_singular} with the highest percentage among these {pie_label_category_plural}" + (f":\n{max_above_mean_label}: {max_above_mean_percentage:.2f}%" if self.detailed_reasoning else "."),
                    },
                ],
                "constraint": f"label of max percentage among {pie_label_category_plural} above mean",
                "answer": max_above_mean_label,
                "mask": {
                    "step_1": target_slice_indices,
                    "step_2": target_slice_indices if self.force_ground else [],
                    "step_3": above_mean_indices,
                    "step_4": max_above_percentage_indices,
                    "answer": max_above_percentage_indices,
                },
            }

            hard_qa_pool["multi_step__threshold__above_mean__sum"] = {
                "question": [
                    f"What is the sum of percentages for {pie_label_category_plural} that have their percentages above the average percentage? Please round to two decimal places.",
                    f"Among {pie_label_category_plural} whose percentages are above the average percentage, what is the sum of their percentages? Please round to two decimal places.",
                    f"What is the total percentage for {pie_label_category_plural} that have their percentages above the mean percentage? Please round to two decimal places.",
                    f"Among {pie_label_category_plural} whose percentages are above the mean percentage, what is the total of their percentages? Please round to two decimal places.",
                ],
                "reasoning": [
                    {
                        "step_1": f"First, I need to read the percentage of each {pie_label_category_singular} in this chart" + (f":\n{read_percent_reason}" if self.detailed_reasoning else "."),
                        "step_2": f"Second, I need to compute the average percentage" + (f":\n{reason_mean_percentage}" if self.detailed_reasoning else "."),
                        "step_3": f"Third, I need to find all the {pie_label_category_plural} whose percentage is higher than {mean_percentage:.2f}%" + (f":\n{', '.join(target_slice_labels[iii] for iii in above_mean_indices)}" if self.detailed_reasoning else "."),
                        "step_4": f"Fourth, I need to calculate the total percentage of these {pie_label_category_plural}" + (f":\n{above_sum_str}" if self.detailed_reasoning else "."),
                    },
                ],
                "constraint": f"sum of percentages for {pie_label_category_plural} above mean",
                "answer": f"{above_sum_percentage:.2f}%",
                "mask": {
                    "step_1": target_slice_indices,
                    "step_2": target_slice_indices if self.force_ground else [],
                    "step_3": above_mean_indices,
                    "step_4": above_mean_indices if self.force_ground else [],
                    "answer": above_mean_indices,
                },
            }

        if below_mean_percentages:
            hard_qa_pool["multi_step__threshold__below_mean__sum"] = {
                "question": [
                    f"What is the sum of percentages for {pie_label_category_plural} that have their percentages below the average percentage? Please round to two decimal places.",
                    f"Among {pie_label_category_plural} whose percentages are below the average percentage, what is the sum of their percentages? Please round to two decimal places.",
                    f"What is the total percentage for {pie_label_category_plural} that have their percentages below the mean percentage? Please round to two decimal places.",
                    f"Among {pie_label_category_plural} whose percentages are below the mean percentage, what is the total of their percentages? Please round to two decimal places.",
                ],
                "reasoning": [
                    {
                        "step_1": f"First, I need to read the percentage of each {pie_label_category_singular} in this chart" + (f":\n{read_percent_reason}" if self.detailed_reasoning else "."),
                        "step_2": f"Second, I need to compute the average percentage" + (f":\n{reason_mean_percentage}" if self.detailed_reasoning else "."),
                        "step_3": f"Third, I need to find all the {pie_label_category_plural} whose percentage is lower than {mean_percentage:.2f}%" + (f":\n{', '.join(target_slice_labels[iii] for iii in below_mean_indices)}" if self.detailed_reasoning else "."),
                        "step_4": f"Fourth, I need to calculate the total percentage of these {pie_label_category_plural}" + (f":\n{below_sum_str}" if self.detailed_reasoning else "."),
                    },
                ],
                "constraint": f"sum of percentages for {pie_label_category_plural} below mean",
                "answer": f"{below_sum_percentage:.2f}%",
                "mask": {
                    "step_1": target_slice_indices,
                    "step_2": target_slice_indices if self.force_ground else [],
                    "step_3": below_mean_indices,
                    "step_4": below_mean_indices if self.force_ground else [],
                    "answer": below_mean_indices,
                },
            }

        if above_mean_percentages and below_mean_percentages:
            hard_qa_pool["multi_step__threshold__mean__sum_diff"] = {
                "question": [
                    f"What is the absolute difference between the total percentage of {pie_label_category_plural} above the average percentage and those below it? Please round to two decimal places.",
                    f"What is the absolute value of the difference between the total percentage for {pie_label_category_plural} above the average and those below the average percentage? Please round to two decimal places.",
                    f"What is the absolute difference between the total percentage of {pie_label_category_plural} above the mean percentage and those below it? Please round to two decimal places.",
                    f"What is the absolute value of the difference between the total percentage for {pie_label_category_plural} above the mean percentage and those below the mean percentage? Please round to two decimal places.",
                ],
                "reasoning": [
                    {
                        "step_1": f"First, I need to read the percentage of each {pie_label_category_singular} in this chart" + (f":\n{read_percent_reason}" if self.detailed_reasoning else "."),
                        "step_2": f"Second, I need to compute the average percentage" + (f":\n{reason_mean_percentage}" if self.detailed_reasoning else "."),
                        "step_3": f"Third, I need to find the first group of {pie_label_category_plural} whose percentage is higher than {mean_percentage:.2f}% and the second group of {pie_label_category_plural} whose percentage is lower than {mean_percentage:.2f}%.",
                        "step_4": f"Fourth, I need to calculate the total percentage of these two groups of {pie_label_category_plural} respectively" + (f":\n* Group 1 (above overall mean): {above_sum_str}\n* Group 2 (below overall mean): {below_sum_str}" if self.detailed_reasoning else "."),
                        "step_5": f"Fifth, I need to calculate the absolute difference between these two groups of {pie_label_category_plural}" + (f":\n|{above_sum_percentage:.2f}% - {below_sum_percentage:.2f}%| = {sum_diff:.2f}%." if self.detailed_reasoning else "."),
                    },
                ],
                "constraint": f"percentage difference between {pie_label_category_plural} above and below mean",
                "answer": f"{sum_diff:.2f}%",
                "mask": {
                    "step_1": target_slice_indices,
                    "step_2": target_slice_indices if self.force_ground else [],
                    "step_3": above_mean_indices + below_mean_indices,
                    "step_4": above_mean_indices + below_mean_indices if self.force_ground else [],
                    "step_5": above_mean_indices + below_mean_indices if self.force_ground else [],
                    "answer": above_mean_indices + below_mean_indices,
                },
            }

        return hard_qa_pool
    