import random
from dataclasses import asdict
from pathlib import Path
import itertools

from otree.api import *
from otree.models.participant import Participant

from utils import bots_utils, csv_utils, tasks_utils
from utils.pageseq_utils import task_page
from utils.scores import calc_outcome, calc_payoff, team_score

DATADIR = Path(__file__).parent.parent / "data"
BOTSDATA = bots_utils.read_bots(DATADIR / "bots.csv")
NAMESDATA = bots_utils.read_names(DATADIR / "names.csv")
TASKDATA = tasks_utils.read(DATADIR / "tasks.csv")

RESULTS = csv_utils.read_csv(DATADIR / "results.csv", {"gender": str, "score": int})


class C(BaseConstants):
    NAME_IN_URL = "main"
    PLAYERS_PER_GROUP = None
    TASKS = ["p1", "p2"]
    NUM_ROUNDS = 4  # start + 2 tasks + final

    RATINGS = {
        "part 1": {
            "default": (3, 3, 3, 4, 4, 4, 4, 5, 5, 5),
            "M": {
                5: 1,
                1: 2,
            },
            "F": {
                1: 1,
                5: 2,
            },
        },
        "part 2": {
            "default": (1, 1, 1, 2, 2, 2, 2, 3, 3, 3),
            "M": {
                19: 5,
                21: 4,
            },
            "F": {
                21: 5,
                19: 4,
            },
        },
    }


def sample_score(bot: bots_utils.Bot):
    results = csv_utils.filter_data(RESULTS, gender=bot.gender)
    return random.choice(results)["score"]


class Subsession(BaseSubsession):
    pass


class Group(BaseGroup):
    pass


class Player(BasePlayer):
    prolific_id = models.StringField(label="Enter you prolific ID")
    gender = models.StringField(label="Select your gender", choices=[("M", "Male"), ("F", "Female"), ("O", "Other")], widget=widgets.RadioSelectHorizontal)
    avatar = models.StringField()

    current_task = models.StringField()
    sequence = models.StringField()  # |-joined list
    expected = models.StringField()  # |-joined list
    answers = models.StringField(initial="")  # |-joined list
    score = models.IntegerField()

    p1_competitor = models.StringField()
    p1_competitor_score = models.IntegerField()
    p1_score = models.IntegerField()
    p1_outcome = models.IntegerField()

    p2_teammate = models.StringField()
    p2_teammate_score = models.IntegerField()
    p2_competitor = models.StringField()
    p2_competitor_score = models.IntegerField()
    p2_score = models.IntegerField()
    p2_outcome = models.IntegerField()

    p3_teammate = models.StringField()
    p3_competitor = models.StringField()
    p3_decision = models.StringField(choices=("compete", "pass"))
    p3_teammate_score1 = models.IntegerField()
    p3_teammate_score2 = models.IntegerField()
    p3_competitor_score = models.IntegerField()
    p3_score = models.IntegerField()
    p3_outcome = models.IntegerField()

    attention = models.BooleanField()

    q1 = models.IntegerField(
        label="""
        How do you see yourself: are you generally a person who is fully prepared to take risks, or do you try to avoid taking risks?
        Please tick a box on the scale, where the value 0 means: 'not at all willing to take risks' and the value 10 means: 'very willing to take risks'.
        """,
        widget=widgets.RadioSelect,
        choices=[
            (0, "Not at all willing to take risks"),
            (1, ""),
            (2, ""),
            (3, ""),
            (4, ""),
            (5, ""),
            (6, ""),
            (7, ""),
            (8, ""),
            (9, ""),
            (10, "Very willing to take risks"),
        ],
    )

    q2 = models.IntegerField(
        label="""It takes me longer to understand mathematics than the average person.""",
        widget=widgets.RadioSelect,
        choices=[
            (1, "Strongly Disagree"),
            (2, "Disagree"),
            (3, "Somewhat Disagree"),
            (4, "Neither Agree nor Disagree"),
            (5, "Somewhat Agree"),
            (6, "Agree"),
            (7, "Strongly Agree"),
        ],
    )

    q3 = models.IntegerField(
        label="""In general, people think men are better at mathematics than women.""",
        widget=widgets.RadioSelect,
        choices=[
            (1, "Strongly Disagree"),
            (2, "Disagree"),
            (3, "Somewhat Disagree"),
            (4, "Neither Agree nor Disagree"),
            (5, "Somewhat Agree"),
            (6, "Agree"),
            (7, "Strongly Agree"),
        ],
    )

    q4 = models.IntegerField(
        label="""Men are inherently better at mathematics than women.""",
        widget=widgets.RadioSelect,
        choices=[
            (1, "Strongly Disagree"),
            (2, "Disagree"),
            (3, "Somewhat Disagree"),
            (4, "Neither Agree nor Disagree"),
            (5, "Somewhat Agree"),
            (6, "Agree"),
            (7, "Strongly Agree"),
        ],
    )


def creating_session(subsession: Subsession):
    for p in subsession.get_players():
        init_player(p, subsession.round_number)


def init_player(player: Player, rnd: int):
    participant = player.participant
    if rnd == 1:
        seq = list(C.TASKS)
        random.shuffle(seq)
        participant.task_sequence = ["start"] + seq + ["end"]
        participant.prolific_id = None
        participant.avatar = None
        participant.seed = hash(participant)

    player.current_task = player.participant.task_sequence[rnd - 1]

    if player.current_task == "p1":
        init_tasks(player, "part 1")

    if player.current_task == "p2":
        init_tasks(player, "part 2")
        bots2 = bots_utils.select(BOTSDATA, section="part 2")
        player.p2_competitor = str(random.choice(bots2))


def init_tasks(player: Player, section: str):
    """load tasks for a section and randomize them"""
    tasks = tasks_utils.select(TASKDATA, section=section)
    random.shuffle(tasks)
    player.sequence = "|".join([t.id for t in tasks])
    player.expected = "|".join([t.correct for t in tasks])


def load_tasks(player: Player):
    """retrieve tasks in the same random order as initialized
    options of each task are shuffled
    """
    ids = player.sequence.split("|")
    tasks = [tasks_utils.select1(TASKDATA, id) for id in ids]
    for t in tasks:
        random.shuffle(t.options)
    return tasks


def calc_results0(player: Player):
    """Player own score = number of matching answers"""
    answers = player.answers.split("|")
    expected = player.expected.split("|")
    player.score = sum(ans == exp for ans, exp in zip(answers, expected))


def calc_results1(player: Player):
    """ "Calc result for part 1 with competitor"""
    calc_results0(player)
    p1_competitor = bots_utils.parse(player.p1_competitor)
    player.p1_competitor_score = sample_score(p1_competitor)
    player.p1_score = player.score
    player.p1_outcome = calc_outcome(player.p1_score, player.p1_competitor_score)
    player.payoff = calc_payoff(player, player.p1_outcome)


def calc_results2(player: Player):
    """Calc result for part 2 with teammate and competitor"""
    calc_results0(player)
    p2_competitor = bots_utils.parse(player.p2_competitor)
    p2_teammate = bots_utils.parse(player.p2_teammate)
    player.p2_competitor_score = sample_score(p2_competitor)
    player.p2_teammate_score = sample_score(p2_teammate)
    player.p2_score = team_score(player.score, player.p2_teammate_score)
    player.p2_outcome = calc_outcome(player.p2_score, player.p2_competitor_score)
    player.payoff = calc_payoff(player, player.p2_outcome)


def all_bots():
    """List all bots with names"""
    names = random.sample(NAMESDATA, k=len(NAMESDATA))
    bots = bots_utils.merge_names(BOTSDATA, names)
    return bots


def list_bots(player: Player, part: str):
    """List bots with ratings accoring to C.RATING"""
    random.seed(player.participant.seed)

    bots = bots_utils.select(all_bots(), section=part)
    random.shuffle(bots)

    ratings_bots = C.RATINGS[part][player.participant.variation]
    ratings_pool = list(C.RATINGS[part]["default"])
    random.shuffle(ratings_pool)

    assert len(bots) == len(ratings_pool) + len(ratings_bots)

    for b in bots:
        if b.id in ratings_bots:
            b.rating = ratings_bots[b.id]
        else:
            b.rating = ratings_pool.pop()
    return bots


# PAGES


@task_page("start")
class Consent(Page):
    template_name = "Consent.html"


@task_page("start")
class ProlificID(Page):
    template_name = "ProlificID.html"
    form_model = "player"
    form_fields = ["prolific_id", "gender"]

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        player.participant.prolific_id = player.prolific_id
        if player.gender in ("M", "F"):
            player.participant.variation = player.gender
        else:
            player.participant.variation = random.choice(["M", "F"])


@task_page("start")
class GeneralInstructions(Page):
    template_name = "GeneralInstructions.html"


@task_page("start")
class Practice(Page):
    template_name = "Practice.html"
    timeout_seconds = 60

    @staticmethod
    def vars_for_template(player: Player):
        tasks = tasks_utils.select(TASKDATA, section="practice")
        return dict(task=asdict(tasks[0]))


@task_page("start")
class PracticeEnd(Page):
    template_name = "PracticeEnd.html"


@task_page("start")
class AvatarChoice(Page):
    template_name = "AvatarChoice.html"
    form_model = "player"
    form_fields = ["avatar"]

    @staticmethod
    def vars_for_template(player: Player):
        avatars = bots_utils.select(BOTSDATA, section="avatar")
        return dict(bots=avatars, bot_field="avatar")

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        player.participant.avatar = player.avatar


#### PART 1


@task_page("p1")
class Instructions1(Page):
    template_name = "Instructions1.html"


@task_page("p1")
class Choice1(Page):
    template_name = "Choice1.html"
    form_model = "player"
    form_fields = ["p1_competitor"]

    @staticmethod
    def vars_for_template(player: Player):
        return {
            "bots": list_bots(player, "part 1"),
            "bot_field": "p1_competitor",
        }


@task_page("p1")
class Start1(Page):
    template_name = "Start.html"


@task_page("p1")
class Match1(Page):
    template_name = "Match.html"
    form_model = "player"
    form_fields = ["answers"]
    timeout_seconds = 150

    @staticmethod
    def vars_for_template(player: Player):
        competitor = bots_utils.parse(player.p1_competitor)
        return {"competitor": competitor, "teammate": None}

    @staticmethod
    def js_vars(player: Player):
        tasks = load_tasks(player)
        return {"tasks": [asdict(t) for t in tasks]}

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        calc_results1(player)


@task_page("p1")
class Results1(Page):
    template_name = "Results1.html"

    @staticmethod
    def vars_for_template(player: Player):
        return {"outcome": player.p1_outcome}


#### PART 2


@task_page("p2")
class Instructions2(Page):
    template_name = "Instructions2.html"


@task_page("p2")
class Choice2(Page):
    template_name = "Choice2.html"
    form_model = "player"
    form_fields = ["p2_teammate"]

    @staticmethod
    def vars_for_template(player: Player):
        return {
            "bots": list_bots(player, "part 2"),
            "bot_field": "p2_teammate",
        }


@task_page("p2")
class Start2(Page):
    template_name = "Start.html"


@task_page("p2")
class Match2(Page):
    template_name = "Match.html"
    form_model = "player"
    form_fields = ["answers"]
    timeout_seconds = 150

    @staticmethod
    def vars_for_template(player: Player):
        teammate = bots_utils.parse(player.p2_teammate)
        return {"competitor": None, "teammate": teammate}

    @staticmethod
    def js_vars(player: Player):
        tasks = load_tasks(player)
        return {"tasks": [asdict(t) for t in tasks]}

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        calc_results2(player)


@task_page("p2")
class Results2(Page):
    template_name = "Results2.html"

    @staticmethod
    def vars_for_template(player: Player):
        return {"outcome": player.p2_outcome}


### END


@task_page("end")
class AttentionCheck(Page):
    template_name = "AttentionCheck.html"
    form_model = "player"
    form_fields = ["attention"]


@task_page("end")
class SurveyIntro1(Page):
    pass


@task_page("end")
class Survey1(Page):
    form_model = "player"
    form_fields = ["q1"]
    template_name = "main/Scale11.html"


@task_page("end")
class SurveyIntro2(Page):
    pass


@task_page("end")
class Survey2(Page):
    form_model = "player"
    form_fields = ["q2"]
    template_name = "main/Scale7.html"


@task_page("end")
class Survey3(Page):
    form_model = "player"
    form_fields = ["q3"]
    template_name = "main/Scale7.html"


@task_page("end")
class Survey4(Page):
    form_model = "player"
    form_fields = ["q4"]
    template_name = "main/Scale7.html"


@task_page("end")
class Debrief(Page):
    template_name = "Debrief.html"


page_sequence = [
    # start
    Consent,
    ProlificID,
    GeneralInstructions,
    Practice,
    PracticeEnd,
    AvatarChoice,
    # p1
    Instructions1,
    Choice1,
    Start1,
    Match1,
    Results1,
    # p2
    Instructions2,
    Choice2,
    Start2,
    Match2,
    Results2,
    # end
    AttentionCheck,
    #
    SurveyIntro1,
    Survey1,
    SurveyIntro2,
    Survey2,
    Survey3,
    Survey4,
    #
    Debrief,
]


def custom_export(players):
    yield [
        "participant.code",
        "participant.label",
        "participant.prolific_id",
        "participant.avatar",
        "participant.avatar_race",
        "participant.avatar_gender",
        "variation",
        "sequence",
        #
        "gender",
        #
        "attention",
        "q.1",
        "q.2",
        "q.3",
        "q.4",
        #
        "p1_competitor",
        "p1_competitor_race",
        "p1_competitor_gender",
        "p1_competitor_rating",
        "p1_score",
        "p1_competitor_score",
        "p1_outcome",
        #
        "p2_teammate",
        "p2_teammate_race",
        "p2_teammate_gender",
        "p2_teammate_rating",
        "p2_competitor",
        "p2_competitor_race",
        "p2_competitor_gender",
        "p2_competitor_rating",
        "p2_score",
        "p2_teammate_score",
        "p2_competitor_score",
        "p2_outcome",
    ]

    # regroup players by participants
    players.sort(key=lambda p: p.participant_id)
    grouped = itertools.groupby(players, key=lambda p: p.participant)

    for participant, p_players in grouped:
        players = {p.current_task: p for p in p_players}

        avatar = bots_utils.parse(participant.avatar)

        rowdata = [
            participant.code,
            participant.label,
            participant.prolific_id,
            avatar,
            avatar.race,
            avatar.gender,
            participant.variation,
            ",".join(filter(lambda t: t in C.TASKS, participant.task_sequence)),
        ]

        if "start" in players:
            player = players["start"]
            rowdata += [
                player.gender,
            ]
        else:
            rowdata += [None]

        if "end" in players:
            player = players["end"]
            rowdata += [
                player.attention,
                player.q1,
                player.q2,
                player.q3,
                player.q4,
            ]
        else:
            rowdata += [None, None, None, None, None]

        if "p1" in players:
            p1_player = players["p1"]
            p1_competitor = bots_utils.parse(p1_player.p1_competitor)
            rowdata += [
                p1_player.p1_competitor,
                p1_competitor.race,
                p1_competitor.gender,
                p1_competitor.rating,
                p1_player.p1_score,
                p1_player.p1_competitor_score,
                p1_player.p1_outcome,
            ]
        else:
            rowdata += [None, None, None, None, None, None, None]

        if "p2" in players:
            p2_player = players["p2"]
            p2_teammate = bots_utils.parse(p2_player.p2_teammate)
            p2_competitor = bots_utils.parse(p2_player.p2_competitor)
            rowdata += [
                p2_player.p2_teammate,
                p2_teammate.race,
                p2_teammate.gender,
                p2_teammate.rating,
                p2_player.p2_competitor,
                p2_competitor.race,
                p2_competitor.gender,
                p2_competitor.rating,
                p2_player.p2_score,
                p2_player.p2_teammate_score,
                p2_player.p2_competitor_score,
                p2_player.p2_outcome,
            ]
        else:
            rowdata += [None, None, None, None, None, None, None, None, None, None, None, None]

        yield rowdata