from experiments.experiment import experiments
from Image_Process.Image_Generator import Image_generator
from Training_Tools.Tools import Tool
from model_data_processing.processing import make_label_list, Read_Image_Root_And_Image_Enhance
from Load_process.LoadData import Load_Data_Prepare
from Calculate_Process.Calculate import Calculate
from merge_class.merge import merge
from model_data_processing.processing_for_cut_image import Cut_Indepentend_Data
from Load_process.LoadData import Loding_Data_Root
from Load_process.file_processing import Process_File
from utils.Stomach_Config import Training_Config, Loading_Config
from Image_Process.Image_Mask_Ground_Truth_Processing import XMLAnnotationProcessor

import time
import torch
 
if __name__ == "__main__":
    # 測試GPU是否可用
    flag = torch.cuda.is_available()
    if not flag:
        print("CUDA不可用\n")
    else:
        print(f"CUDA可用，數量為{torch.cuda.device_count()}\n")

    # 测试GPU是否可用
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"使用设备: {device}")
    print(f"GPU: {torch.cuda.get_device_name(0)}")

    tool = Tool()
    Status = 1

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

    Prepare = Load_Data_Prepare()
    Indepentend = Cut_Indepentend_Data(Loading_Config["Train_Data_Root"], Loading_Config["Training_Labels"])
    Merge = merge()
    Calculate_Tool = Calculate()
    file = Process_File()

    Train_Size = 0
 
    # 讀取資料
    # 將測試資料獨立出來
    test_size = 0.2
    Indepentend.IndependentData_main(Loading_Config["Test_Data_Root"], test_size)

    # 創建處理切割的部分
    if not file.JudgeRoot_MakeDir(Loading_Config['Annotation_Training_Root']) and not file.JudgeRoot_MakeDir(Loading_Config['Annotation_Testing_Root']):
        processor_train = XMLAnnotationProcessor(
            dataset_root = Loading_Config["Train_Data_Root"],
        )
        processor_test = XMLAnnotationProcessor(
            dataset_root = Loading_Config["Test_Data_Root"],
        )
        
        # 設定自訂樣式（可選）
        processor_train.set_drawing_style(
            box_color=(0, 0, 255),  # 紅色邊界框
            text_color=(255, 255, 255),  # 白色文字
            box_thickness=3,
            font_scale=0.7
        )
        processor_test.set_drawing_style(
            box_color=(0, 0, 255),  # 紅色邊界框
            text_color=(255, 255, 255),  # 白色文字
            box_thickness=3,
            font_scale=0.7
        )
        
        print("XML標註處理器已準備就绪")
        for Label_Iamge_List in Loading_Config["Label_Image_Labels"]:
            if Label_Iamge_List == "CA":
                Label = "stomach_cancer_Crop"
            else:
                Label = "Have_Question_Crop"
            training_results = processor_train.process_multiple_xml(f"../Label_Image/{Label_Iamge_List}", Loading_Config['Annotation_Training_Root'], Label)
            testing_results = processor_test.process_multiple_xml(f"../Label_Image/{Label_Iamge_List}", Loading_Config['Annotation_Testing_Root'], Label)


    else:
        print("Training and Testing annoation is exist!!!!")

    # 讀取切割完成後的檔案
    print("Mask Ground truth is Finished\n")
    Mask_load = Loding_Data_Root(Loading_Config["XML_Loading_Label"], Loading_Config['Annotation_Training_Root'], "")
    Mask_Data_Dict_Data = Mask_load.process_main(False)

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

    # 讀取原始資料集
    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)

    # 製作資料增強資料
    if not file.Judge_File_Exist(Loading_Config['Image enhance processing save root']):
        for label in Loading_Config["Training_Labels"]:
            Read_Image_Root_And_Image_Enhance(Data_Dict_Data[label], f"{Loading_Config['Image enhance processing save root']}/{label}")

    tmp_load = Loding_Data_Root(Loading_Config["Training_Labels"], Loading_Config['Image enhance processing save root'], Loading_Config["ImageGenerator_Data_Root"])
    Data_Dict_Data = tmp_load.process_main(False)

    Total_Size_List = []
    print("前處理後資料集總數")
    for label in Loading_Config["Training_Labels"]:
        Train_Size += len(Data_Dict_Data[label])
        Total_Size_List.append(len(Data_Dict_Data[label]))
        print(f"Labels: {label}, 總數為: {len(Data_Dict_Data[label])}")
        
    print("總共有 " + str(Train_Size) + " 筆資料")
    print(f"要被Mask的資料共有{int(Total_Size_List[0]) + int(Total_Size_List[2])}筆")

    # 做出跟資料相同數量的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())

    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

    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

    Training_Mask_Data = Merge.merge_all_image_data(Mask_Data_Dict_Data[Mask_Keys[0]], Mask_Data_Dict_Data[Mask_Keys[1]]) # 將訓練資料合併成一個list
    for i in range(2, len(Mask_Keys)):
        Training_Mask_Data = Merge.merge_all_image_data(Training_Mask_Data, Mask_Data_Dict_Data[Mask_Keys[i]]) # 將訓練資料合併成一個list

    experiment = experiments(
        Xception_Training_Data=Training_Data, 
        Xception_Training_Label=Training_Label, 
        Xception_Training_Mask_Data=Training_Mask_Data,
        tools=tool, 
        Number_Of_Classes=Label_Length, 
        status=Status
        )
    start = time.time()
    experiment.processing_main() # 執行訓練方法
    end = time.time()
    print(f"\n\n\n訓練時間：{end - start} 秒\n\n")