import os
import shutil

from PIL import Image
from torchvision.transforms.functional import to_pil_image
from Load_process.file_processing import Process_File
from utils.Stomach_Config import Image_Enhance, Training_Config, Loading_Config
from Training_Tools.PreProcess import Training_Precesses
from Load_process.LoadData import Loding_Data_Root
from Training_Tools.Tools import Tool
from merge_class.merge import merge
from Load_process.LoadData import Load_Data_Prepare

def make_label_list(length, content):
    '''製作label的列表'''
    label_list = []
    for i in range(length):
        label_list.append(content)
    return label_list

def Read_Image_Root_And_Image_Enhance(Image_Roots, Save_Root):
    '''讀取Image Root'''
    i = 0
    file = Process_File()
    for Image_Root in Image_Roots:
        try:
            Images = Image.open(Image_Root).convert("RGB")

            Images = Image_Enhance["gamma"](Images, Image_Enhance["Gamma_Value"])
            for j in range(5):
                Images = Image_Enhance["Median"](Images)
            Images = Image_Enhance["Shapen"](Images)

            file.JudgeRoot_MakeDir(Save_Root)
            # 使用原始檔名而不是索引編號
            original_filename = os.path.basename(Image_Root)
            path = file.Make_Save_Root(original_filename, Save_Root)
            
            Images.save(path)
            # Image = cv2.imread(Image_Root, cv2.IMREAD_COLOR) # 讀檔(彩色)
            # Image = cv2.cvtColor(Image, cv2.COLOR_BGR2RGB)
        except Exception as e:
            print(e)

        i += 1
        print("已處理 " + str(i) + " 張圖片")

    pass

def Image_Enhance_Training_Data(Training_Loader, Save_Root):
    '''讀取Image Root'''
    # 改成只要在Training時進行影像前處理
    Images, Labels, File_Names, File_Classes = [], [], [], []
    Label_Length = len(Loading_Config["Training_Labels"])
    Process = Training_Precesses(ImageSize = Training_Config["Image_Size"])
    file = Process_File()
    Prepare = Load_Data_Prepare()
    tool = Tool()
    Merge = merge()

    if file.Judge_File_Exist(Save_Root):
        shutil.rmtree(Save_Root)
        print("檔案已存在，刪除檔案")

    # 取得One-hot encording 的資料
    tool.Set_OneHotEncording(Loading_Config["Training_Labels"])
    Encording_Label = tool.Get_OneHot_Encording_Label()

    for inputs, labels, File_Name, File_Class in Training_Loader:
        for i in range(inputs.shape[0]):
            Permute_Image = inputs[i].permute(1, 2, 0)
            Permute_Image *= 255
            Permute_Image = Image_Enhance["gamma"](Permute_Image, Image_Enhance["Gamma_Value"])
            for j in range(5):
                Permute_Image = Image_Enhance["Median"](Permute_Image)
                
            Images = Image_Enhance["Shapen"](Permute_Image)
            
            Save = file.Make_Save_Root(File_Class[i], Save_Root)            
            file.JudgeRoot_MakeDir(Save)

            path = file.Make_Save_Root(File_Name[i], Save)
            Images.save(path)

    load = Loding_Data_Root(Loading_Config["Training_Labels"], Loading_Config["Train_Data_Root"], Loading_Config["ImageGenerator_Data_Root"])
    Data_Dict_Data = load.process_main(False)

    Total_Size_List = []
    for label in Loading_Config["Training_Labels"]:
        Total_Size_List.append(len(Data_Dict_Data[label]))

    # 做出跟資料相同數量的Label
    Classes = []
    i = 0
    for encording in Encording_Label:
        Classes.append(make_label_list(Total_Size_List[i], encording))
        i += 1

    # 將資料做成Dict的資料型態
    Prepare.Set_Final_Dict_Data(Loading_Config["Training_Labels"], Data_Dict_Data, Classes, Label_Length)
    Final_Dict_Data = Prepare.Get_Final_Data_Dict()
    keys = list(Final_Dict_Data.keys())
    # Mask_Keys = list(Mask_Data_Dict_Data.keys())

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

    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):
        Label = Merge.merge_all_image_data(Label, Final_Dict_Data[keys[i]]) # 將訓練資料合併成一個list

    DataSet = Process.Setting_DataSet(
        Datas = Data, 
        Labels = Label,
        Mask_List = None, 
        transform = "Transform"
    )
    Dataloader = Process.Dataloader_Sampler(DataSet, Training_Config["Train_Batch_Size"], True)
    return Dataloader