import argparse
import os
import sys
import torch
import torch.backends
from exp.exp_long_term_forecasting import Exp_Long_Term_Forecast
from exp.exp_imputation import Exp_Imputation
from exp.exp_short_term_forecasting import Exp_Short_Term_Forecast
from exp.exp_anomaly_detection import Exp_Anomaly_Detection
from exp.exp_classification import Exp_Classification
from exp.exp_zero_shot_forecasting import Exp_Zero_Shot_Forecast
from utils.print_args import print_args
import random
import numpy as np
import threading
import queue
import tkinter as tk
from tkinter import ttk, scrolledtext, messagebox
import matplotlib
matplotlib.use('TkAgg')
from matplotlib.figure import Figure
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
import re
import io
from contextlib import redirect_stdout

class QueueWriter:
    def __init__(self, queue):
        self.queue = queue
    
    def write(self, text):
        self.queue.put(text)
    
    def flush(self):
        pass

class TrainingThread(threading.Thread):
    def __init__(self, args, log_queue, loss_queue, dataset_info_queue):
        super().__init__()
        self.args = args
        self.log_queue = log_queue
        self.loss_queue = loss_queue
        self.dataset_info_queue = dataset_info_queue
        self.daemon = True
        
    def run(self):
        # 重定向stdout到队列
        original_stdout = sys.stdout
        queue_writer = QueueWriter(self.log_queue)
        sys.stdout = queue_writer
        
        try:
            fix_seed = 2021
            random.seed(fix_seed)
            torch.manual_seed(fix_seed)
            np.random.seed(fix_seed)
            
            # 设备设置
            if torch.cuda.is_available() and self.args.use_gpu:
                self.args.device = torch.device('cuda:{}'.format(self.args.gpu))
            else:
                if hasattr(torch.backends, "mps"):
                    self.args.device = torch.device("mps") if torch.backends.mps.is_available() else torch.device("cpu")
                else:
                    self.args.device = torch.device("cpu")
            
            if self.args.use_gpu and self.args.use_multi_gpu:
                self.args.devices = self.args.devices.replace(' ', '')
                device_ids = self.args.devices.split(',')
                self.args.device_ids = [int(id_) for id_ in device_ids]
                self.args.gpu = self.args.device_ids[0]
            
            # 选择实验类
            if self.args.task_name == 'long_term_forecast':
                Exp = Exp_Long_Term_Forecast
            elif self.args.task_name == 'short_term_forecast':
                Exp = Exp_Short_Term_Forecast
            elif self.args.task_name == 'imputation':
                Exp = Exp_Imputation
            elif self.args.task_name == 'anomaly_detection':
                Exp = Exp_Anomaly_Detection
            elif self.args.task_name == 'classification':
                Exp = Exp_Classification
            elif self.args.task_name == 'zero_shot_forecast':
                Exp = Exp_Zero_Shot_Forecast
            else:
                Exp = Exp_Long_Term_Forecast
            
            # 创建实验实例
            exp = Exp(self.args)
            
            # 获取数据集信息
            try:
                train_data, train_loader = exp._get_data(flag='train')
                test_data, test_loader = exp._get_data(flag='test')
                val_data, val_loader = exp._get_data(flag='val')
                
                train_size = len(train_data)
                test_size = len(test_data)
                val_size = len(val_data) if val_data else 0
                
                dataset_info = {
                    'train_size': train_size,
                    'test_size': test_size,
                    'val_size': val_size,
                    'train_batches': len(train_loader),
                    'test_batches': len(test_loader),
                    'val_batches': len(val_loader) if val_loader else 0
                }
                self.dataset_info_queue.put(dataset_info)
            except Exception as e:
                self.log_queue.put(f"获取数据集信息失败: {str(e)}\n")
            
            if self.args.is_training:
                for ii in range(self.args.itr):
                    setting = '{}_{}_{}_{}_ft{}_sl{}_ll{}_pl{}_dm{}_nh{}_el{}_dl{}_df{}_expand{}_dc{}_fc{}_eb{}_dt{}_{}_{}'.format(
                        self.args.task_name,
                        self.args.model_id,
                        self.args.model,
                        self.args.data,
                        self.args.features,
                        self.args.seq_len,
                        self.args.label_len,
                        self.args.pred_len,
                        self.args.d_model,
                        self.args.n_heads,
                        self.args.e_layers,
                        self.args.d_layers,
                        self.args.d_ff,
                        self.args.expand,
                        self.args.d_conv,
                        self.args.factor,
                        self.args.embed,
                        self.args.distil,
                        self.args.des, ii)
                    
                    print(f'开始训练: {setting}')
                    
                    # 训练过程
                    exp.train(setting)
                    
                    print(f'开始测试: {setting}')
                    exp.test(setting)
                    
                    if self.args.gpu_type == 'mps':
                        torch.backends.mps.empty_cache()
                    elif self.args.gpu_type == 'cuda':
                        torch.cuda.empty_cache()
            
            print("训练完成！")
            
        except Exception as e:
            self.log_queue.put(f"训练过程中出错: {str(e)}\n")
            import traceback
            self.log_queue.put(traceback.format_exc() + "\n")
        finally:
            sys.stdout = original_stdout


class TrainingGUI:
    def __init__(self, root):
        self.root = root
        self.root.title("时间序列预测模型训练界面")
        self.root.geometry("1400x900")
        
        # 训练线程
        self.training_thread = None
        self.is_training = False
        
        # 队列
        self.log_queue = queue.Queue()
        self.loss_queue = queue.Queue()
        self.dataset_info_queue = queue.Queue()
        
        # Loss数据
        self.train_losses = []
        self.val_losses = []
        self.test_losses = []
        self.epochs = []
        
        self.create_widgets()
        self.process_queue()
        
    def create_widgets(self):
        # 创建主框架
        main_frame = ttk.Frame(self.root, padding="10")
        main_frame.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
        
        self.root.columnconfigure(0, weight=1)
        self.root.rowconfigure(0, weight=1)
        main_frame.columnconfigure(1, weight=1)
        main_frame.rowconfigure(1, weight=1)
        
        # 左侧参数面板
        left_frame = ttk.Frame(main_frame)
        left_frame.grid(row=0, column=0, rowspan=2, sticky=(tk.W, tk.E, tk.N, tk.S), padx=(0, 10))
        left_frame.columnconfigure(1, weight=1)
        
        # 创建滚动框架
        canvas = tk.Canvas(left_frame, width=400, height=700)
        scrollbar = ttk.Scrollbar(left_frame, orient="vertical", command=canvas.yview)
        scrollable_frame = ttk.Frame(canvas)
        
        scrollable_frame.bind(
            "<Configure>",
            lambda e: canvas.configure(scrollregion=canvas.bbox("all"))
        )
        
        canvas.create_window((0, 0), window=scrollable_frame, anchor="nw")
        canvas.configure(yscrollcommand=scrollbar.set)
        
        canvas.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
        scrollbar.grid(row=0, column=1, sticky=(tk.N, tk.S))
        left_frame.rowconfigure(0, weight=1)
        
        # 参数输入
        self.params = {}
        row = 0
        
        # 基本配置
        ttk.Label(scrollable_frame, text="基本配置", font=('Arial', 12, 'bold')).grid(row=row, column=0, columnspan=2, pady=10, sticky=tk.W)
        row += 1
        
        self.create_param_row(scrollable_frame, "任务类型 (task_name)", "long_term_forecast", row)
        row += 1
        self.create_param_row(scrollable_frame, "模型 (model)", "TimesNet", row)
        row += 1
        self.create_param_row(scrollable_frame, "模型ID (model_id)", "test", row)
        row += 1
        
        # 数据配置
        ttk.Label(scrollable_frame, text="数据配置", font=('Arial', 12, 'bold')).grid(row=row, column=0, columnspan=2, pady=10, sticky=tk.W)
        row += 1
        
        self.create_param_row(scrollable_frame, "数据集 (data)", "ETTh1", row)
        row += 1
        self.create_param_row(scrollable_frame, "数据路径 (root_path)", "./dataset/ETT-small/", row)
        row += 1
        self.create_param_row(scrollable_frame, "数据文件 (data_path)", "ETTh1.csv", row)
        row += 1
        self.create_param_row(scrollable_frame, "特征类型 (features)", "M", row)
        row += 1
        self.create_param_row(scrollable_frame, "频率 (freq)", "h", row)
        row += 1
        
        # 预测配置
        ttk.Label(scrollable_frame, text="预测配置", font=('Arial', 12, 'bold')).grid(row=row, column=0, columnspan=2, pady=10, sticky=tk.W)
        row += 1
        
        self.create_param_row(scrollable_frame, "序列长度 (seq_len)", "96", row, int)
        row += 1
        self.create_param_row(scrollable_frame, "标签长度 (label_len)", "48", row, int)
        row += 1
        self.create_param_row(scrollable_frame, "预测长度 (pred_len)", "96", row, int)
        row += 1
        
        # 模型配置
        ttk.Label(scrollable_frame, text="模型配置", font=('Arial', 12, 'bold')).grid(row=row, column=0, columnspan=2, pady=10, sticky=tk.W)
        row += 1
        
        self.create_param_row(scrollable_frame, "模型维度 (d_model)", "512", row, int)
        row += 1
        self.create_param_row(scrollable_frame, "注意力头数 (n_heads)", "8", row, int)
        row += 1
        self.create_param_row(scrollable_frame, "编码器层数 (e_layers)", "2", row, int)
        row += 1
        self.create_param_row(scrollable_frame, "解码器层数 (d_layers)", "1", row, int)
        row += 1
        self.create_param_row(scrollable_frame, "FFN维度 (d_ff)", "2048", row, int)
        row += 1
        self.create_param_row(scrollable_frame, "Dropout", "0.1", row, float)
        row += 1
        
        # 训练配置
        ttk.Label(scrollable_frame, text="训练配置", font=('Arial', 12, 'bold')).grid(row=row, column=0, columnspan=2, pady=10, sticky=tk.W)
        row += 1
        
        self.create_param_row(scrollable_frame, "训练轮数 (train_epochs)", "10", row, int)
        row += 1
        self.create_param_row(scrollable_frame, "批次大小 (batch_size)", "32", row, int)
        row += 1
        self.create_param_row(scrollable_frame, "学习率 (learning_rate)", "0.0001", row, float)
        row += 1
        self.create_param_row(scrollable_frame, "早停耐心值 (patience)", "3", row, int)
        row += 1
        self.create_param_row(scrollable_frame, "实验次数 (itr)", "1", row, int)
        row += 1
        
        # GPU配置
        ttk.Label(scrollable_frame, text="GPU配置", font=('Arial', 12, 'bold')).grid(row=row, column=0, columnspan=2, pady=10, sticky=tk.W)
        row += 1
        
        self.create_param_row(scrollable_frame, "使用GPU (use_gpu)", "True", row, bool)
        row += 1
        self.create_param_row(scrollable_frame, "GPU设备 (gpu)", "0", row, int)
        row += 1
        
        # 按钮
        button_frame = ttk.Frame(scrollable_frame)
        button_frame.grid(row=row, column=0, columnspan=2, pady=20)
        row += 1
        
        self.train_button = ttk.Button(button_frame, text="开始训练", command=self.start_training, width=20)
        self.train_button.pack(side=tk.LEFT, padx=5)
        
        self.stop_button = ttk.Button(button_frame, text="停止训练", command=self.stop_training, width=20, state=tk.DISABLED)
        self.stop_button.pack(side=tk.LEFT, padx=5)
        
        # 右侧显示面板
        right_frame = ttk.Frame(main_frame)
        right_frame.grid(row=0, column=1, sticky=(tk.W, tk.E, tk.N, tk.S))
        right_frame.columnconfigure(0, weight=1)
        right_frame.rowconfigure(1, weight=1)
        
        # 数据集信息
        info_frame = ttk.LabelFrame(right_frame, text="数据集信息", padding="10")
        info_frame.grid(row=0, column=0, sticky=(tk.W, tk.E), pady=(0, 10))
        info_frame.columnconfigure(0, weight=1)
        
        self.dataset_info_text = tk.Text(info_frame, height=6, wrap=tk.WORD)
        self.dataset_info_text.pack(fill=tk.BOTH, expand=True)
        self.dataset_info_text.insert("1.0", "等待训练开始...\n")
        self.dataset_info_text.config(state=tk.DISABLED)
        
        # 创建标签页
        notebook = ttk.Notebook(right_frame)
        notebook.grid(row=1, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
        
        # 日志标签页
        log_frame = ttk.Frame(notebook)
        notebook.add(log_frame, text="训练日志")
        log_frame.columnconfigure(0, weight=1)
        log_frame.rowconfigure(0, weight=1)
        
        self.log_text = scrolledtext.ScrolledText(log_frame, wrap=tk.WORD, height=20)
        self.log_text.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
        
        # Loss可视化标签页
        loss_frame = ttk.Frame(notebook)
        notebook.add(loss_frame, text="Loss可视化")
        loss_frame.columnconfigure(0, weight=1)
        loss_frame.rowconfigure(0, weight=1)
        
        self.fig = Figure(figsize=(10, 6), dpi=100)
        self.ax = self.fig.add_subplot(111)
        self.ax.set_xlabel('Epoch')
        self.ax.set_ylabel('Loss')
        self.ax.set_title('训练过程Loss曲线')
        self.ax.grid(True)
        
        self.canvas = FigureCanvasTkAgg(self.fig, loss_frame)
        self.canvas.get_tk_widget().grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
        
    def create_param_row(self, parent, label, default, row, param_type=str):
        ttk.Label(parent, text=label, width=20).grid(row=row, column=0, sticky=tk.W, padx=5, pady=2)
        
        entry = ttk.Entry(parent, width=25)
        entry.insert(0, str(default))
        entry.grid(row=row, column=1, sticky=(tk.W, tk.E), padx=5, pady=2)
        
        # 提取参数名：如果标签包含括号，提取括号内的内容；否则使用标签的小写形式
        if '(' in label and ')' in label:
            param_name = label.split('(')[1].split(')')[0]
        else:
            # 没有括号的标签，使用标签的小写形式作为参数名
            param_name = label.lower()
        
        self.params[param_name] = (entry, param_type)
        
    def get_args(self):
        args = argparse.Namespace()
        
        # 基本配置
        args.task_name = self.params['task_name'][0].get()
        args.model = self.params['model'][0].get()
        args.model_id = self.params['model_id'][0].get()
        args.is_training = 1
        
        # 数据配置
        args.data = self.params['data'][0].get()
        args.root_path = self.params['root_path'][0].get()
        args.data_path = self.params['data_path'][0].get()
        args.features = self.params['features'][0].get()
        args.target = 'OT'
        args.freq = self.params['freq'][0].get()
        args.checkpoints = './checkpoints/'
        
        # 预测配置
        args.seq_len = int(self.params['seq_len'][0].get())
        args.label_len = int(self.params['label_len'][0].get())
        args.pred_len = int(self.params['pred_len'][0].get())
        args.seasonal_patterns = 'Monthly'
        args.inverse = False
        
        # 输入任务
        args.mask_rate = 0.25
        
        # 异常检测任务
        args.anomaly_ratio = 0.25
        
        # 模型配置
        args.expand = 2
        args.d_conv = 4
        args.top_k = 5
        args.num_kernels = 6
        args.enc_in = 7
        args.dec_in = 7
        args.c_out = 7
        args.d_model = int(self.params['d_model'][0].get())
        args.n_heads = int(self.params['n_heads'][0].get())
        args.e_layers = int(self.params['e_layers'][0].get())
        args.d_layers = int(self.params['d_layers'][0].get())
        args.d_ff = int(self.params['d_ff'][0].get())
        args.moving_avg = 25
        args.factor = 1
        args.distil = True
        args.dropout = float(self.params['dropout'][0].get())
        args.embed = 'timeF'
        args.activation = 'gelu'
        args.channel_independence = 1
        args.decomp_method = 'moving_avg'
        args.use_norm = 1
        args.down_sampling_layers = 0
        args.down_sampling_window = 1
        args.down_sampling_method = None
        args.seg_len = 96
        
        # 优化配置
        args.num_workers = 10
        args.itr = int(self.params['itr'][0].get())
        args.train_epochs = int(self.params['train_epochs'][0].get())
        args.batch_size = int(self.params['batch_size'][0].get())
        args.patience = int(self.params['patience'][0].get())
        args.learning_rate = float(self.params['learning_rate'][0].get())
        args.des = 'test'
        args.loss = 'MSE'
        args.lradj = 'type1'
        args.use_amp = False
        
        # GPU配置
        use_gpu_str = self.params['use_gpu'][0].get().lower()
        args.use_gpu = use_gpu_str in ['true', '1', 'yes']
        args.gpu = int(self.params['gpu'][0].get())
        args.gpu_type = 'cuda'
        args.use_multi_gpu = False
        args.devices = '0,1,2,3'
        
        # 其他参数
        args.p_hidden_dims = [128, 128]
        args.p_hidden_layers = 2
        args.use_dtw = False
        args.augmentation_ratio = 0
        args.seed = 2
        args.jitter = False
        args.scaling = False
        args.permutation = False
        args.randompermutation = False
        args.magwarp = False
        args.timewarp = False
        args.windowslice = False
        args.windowwarp = False
        args.rotation = False
        args.spawner = False
        args.dtwwarp = False
        args.shapedtwwarp = False
        args.wdba = False
        args.discdtw = False
        args.discsdtw = False
        args.extra_tag = ""
        args.patch_len = 16
        args.node_dim = 10
        args.gcn_depth = 2
        args.gcn_dropout = 0.3
        args.propalpha = 0.3
        args.conv_channel = 32
        args.skip_channel = 32
        args.individual = False
        args.alpha = 0.1
        args.top_p = 0.5
        args.pos = 1
        
        return args
    
    def start_training(self):
        if self.is_training:
            messagebox.showwarning("警告", "训练已在进行中！")
            return
        
        try:
            args = self.get_args()
            self.is_training = True
            self.train_button.config(state=tk.DISABLED)
            self.stop_button.config(state=tk.NORMAL)
            
            # 清空日志和loss
            self.log_text.delete("1.0", tk.END)
            self.train_losses = []
            self.val_losses = []
            self.test_losses = []
            self.epochs = []
            self.ax.clear()
            self.ax.set_xlabel('Epoch')
            self.ax.set_ylabel('Loss')
            self.ax.set_title('训练过程Loss曲线')
            self.ax.grid(True)
            self.canvas.draw()
            
            # 启动训练线程
            self.training_thread = TrainingThread(args, self.log_queue, self.loss_queue, self.dataset_info_queue)
            self.training_thread.start()
            
            self.log_text.insert(tk.END, "训练线程已启动...\n")
            self.log_text.see(tk.END)
            
        except Exception as e:
            messagebox.showerror("错误", f"启动训练失败: {str(e)}")
            self.is_training = False
            self.train_button.config(state=tk.NORMAL)
            self.stop_button.config(state=tk.DISABLED)
    
    def stop_training(self):
        # 注意：这个实现可能无法立即停止训练
        # 更好的实现需要修改训练代码以支持中断
        if self.training_thread and self.training_thread.is_alive():
            messagebox.showwarning("警告", "训练无法立即停止，请等待当前epoch完成")
            # 可以设置一个标志位，让训练代码检查
    
    def process_queue(self):
        # 处理日志队列
        try:
            while True:
                msg = self.log_queue.get_nowait()
                self.log_text.insert(tk.END, msg)
                self.log_text.see(tk.END)
                
                # 检测训练完成
                if "训练完成" in msg or "训练过程中出错" in msg:
                    if self.training_thread and not self.training_thread.is_alive():
                        self.is_training = False
                        self.train_button.config(state=tk.NORMAL)
                        self.stop_button.config(state=tk.DISABLED)
                
                # 从日志中提取loss信息
                # 匹配格式: "Epoch: {0}, Steps: {1} | Train Loss: {2:.7f} Vali Loss: {3:.7f} Test Loss: {4:.7f}"
                loss_pattern = r'Epoch:\s+(\d+).*?Train Loss:\s+([\d.]+).*?Vali Loss:\s+([\d.]+).*?Test Loss:\s+([\d.]+)'
                match = re.search(loss_pattern, msg)
                if match:
                    epoch = int(match.group(1))
                    train_loss = float(match.group(2))
                    val_loss = float(match.group(3))
                    test_loss = float(match.group(4))
                    
                    self.loss_queue.put({
                        'epoch': epoch,
                        'train_loss': train_loss,
                        'val_loss': val_loss,
                        'test_loss': test_loss
                    })
        except queue.Empty:
            pass
        
        # 检查训练线程是否结束
        if self.is_training and self.training_thread and not self.training_thread.is_alive():
            self.is_training = False
            self.train_button.config(state=tk.NORMAL)
            self.stop_button.config(state=tk.DISABLED)
        
        # 处理数据集信息队列
        try:
            while True:
                info = self.dataset_info_queue.get_nowait()
                self.dataset_info_text.config(state=tk.NORMAL)
                self.dataset_info_text.delete("1.0", tk.END)
                self.dataset_info_text.insert("1.0", 
                    f"训练集大小: {info['train_size']}\n"
                    f"测试集大小: {info['test_size']}\n"
                    f"验证集大小: {info['val_size']}\n"
                    f"训练批次数: {info['train_batches']}\n"
                    f"测试批次数: {info['test_batches']}\n"
                    f"验证批次数: {info['val_batches']}\n"
                )
                self.dataset_info_text.config(state=tk.DISABLED)
        except queue.Empty:
            pass
        
        # 处理loss队列
        try:
            while True:
                loss_data = self.loss_queue.get_nowait()
                epoch = loss_data.get('epoch', 0)
                train_loss = loss_data.get('train_loss', 0)
                val_loss = loss_data.get('val_loss', 0)
                test_loss = loss_data.get('test_loss', 0)
                
                self.epochs.append(epoch)
                self.train_losses.append(train_loss)
                self.val_losses.append(val_loss)
                self.test_losses.append(test_loss)
                
                self.update_loss_plot()
        except queue.Empty:
            pass
        
        self.root.after(100, self.process_queue)
    
    def update_loss_plot(self):
        if len(self.epochs) > 0:
            self.ax.clear()
            self.ax.set_xlabel('Epoch')
            self.ax.set_ylabel('Loss')
            self.ax.set_title('训练过程Loss曲线')
            self.ax.grid(True)
            
            if len(self.train_losses) > 0:
                self.ax.plot(self.epochs, self.train_losses, label='Train Loss', marker='o')
            if len(self.val_losses) > 0:
                self.ax.plot(self.epochs, self.val_losses, label='Val Loss', marker='s')
            if len(self.test_losses) > 0:
                self.ax.plot(self.epochs, self.test_losses, label='Test Loss', marker='^')
            
            self.ax.legend()
            self.canvas.draw()


def main():
    root = tk.Tk()
    app = TrainingGUI(root)
    root.mainloop()


if __name__ == '__main__':
    main()