import os
from pyexpat import features
import sys
import json
import torch
import optuna
import joblib

weather_utils_path = os.path.abspath(
    os.path.join(os.path.dirname(__file__), "..")
)
if weather_utils_path not in sys.path:
    sys.path.insert(0, weather_utils_path)
    
from model.DLinear import DLinearModel
from model.GRU import GRUModel
from model.TSMixer import TSMixer
from model.TCN import TemporalConvNet
from model.Mamba import Mamba
from model.CNN_LSTM import CNNLSTM
from model.TransformerEn import TransformerEncoderForecaster


class SolorFeature:
    def __init__(self) -> None:  
        self.features_dict = {
                'v24': ['ssrd','tcc','lcc','mcc','hcc','2t','2d','10u','10v']
            }
    def get_feature(self,release):
        return self.features_dict.get(release)

class ModelLoading:
    """
    ModelLoading类用于根据传入的模型名称和Optuna调参trial对象,自动配置并实例化深度学习模型(如DLinear、GRU、TSMixer(以及相关的超参数设置。
    该类支持模型的加载训练调优模式和预测权重加载模式,便于模型选择、超参数搜索及后续的推理流程。
    """
    def __init__(self, 
                trial,
                projs = 'OL_JIANGMEN',
                model_name = 'Dlinear',
                release = 'None',
                input_seq = 288,
                saved_model_path = 'None',
                predict_mode = False,
                capacity = '4050'
    ) -> None:
        """
            参数说明:
                trial: Optuna的trial对象或占位符(None(,用以搜索和记录超参数组合。如在预测或确定性训练场景下可设为None。
                model_name: 选择要加载/调优的模型名称。支持'Dlinear', 'GRU', 'TSMixer'等(默认'Dlinear'）。
                input_seq: 输入序列长度,通常为一天内的采样步数等,用于确定模型的输入和输出时序长度(默认288)。
                model_path: 预测模式下(指定要加载预训练权重的模型文件路径(默认'None')。
                predict_mode: 是否为预测模式(如果为True,则跳过超参数采样只加载已有权重的最佳模型(默认False)。
        """
    
        self.trial = trial
        self.projs = projs
        self.release = release
        self.features = SolorFeature().get_feature(self.release)
        self.input_seq = input_seq
        self.model_name = model_name
        self.saved_model_path = saved_model_path
        self.predict_mode = predict_mode
        self.capacity = capacity
        
        if self.release != 'None':
            os.makedirs(os.path.join(self.saved_model_path, self.projs, self.model_name, str(self.input_seq), self.release), exist_ok=True)
            self.saved_model_root = os.path.join(self.saved_model_path, self.projs, self.model_name, str(self.input_seq), self.release)
        else:
            os.makedirs(os.path.join(self.saved_model_path, self.projs, self.model_name, str(self.input_seq)), exist_ok=True)
            self.saved_model_root = os.path.join(self.saved_model_path, self.projs, self.model_name, str(self.input_seq))

        # 训练模式下，进行超参数搜索
        if isinstance(self.trial, optuna.trial.Trial):
            if self.model_name == 'Dlinear':
                kernel_size = self.trial.suggest_categorical("kernel_size", [5, 9, 15, 25])  
                lr = self.trial.suggest_float("learning_rate", 1e-5, 1e-2, log=True)
                batch_size = self.trial.suggest_categorical("batch_size", [4, 8, 16, 32, 64, 128])
                self.model_config = {
                    'model': DLinearModel(
                                seq_len = self.input_seq,
                                pred_len = self.input_seq,
                                individual = True,
                                channels = len(self.features),
                                channels_outs = 1,
                                kernel_size = kernel_size
                    ),
                    'lr': lr,
                    'batch_size': batch_size
                }
            elif self.model_name == 'GRU':
                hidden_channels = self.trial.suggest_categorical("hidden_channels", [4,16,32,64,128,256,512])
                layers = self.trial.suggest_categorical("layers", [1,2,3,4,5,6,7,8,9,10])
                lr = self.trial.suggest_float("learning_rate", 1e-5, 1e-2, log=True)
                batch_size = self.trial.suggest_categorical("batch_size", [4, 8, 16, 32, 64, 128])
                self.model_config = {
                    'model': GRUModel(
                        input_channels = len(self.features),
                        output_channels = 1,
                        input_len = self.input_seq,
                        pred_len = self.input_seq,
                        hidden_channels = hidden_channels,
                        layers = layers
                    ),
                    'lr': lr,
                    'batch_size': batch_size
                }
            elif self.model_name == 'TSMixer':
                 num_blocks = self.trial.suggest_categorical("num_blocks", [2,4,6,8,10,12,14,16])
                 dropout_rate = self.trial.suggest_categorical("dropout_rate", [0.1, 0.2, 0.3, 0.4, 0.5, 0.6])
                 ff_dim = self.trial.suggest_categorical("ff_dim", [8,16,32,64,128,256])
                 lr = self.trial.suggest_float("learning_rate", 1e-5, 1e-2, log=True)
                 batch_size = self.trial.suggest_categorical("batch_size", [4, 8, 16, 32, 64, 128])
                 self.model_config = {
                    'model': TSMixer(
                            sequence_length = self.input_seq,
                            prediction_length = self.input_seq,
                            input_channels = len(self.features),
                            output_channels = 1,
                            activation_fn = 'relu',
                            num_blocks = num_blocks,
                            dropout_rate =  dropout_rate,
                            ff_dim = ff_dim,
                            normalize_before = True,
                            norm_type = 'batch',
                    ),
                    'lr': lr,
                    'batch_size':batch_size
                }
            elif self.model_name == 'TCN':
                 num_channels_options = [
                     (8, 1),
                     (16, 8, 1),
                     (32, 16, 8, 1),
                     (64, 32, 16, 8, 1),
                     (128, 64, 32, 16, 8, 1)
                 ]
                 num_channels = self.trial.suggest_categorical("num_channels", num_channels_options)
                 kernel_size = self.trial.suggest_categorical("kernel_size", [2,4,6,8,10,12,14,16])
                 dropout = self.trial.suggest_categorical("dropout", [0.1, 0.2, 0.3, 0.4, 0.5, 0.6])
                 lr = self.trial.suggest_float("learning_rate", 1e-5, 1e-2, log=True)
                 batch_size = self.trial.suggest_categorical("batch_size", [4, 8, 16, 32, 64, 128])
                 self.model_config = {
                    'model': TemporalConvNet(
                                num_inputs= len(self.features), 
                                num_channels= num_channels, 
                                kernel_size= kernel_size, 
                                dropout= dropout
                    ),
                    'lr': lr,
                    'batch_size':batch_size
                }
            elif self.model_name == 'Mamba':
                 d_model = self.trial.suggest_categorical("d_model", [4, 8, 16, 32, 64, 128, 256])
                 d_conv = self.trial.suggest_categorical("d_conv", [4,6,8,10,12,14,16])
                 n_layer = self.trial.suggest_categorical("n_layer", [1,2,3,4,5,6])
                 d_state = self.trial.suggest_categorical("d_state", [2,4,8,16,32])
                 expand = self.trial.suggest_categorical("expand", [1,2,3,4])
                 lr = self.trial.suggest_float("learning_rate", 1e-5, 1e-2, log=True)
                 batch_size = self.trial.suggest_categorical("batch_size", [4, 8, 16, 32, 64, 128])
                 self.model_config = {
                    'model': Mamba(
                                d_model = d_model,
                                n_layer = n_layer,
                                input_dim = len(self.features),
                                output_dim = 1,
                                d_state = d_state,
                                expand = expand,
                                dt_rank = 'auto',
                                d_conv = d_conv
                    ),
                    'lr': lr,
                    'batch_size':batch_size
                }
            elif self.model_name == 'CNNLSTM':
                cnn_channels = self.trial.suggest_categorical("cnn_channels", [16, 32, 64, 128])
                lstm_hidden = self.trial.suggest_categorical("lstm_hidden", [16, 32, 64, 128])
                lstm_layers = self.trial.suggest_categorical("lstm_layers", [1, 2, 3, 4])
                dropout = self.trial.suggest_categorical("dropout", [0.1, 0.2, 0.3, 0.4, 0.5, 0.6])
                lr = self.trial.suggest_float("learning_rate", 1e-5, 1e-2, log=True)
                batch_size = self.trial.suggest_categorical("batch_size", [4, 8, 16, 32, 64, 128])
                self.model_config = {
                    'model': CNNLSTM(
                        n_features = len(self.features),
                        pred_len = self.input_seq,
                        cnn_channels = cnn_channels,
                        lstm_hidden = lstm_hidden,
                        lstm_layers = lstm_layers,
                        dropout = dropout
                    ),
                    'lr': lr,
                    'batch_size': batch_size
                }
            elif self.model_name == 'TransformerEn':
                d_model = self.trial.suggest_categorical("d_model", [16, 32, 64, 128])
                n_heads = self.trial.suggest_categorical("n_heads", [1, 2, 4, 8])
                num_layers = self.trial.suggest_categorical("num_layers", [1, 2, 3, 4])
                dim_feedforward = self.trial.suggest_categorical("dim_feedforward", [16, 32, 64, 128])
                dropout = self.trial.suggest_categorical("dropout", [0.1, 0.2, 0.3, 0.4, 0.5, 0.6])
                lr = self.trial.suggest_float("learning_rate", 1e-5, 1e-2, log=True)
                batch_size = self.trial.suggest_categorical("batch_size", [4, 8, 16, 32, 64, 128])
                pooling = self.trial.suggest_categorical("pooling", ["last", "mean"])
                self.model_config = {
                    'model': TransformerEncoderForecaster(
                        n_features = len(self.features),
                        pred_len = self.input_seq,
                        d_model = d_model,
                        n_heads = n_heads,
                        num_layers = num_layers,
                        dim_feedforward = dim_feedforward,
                        dropout = dropout,
                        pooling = pooling
                    ),
                    'lr': lr,
                    'batch_size': batch_size
                }
            else:
                raise ValueError(f"Unsupported model_name: {self.model_name}")
        else:
            # 加载最优参数,只为当前模型加载对应参数
            # 获取最佳参数路径
            model_param_path = os.path.join(self.saved_model_root,'best_tpe_params.json')
            if os.path.exists(model_param_path):
                # 加载最佳参数
                with open(model_param_path, 'r', encoding='utf-8') as f:
                    self.model_param_dict = json.load(f)
                best_params = self.model_param_dict['best_params']
            else:
                raise FileNotFoundError(f"Model parameter file not found:{self.projs}-{self.release}-{self.model_name}")

            if self.model_name == 'Dlinear':
                self.model_config = {
                    'model': DLinearModel(
                                seq_len = self.input_seq,
                                pred_len = self.input_seq,
                                individual = True,
                                channels = len(self.features),
                                channels_outs = 1,
                                kernel_size = best_params['kernel_size']
                    ),
                    'lr': best_params['learning_rate'],
                    'batch_size': best_params['batch_size']
                }
            elif self.model_name == 'GRU':
                self.model_config = {
                    'model': GRUModel(
                        input_channels = len(self.features),
                        output_channels = 1,
                        input_len = self.input_seq,
                        pred_len = self.input_seq,
                        hidden_channels = best_params['hidden_channels'],
                        layers = best_params['layers']
                    ),
                    'lr': best_params['learning_rate'],
                    'batch_size': best_params['batch_size']
                }
            elif self.model_name == 'TSMixer':
                 self.model_config = {
                    'model': TSMixer(
                            sequence_length = self.input_seq,
                            prediction_length = self.input_seq,
                            input_channels = len(self.features),
                            output_channels = 1,
                            activation_fn = 'relu',
                            num_blocks = best_params['num_blocks'],
                            dropout_rate =  best_params['dropout_rate'],
                            ff_dim = best_params['ff_dim'],
                            normalize_before = True,
                            norm_type = 'batch',
                    ),
                    'lr': best_params['learning_rate'],
                    'batch_size': best_params['batch_size']
                }
            elif self.model_name == 'TCN':
                 self.model_config = {
                    'model': TemporalConvNet(
                                num_inputs= len(self.features), 
                                num_channels= best_params["num_channels"], 
                                kernel_size= best_params["kernel_size"], 
                                dropout= best_params["dropout"]
                    ),
                    'lr': best_params['learning_rate'],
                    'batch_size': best_params['batch_size']
                }
            elif self.model_name == 'Mamba':
                 self.model_config = {
                    'model': Mamba(
                                d_model = best_params["d_model"],
                                n_layer = best_params["n_layer"],
                                input_dim = len(self.features),
                                output_dim = 1,
                                d_state = best_params["d_state"],
                                expand = best_params["expand"],
                                d_conv = best_params["d_conv"]
                    ),
                    'lr': best_params['learning_rate'],
                    'batch_size': best_params['batch_size']
                }
            elif self.model_name == 'CNNLSTM':
                self.model_config = {
                    'model': CNNLSTM(
                        n_features = len(self.features),
                        pred_len = self.input_seq,
                        cnn_channels = best_params['cnn_channels'],
                        lstm_hidden = best_params['lstm_hidden'],
                        lstm_layers = best_params['lstm_layers'],
                        dropout = best_params['dropout']
                    ),
                    'lr': best_params['learning_rate'],
                    'batch_size': best_params['batch_size']
                }
            elif self.model_name == 'TransformerEn':
                self.model_config = {
                    'model': TransformerEncoderForecaster(
                        n_features = len(self.features),
                        pred_len = self.input_seq,
                        d_model = best_params['d_model'],
                        n_heads = best_params['n_heads'],
                        num_layers = best_params['num_layers'],
                        dim_feedforward = best_params['dim_feedforward'],
                        dropout = best_params['dropout'],
                        pooling = best_params['pooling']
                    ),
                    'lr': best_params['learning_rate'],
                    'batch_size': best_params['batch_size']
                }
            else:
                raise ValueError(f"Unsupported model_name: {self.model_name}")
        
        # 启动预测模式
        if self.predict_mode:
            print(f'{self.model_name}--进入预测模式！！！')
            # 判断当前是否有可用的GPU,如果有则使用cuda,否则使用cpu
            # 加载最佳模型
            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
            self.model = self.get_model()
            model_weigth_file = os.path.join(self.saved_model_root, 'best_model.pth') 
            self.model.load_state_dict(torch.load(model_weigth_file, map_location=device))
            self.model.eval()

            # 加载归一化器
            scaler_path = os.path.join(self.saved_model_root, "scaler.pkl")
            if not os.path.exists(scaler_path):
                raise FileNotFoundError(f"Scaler file not found: {self.proj}-{self.model_name}-{self.pre_seq}-{self.release}")
            self.scaler = joblib.load(scaler_path)

    def get_model(self):
        return self.model_config.get('model')

    def get_lr(self):
        return self.model_config.get('lr')
        
    def get_batch_size(self):
        return self.model_config.get('batch_size')

    def predict(self, features_in):
        """
        功能:
            用于对输入的特征进行模型预测,输出预测结果。
        参数:
            features_in: 输入特征数据,格式应为模型可接受的数据类型(如numpy数组、Tensor等),
                        shape = (batch_size, seq_len, feature_dim)。
        返回:
            预测结果shape = (batch_size, seq_len, 1)。
        """
        if not self.predict_mode:
            raise RuntimeError("Model is not in prediction mode. Please set predict=True in ModelLoading to use predict().")
        else:
            x_scaled = self.scaler_fn(features_in)
            return self.inverse_transform(self.model(x_scaled),self.capacity)

    def inverse_transform(self, y, capacity):
        """
        对模型输出进行反归一化，以恢复到原始的功率尺度。

        参数:
            y: 模型输出,归一化后的数据(如Tensor或ndarray)。
            capacity: 电站容量（标量），用于恢复为实际功率。

        返回:
            反归一化后的功率预测值。
        """
        return y * capacity
   

    def scaler_fn(self, x):
        """
        对输入特征 x 进行归一化处理。
        参数:
            x: 可以为 torch.Tensor 或 numpy.ndarray 类型的输入特征，
               通常形状为 (1, seq_len, feature_dim)。

        返回:
            归一化后的数据,返回torch.Tensor;
        """
        # 如果传入的是torch张量，则转为numpy数组进行归一化
        B, L, M = x.shape
        if isinstance(x, torch.Tensor):
            x_np = x.detach().cpu().numpy()  # (B, L, M)
            x_np_flat = x_np.reshape(-1, M)  # (B*L, M)
            x_scaled_flat = self.scaler.transform(x_np_flat)  # (B*L, M)
            x_scaled = x_scaled_flat.reshape(B, L, M)  # (B, L, M)
            return torch.from_numpy(x_scaled).to(x.device).type(x.dtype)
            
        else:
            x_flat = x.reshape(-1, M)
            x_scaled_flat = self.scaler.transform(x_flat)
            x_scaled = x_scaled_flat.reshape(B, L, M)
            return x_scaled