from otree.api import (
    models, widgets, BaseConstants, BaseSubsession, BaseGroup, BasePlayer,
    Currency as c, currency_range
)

author = 'Elisa Macchi'

doc = """ Dictator or Trust or Public Good Game with Advisor """

import random
import itertools
import pandas as pd


class Constants(BaseConstants):

    name_in_url = 'games'
    num_rounds = 1
    pilot = 'yes'
    n_players = 3

    n_questions = 4

    # Payoffs:
    min_bonus = c(0)
    bonus = c(12)

    # Image info read from csv:

    img_f = ['1_thin.jpg', '2_thin.jpg', '3_thin.jpg']
    img_o = ['1_original.jpg', '2_original.jpg', '3_original.jpg']
    img_t = ['1_fat.jpg', '2_fat.jpg', '3_fat.jpg']

    age = [31, 40, 52]
    gender = [1, 0, 1]

    # Endowments
    endowment = c(4)

    # Trust game:
    multiplier = 2

    # Public Good Game:

    players_per_group = 3
    pg_multiplier = 2


class Subsession(BaseSubsession):

    def creating_session(self):

        # Random Game Choice

        games_list = itertools.cycle([1, 2, 3])

        # randomize to treatments
        treatments = itertools.cycle(['original', 'thin', 'fat'])

        for p in self.get_players():

            p.game = next(games_list)

            p.suggested_contribution = random.choice([0, Constants.endowment])
            p.participant.vars['suggested_contribution'] = p.suggested_contribution

            p.treatment = next(treatments)

            if p.treatment == "fat":
                imgs = itertools.cycle(Constants.img_f)
                imgs_age = itertools.cycle(Constants.age)
                imgs_gen = itertools.cycle(Constants.gender)

                p.participant.vars['images'] = next(imgs)
                p.image_url = next(imgs)
                p.image_age = next(imgs_age)
                p.image_gender = next(imgs_gen)

            if p.treatment == "original":
                imgs = itertools.cycle(Constants.img_o)
                imgs_age = itertools.cycle(Constants.age)
                imgs_gen = itertools.cycle(Constants.gender)

                p.participant.vars['images'] = next(imgs)
                p.image_url = next(imgs)
                p.image_age = next(imgs_age)
                p.image_gender = next(imgs_gen)

            if p.treatment == "thin":
                imgs = itertools.cycle(Constants.img_t)
                imgs_age = itertools.cycle(Constants.age)
                imgs_gen = itertools.cycle(Constants.gender)

                p.participant.vars['images'] = next(imgs)
                p.image_url = next(imgs)
                p.image_age = next(imgs_age)
                p.image_gender = next(imgs_gen)

        def vars_for_admin_report(self):
            contributions = [p.contribution for p in self.get_players() if p.contribution != None]
            if contributions:
                return {
                    'avg_contribution': sum(contributions) / len(contributions),
                    'min_contribution': min(contributions),
                    'max_contribution': max(contributions),
                }
            else:
                return {
                    'avg_contribution': '(no data)',
                    'min_contribution': '(no data)',
                    'max_contribution': '(no data)',
                }


class Group(BaseGroup):

    total_contribution = models.CurrencyField()
    individual_share = models.CurrencyField()

    def set_payoffs(self):

        # Payoff Allocation:

        players = self.get_players()

        for p in players:
            game_that_counts = p.game
            p.participant.vars['game_that_counts'] = game_that_counts

        # Trust Game Amounts
        img_list = Constants.img_f + Constants.img_t
        exp_contribution = [0] * len(img_list)
        avg_contribution = [0] * len(img_list)

        for i in range(len(img_list)):
            count = 0
            for p in players:
                if p.game == 2:
                    if img_list[i] == p.image_url:
                        count += 1
                        exp_contribution[i] += p.beliefs_sent_back_amount
            if count > 0:
                avg_contribution[i] = exp_contribution[i] / count

        for p in players:
            if p.game == 2:
                for i in range(len(img_list)):
                    if p.image_url == img_list[i]:
                        p.sent_back_amount = avg_contribution[i]

        self.total_contribution = 0
        for p in players:
            if p.game == 3:
                self.total_contribution += p.contribution

        self.individual_share = self.total_contribution * Constants.multiplier / Constants.players_per_group

        for p in players:
            if p.game == 1:
                p.payoff = Constants.endowment - p.kept
            elif p.game == 2:
                p.payoff = Constants.endowment - p.sent_amount + p.sent_back_amount
            elif p.game == 3:
                p.payoff = (Constants.endowment - p.contribution) + self.individual_share


class Player(BasePlayer):

    # Treatment:

    game = models.IntegerField()
    treatment = models.StringField()
    image_url = models.StringField()


    # Images Characteristics

    image_age = models.IntegerField()
    image_gender = models.BooleanField()

    # Dictator Game:

    kept = models.CurrencyField(choices=currency_range(0, Constants.endowment, c(1)),
                                doc="""dictator's choice""")

    # Trust Game
    sent_amount = models.CurrencyField(choices=currency_range(0, Constants.endowment, c(1)),
                                       doc="""Amount sent by P1""")

    beliefs_sent_back_amount = models.CurrencyField(min =0, max= 8, doc="""Beliefs on amount sent by P2""")

    sent_back_amount = models.CurrencyField()

    implicit_rate = models.FloatField()
    image_return = models.FloatField()

    payoff_trust = models.CurrencyField()

    # Public Good Game

    suggested_contribution = models.CurrencyField(
        doc="""Suggested contribution by advisor """,
    )

    contribution = models.CurrencyField(
        choices=currency_range(0, Constants.endowment, c(1)),
        doc="""Amount contributed by the player""",
    )

    # Winner

    group_b1 = models.IntegerField()
    group_b2 = models.IntegerField()
    group_b3 = models.IntegerField()

    is_winner = models.BooleanField()
    is_selected = models.BooleanField()
    is_selected_trust = models.BooleanField()