from otree.api import *
import itertools
import random

doc = """ Stock Experiences """


class C(BaseConstants):
    NAME_IN_URL = 'stock5'
    PLAYERS_PER_GROUP = None
    NUM_ROUNDS = 11

class Subsession(BaseSubsession):
    pass


class Group(BaseGroup):
    pass


class Player(BasePlayer):

    purchase_price = models.FloatField()
    price = models.FloatField()
    period_return = models.FloatField()
    inv_return = models.FloatField()

    invest = models.IntegerField(min=0, max=10000)
    quality = models.FloatField(min=-50, max=50)
    value = models.FloatField(min=-50, max=50)
    volatility = models.FloatField(min=-50, max=50)
    exp_return = models.FloatField(min=-50, max=50)

# FUNCTIONS
def creating_session(subsession):
    if subsession.round_number != 1:
        return

    for player in subsession.get_players():
        participant = player.participant
        for round_number in range(1, C.NUM_ROUNDS + 1):
            round_player = player.in_round(round_number)
            round_player.price = participant.path5[round_number - 1]
            round_player.purchase_price = participant.path5[0]

            if round_number == 1:
                round_player.period_return = 0
                round_player.inv_return = 0
            else:
                prev_price = participant.path5[round_number - 2]
                round_player.period_return = (round_player.price - prev_price) / prev_price
                round_player.inv_return = (round_player.price - round_player.purchase_price) / round_player.purchase_price
# PAGES
class Start(Page):
    form_model = "player"

    @staticmethod
    def is_displayed(player):
        return player.round_number==1


class PriceDevelopment(Page):

    @staticmethod
    def vars_for_template(player: Player):

        player.price = player.participant.path5[player.round_number - 1]

        if player.round_number == 1:
            player.purchase_price = player.price
            player.period_return = 0
            player.inv_return = 0
        else:
            prev_player = player.in_round(player.round_number - 1)
            round1_player = player.in_round(1)  # get purchase_price from round 1
            player.period_return = (player.price - prev_player.price ) / prev_player.price
            player.inv_return = (player.price - round1_player.purchase_price ) / round1_player.purchase_price
            player.purchase_price = round1_player.purchase_price


        return dict(
            round=player.round_number-1,
            price=f"{player.price:.2f}",
            purchase_price=f"{player.purchase_price:.2f}",
            period_return=f"{player.period_return * 100:.2f}",
            inv_return=f"{player.inv_return * 100:.2f}",
            period_return_raw = player.period_return,
            inv_return_raw = player.inv_return,
        )

    @staticmethod
    def is_displayed(player: Player):
        return player.participant.treatment == "experience"


class PriceSummary(Page):

    @staticmethod
    def vars_for_template(player: Player):
        purchase_price = player.participant.path5[0]
        rows = []
        prev_price = purchase_price

        for r in range(2, player.round_number + 1):  # skip round 1 (allocation only)
            price = player.participant.path5[r - 1]
            period_return = (price - prev_price) / prev_price
            inv_return = (price - purchase_price) / purchase_price

            rows.append(dict(
                round=r - 1,  # matches your existing "period" numbering
                price=f"{price:.2f}",
                period_return=f"{period_return * 100:.2f}",
                inv_return=f"{inv_return * 100:.2f}",
                period_return_raw=period_return,
                inv_return_raw=inv_return,
            ))
            prev_price = price #assig price to previous price for next round

        return dict(
            round=player.round_number - 1,
            purchase_price=f"{purchase_price:.2f}",
            rows=rows,
        )

    @staticmethod
    def is_displayed(player: Player):
        return (player.round_number == 1 or player.round_number == C.NUM_ROUNDS) and player.participant.treatment != "experience"


class Evaluation(Page):
    form_model = 'player'
    form_fields = ["exp_return"]

    @staticmethod
    def vars_for_template(player: Player):
        return dict(
            round=player.round_number-1,
        )

    @staticmethod
    def is_displayed(player):
        return player.round_number ==11 and player.participant.incentive=="beliefs"

class InvestmentBeliefs(Page):
    form_model = 'player'
    form_fields = ["invest"]

    @staticmethod
    def vars_for_template(player: Player):
        return dict(
            round=player.round_number - 1,
        )

    @staticmethod
    def is_displayed(player):
        return player.round_number ==11 and player.participant.incentive=="beliefs"

class Investment(Page):
    form_model = 'player'
    form_fields = ["invest"]

    @staticmethod
    def vars_for_template(player: Player):
        return dict(
             round=player.round_number - 1,
        )

    @staticmethod
    def is_displayed(player):
        return player.round_number ==11 and player.participant.incentive=="choice"

class EvaluationChoice(Page):
    form_model = 'player'
    form_fields = ["exp_return"]

    @staticmethod
    def vars_for_template(player: Player):
        return dict(
             round=player.round_number - 1,
        )

    @staticmethod
    def is_displayed(player):
        return player.round_number ==11 and player.participant.incentive=="choice"

class FurtherQuestions(Page):
    form_model = 'player'
    form_fields = ["quality", "value", "volatility"]

    @staticmethod
    def vars_for_template(player: Player):
        return dict(
             round=player.round_number - 1,
        )

    @staticmethod
    def is_displayed(player):
        return player.round_number ==11

    @staticmethod
    def before_next_page(player: Player, timeout_happened):
        if player.participant.bonusmarket == 5:
            all_players = player.in_all_rounds()
            period_returns = [p.period_return for p in all_players if p.round_number > 1]
            avg_return = sum(period_returns) / len(period_returns)
            random_return = random.choice(period_returns)

            if player.participant.incentive == "choice":
                player.participant.finalwealth = player.invest * (1 + random_return) + (10000 - player.invest)
                player.participant.bonus = player.participant.finalwealth / 10000
            else:
                if player.exp_return <= (avg_return + 0.02) * 100 and player.exp_return >= (avg_return - 0.02) * 100:
                    player.participant.finalwealth = 0
                    player.participant.bonus = 1.5
                else:
                    player.participant.finalwealth = 0
                    player.participant.bonus = 0


page_sequence = [Start, PriceDevelopment, PriceSummary, Evaluation, InvestmentBeliefs, Investment, EvaluationChoice, FurtherQuestions]