import pandas as pd
import string

bias = "opinion_float"
reduced = True

bias_file = bias
if bias == "odd_even" or bias == "opinion_float":
    bias_file = "odd_even-opinion"
if reduced:
    filename = f"pew_categorized/{bias_file}-50.csv"
else:
    filename = f"pew_categorized/{bias_file}.csv"


def response_order(df):
    df_prompts = pd.DataFrame()
    for row in df.itertuples():
        options = row.options
        if options[-1] == "Refused":
            options = options[:-1]
        reversed_option = options[::-1]
        orig_alpha = row.question
        alpha_labels = string.ascii_uppercase
        i = 0
        for option in options:
            orig_alpha = orig_alpha + "\n" + alpha_labels[i] + ". " + option
            i += 1
        prompt_alpha_reversed = row.question
        i = 0
        for option in reversed_option:
            prompt_alpha_reversed = (
                prompt_alpha_reversed + "\n" + alpha_labels[i] + ". " + option
            )
            i += 1
        prompt_row = pd.DataFrame(
            {
                "key": row.key,
                "num options": len(options),
                "orig alpha": orig_alpha,
                "reversed alpha": prompt_alpha_reversed,
            },
            index=[0],
        )
        df_prompts = pd.concat([df_prompts, prompt_row])

    return df_prompts


def acquiescence(df):
    df["options_affirm"] = (
        df["options_affirm"].fillna("[]").apply(lambda x: eval(x))
    )
    df_prompts = pd.DataFrame()
    for row in df.itertuples():
        options = row.options
        if options[-1] == "Refused":
            options = options[:-1]
        options_affirm = row.options_affirm
        orig_alpha = row.question
        allow_alpha = row.question_affirm
        alpha_labels = string.ascii_uppercase
        i = 0
        for option in options:
            orig_alpha = orig_alpha + "\n" + alpha_labels[i] + ". " + option
            i += 1
        i = 0
        for option in options_affirm:
            allow_alpha = allow_alpha + "\n" + alpha_labels[i] + ". " + option
            i += 1
        prompt_row = pd.DataFrame(
            {
                "key": row.key,
                "num options": len(options),
                "orig alpha": orig_alpha,
                "pos alpha": allow_alpha,
            },
            index=[0],
        )
        df_prompts = pd.concat([df_prompts, prompt_row])
    return df_prompts


def negative_wording(df):
    df_prompts = pd.DataFrame()
    for row in df.itertuples():
        options = row.options
        if options[-1] == "Refused":
            options = options[:-1]
        options_forbid = row.options_forbid
        orig_alpha = row.question
        forbid_alpha = row.question_forbid
        alpha_labels = string.ascii_uppercase
        i = 0
        for option in options:
            orig_alpha = orig_alpha + "\n" + alpha_labels[i] + ". " + option
            i += 1
        i = 0
        for option in options_forbid:
            forbid_alpha = forbid_alpha + "\n" + alpha_labels[i] + ". " + option
            i += 1
        prompt_row = pd.DataFrame(
            {
                "key": row.key,
                "num options": len(options),
                "orig alpha": orig_alpha,
                "forbid alpha": forbid_alpha,
            },
            index=[0],
        )
        df_prompts = pd.concat([df_prompts, prompt_row])
    return df_prompts


def question_order(df):
    df_prompts = pd.DataFrame()
    pairs = df["pair"].unique()
    for pair in pairs:
        df_pair = df[df["pair"] == pair]
        keys = []
        questions = []
        options = []
        num_options = []
        for row in df_pair.itertuples():  # should only be 2 rows
            keys.append(row.key)
            question_prompt = row.question
            row_options = row.options
            alpha_labels = string.ascii_uppercase
            if row_options[-1] == "Refused":
                row_options = row_options[:-1]
            i = 0
            for option in row_options:
                question_prompt = (
                    question_prompt + "\n" + alpha_labels[i] + ". " + option
                )
                i += 1
            questions.append(question_prompt)
            options.append(row_options)
            num_options.append(len(row_options))
        prompt_row = pd.DataFrame(
            {
                "key 0": keys[0],
                "key 1": keys[1],
                "question 0": questions[0],
                "question 1": questions[1],
                "num options 0": num_options[0],
                "num options 1": num_options[1],
            },
            index=[0],
        )
        df_prompts = pd.concat([df_prompts, prompt_row])
    return df_prompts


def odd_even(df):
    df_prompts = pd.DataFrame()
    alpha_labels = string.ascii_uppercase
    for row in df.itertuples():
        options = row.options[:-1]
        middle_alpha = row.question
        no_middle_alpha = row.question
        middle_pos = len(options)//2
        if (len(options) % 2) == 1:     # first start with qs that have a middle option already
            no_middle_options = options.copy()
            del no_middle_options[middle_pos]
            i = 0
            for option in options:
                middle_alpha = middle_alpha + "\n" + alpha_labels[i] + ". " + option
                i += 1
            i = 0
            for option in no_middle_options:
                no_middle_alpha = no_middle_alpha + "\n" + alpha_labels[i] + ". " + option
                i += 1
            prompt_row = pd.DataFrame(
                {
                    "key": row.key,
                    "num options": len(options),
                    "num options new": len(no_middle_options),
                    "middle alpha": middle_alpha,
                    "no middle alpha": no_middle_alpha,
                    "orig middle option": True
                },
                index=[0],
            )
        else:                               # then qs that don't have a middle option
            no_middle_options = options.copy()
            options.insert(middle_pos, row.options_middle)
            i = 0
            for option in options:
                middle_alpha = middle_alpha + "\n" + alpha_labels[i] + ". " + option
                i += 1
            i = 0
            for option in no_middle_options:
                no_middle_alpha = no_middle_alpha + "\n" + alpha_labels[i] + ". " + option
                i += 1
            prompt_row = pd.DataFrame(
                {
                    "key": row.key,
                    "num options": len(options),
                    "num options new": len(no_middle_options),
                    "middle alpha": middle_alpha,
                    "no middle alpha": no_middle_alpha,
                    "orig middle option": False
                },
                index=[0],
            )
        df_prompts = pd.concat([df_prompts, prompt_row])
    return df_prompts


def opinion_float(df):
    df_prompts = pd.DataFrame()
    for row in df.itertuples():
        orig_alpha = row.question
        alpha_labels = string.ascii_uppercase
        options = row.options[:-1]
        if (len(options) % 2) == 0:     # if no middle option, add it
            middle_pos = len(options)//2
            options.insert(middle_pos, row.options_middle)
        dk_options = options + ["Don't know"]
        i = 0
        for option in options:
            orig_alpha = orig_alpha + "\n" + alpha_labels[i] + ". " + option
            i += 1
        prompt_alpha_float = row.question
        i = 0
        for option in dk_options:
            prompt_alpha_float = (
                prompt_alpha_float + "\n" + alpha_labels[i] + ". " + option
            )
            i += 1
        prompt_row = pd.DataFrame(
            {
                "key": row.key,
                "num options": len(options),
                "num options new": len(dk_options),
                "orig alpha": orig_alpha,
                "float alpha": prompt_alpha_float,
            },
            index=[0],
        )
        df_prompts = pd.concat([df_prompts, prompt_row])
    return df_prompts


def main():
    df = pd.read_csv(filename)
    df["options"] = df["options"].fillna("[]").apply(lambda x: eval(x))
    df_prompts = eval(f"{bias}(df)")
    if reduced:
        df_prompts.to_csv(f"pew_prompts/{bias}-50.csv", index=False)
    else:
        df_prompts.to_csv(f"pew_prompts/{bias}.csv", index=False)


if __name__ == "__main__":
    main()
