from otree.api import (
    Currency,
    cu,
    currency_range,
    models,
    widgets,
    BaseConstants,
    BaseSubsession,
    BaseGroup,
    BasePlayer,
    ExtraModel,
    WaitPage,
    Page,
    read_csv,
    url_of_static_file,
)

import units
import shared_out

doc = ''
class C(BaseConstants):
    # built-in constants
    NAME_IN_URL = 'irrigation'
    PLAYERS_PER_GROUP = 4
    NUM_ROUNDS = 9
    # user-defined constants
    INITIAL_BALANCE = 5
    MAX_INVESTMENT = 3
    MAX_SPENDING = 5
    MIN_INVESTMENT_FOR_WATER = 2
    INVESTMENT_THRESHOLD_MAX = 8
    MAX_GROUNDWATER = 12
    WATER_MULTIPLIER = 2
    WATER_OFFSET = 4
class Subsession(BaseSubsession):
    surface_water = models.IntegerField()
class Group(BaseGroup):
    total_investment = models.IntegerField()
    groundwater_produced = models.IntegerField()
    total_water_available = models.IntegerField()
    water_outflow = models.IntegerField()
class Player(BasePlayer):
    investment = models.IntegerField(initial=0, label='Your investment in groundwater (0-3 tokens)', max=C.MAX_INVESTMENT, min=0)
    card_choice = models.StringField(choices=['None', 'A1', 'A2', 'A3', 'A4', 'A5'], initial='None', label='Choose your activity card')
    water_received = models.IntegerField()
    card_success = models.BooleanField()
    round_payoff = models.IntegerField()
    cumulative_balance = models.IntegerField()
    water_input_needed = models.IntegerField()
    water_output_released = models.IntegerField()
    card_cost = models.IntegerField()
    card_output = models.IntegerField()
    selected_round_for_payment = models.IntegerField()
# built-in hook function(s) (called automatically by oTree)
# <hook-functions>
def creating_session(subsession: Subsession):
    
    for player in subsession.get_players():
        player.cumulative_balance = C.INITIAL_BALANCE
    subsession.surface_water = subsession.session.config['surface_water']
    
# </hook-functions>
# the below function(s) are user-defined, not called by oTree
# <helper-functions>
def calculate_groundwater(total_investment):
    
    if total_investment < C.MIN_INVESTMENT_FOR_WATER:
        return 0
    elif total_investment <= C.INVESTMENT_THRESHOLD_MAX:
        return C.WATER_MULTIPLIER * total_investment - C.WATER_OFFSET
    else:
        return C.MAX_GROUNDWATER
    
def get_card_properties(card_name):
    
    return card_data()[card_name]
    
def card_data():
    
    return {
        "A1": {"water_input": 0, "cost": 2, "water_output": 2, "financial_output": 1},
        "A2": {"water_input": 0, "cost": 1, "water_output": 1, "financial_output": 1},
        "A3": {"water_input": 2, "cost": 1, "water_output": 1, "financial_output": 3},
        "A4": {"water_input": 2, "cost": 1, "water_output": 0, "financial_output": 4},
        "A5": {"water_input": 3, "cost": 2, "water_output": 0, "financial_output": 6},
        "None": {"water_input": 0, "cost": 0, "water_output": 0, "financial_output": 0}
    }
    
def process_water_flow(group: Group):
    
    water_available = group.total_water_available
    for player in group.get_players():
        card_props = get_card_properties(player.card_choice)
        player.water_input_needed = card_props['water_input']
        player.card_cost = card_props['cost']
        player.card_output = card_props['financial_output']
        player.water_output_released = card_props['water_output']
    
        if water_available >= player.water_input_needed:
            player.card_success = True
            water_available -= player.water_input_needed
            water_available += player.water_output_released
            player.water_received = player.water_input_needed
        else:
            player.card_success = False
            player.water_received = water_available
    
        previous_balance = player.in_round(player.round_number - 1).cumulative_balance if player.round_number > 1 else C.INITIAL_BALANCE
    
        if player.card_success:
            player.round_payoff = player.card_output - player.card_cost - player.investment
        else:
            player.round_payoff = 0 - player.card_cost - player.investment
    
        player.cumulative_balance = previous_balance + player.round_payoff
    
    group.water_outflow = water_available
    
# </helper-functions>
class Investment(Page):
    form_model = 'player'
    form_fields = ['investment']
    @staticmethod
    def vars_for_template(player: Player):
        
        return dict(
            round_number=player.round_number,
            position=player.id_in_group,
            balance=player.in_round(player.round_number - 1).cumulative_balance if player.round_number > 1 else C.INITIAL_BALANCE
        )
        
class AfterInvestment(WaitPage):
    wait_for_all_groups = True
    @staticmethod
    def after_all_players_arrive(subsession: Subsession):
        
        for group in subsession.get_groups():
            group.total_investment = sum([p.investment for p in group.get_players()])
            group.groundwater_produced = calculate_groundwater(group.total_investment)
            group.total_water_available = subsession.surface_water + group.groundwater_produced
        
class CardChoice(Page):
    form_model = 'player'
    form_fields = ['card_choice']
    @staticmethod
    def vars_for_template(player: Player):
        
        cards = card_data()
        return dict(
            position=player.id_in_group,
            cards=cards,
            groundwater=player.group.groundwater_produced,
            total_water=player.group.total_water_available
        )
        
    @staticmethod
    def error_message(player: Player, values):
        
        card_choice = values['card_choice']
        card_props = get_card_properties(card_choice)
        previous_balance = player.in_round(player.round_number - 1).cumulative_balance if player.round_number > 1 else C.INITIAL_BALANCE
        total_cost = player.investment + card_props['cost']
        if total_cost > C.MAX_SPENDING:
            return f"Total spending ({total_cost} tokens) exceeds maximum of {C.MAX_SPENDING} tokens"
        if total_cost > previous_balance:
            return f"Insufficient balance. You have {previous_balance} tokens but need {total_cost} tokens"
        
class AfterCardChoice(WaitPage):
    wait_for_all_groups = True
    @staticmethod
    def after_all_players_arrive(subsession: Subsession):
        
        for group in subsession.get_groups():
            process_water_flow(group)
        
class Results(Page):
    form_model = 'player'
    @staticmethod
    def is_displayed(player: Player):
        
        import random
        if player.round_number == C.NUM_ROUNDS:
            player.selected_round_for_payment = random.randint(1, C.NUM_ROUNDS)
            player.participant.payoff = player.in_round(player.selected_round_for_payment).round_payoff
        return True
        
    @staticmethod
    def vars_for_template(player: Player):
        
        return dict(
            position=player.id_in_group,
            round_number=player.round_number,
            total_rounds=C.NUM_ROUNDS
        )
        
page_sequence = [Investment, AfterInvestment, CardChoice, AfterCardChoice, Results]