from otree.api import *

# Import risk params
from input.riskParams import expParts ## MARK: Only keep this!
from random import randint
import numpy as np
import matplotlib as plt
from matplotlib import cm
import json # Save dict as string

doc = """
Your app description
"""

# MARK: Change study number here to make pages visible
thisStudyId = "priceList"


# Function to make feedback field
def make_field(label):
    return models.LongStringField(
        blank=True,
        label=label,
    )

class Constants(BaseConstants):
    name_in_url = 'mpl'
    players_per_group = None
    num_rounds = 1


class Subsession(BaseSubsession):
    pass


class Group(BaseGroup):
    pass


class Player(BasePlayer):
    # Select fixed amount
    randomStudy = models.StringField() # Save name of randomly selected study
    study = models.IntegerField()
    compQuestion_risk_1 = models.StringField()
    compQuestion_risk_2 = models.StringField()
    
    # Safe info about instruction choices
    choice_Instructions = models.IntegerField() # Save shell selected during instructions
    outcome_Instructions = models.StringField() # Save risk outcome during instructions
    multiSwitch_Instructions = models.IntegerField(
        initial=0
        )
    switches_Instructions = models.LongStringField()

    
    # Select color to determine type
    colorSelect=models.IntegerField()

    # Save key variable
    choice_Study2 = models.IntegerField()

    # Save inconsistent switches
    multiSwitch = models.IntegerField(
        initial=0
        )
    inconsistentSwitches = models.LongStringField()
    
    # Internal Feedback
    feedback_colorSelect = make_field("Feedback commentary:")
    feedback_generalTask = make_field("Feedback commentary:")
    feedback_option1 = make_field("Feedback commentary:")
    feedback_option2 = make_field("Feedback commentary:")
    feedback_attentionScreen = make_field("Feedback commentary:")
    feedback_task = make_field("Feedback commentary:")


# PAGES
class taskExplain(Page):
    form_model = "player"
    form_fields = [
        "feedback_generalTask",
        ]

    @staticmethod
    def is_displayed(player):
        # Only in the first round and correct study
        return (player.round_number == 1) and (player.participant.studyName == thisStudyId)
    
    @staticmethod
    def js_vars(player: Player):
        return dict(
            testing = player.session.config["testing"],
        )
    
    @staticmethod
    def vars_for_template(player: Player):
        return dict(
            testing = player.session.config["testing"]
        ) 

class option1(Page):
    form_model = "player"
    form_fields = [
        "feedback_option1",
        "choice_Instructions",
        "outcome_Instructions",
        ]
    
    @staticmethod
    def is_displayed(player: Player):
        # Only in the first round and correct study
        return (player.round_number == 1) and (player.participant.studyName == thisStudyId)
    
    @staticmethod
    def js_vars(player: Player):
        return dict(
            #payoffs = ['🍐', '🍒'],
            payoffs =  ['£1.10', '£0.90'],
            treatment = player.participant.treatment
        ) 
    
    @staticmethod
    def vars_for_template(player: Player):
        return dict(
            testing = player.session.config["testing"],
        ) 

class option2(Page):
    form_model = "player"
    form_fields = [
        "feedback_option2"
        ]

    @staticmethod
    def live_method(player, data):
        player.multiSwitch_Instructions += 1
        player.switches_Instructions = json.dumps(data)


    @staticmethod
    def is_displayed(player: Player):
        # Only in the first round and correct study
        return (player.round_number == 1) and (player.participant.studyName == thisStudyId)
    
    @staticmethod
    def js_vars(player: Player):
        fixedAmount = [0.50, 1.50]
        increment = 0.10
        rows = int((max(fixedAmount) - min(fixedAmount))/increment) + 1
        allAmounts = np.linspace(min(fixedAmount),max(fixedAmount), num=rows, endpoint=True ).tolist()

        amount = []
        for index, element in enumerate(allAmounts):
            amount.append([index, '%.2f' % element])

        return dict(
            payoffs =  ['£1.10', '£0.90'],
            amounts = amount,
        ) 

    @staticmethod
    def vars_for_template(player: Player):
        fixedAmount = [0.50, 1.50]
        increment = 0.10
        rows = int((max(fixedAmount) - min(fixedAmount))/increment) + 1
        allAmounts = np.linspace(min(fixedAmount),max(fixedAmount), num=rows, endpoint=True ).tolist()

        amount = []
        for index, element in enumerate(allAmounts):
            amount.append([index, '%.2f' % element])

        return dict(
            amounts = amount,
            numOptions = amount[-1][0] + 1,
            testing = player.session.config["testing"],
        ) 

class colorSelect(Page):
    form_model = "player"
    form_fields = [
        "feedback_colorSelect",
        "colorSelect",
        ]

    @staticmethod
    def is_displayed(player):
        # Only in the first round and correct study
        return (player.round_number == 1) and (player.participant.studyName == thisStudyId)
    
    @staticmethod
    def js_vars(player: Player):
        return dict(
            testing = player.session.config["testing"],
        )
    
    def error_message(player, values):
        if values["colorSelect"] == -9999:
            msg = "Please select a color before proceeding."
            return msg
    
    @staticmethod
    def vars_for_template(player: Player):
        colorScheme = [ [idx,val[0]] for idx, val in player.participant.typeMatching.items()] 
        print(colorScheme)
        return dict(
            testing = player.session.config["testing"],
            numColors = len(colorScheme),
            colorScheme = colorScheme, 
            treatment = player.participant.treatment
        ) 


    @staticmethod
    def before_next_page(player, timeout_happened):
        player.participant.type = player.colorSelect

        tmpData = player.participant.typeMatching[player.colorSelect]
        player.participant.color= tmpData[0]
        

class attentionScreen(Page):
    form_model = "player"
    form_fields = [
        "feedback_attentionScreen"
        ]
    
    @staticmethod
    def is_displayed(player: Player):
        # Only in the first round and correct study
        return (player.round_number == 1) and (player.participant.studyName == thisStudyId)
    
    @staticmethod
    def vars_for_template(player: Player):
        return dict(
            testing = player.session.config["testing"],
        ) 

class task(Page):
    form_model = "player"
    form_fields = [
        "feedback_task", 
        "choice_Study2"
        ]
    
    @staticmethod
    def live_method(player, data):
        player.multiSwitch += 1
        player.inconsistentSwitches = json.dumps(data) 
    
    def error_message(player, values):
        if values["choice_Study2"] == -9999:
            msg = "Please select a switching row before proceeding."
            return msg

    @staticmethod
    def is_displayed(player: Player):
        # Only in the first round and correct study
        return (player.round_number == 1) and (player.participant.studyName == thisStudyId)
    
    @staticmethod
    def js_vars(player: Player):
        studyKey = list(expParts.keys())[player.participant.study]
        studyProps = expParts[studyKey]

        payoffs = []
        probabilities = []

        for elem in studyProps[2]:
            payoffs.append(elem[0])
            probabilities.append(elem[1])
        
        allAmounts = [val[2] for val in player.participant.typeMatching.values()]
        allAmounts.sort()

        amount = []
        for index, element in enumerate(allAmounts):
            amount.append([index, '%.2f' % element])

        return dict(
            payoffs = payoffs,
            probabilities = probabilities,    
            fixedAmounts = amount,       
        ) 
    
    @staticmethod
    def vars_for_template(player: Player):
        allAmounts = [val[2] for val in player.participant.typeMatching.values()]
        allAmounts.sort()
        
        amount = []
        for index, element in enumerate(allAmounts):
            amount.append([index, '%.2f' % element])

        return dict(
            amounts = amount,
            numOptions = len(amount),
            testing = player.session.config["testing"],
        ) 
    
    @staticmethod
    def app_after_this_page(player, upcoming_apps):
        # Always go to risk resolution
        return 'riskResolution'
    
    
    @staticmethod
    def before_next_page(player, timeout_happened):
        player.participant.singleChoice = player.choice_Study2


page_sequence = [
    taskExplain,
    option1,
    option2,
    colorSelect, 
    attentionScreen,
    task
    ]