import xlrd
import os

def transition(label):

    nl = []
    
    for l in label:
        
        if l == 0:
            nl.append(0)
        elif l == 2:
            nl.append(1)
        elif l == 1:
            nl.append(2)
    
    return nl


def form_data(segment, time, predLabel, trueLabel, safety, comfort, otherFeelings, interest, enjoyment, surprise, fear, tension, satisfaction, mode=0):
    
    allSegment = sum(segment, [])
    allPredLabel = sum(predLabel, [])
    allTrueLabel = sum(trueLabel, [])
    
    allSafety = sum(safety, [])
    allComfort = sum(comfort, [])
    
    allOtherFeelings = sum(otherFeelings, [])
    allEnjoyment = [i for i in sum([i[1:] for i in enjoyment], [])]
    allInterest = [i for i in sum([i[1:] for i in interest], [])]
    allSurprise = [i for i in sum([i[1:] for i in surprise], [])]
    allFear = [i for i in sum([i[1:] for i in fear], [])]
    allTension = [i for i in sum([i[1:] for i in tension], [])]
    allSatisfaction = [i for i in sum([i[1:] for i in satisfaction], [])]
    
    baseEnjoyment = [i for i in sum([[(i, item[0]) for i in item[1:]] for item in enjoyment], [])]
    baseInterest = [i for i in sum([[(i, item[0]) for i in item[1:]] for item in interest], [])]
    baseSurprise = [i for i in sum([[(i, item[0]) for i in item[1:]] for item in surprise], [])]
    baseFear = [i for i in sum([[(i, item[0]) for i in item[1:]] for item in fear], [])]
    baseTension = [i for i in sum([[(i, item[0]) for i in item[1:]] for item in tension], [])]
    baseSatisfaction = [i for i in sum([[(i, item[0]) for i in item[1:]] for item in satisfaction], [])]
    allAffect = [[[allEnjoyment[i], allInterest[i], allSurprise[i], allFear[i], allTension[i], allSatisfaction[i]], [baseEnjoyment[i][1], baseInterest[i][1], baseSurprise[i][1], baseFear[i][1], baseTension[i][1], baseSatisfaction[i][1]]] for i in range(len(allPredLabel))]
    
    if mode != 0:
        return [allAffect[i] for i, j in enumerate(allSegment) if j == mode], [allOtherFeelings[i] for i, j in enumerate(allSegment) if j == mode], [allSafety[i] for i, j in enumerate(allSegment) if j == mode], [allComfort[i] for i, j in enumerate(allSegment) if j == mode], transition([allPredLabel[i] for i, j in enumerate(allSegment) if j == mode]), transition([allTrueLabel[i] for i, j in enumerate(allSegment) if j == mode])
    else:
        return allSegment, allAffect, allOtherFeelings, transition(allPredLabel), transition(allTrueLabel), len(allPredLabel), len(predLabel)

    
def load_data(mode=0):
    
    try:
        dataBook = xlrd.open_workbook('./data/xls_data/data.xls')
    except:
        dataBook = xlrd.open_workbook('/'.join(os.getcwd().split('/')[:-1])+'/data/xls_data/data.xls')
    
    dataSheet = dataBook.sheet_by_index(0)
    segment = []
    time = []
    predLabel, trueLabel = [], [] 
    safety, comfort = [], [] 
    otherFeelings = []
    interest, enjoyment, surprise, fear, tension, satisfaction = [], [], [], [], [], []
    i = 1

    while i < dataSheet.nrows:

        seg = []
        pl, tl = [], []
        s, c = [], []
        of = []
        if dataSheet.cell_value(i, 4) != 0:
            t = []
            ei, ee, esur, ef, et, esa = [], [], [], [], [], []
            t.append(dataSheet.cell_value(i, 1))
            ee.append(dataSheet.cell_value(i, 10))
            ei.append(dataSheet.cell_value(i, 11))
            esur.append(dataSheet.cell_value(i, 12))
            ef.append(dataSheet.cell_value(i, 13))
            et.append(dataSheet.cell_value(i, 14))
            esa.append(dataSheet.cell_value(i, 15))
            i += 1
            if i == dataSheet.nrows:
                break
        else:
            while dataSheet.cell_value(i, 4) == 0:
                
                if dataSheet.cell_value(i, 3) == '赛段一':
                    seg.append(1)
                elif dataSheet.cell_value(i, 3) == '赛段二':
                    seg.append(2)
                elif dataSheet.cell_value(i, 3) == '赛段三':
                    seg.append(3)
                else:
                    print('Error!')
                
                t.append(dataSheet.cell_value(i, 1))
                pl.append(int(dataSheet.cell_value(i, 5)))
                tl.append(int(dataSheet.cell_value(i, 6)))
                s.append(dataSheet.cell_value(i, 7))
                c.append(dataSheet.cell_value(i, 8))
                of.append(dataSheet.cell_value(i, 9))
                ee.append(dataSheet.cell_value(i, 10))
                ei.append(dataSheet.cell_value(i, 11))
                esur.append(dataSheet.cell_value(i, 12))
                ef.append(dataSheet.cell_value(i, 13))
                et.append(dataSheet.cell_value(i, 14))
                esa.append(dataSheet.cell_value(i, 15))         
                i += 1
                if i == dataSheet.nrows:
                    break
                    
            segment.append(seg)
            time.append(t)
            predLabel.append(pl); trueLabel.append(tl)
            safety.append(s); comfort.append(c)
            otherFeelings.append(of)
            enjoyment.append(ee), interest.append(ei), surprise.append(esur), fear.append(ef), tension.append(et), satisfaction.append(esa)
    
    return form_data(segment, time, predLabel, trueLabel, safety, comfort, otherFeelings, interest, enjoyment, surprise, fear, tension, satisfaction, mode=mode)