from Load_process.LoadData import Loding_Data_Root
from utils.Stomach_Config import Loading_Config
from merge_class.merge import merge
from Load_process.LoadData import Load_Data_Prepare
from Training_Tools.Tools import Tool

class Training_Preprocessing():
    def __init__(self):
        pass

    def Process_Main(self, Loading_Data_Roots, Labels, Have_Label_Lists = True, Generator_Roots = None):
        """
        讀取資料成 Dict 資料型態，並計算各資料筆數與總共資料筆數
        回傳Loading_Data_Dict_Data
        """
        Loading_Data_Dict_Data, Total_Size_Lists = self.Loading_Data_Dicts(Loading_Data_Roots, Labels, Generator_Roots)

        # 做出跟資料相同數量的Label
        # 取得One-hot encording 的資料
        tool = Tool()
        Label_Length = len(Labels)

        if Have_Label_Lists:
            tool.Set_OneHotEncording(Loading_Config["Training_Labels"])
            Encording_Label = tool.Get_OneHot_Encording_Label()
            
            Classes = self.make_label_list(Encording_Label, Total_Size_Lists)
            return self.Merge_Total_Data(Loading_Data_Dict_Data, Classes, Label_Length)

        return self.Merge_Total_Data(Loading_Data_Dict_Data, Classes, Label_Length)

    def Loading_Data_Dicts(self, Loading_Data_Roots, Labels, Generator_Roots):
        """
        讀取資料成 Dict 資料型態，並計算各資料筆數與總共資料筆數
        
        回傳Loading_Data_Dict_Data
        """
        Load_Data = Loding_Data_Root(Loading_Data_Roots, Labels, Generator_Roots)
        Loading_Data_Dict_Data = Load_Data.process_main(False)
        Total_Size_Lists = []

        print("讀取的資料集總數")
        for label in Labels:
            Train_Size += len(Loading_Data_Dict_Data[label])
            Total_Size_Lists.append(len(Loading_Data_Dict_Data[label]))
            print(f"Labels: {label}, 總數為: {len(Loading_Data_Dict_Data[label])}")
            
        print("總共有 " + str(Train_Size) + " 筆資料")

        return Loading_Data_Dict_Data, Total_Size_Lists

    def make_label_list(Contents, Lengths):
        '''製作label的列表'''
        Classes = []
        for Content in Contents:
            label_list = []
            for i in range(Lengths):
                label_list.append(Content)
            Classes.append(label_list)

        return label_list

    def Merge_Total_Data(self, Data_Dict, Classes, Label_Length):
        """
        將所有的資料合併成一個list
        """
        Prepare = Load_Data_Prepare()
        Merge = merge()

        # 將資料做成Dict的資料型態
        if Classes == None:
            Classes = [[], [], []]
        Prepare.Set_Final_Dict_Data(Loading_Config["Training_Labels"], Data_Dict, Classes, Label_Length)
        Final_Dict_Data = Prepare.Get_Final_Data_Dict()
        keys = Loading_Config["Training_Labels"]

        Training_Data = Merge.merge_all_image_data(Final_Dict_Data[keys[0]], Final_Dict_Data[keys[1]]) # 將訓練資料合併成一個list
        for i in range(2, Label_Length):
            Training_Data = Merge.merge_all_image_data(Training_Data, Final_Dict_Data[keys[i]]) # 將訓練資料合併成一個list

        if Classes is not [[], [], []]:
            Training_Label = Merge.merge_all_image_data(Final_Dict_Data[keys[Label_Length]], Final_Dict_Data[keys[Label_Length + 1]]) #將訓練資料的label合併成一個label的list
            for i in range(Label_Length + 2, 2 * Label_Length):
                Training_Label = Merge.merge_all_image_data(Training_Label, Final_Dict_Data[keys[i]]) # 將訓練資料合併成一個list
                return Training_Data, Training_Label

        return Training_Data
