import statistics
import numpy as np

class Prediction:
    def __init__(self, ilevels, levels, betap, deltaP, penalisation, gradeT):
        #The level has the note and the time for each question
        #Test data
        self.ilevels=ilevels
        self.levels=levels
        #Initialization of beta distributions for each level
        self.betap=betap
        self.INLevel=[]
        self.NLevel=[]

        self.maxGrade=10
        self.factor=10
        
        self.deltaPlus=deltaP
        self.gradeThreshold=gradeT
        self.timePenalisation=penalisation

    def Structure(self, base, total, questions):
        v1=[]
        v2=[]
        elem=1
        #for i in range(1,150,2):
        for i in range(1,total,2):
            #print(i," ",dr.tolist()[0+i:2+i])
            v1.append(base.tolist()[0+i:2+i])
            #if (elem % 15) == 0:
            if (elem % questions) == 0:
                v2.append(v1)
                v1=[]
            elem+=1
        #In this case, v2 is the last of all rows
        return v2

    def CalculateGradePenalization(self):
        self.ilevels=self.Structure(self.ilevels,30,3)
        self.levels=self.Structure(self.levels,150,15)
        #Calculate the note with penalization for time used
        self.INLevel=[[ilist[0]-((ilist[0])*((ilist[1]/60)*self.timePenalisation)) for ilist in lev] for lev in self.ilevels]
        self.NLevel=[[ilist[0]-((ilist[0])*((ilist[1]/60)*self.timePenalisation)) for ilist in lev] for lev in self.levels]
#Generalized Thompson sampling
    def Calculate1(self, NLevel):

        r=1
        maxGrade=10
        #Here start the Multi armed bandits for choose the best level. Thompson Sampling...
        for i in range(min(len(self.levels[0]),len(self.levels[1]),len(self.levels[2]))):
            if i==0:
                IRLevel=0
            else:
                #Take a sample for all calculated posterior distributions
                Rlevel=[np.random.beta(self.betap[0][0],self.betap[0][1]),
                        np.random.beta(self.betap[1][0],self.betap[1][1]),
                        np.random.beta(self.betap[2][0],self.betap[2][1])]
                IRLevel=max(enumerate(Rlevel),key=lambda x: x[1])[0]
                print(Rlevel, self.betap)
                print("Mean 1: ", self.betap[0][0]/(self.betap[0][0]+self.betap[0][1]))
                print("Mean 2: ", self.betap[1][0]/(self.betap[1][0]+self.betap[1][1]))
                print("Mean 3: ", self.betap[2][0]/(self.betap[2][0]+self.betap[2][1]))
                print(IRLevel,NLevel[IRLevel][i])
                #Rewards for succes or not of machine results, here the rewards are inversed because we want to rest into the machine with less success. The reward is normalized value between 0 and 1
                #NLevel[IRLevel][i]=NLevel[IRLevel][i]/self.maxGrade
                grade=NLevel[IRLevel][i]
                deltaMu=(maxGrade-(2*grade))/10

                print("DeltaMu: ",deltaMu)

                #Change the value of beta parameter index 0
                sumAlphaBeta=self.betap[IRLevel][0]+self.betap[IRLevel][1]
                self.betap[IRLevel][0]=round((self.betap[IRLevel][0]*r)+(deltaMu*(sumAlphaBeta)),0)
                #Control the limits of value for first parameter
                self.betap[IRLevel][0]=max(1, self.betap[IRLevel][0])

                #Change the value of beta parameter index 1
                self.betap[IRLevel][1]=round((sumAlphaBeta*r)-self.betap[IRLevel][0],0)
                #Control the limits of value for second parameter
                self.betap[IRLevel][1]=max(1, self.betap[IRLevel][1])

                print(self.betap[IRLevel][0], self.betap[IRLevel][1])
                print(Rlevel, self.betap)

#Bernoulli Thompson Sampling
    def UpdateBeta(self, grade, level):
        if grade >= self.gradeThreshold:
            #Change the value of beta parameter index 1
            self.betap[level][1]+=self.deltaPlus
            #Correlated Thompson Sampling
            if level>0:
                self.betap[level-1][1]+=self.deltaPlus/2
            if level<len(self.betap)-1:
                self.betap[level+1][0]+=self.deltaPlus/2

        else:
            #Change the value of alpha parameter index 0
            self.betap[level][0]+=self.deltaPlus
            #Correlated Thompson Sampling
            if level>0:
                self.betap[level-1][0]+=self.deltaPlus/2
            if level<len(self.betap)-1:
                self.betap[level+1][1]+=self.deltaPlus/2

    def InitializeBeta(self):
        c=0
        for itemc in self.INLevel:
            for i in range(len(itemc)):
                #print(itemc[i])
                self.UpdateBeta(itemc[i],c)
            c+=1

    def Calculate(self):
        self.InitializeBeta()
        NLevel=self.NLevel
        file1=open('results_slevel.csv','a+')
        file2=open('results_sgrade.csv','a+')
        #Here start the Multi armed bandits for choose the best level. Thompson Sampling...
        #print("NLevel Vector: ",NLevel)
        for i in range(len(NLevel[0])):
            #if i==0:
            #    IRLevel=0
            #else:
            #Take a sample for all calculated posterior distributions
            Rlevel=[np.random.beta(p[0],p[1]) for p in self.betap]
            #Take the max probability value
            #IRLevel=max(enumerate(Rlevel),key=lambda x: x[1])[0]
            IRLevel=max( (v, i) for i, v in enumerate(Rlevel) )[1]
            #print(Rlevel, self.betap, [p[0]/(p[0]+p[1]) for p in self.betap])
            print("Stochastic ",i," ",IRLevel)
            #print(NLevel[IRLevel][i])
            file1.write(str(IRLevel)+" ")
            file2.write(str(IRLevel)+" "+str(NLevel[IRLevel][i])+" ")
            #Rewards for succes or not of machine results, here the rewards are inversed because we want to rest into the machine with less success
            self.UpdateBeta(NLevel[IRLevel][i],IRLevel)
        file1.write("\n")
        file2.write("\n")
        file1.close()
        file2.close()

    def CalculateSW(self):
        file1=open('results_dlevel.csv','a+')
        file2=open('results_dgrade.csv','a+')
        #file2=open('results_dtm.csv','a+')
        IRLevel=0
        NIRLevel=0
        step=0
        mc0=0
        mc1=0
        mc2=0
        mc3=0
        mc4=0
        #print("INLevel: ",INLevel)
        #for i in range(len(NLevel[0])):

        Level=self.INLevel
        clevel=[0,0,0,0,0]
        vindex=[0,0,0,0,0]

        for i in range(15):
            if i>0:
                #print("Value ",Level)
                Level[NIRLevel].append(self.NLevel[NIRLevel][vindex[NIRLevel]])
                Level[NIRLevel].pop(0)
                #print("Value ",Level)
                vindex[NIRLevel]+=1
                #vindex+=1
            mc0=0
            mc1=0
            mc2=0
            mc3=0
            mc4=0
            mc=0
            #print("Level Vector ",Level)
            for IRL in range(4):
                if IRL == 0:
                    mc0=Level[IRL][2]+Level[IRL][1]+Level[IRL][0]
                    mc0=mc0/3
                    mc0=10*mc0/5
                if IRL == 1:
                    mc1=Level[IRL][2]+Level[IRL][1]+Level[IRL][0]
                    mc1=mc1/3
                if IRL == 2:
                    mc2=Level[IRL][2]+Level[IRL][1]+Level[IRL][0]
                    mc2=mc2/3
                if IRL == 3:
                    mc3=Level[IRL][2]+Level[IRL][1]+Level[IRL][0]
                    mc3=mc3/3
                if IRL == 4:
                    mc4=Level[IRL][2]+Level[IRL][1]+Level[IRL][0]
                    mc4=mc4/3

            #print(mc0," ",mc1," ",mc2,"",mc3)
            mc1=max(mc0+((10/5)*mc1), (20/5)*mc1)
            mc2=max(mc1+((10/5)*mc2), (30/5)*mc2)
            mc3=max(mc2+((10/5)*mc3), (40/5)*mc3)
            mc=max(mc3+((10/5)*mc4), (50/5)*mc4)
            #print(mc0," ",mc1," ",mc2," ",mc3," ",mc)
            #print(mc)
            #file2.write(str(mc)+" ")
            if mc >= 0 and mc <= 15:
                NIRLevel=0
            elif mc >= 16 and mc <= 25:
                NIRLevel=1
            elif mc >= 26 and mc <= 35:
                NIRLevel=2
            elif mc >= 36 and mc <= 42:
                NIRLevel=2
            elif mc >= 43 and mc <= 50:
                NIRLevel=3
            elif mc >= 51 and mc <= 75:
                NIRLevel=3
            elif mc >= 76 and mc <= 100:
                NIRLevel=4
            print("Deterministic ",i," ",NIRLevel)
            if NIRLevel != IRLevel:
                IRLevel=NIRLevel
            file1.write(str(IRLevel)+" ")
            file2.write(str(IRLevel)+" "+str(self.NLevel[IRLevel][vindex[IRLevel]])+" ")
        file1.write("\n")
        file2.write("\n")
        file1.close()
        file2.close()