# -*- coding: utf-8 -*-
"""
oTree app: new_decision_making_experiment
Between-subjects design – every participant completes ONE round, assigned to either:

  Control   – grid shows payoffs only (no initials shown), participant types initials.
  AvgZones  – grid shows initials + payoffs, participant types initials.

Bonus = the single round payoff.
"""

from datetime import datetime
import random, json
from otree.api import *

# ─────────────────────────────────────────────────────────────
# 1.  HARD-CODED CELL DEFINITIONS
# ─────────────────────────────────────────────────────────────
GRID_ZONE_A = [
    (1,  18.7, 'KF'), (2, 19.2, 'DL'), (3, 18.9, 'TJ'),
    (4,  19.4, 'QS'), (5, 19.1, 'BP'), (6, 18.6, 'ZM'),
    (7,  19.3, 'HR'), (8, 18.6, 'VC'), (9, 18.8, 'NW'),
    (10, 19.2, 'GL'), (11, 19.5, 'EX'), (12, 18.7, 'UY'),
    (14, 19.0, 'JB'), (15, 18.9, 'SO'), (16, 19.8, 'PK'),
    (17, 18.8, 'RA'), (18, 19.2, 'LM'), (19, 19.3, 'FD'),
    (20, 18.9, 'QT'), (21, 18.1, 'CX'),
]

GRID_ZONE_B = [
    (1,  13.9, 'HV'), (2, 14.2, 'AB'), (3, 14.1, 'PO'),
    (4,  14.3, 'MI'), (5, 14.0, 'RU'), (6, 13.4, 'CY'),
    (7,  13.3, 'TZ'), (8, 13.4, 'EG'), (9, 14.5, 'KB'),
    (10, 114.0,'DS'), (11, 14.3, 'LN'), (12, 14.2, 'WI'),
    (14, 14.1, 'QF'), (15, 14.4, 'SM'), (16, 14.2, 'JO'),
    (17, 14.0, 'VB'), (18, 14.3, 'CP'), (19, 13.8, 'UA'),
    (20, 13.4, 'NK'), (21, 14.2, 'GD'),
]

# ─────────────────────────────────────────────────────────────
# 2.  LOOK-UP TABLE  tag → zone
# ─────────────────────────────────────────────────────────────
TAG2ZONE = {t: 'A' for _, _, t in GRID_ZONE_A if t != '--'}
TAG2ZONE.update({t: 'B' for _, _, t in GRID_ZONE_B if t != '--'})
# Control condition uses fruit words instead of initials
TAG2ZONE['MANGO']   = 'A'   # Mango → Zone A (the "safe" zone)
TAG2ZONE['AVOCADO'] = 'B'   # Avocado → Zone B (the "risky" zone)

# ─────────────────────────────────────────────────────────────
# 3.  GRID BUILDER
# ─────────────────────────────────────────────────────────────
def shuffle_zone(zone_constant):
    entries = [(pay, ini) for idx, pay, ini in zone_constant if ini != '--']
    random.shuffle(entries)
    return [(i+1, pay, ini) for i, (pay, ini) in enumerate(entries)]

# ─────────────────────────────────────────────────────────────
# 4.  PAYOFF HELPER
# ─────────────────────────────────────────────────────────────
def assign_payoff(player, tag):
    """Compute and store payoff based on typed initials."""
    zone = TAG2ZONE[tag]
    player.raw_choice_val = tag
    player.choice_val     = zone   # 'A' or 'B'

    if zone == 'B':
        player.points_this_round = (
            114.0 if random.random() <= 0.05 else 14.0
        )
    else:
        player.points_this_round = 19.0

    player.total_points = player.points_this_round

# ─────────────────────────────────────────────────────────────
# 5.  MODELS
# ─────────────────────────────────────────────────────────────
class C(BaseConstants):
    NAME_IN_URL       = 'new_decision_making_experiment'
    PLAYERS_PER_GROUP = None
    NUM_ROUNDS        = 1

    base_pay = 20      # 20 pence base pay


class Subsession(BaseSubsession):
    pass


class Group(BaseGroup):
    pass


class Player(BasePlayer):
    # demographics
    gender     = models.StringField(choices=['Male', 'Female', 'Other'])
    age        = models.IntegerField(min=15, max=90)
    connect_id = models.StringField()

    # condition label for export
    condition           = models.StringField()
    description_version = models.StringField(doc="1=Safer/Risky framing, 2=Boring/Better framing")
    zone_order          = models.StringField(doc="AB = Zone A on top, BA = Zone B on top")

    # choice tracking
    chosen_initials = models.StringField(blank=True)
    raw_choice_val  = models.StringField()
    choice_val      = models.StringField()

    points_this_round = models.FloatField(initial=0)
    total_points      = models.FloatField(initial=0)

    # payoff page
    total_payoff   = models.FloatField(initial=0)
    received_bonus = models.BooleanField(initial=False)

    # sampling / click tracking
    sampled_cells      = models.LongStringField(initial='[]', doc="JSON list of revealed cells: {tag, zone, payoff, timestamp, order}")
    num_sampled_total  = models.IntegerField(initial=0, doc="Total cells revealed")
    num_sampled_mango  = models.IntegerField(initial=0, doc="Mango (A) cells revealed")
    num_sampled_avocado= models.IntegerField(initial=0, doc="Avocado (B) cells revealed")
    saw_114            = models.BooleanField(initial=False, doc="1 if participant revealed the DS (114) cell")

    # timing
    start_time = models.StringField()
    end_time   = models.StringField()
    total_time = models.FloatField()

# ─────────────────────────────────────────────────────────────
# 6.  SESSION INITIALISATION
# ─────────────────────────────────────────────────────────────
def creating_session(subsession: Subsession):
    participants = subsession.session.get_participants()
    for i, pp in enumerate(participants):
        pp.vars['grid_a']              = shuffle_zone(GRID_ZONE_A)
        pp.vars['grid_b']              = shuffle_zone(GRID_ZONE_B)
        pp.vars['zone_order']          = 'AB' if random.random() < 0.5 else 'BA'
        #pp.vars['description_version'] = '1' if i % 3 == 0 else '2'  # guaranteed 50/50
        ## Cycle through three conditions for even balance
        #condition_cycle = ['control', 'avg_zones', 'description_only']
        #pp.vars['condition'] = condition_cycle[i % 3]
        pp.vars['condition'] = 'avg_zones'
    for p in subsession.get_players():
        p.condition = p.participant.vars['condition']

# ─────────────────────────────────────────────────────────────
# 7.  PAGES
# ─────────────────────────────────────────────────────────────
class Intro(Page):
    form_model  = 'player'
    form_fields = ['gender', 'age', 'connect_id']

    @staticmethod
    def before_next_page(player, timeout_happened):
        now = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        player.start_time = now
        player.participant.vars['start_time'] = now


class Instructions(Page):
    pass


# ─── Control condition ───────────────────────────────────────
class Control(Page):
    form_model  = 'player'
    form_fields = ['chosen_initials']

    @staticmethod
    def is_displayed(player):
        return player.participant.vars.get('condition') == 'control'

    @staticmethod
    def vars_for_template(player):
        zone_order = player.participant.vars.get('zone_order', 'AB')
        # Pass payoffs + initials (initials used server-side for validation only;
        # the template will NOT render the tag divs)
        grid_a = [(idx, f"{pay:.1f}", ini) for idx, pay, ini in player.participant.vars['grid_a']]
        grid_b = [(idx, f"{pay:.1f}", ini) for idx, pay, ini in player.participant.vars['grid_b']]
        return dict(
            zone_order            = zone_order,
            description_version   = player.participant.vars.get('description_version', '1'),
            values_with_indices_a = grid_a,
            values_with_indices_b = grid_b,
        )

    @staticmethod
    def live_method(player, data):
        """Receive a cell-reveal event from the browser and log it."""
        tag     = (data.get('tag') or '').upper()
        zone    = data.get('zone', '')
        payoff  = data.get('payoff', '')
        ts      = data.get('timestamp', '')
        order   = data.get('order', 0)

        log = json.loads(player.sampled_cells or '[]')
        # Only record each cell once
        already = any(e['tag'] == tag for e in log)
        if not already:
            log.append({'tag': tag, 'zone': zone, 'payoff': payoff,
                        'timestamp': ts, 'order': order})
            player.sampled_cells       = json.dumps(log)
            player.num_sampled_total   = len(log)
            player.num_sampled_mango   = sum(1 for e in log if e['zone'] == 'A')
            player.num_sampled_avocado = sum(1 for e in log if e['zone'] == 'B')
            if tag == 'DS':
                player.saw_114 = True
        return {}

    @staticmethod
    def error_message(player, values):
        tag = (values.get('chosen_initials') or '').strip().upper()
        if tag not in ('MANGO', 'AVOCADO'):
            return "Please type Mango or Avocado to select a zone."

    @staticmethod
    def before_next_page(player, timeout_happened):
        player.zone_order          = player.participant.vars.get('zone_order', 'AB')
        player.description_version = player.participant.vars.get('description_version', '1')
        player.chosen_initials     = player.chosen_initials.upper()
        assign_payoff(player, player.chosen_initials)


# ─── AvgZones condition ──────────────────────────────────────
class AvgZones(Page):
    form_model  = 'player'
    form_fields = ['chosen_initials']

    @staticmethod
    def is_displayed(player):
        return player.participant.vars.get('condition') == 'avg_zones'

    @staticmethod
    def vars_for_template(player):
        zone_order = player.participant.vars.get('zone_order', 'AB')
        grid_a = [(idx, f"{pay:.1f}", ini) for idx, pay, ini in player.participant.vars['grid_a']]
        grid_b = [(idx, f"{pay:.1f}", ini) for idx, pay, ini in player.participant.vars['grid_b']]
        return dict(
            values_with_indices_a = grid_a,
            values_with_indices_b = grid_b,
            zone_order            = zone_order,
            description_version   = player.participant.vars.get('description_version', '1'),
        )

    @staticmethod
    def live_method(player, data):
        """Receive a cell-reveal event from the browser and log it."""
        tag     = (data.get('tag') or '').upper()
        zone    = data.get('zone', '')
        payoff  = data.get('payoff', '')
        ts      = data.get('timestamp', '')
        order   = data.get('order', 0)

        log = json.loads(player.sampled_cells or '[]')
        already = any(e['tag'] == tag for e in log)
        if not already:
            log.append({'tag': tag, 'zone': zone, 'payoff': payoff,
                        'timestamp': ts, 'order': order})
            player.sampled_cells       = json.dumps(log)
            player.num_sampled_total   = len(log)
            player.num_sampled_mango   = sum(1 for e in log if e['zone'] == 'A')
            player.num_sampled_avocado = sum(1 for e in log if e['zone'] == 'B')
            if tag == 'DS':
                player.saw_114 = True
        return {}

    @staticmethod
    def error_message(player, values):
        tag = (values.get('chosen_initials') or '').strip().upper()
        if not tag or len(tag) != 2:
            return "Please type the two-letter initials of the investor you wish to imitate."
        if tag not in TAG2ZONE:
            return "The initials you entered were not recognised. Please type valid two-letter initials from the grid."

    @staticmethod
    def before_next_page(player, timeout_happened):
        player.zone_order          = player.participant.vars.get('zone_order', 'AB')
        player.description_version = player.participant.vars.get('description_version', '1')
        player.chosen_initials     = player.chosen_initials.upper()
        assign_payoff(player, player.chosen_initials)


# ─── DescriptionOnly condition ──────────────────────────────
class DescriptionOnly(Page):
    form_model  = 'player'
    form_fields = ['chosen_initials']

    @staticmethod
    def is_displayed(player):
        return player.participant.vars.get('condition') == 'description_only'

    @staticmethod
    def error_message(player, values):
        tag = (values.get('chosen_initials') or '').strip().upper()
        if tag not in ('MANGO', 'AVOCADO'):
            return "Please type Mango or Avocado to select a zone."

    @staticmethod
    def before_next_page(player, timeout_happened):
        player.zone_order          = player.participant.vars.get('zone_order', 'AB')
        player.description_version = player.participant.vars.get('description_version', '1')
        player.chosen_initials     = player.chosen_initials.upper()
        assign_payoff(player, player.chosen_initials)


# ─── Shared pages ────────────────────────────────────────────
class Results(Page):
    @staticmethod
    def vars_for_template(player):
        return dict(total_points=f"{player.total_points:.0f}")


class Payoff(Page):
    @staticmethod
    def vars_for_template(player):
        bonus_amount = player.total_points

        end_time_str = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        player.end_time = end_time_str

        start_time_str = player.participant.vars.get('start_time') or player.field_maybe_none('start_time')
        if start_time_str:
            start_dt = datetime.strptime(start_time_str, '%Y-%m-%d %H:%M:%S')
            end_dt   = datetime.strptime(end_time_str,   '%Y-%m-%d %H:%M:%S')
            player.total_time = (end_dt - start_dt).total_seconds() / 60
        else:
            player.total_time = 0.0

        player.total_payoff   = C.base_pay + bonus_amount
        player.received_bonus = True

        # Prolific completion link based on bonus amount
        if abs(bonus_amount - 114.0) < 0.5:
            submission_link = "https://app.prolific.com/submissions/complete?cc=CLUQD9DX"
        elif abs(bonus_amount - 19.0) < 0.5:
            submission_link = "https://app.prolific.com/submissions/complete?cc=C184QS4S"
        else:
            submission_link = "https://app.prolific.com/submissions/complete?cc=C1HOCDRP"

        return dict(
            total_payoff    = f"£{player.total_payoff/100:.2f}",
            received_bonus  = player.received_bonus,
            bonus_amount    = f"£{bonus_amount/100:.2f}",
            final_points    = bonus_amount,
            start_time      = start_time_str,
            end_time        = end_time_str,
            total_time      = player.total_time,
            submission_link = submission_link,
        )

# ─────────────────────────────────────────────────────────────
# 8.  PAGE SEQUENCE
# ─────────────────────────────────────────────────────────────
page_sequence = [
    Intro,             # demographics
    Instructions,
    Control,           # control condition only – payoffs only, no initials shown
    AvgZones,          # avg_zones condition only – initials + payoffs
    DescriptionOnly,   # description_only condition – text description, no grid
    Results,
    Payoff,
]