from otree.api import *
import json
from decimal import Decimal, ROUND_HALF_UP
from datetime import datetime, timezone
from .lexicon_en import Lexicon as LexiconEN
from .lexicon_fr import Lexicon as LexiconFR

doc = """
Exit questionnaire and payoff
"""

# Helper function language access
def get_language_and_lexicon(player, strict=True):
    from .lexicon_fr import Lexicon as LexiconFR
    from .lexicon_en import Lexicon as LexiconEN
    lang = (
        player.participant.vars['language']
        if strict
        else player.participant.vars.get('language') or player.field_maybe_none('language')
    )
    if lang not in ['fr', 'en']:
        raise ValueError(f"Unsupported or missing language: {lang}")
    if lang == 'fr':
        return LexiconFR, 'fr'
    else:
        return LexiconEN, 'en'

class C(BaseConstants):
    NAME_IN_URL = 'questionnaire'
    PLAYERS_PER_GROUP = None
    NUM_ROUNDS = 1

class Subsession(BaseSubsession):
    pass

class Group(BaseGroup):
    pass


def utc_now_iso():
    return datetime.now(timezone.utc).isoformat()


def append_page_timestamp(player, field_name, page_name):
    raw = player.field_maybe_none(field_name) or '{}'
    try:
        payload = json.loads(raw)
    except json.JSONDecodeError:
        payload = {}
    if field_name == 'page_shown_at_utc' and page_name in payload:
        return
    payload[page_name] = utc_now_iso()
    setattr(player, field_name, json.dumps(payload, ensure_ascii=True))

class Player(BasePlayer):
    gender = models.IntegerField()
    sex_at_birth = models.IntegerField()
    eye_color = models.IntegerField()
    wear_glasses_or_contacts = models.IntegerField()
    vision_problems = models.StringField()
    age = models.IntegerField() # don't set min/max, because it is validated in the page
    education = models.IntegerField()
    work = models.IntegerField()
    income = models.IntegerField()
    finlit1 = models.IntegerField()
    finlit2 = models.IntegerField()
    finlit3 = models.IntegerField()
    experience = models.IntegerField()
    received_advice = models.IntegerField()
    robo_curious = models.IntegerField(blank=True)
    robo_satisfied = models.IntegerField(blank=True)
    trust_friends = models.IntegerField()
    trust_influencers = models.IntegerField()
    trust_advisors = models.IntegerField()
    trust_robos = models.IntegerField()
    difficulty_forecast = models.IntegerField()
    difficulty_allocation = models.IntegerField()
    comfortable_allocation = models.IntegerField()
    advice_helpful = models.IntegerField()
    trust_technology = models.IntegerField()
    regret_decision = models.IntegerField()
    resist_blame = models.IntegerField()
    comments = models.StringField(blank=True)
    played_lottery = models.BooleanField()
    lottery_allocation = models.FloatField()
    lottery_outcome = models.StringField()  # 'up' or 'down'
    lottery_payoff = models.FloatField()
    selected_round = models.IntegerField()
    forecast_value = models.FloatField()
    realized_return = models.FloatField()
    forecast_bonus = models.BooleanField()
    final_tokens = models.FloatField()
    tokens_to_cad = models.FloatField()  # This is the conversion of final tokens to CAD
    final_cad = models.FloatField()
    invest_payoff = models.FloatField()
    investment = models.FloatField()
    page_shown_at_utc = models.LongStringField(blank=True)
    page_submitted_at_utc = models.LongStringField(blank=True)

def calculate_final_payoff(player: Player):
    import random

    pvars = player.participant.vars
    player.played_lottery = random.randint(1, 3) == 1
    lottery_payoff = 0

    # --- Lottery logic ---
    if player.played_lottery:
        investment = pvars.get('decisions', [{}])[0].get('allocation', 0)
        endowment = player.session.config['endowment_lottery']
        mu = player.session.config['mu']
        sigma = player.session.config['sigma']

        draw_up = random.random() < 0.5
        player.lottery_outcome = 'up' if draw_up else 'down'
        ret = 1 + mu + sigma if draw_up else 1 + mu - sigma

        lottery_payoff = investment * ret + (endowment - investment)
        player.lottery_payoff = lottery_payoff
    else:
        player.lottery_outcome = 'none'
        player.lottery_payoff = 0

    # --- Investment round logic ---
    selected = random.randint(1, 3)
    player.selected_round = selected
    invest_player = [
        p for p in player.participant.get_players()
        if p.__class__.__module__.startswith('investment_decisions_real') and p.round_number == selected
    ][0]

    invest_payoff = invest_player.payoff_this_round
    forecast = invest_player.field_maybe_none('return_expectation')
    realized = invest_player.realized_return * 100
  
    player.forecast_value = forecast
    player.realized_return = realized
    player.investment = invest_player.portfolio_allocation
    player.invest_payoff = invest_player.payoff_this_round

    print(forecast)
    print(realized)

    if player.field_maybe_none('forecast_value') is not None:
        if abs(forecast - realized) < 1:
            player.forecast_bonus = True
            invest_payoff += 50
        else:
            player.forecast_bonus = False
    else:
        player.forecast_bonus = False

    # --- Final conversion ---
    player.final_tokens = invest_payoff + lottery_payoff
    show_up_fee = Decimal(str(float(player.session.config['participation_fee'])))
    conversion_rate = Decimal(str(float(player.session.config['real_world_currency_per_point'])))
    final_tokens_dec = Decimal(str(player.final_tokens))

    # Canonical monetary rounding: 2 decimals, half-up, used by both UI and oTree payment fields.
    tokens_to_cad_dec = (final_tokens_dec * conversion_rate).quantize(Decimal('0.01'), rounding=ROUND_HALF_UP)
    final_cad_dec = (show_up_fee + tokens_to_cad_dec).quantize(Decimal('0.01'), rounding=ROUND_HALF_UP)

    player.tokens_to_cad = float(tokens_to_cad_dec)
    player.final_cad = float(final_cad_dec)

    # Keep oTree's built-in payoff aligned to the rounded bonus shown to participants.
    if conversion_rate != 0:
        rounded_payoff_points = tokens_to_cad_dec / conversion_rate
        player.payoff = cu(float(rounded_payoff_points))
    else:
        player.payoff = cu(0)

class Page1(Page):
    form_model = 'player'
    form_fields = ['gender', 'sex_at_birth', 'eye_color', 'wear_glasses_or_contacts', 'vision_problems', 'age', 'education', 'work', 'income']
    
    @staticmethod
    def vars_for_template(player):
        append_page_timestamp(player, 'page_shown_at_utc', Page1.__name__)
        from . import get_language_and_lexicon
        Lexicon, lang = get_language_and_lexicon(player)

        gender_choices = [
            (1, Lexicon.gender_choice_1),
            (2, Lexicon.gender_choice_2),
            (3, Lexicon.gender_choice_3),
            (4, Lexicon.gender_choice_4),
        ]
        sex_at_birth_choices = [
            (1, Lexicon.sex_at_birth_choice_1),
            (2, Lexicon.sex_at_birth_choice_2),
            (3, Lexicon.sex_at_birth_choice_3),
            (4, Lexicon.sex_at_birth_choice_4),
        ]
        eye_color_choices = [
            (1, Lexicon.eye_color_choice_1),
            (2, Lexicon.eye_color_choice_2),
        ]
        wear_glasses_or_contacts_choices = [
            (1, Lexicon.wear_glasses_or_contacts_choice_1),
            (2, Lexicon.wear_glasses_or_contacts_choice_2),
            (3, Lexicon.wear_glasses_or_contacts_choice_3),
        ]
        vision_problems_choices = [
            (1, Lexicon.vision_problems_choice_1),
            (2, Lexicon.vision_problems_choice_2),
            (3, Lexicon.vision_problems_choice_3),
            (4, Lexicon.vision_problems_choice_4),
            (5, Lexicon.vision_problems_choice_5),
            (6, Lexicon.vision_problems_choice_6),
        ]
        selected_vision_problems = [
            int(s) for s in (player.field_maybe_none('vision_problems') or '').split(',') if s.isdigit()
        ]
        vision_problems_serialized = player.field_maybe_none('vision_problems') or ''
        education_choices = [
            (1, Lexicon.education_choice_1),
            (2, Lexicon.education_choice_2),
            (3, Lexicon.education_choice_3),
            (4, Lexicon.education_choice_4),
            (5, Lexicon.education_choice_5),
        ]

        work_choices = [
            (1, Lexicon.work_choice_1),
            (2, Lexicon.work_choice_2),
            (3, Lexicon.work_choice_3),
            (4, Lexicon.work_choice_4),
            (5, Lexicon.work_choice_5),
        ]
        income_choices = [
            (1, Lexicon.income_choice_1),
            (2, Lexicon.income_choice_2),
            (3, Lexicon.income_choice_3),
            (4, Lexicon.income_choice_4),
            (5, Lexicon.income_choice_5),
            (6, Lexicon.income_choice_6),
        ]

        return dict(
            Lexicon=Lexicon,
            lang=lang,
            gender_choices=gender_choices,
            sex_at_birth_choices=sex_at_birth_choices,
            eye_color_choices=eye_color_choices,
            wear_glasses_or_contacts_choices=wear_glasses_or_contacts_choices,
            vision_problems_choices=vision_problems_choices,
            selected_vision_problems=selected_vision_problems,
            vision_problems_serialized=vision_problems_serialized,
            education_choices=education_choices,
            work_choices=work_choices,
            income_choices=income_choices,
        )   

    @staticmethod
    def before_next_page(player, timeout_happened):
        append_page_timestamp(player, 'page_submitted_at_utc', Page1.__name__)


    def error_message(player, values):
        from . import get_language_and_lexicon
        Lexicon, _ = get_language_and_lexicon(player)

        errors = dict()

        # Age validation
        age = values.get('age')
        if age is not None and not (18 <= age <= 99):
            errors['age'] = Lexicon.age_out_of_range

        # Vision problems is a multi-select serialized as a comma-separated string.
        if not values.get('vision_problems'):
            errors['vision_problems'] = Lexicon.error_incomplete

        return errors if errors else None
    
class Page2(Page):
    form_model = 'player'
    form_fields = ['finlit1','finlit2','finlit3']
    
    @staticmethod
    def vars_for_template(player):
        append_page_timestamp(player, 'page_shown_at_utc', Page2.__name__)
        from . import get_language_and_lexicon
        Lexicon, lang = get_language_and_lexicon(player)
        finlit1_choices = [
            (1, Lexicon.finlit1_choice_1),
            (2, Lexicon.finlit1_choice_2),
            (3, Lexicon.finlit1_choice_3),
            (4, Lexicon.finlit1_choice_4),
            (5, Lexicon.finlit1_choice_5),
        ]
        finlit2_choices = [
            (1, Lexicon.finlit2_choice_1),
            (2, Lexicon.finlit2_choice_2),
            (3, Lexicon.finlit2_choice_3),
            (4, Lexicon.finlit2_choice_4),
            (5, Lexicon.finlit2_choice_5),
        ]
        finlit3_choices = [
            (1, Lexicon.finlit3_choice_1),
            (2, Lexicon.finlit3_choice_2),
            (3, Lexicon.finlit3_choice_3),
            (4, Lexicon.finlit3_choice_4),
        ]

        return dict(
            Lexicon=Lexicon,
            lang=lang,
            finlit1_choices=finlit1_choices,
            finlit2_choices=finlit2_choices,
            finlit3_choices=finlit3_choices,
        )

    @staticmethod
    def before_next_page(player, timeout_happened):
        append_page_timestamp(player, 'page_submitted_at_utc', Page2.__name__)
    
class Page3(Page):
    form_model = 'player'
    form_fields = ['experience','received_advice','robo_curious','robo_satisfied','trust_friends', 'trust_influencers', 'trust_advisors', 'trust_robos']
    
    # need a custom error message if dynamic follow-up questions aren't answered
    def error_message(player, values):
        from . import get_language_and_lexicon
        Lexicon, _ = get_language_and_lexicon(player)

        if values['received_advice'] == 2 and not values.get('robo_satisfied'):
            return Lexicon.error_incomplete

        if values['received_advice'] in [1, 3] and not values.get('robo_curious'):
            return Lexicon.error_incomplete
        
    @staticmethod
    def vars_for_template(player):
        append_page_timestamp(player, 'page_shown_at_utc', Page3.__name__)
        from . import get_language_and_lexicon
        Lexicon, lang = get_language_and_lexicon(player)
        experience_choices = [
            (1, Lexicon.experience_choice_1),
            (2, Lexicon.experience_choice_2),
            (3, Lexicon.experience_choice_3),
            ]
        received_advice_choices = [
            (1, Lexicon.received_advice_choice_1),
            (2, Lexicon.received_advice_choice_2),
            (3, Lexicon.received_advice_choice_3),
            (4, Lexicon.received_advice_choice_4),
        ]
        robo_curious_choices = [
            (1, Lexicon.robo_curious_choice_1),
            (2, Lexicon.robo_curious_choice_2),
            (3, Lexicon.robo_curious_choice_3),
        ]
        robo_satisfied_choices = [
            (1, Lexicon.robo_satisfied_choice_1),
            (2, Lexicon.robo_satisfied_choice_2),
            (3, Lexicon.robo_satisfied_choice_3),
            (4, Lexicon.robo_satisfied_choice_4),
        ]

        rows = []
        for field_name, label in Lexicon.trust_sources:
            value = player.field_maybe_none(field_name)
            rows.append(dict(
                field_name=field_name,
                label=label,
                value=value
            ))

        return dict(
            Lexicon=Lexicon,
            lang=lang,
            experience_choices=experience_choices,
            received_advice_choices=received_advice_choices,
            robo_curious_choices=robo_curious_choices,
            robo_satisfied_choices=robo_satisfied_choices,
            trust_question=Lexicon.trust_question,
            trust_sources=Lexicon.trust_sources,
            likert_labels=Lexicon.likert_labels,
            rows=rows,
        )
    

    def before_next_page(player,timeout_happened):
        append_page_timestamp(player, 'page_submitted_at_utc', Page3.__name__)
        if player.received_advice is None:
            return  # skip logic if form didn’t validate

        if player.received_advice == 2:  # Option 2 → robo_satisfied
            player.robo_curious = None
        elif player.received_advice in [1, 3]:  # Options 1 or 3 → robo_curious
            player.robo_satisfied = None

class Page4(Page):
    form_model = 'player'
    form_fields = ['difficulty_forecast', 'difficulty_allocation', 'comfortable_allocation','advice_helpful','trust_technology','regret_decision','resist_blame','comments']
    
    @staticmethod
    def vars_for_template(player):
        append_page_timestamp(player, 'page_shown_at_utc', Page4.__name__)
        from . import get_language_and_lexicon
        Lexicon, lang = get_language_and_lexicon(player)
        difficulty_forecast_choices = [
            (1, Lexicon.difficulty_forecast_choice_1),
            (2, Lexicon.difficulty_forecast_choice_2),
            (3, Lexicon.difficulty_forecast_choice_3),
            (4, Lexicon.difficulty_forecast_choice_4),
            (5, Lexicon.difficulty_forecast_choice_5),
        ]
        difficulty_allocation_choices = [
            (1, Lexicon.difficulty_allocation_choice_1),
            (2, Lexicon.difficulty_allocation_choice_2),
            (3, Lexicon.difficulty_allocation_choice_3),
            (4, Lexicon.difficulty_allocation_choice_4),
            (5, Lexicon.difficulty_allocation_choice_5),
        ]

        comfortable_allocation_choices = [
            (1, Lexicon.comfortable_allocation_choice_1),
            (2, Lexicon.comfortable_allocation_choice_2),
            (3, Lexicon.comfortable_allocation_choice_3),
            (4, Lexicon.comfortable_allocation_choice_4),
            (5, Lexicon.comfortable_allocation_choice_5),
            (6, Lexicon.comfortable_allocation_choice_6),
        ]
        advice_helpful_choices = [
            (1, Lexicon.advice_helpful_choice_1),
            (2, Lexicon.advice_helpful_choice_2),
            (3, Lexicon.advice_helpful_choice_3),
            (4, Lexicon.advice_helpful_choice_4),
            (5, Lexicon.advice_helpful_choice_5),
        ]
        trust_technology_choices = [
            (1, Lexicon.trust_technology_choice_1),
            (2, Lexicon.trust_technology_choice_2),
            (3, Lexicon.trust_technology_choice_3),
            (4, Lexicon.trust_technology_choice_4),
            (5, Lexicon.trust_technology_choice_5),
            (6, Lexicon.trust_technology_choice_6),
        ] 
        regret_decision_choices = [
            (1, Lexicon.regret_decision_choice_1),
            (2, Lexicon.regret_decision_choice_2),
            (3, Lexicon.regret_decision_choice_3),
            (4, Lexicon.regret_decision_choice_4),
            (5, Lexicon.regret_decision_choice_5),
            (6, Lexicon.regret_decision_choice_6),
            (7, Lexicon.regret_decision_choice_7),
        ]
        resist_blame_choices = [
            (1, Lexicon.resist_blame_choice_1),
            (2, Lexicon.resist_blame_choice_2),
            (3, Lexicon.resist_blame_choice_3),
            (4, Lexicon.resist_blame_choice_4),
            (5, Lexicon.resist_blame_choice_5),
            (6, Lexicon.resist_blame_choice_6),
            (7, Lexicon.resist_blame_choice_7),
        ]

        return dict(
            Lexicon=Lexicon,
            lang=lang,
            difficulty_forecast_choices=difficulty_forecast_choices,
            difficulty_allocation_choices=difficulty_allocation_choices,
            comfortable_allocation_choices=comfortable_allocation_choices,
            advice_helpful_choices=advice_helpful_choices,
            trust_technology_choices=trust_technology_choices,
            regret_decision_choices=regret_decision_choices,
            resist_blame_choices=resist_blame_choices,
        )   

    @staticmethod
    def before_next_page(player, timeout_happened):
        append_page_timestamp(player, 'page_submitted_at_utc', Page4.__name__)

class Payoff(Page):
    
    def vars_for_template(player):
        append_page_timestamp(player, 'page_shown_at_utc', Payoff.__name__)
        Lexicon, lang = get_language_and_lexicon(player)
        # Freeze random payoff draws after first computation for this participant.
        if player.field_maybe_none('final_tokens') is None:
            calculate_final_payoff(player)
        if player.lottery_outcome == 'up':
            raw_return = player.session.config['mu'] + player.session.config['sigma']
        elif player.lottery_outcome == 'down':
            raw_return = player.session.config['mu'] - player.session.config['sigma']
        else:
            raw_return = 0
        lottery_return_signed = f"{raw_return * 100:+.0f}"  # e.g. "+27"
        investment_return_signed = f"{player.realized_return:+.1f}"  # e.g. "+27.1"    
        
        return dict(
            Lexicon=Lexicon,
            lang=lang,
            show_up_fee = player.session.config['participation_fee'],
            played_lottery=player.played_lottery,
            lottery_allocation= player.participant.vars.get('decisions', [{}])[0].get('allocation', 0), 
            lottery_tokens=player.lottery_payoff,
            lottery_outcome=player.lottery_outcome,
            lottery_return_signed=lottery_return_signed,
            selected_round=player.selected_round,
            forecast_value=player.field_maybe_none('forecast_value'),
            realized_return=player.realized_return,
            forecast_bonus=player.forecast_bonus,
            investment_allocation=player.investment,
            investment_return_signed=investment_return_signed,
            invest_payoff=player.invest_payoff,
            total_tokens=player.final_tokens,
            tokens_to_cad=player.tokens_to_cad,
            final_cad=player.final_cad,
        )

    @staticmethod
    def before_next_page(player, timeout_happened):
        # Defensive: ensure final payoff exists even if template was not fully rendered.
        if player.field_maybe_none('final_tokens') is None:
            calculate_final_payoff(player)
        append_page_timestamp(player, 'page_submitted_at_utc', Payoff.__name__)
            
page_sequence = [Page4, Page2, Page3, Page1, Payoff]