from otree.api import *
# import libraries
import matplotlib as plt
from matplotlib import cm

import numpy as np

# Import skripts to write PDFs
import input.createBonusFiles as bf #Skript for bonus files
import input.createVerification as vf #Skript for verification file

# Import information about risks
from input.riskParams import expParts ## MARK: Only keep this!

doc = """
Your app description
"""

# Make textfield for feedback
def make_field(label):
    return models.LongStringField(
        blank=True,
        label=label,
    )

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

class Subsession(BaseSubsession):
    pass


class Group(BaseGroup):
    pass


class Player(BasePlayer):   
    # Select study during testing
    randomStudy = models.StringField() # Save name of randomly selected study
    study = models.IntegerField(
        widget=widgets.RadioSelect
        )

    # Prolific ID
    prolID = models.StringField(
        label="<b>Please enter your Prolific ID</b> before we start. This allows us to approve your submission and transfer your payment to your account:",
        )
    
    # Password
    password_bonus = models.LongStringField()
    password_verification = models.LongStringField()
    
    # Privacy Policy
    consent_privacy = models.BooleanField(
        label="I have read the privacy policy and consent with it:",
        choices=[[True, "Yes"], [False, "No"]],
    )

    # Verification PDF
    riskResolution = models.LongStringField() # Shuffled positions for risk resolutions
    verification_download_bool = models.BooleanField(initial=False)
    verification_popup = models.BooleanField(initial=False)

    # Bonus Files
    #option2 = models.LongStringField()
    bonusFileDict = models.LongStringField()
    bfDownload = models.BooleanField(
        initial=False,
    )
    chosenBF = models.StringField(initial="None")


    # Select treatment for Testing
    # MARK: Must be changed to automatic assignmeent 
    # Balanced Treatments: https://otree.readthedocs.io/en/latest/treatments.html
    selectTreatment = models.IntegerField(
        label="For testing:",
        choices=[
            [0, "Verification"],
            [1, "No verification"],
        ],
        blank=True,
        initial=0,
    )

    # Programming feedback
    feedback_welcomePage = make_field("Feedback commentary:")
    feedback_exitPage = make_field("Feedback commentary:")
    feedback_noDeception = make_field("Feedback commentary:")
    feedback_dwnldVerification = make_field("Feedback commentary:")
    feedback_dwnldBonusFiles = make_field("Feedback commentary:")
    feedback_generalTask = make_field("Feedback commentary:")


# FUNCTIONS
def study_choices(player):
    availStudy = []
    for idx, key in enumerate(expParts.keys()):
        #desc = key + ": " + expParts[key][0]
        availStudy.append([idx, expParts[key][0]])
    return availStudy

def create_bonusFiles(player):
    print("BonusFile")
     

def creating_session(subsession):
    import itertools
    studyList = list(expParts.keys())
    treatmentList = [0,1]

    studyAssignment = []
    for study in studyList:
        for treatment in treatmentList:
            studyAssignment.append([study,treatment])

    if subsession.round_number == 1:
        selectStudy = itertools.cycle(studyAssignment)
        for p in subsession.get_players():
            # Set initial choice for study selection
            thisStudy = next(selectStudy)
            if p.session.config["testing"]:
                #p.study = len(expParts.keys()) - 1
                p.study = 2 ## Price List
            elif p.session.config["study"] == "Study 1":
                # Randomly assign study only if Study 1
                p.randomStudy = thisStudy[0]
                p.study = list(expParts.keys()).index(p.randomStudy)
            else:
                p.study = list(expParts.keys()).index(p.session.config["study"])
            #p.selectTreatment = thisStudy[1]
            p.session.participantCounter = 0 # Initialize with zero; Only used for substudies Study 1a - 1c
            p.session.choicePageReached_verify = 0 # Initialize with zero; Only used for substudies Study 1a - 1c
            p.session.choicePageReached_noVerify = 0 # Initialize with zero; Only used for substudies Study 1a - 1c


# PAGES
class taskSelection_forTest(Page):
    form_model = "player"
    form_fields = [
        "study"
        ]
    
    @staticmethod
    def is_displayed(player):
        return player.session.config["testing"]


class welcomePage(Page):
    form_model = "player"
    form_fields = [
        "prolID",
        "consent_privacy",
        "selectTreatment",
        "feedback_welcomePage"
        ]
    
    @staticmethod
    def js_vars(player: Player):         
        return dict(
            testing = player.session.config["testing"],
            )
   
    @staticmethod
    def vars_for_template(player: Player):
        attentionChecks = False

        if player.study == 2:
            attentionChecks = True

        return dict(
            attentionChecks = attentionChecks,
            testing = player.session.config["testing"],
            )

    @staticmethod
    def prolID_error_message(player, value):
        if not value:
            return 'Please enter your Prolific ID.'

    @staticmethod
    def app_after_this_page(player, upcoming_apps):
        if player.selectTreatment == 1 and player.consent_privacy:
                return player.field_display('study') # Gives name of study ("binaryChoice", "priceList", etc.)
    
    @staticmethod
    def before_next_page(player, timeout_happened):
        if player.consent_privacy:
            # Assign Treatment
            player.session.participantCounter += 1
            if not player.session.config["testing"]:                
                player.selectTreatment = player.session.participantCounter % 2 

            # Save variables
            player.participant.study = player.study
            player.participant.studyName = player.field_display('study')
            player.participant.treatment = player.selectTreatment
            player.participant.color = "#ffffff00" # hex color saved later - transparent here
            player.participant.shape = "d-none" # shape saved later - hidden here
            player.participant.typeMatching = [] # Initialize
            
            # Generate random data for verification file
            studyKey = list(expParts.keys())[player.study]
            data = expParts[studyKey]
            _, resolutionType, risk, fixedAmount  = data

            # Generate colorbar if study is study 2
            if studyKey == "Study 2":
                increment = 0.1
                n_colors = int((max(fixedAmount) - min(fixedAmount))/increment) + 1

                allAmounts_tmp = np.linspace(min(fixedAmount),max(fixedAmount), num=n_colors, endpoint=True ).tolist()
                allAmounts = [round(amount, 2) for amount in allAmounts_tmp] 

                if n_colors <= 20:
                    colors = cm.tab20(np.linspace(0, 1, n_colors)).tolist()
                elif 20 < n_colors <= 40:
                    tmpPositions = np.linspace(0, 1, 20)
                    colors = cm.tab20(tmpPositions).tolist()
                    for i in range(n_colors-20):
                        #Append from different colormap
                        colors.append(cm.gist_rainbow(tmpPositions[i]))

                hexcolor = [plt.colors.to_hex(rgba) for rgba in colors]
                
                # Dictionary with {"Num" : [hexCol, newRow, payOff]} 
                player.participant.typeMatching = vf.shuffleColorTypes(hexcolor, allAmounts)
            
            elif studyKey == "Study 2a":
                increment = 0.1
                n_amounts = int((max(fixedAmount) - min(fixedAmount))/increment) + 1

                allAmounts_tmp = np.linspace(min(fixedAmount),max(fixedAmount), num=n_amounts, endpoint=True ).tolist()
                allAmounts = [round(amount, 2) for amount in allAmounts_tmp]

                player.participant.typeMatching = vf.shuffleAmounts(allAmounts)           
                
            elif studyKey == "Study 3" or studyKey == "Study 3a":
                shapes = ["square", "circle", "triangle", "pentagon", "trapezoid"]
                n_colors = int(101 / len(shapes)) + 1 
                
                if n_colors <= 20:
                    colors = cm.tab20(np.linspace(0, 1, n_colors)).tolist()
                elif 20 < n_colors <= 40:
                    tmpPositions = np.linspace(0, 1, 20)
                    colors = cm.tab20(tmpPositions).tolist()
                    for i in range(n_colors-20):
                        #Append from different colormap
                        colors.append(cm.gist_rainbow(tmpPositions[i]))

                hexcolor = [plt.colors.to_hex(rgba) for rgba in colors]

                player.participant.typeMatching = vf.matchingProbabilities(hexcolor, shapes)
                
            try:
                # Ensure that payoffs only shuffled once 
                print(player.participant.shuffledPayoff)
            except:
                # Save shuffled outcomes for risk resolution
                shuffledPayoff = vf.shuffleOutcomes(risk, resolutionType)
                player.participant.shuffledPayoff = shuffledPayoff
                player.riskResolution = str(shuffledPayoff)
            
                if player.session.config['testing']:
                    player.password_verification = "Loewe50"
                else:
                    player.password_verification = bf.createPassword(5)
                
                player.participant.password_verification = player.password_verification


class exitPage(Page):
    form_model = "player"
    form_fields = [
        "feedback_exitPage"
        ]

    @staticmethod
    def is_displayed(player):
        return player.consent_privacy == False
    
    @staticmethod
    def vars_for_template(player: Player):
        return dict(
            testing = player.session.config["testing"],
        )

class dwnldVerification(Page):
    form_model = "player"
    form_fields = [
        "verification_download_bool",
        "verification_popup",
        "feedback_dwnldVerification"
        ]

    @staticmethod
    def is_displayed(player):
        # Write verification pdf
        studyKey = list(expParts.keys())[player.study]
        data = expParts[studyKey]
        studyName, _, _, _  = data

        vf.createVerificationPDF(
            studyName,
            player.participant.shuffledPayoff,
            player.password_verification,
            player.participant.id_in_session,
            player.participant.typeMatching            
        )        

        # Only displayed if verification treatment
        showPage = False
        if player.selectTreatment == 0 and player.consent_privacy:
            showPage = True
                 
        return showPage
    
    @staticmethod
    def js_vars(player: Player):
        minTime = 0
        if player.study == 2:
            minTime = player.session.config["minTimeOnPage"],

        return dict(
            testing = player.session.config["testing"],
            minTime = minTime,
        )
    
    @staticmethod
    def vars_for_template(player: Player):
        id = player.participant.id_in_session
        folder = "bonusFiles/" + str(id) + "/"
        return dict(
            download_link= folder + "Verification.pdf",
            testing = player.session.config["testing"],
            
        )
    
    @staticmethod
    def app_after_this_page(player, upcoming_apps):
        return player.field_display('study') # Gives name of study ("binaryChoice", "priceList", etc.)

'''
class dwnldBonusFiles(Page):
    form_model = "player"
    form_fields = [
        "bfDownload", 
        "chosenBF",
        "feedback_dwnldBonusFiles"
        ]

    @staticmethod
    def is_displayed(player):
        # Only shown in study 2 with 4 choice tasks
        # Generate bonus files
        bf.rewriteAllBonusFiles(
            player.participant.bonusFileDict,
            player.participant.password_bonus, 
            player.participant.id_in_session)
        
        return player.consent_privacy == True and player.study == 2
    
    def live_method(player, data):
        print(data)
        player.chosenBF = data
        return {player.id_in_group: data}
    
    @staticmethod
    def vars_for_template(player: Player):
        download_links = []

        id = player.participant.id_in_session
        folder = "bonusFiles/" + str(id) + "/"

        for file in player.participant.bonusFileDict:
            download_links.append(
                [file, f"BonusFile{file}", folder + f"BonusFile{file}.pdf"]
            )

        return dict(
            download_links=download_links,
            testing = player.session.config["testing"],
            )

    def error_message(player: Player, values):
        if values["bfDownload"] == False:
            msg = "Please download a bonus file."
        elif values["chosenBF"] == "":
            msg = "Please decide for a bonus file and save it to a convenient location."
        else:
            msg = ""
        return msg

    @staticmethod
    def js_vars(player:Player):
        return dict(
            min_time_on_page = player.session.config["min_time_instructions"],
            chosenBF = player.field_maybe_none('chosenBF'),
            testing = player.session.config["testing"],
            )
    
    def before_next_page(player, timeout_happened):
        player.participant.chosenBF = player.chosenBF

        user_chosen_decision_no = int(
            player.participant.chosenBF.replace("BonusFile", "")
        )

        player.participant.bfNum = user_chosen_decision_no

'''

page_sequence = [
    taskSelection_forTest,
    welcomePage,
    exitPage,
    dwnldVerification,
    #dwnldBonusFiles, # Belongs to specific app
    #generalTask, # Belongs to specific app
    ]