import torch
import torch.nn as nn
import torch.nn.functional as F
from ns_layers.SelfAttention_Family import AttentionLayer, FullAttention
from layers.Embed import DataEmbedding, DataEmbedding_inverted, PositionalEmbedding
import numpy as np
import math
class FlattenHead(nn.Module):
    def __init__(self, n_vars, nf, target_window, head_dropout=0):
        super().__init__()
        self.n_vars = n_vars
        self.flatten = nn.Flatten(start_dim=-2)
        self.linear = nn.Linear(nf, target_window)
        #self.dropout = nn.Dropout(head_dropout)

    def forward(self, x):  # x: [bs x nvars x d_model x patch_num]
        x = self.flatten(x)
        x = self.linear(x)
        #x = self.dropout(x)

        return x


class EnEmbedding(nn.Module):
    def __init__(self, n_vars, d_model, patch_len, dropout):
        super(EnEmbedding, self).__init__()
        # Patching
        self.patch_len = patch_len

        self.value_embedding = nn.Linear(patch_len, d_model, bias=False)
        self.glb_token = nn.Parameter(torch.randn(1, n_vars, 1, d_model))
        self.position_embedding = PositionalEmbedding(d_model)

        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None, tt=None, **args):
        # do patching
        n_vars = x.shape[1]
        glb = self.glb_token.repeat((x.shape[0], 1, 1, 1))

        x = x.unfold(dimension=-1, size=self.patch_len, step=self.patch_len)
        x = torch.reshape(x, (x.shape[0] * x.shape[1], x.shape[2], x.shape[3]))
        # Input encoding
        x = self.value_embedding(x) + self.position_embedding(x)
        x = torch.reshape(x, (-1, n_vars, x.shape[-2], x.shape[-1]))
        x = torch.cat([x, glb], dim=2)
        x = torch.reshape(x, (x.shape[0] * x.shape[1], x.shape[2], x.shape[3]))
        return self.dropout(x), n_vars
    
class EntropyAwareEmbedding(nn.Module):
    def __init__(self, n_vars, d_model, patch_len, patch_num,  dropout, target_len):
        super(EntropyAwareEmbedding, self).__init__()
        # Patching
        self.patch_len = patch_len
        self.patch_num = patch_num
        self.target_len = target_len
        self.te_dim = 8
        self.groups = 16
        self.perproject = nn.Linear(1, self.te_dim)
        self.device = torch.device('cuda:0')
        self.glb_token = nn.Parameter(torch.randn(1, n_vars, 1, d_model))
        self.position_embedding = PositionalEmbedding(self.te_dim)
        self.dropout = nn.Dropout(dropout)
        self.tau = nn.Parameter(torch.zeros(size=(n_vars, self.groups)))
        self.boundarynet = nn.Sequential(
            nn.Linear(target_len*3, patch_num*8),
            nn.ELU(),
            nn.Linear(patch_num*8, patch_num*2)
        )
        # self.boundarynet_r = nn.Sequential(
        #     nn.Linear(target_len*3, target_len*2),
        #     nn.ELU(),
        #     nn.Linear(target_len*2, patch_num)
        # )
        patch_centers = torch.linspace(patch_len / 2, 
                                            patch_num * patch_len - patch_len / 2, 
                                            patch_num, 
                                            device=self.device, requires_grad=False)
        self.register_buffer("patch_centers", patch_centers)
        
        base_left_boundaries = (patch_centers - patch_len / 2).unsqueeze(0).unsqueeze(-1)  # [1,  P, 1]
        self.register_buffer("base_left_boundaries", base_left_boundaries)
        self.register_buffer("time", torch.arange(0, self.target_len, 1, device=self.device))  # [ patch_num*patch_len]
        
        self.project = nn.Linear(self.te_dim*self.groups, d_model)

        self.te_sacle_list = nn.ModuleList([ nn.Linear(1, 1) for i in range(4) ])
        self.te_periodic_list = nn.ModuleList([nn.Linear(1, (self.te_dim-4)//4 ) for i in range(4)])


    


    def forward(self, x, mask=None, tt=None):
        # do patching
        x # [B, n_vars, L]
        if mask is None:
            mask = torch.ones_like(x)  # [B, n_vars, L]
        mask # [B, n_vars, L]
        n_vars = x.shape[1]
   
        #te = self.LearnableTE(tt).unsqueeze(1).repeat(1, n_vars, 1, 1) # [B, N,L, TE]
      
        x_with_te = self.perproject(x.unsqueeze(-1)) + self.position_embedding(x) 
        # x_with_te = x.unsqueeze(-1)
        glb = self.glb_token.repeat((x.shape[0], 1, 1, 1))
        # x = self.interpolate_batch_vectorized_linear(x.reshape(-1, x.shape[-1]), mask.reshape(-1, mask.shape[-1]))
        x = x.reshape(-1,  x.shape[-1]) #[B*N, L]
        # x = x.unfold(dimension=-1, size=self.patch_len, step=self.patch_len)
        # x = torch.reshape(x, (x.shape[0] * x.shape[1], x.shape[2], x.shape[3]))   # [B*n_vars, num_patches, patch_len]

        SampEn, counts = self.compute_SampleEntropy_optimized(x, m=2, r=0.2, mask=mask)  # [B*n_vars, num_patches]
        counts = counts/192
        counts = counts.reshape(-1, counts.shape[-2]*2)
        x = torch.concat([x, counts], dim=-1)  # [B*n_vars, num_patches, patch_len*3]

        shifts = self.boundarynet(x).reshape(-1, self.patch_num, 2)  # [B*n_vars, num_patches*2, 1]
        shifts_left = shifts[..., 0:1]


        tau_vars = F.softplus(self.tau).unsqueeze(-1).unsqueeze(-1) + 1e-3  # [n_vars, g, 1, 1]
        shifts_right = shifts[..., 1:2]
     
        x =  self.patch_aggregation(shifts_left, 
                                    shifts_right, 
                                    tau_vars, 
                                    x_with_te.reshape(-1, x_with_te.shape[-2], x_with_te.shape[-1]), 
                                    mask.reshape(-1, mask.shape[-1]), 
                                    glb, 
                                    n_vars)
        return x, n_vars


    def LearnableTE(self, tt):
        
        B, L, te = tt.shape
      
        te_emb = torch.zeros(size=(B, L, self.te_dim), device=self.device )
        for i,( te_scale, te_periodic) in enumerate(zip(self.te_sacle_list, self.te_periodic_list)):
            out2 = torch.sin(te_periodic(tt[..., i:i+1])) # [..., D_te - 1]
            te_emb[..., i*4:i*4+1] = te_scale(tt[..., i:i+1])
            te_emb[..., i*4+1:i*4+5]
        return te_emb # [..., D_te]
    
    def patch_aggregation(self, shift_left, shift_right, tau, x, mask, glb, n_vars):
        
        adjusted_left = self.base_left_boundaries + shift_left  # [B*n_vars, num_patches, 1]
        adjusted_right = adjusted_left + shift_right  # [B*n_vars, num_patches, 1]
        adjusted_left = adjusted_left.unsqueeze(1)
        adjusted_right = adjusted_right.unsqueeze(1)
        weights_raw = torch.sigmoid((adjusted_right - self.time) / tau) * \
                      torch.sigmoid((self.time - adjusted_left) / tau)  # [B*n_vars, g,  num_patches, patch_num*patch_len]
       
        weights = weights_raw * mask.reshape(-1, 1, 1, mask.shape[-1])  # [B*n_vars, g, num_patches, patch_num*patch_len]
        x = x.unsqueeze(1)
        x = weights @ x  # [B*n_vars, g, num_patches, te_dim]
        x = x / torch.sum(weights, dim=-1, keepdim=True)
        x = x.permute(0, 2, 1, 3)
       
        x = torch.reshape(x, (-1, n_vars, x.shape[1], x.shape[2]*x.shape[3]))
        x = self.project(x)
        x = torch.cat([x, glb], dim=2)
        x = x.reshape((x.shape[0] * x.shape[1], x.shape[2], x.shape[3]))
        return self.dropout(x)


    def compute_SampleEntropy_optimized(self, x, m=2, r=0.2, mask=None):
        """
        
        """
        def _phi_optimized(m):
            shapes = x.shape
            N = shapes[-1]
            L = N - m + 1  

            x_m = x.unfold(dimension=-1, size=m, step=1)  # shape: (..., L, m)
            

            x_m_expanded1 = x_m.unsqueeze(-3)  # shape: (..., 1, L, m)
            x_m_expanded2 = x_m.unsqueeze(-2)  # shape: (..., L, 1, m)

            diff = torch.abs(x_m_expanded1 - x_m_expanded2)
            max_diff = diff.max(dim=-1)[0]  # shape: (..., L, L)

            matches = max_diff <= r

            mask_matrix = ~torch.eye(L, dtype=torch.bool, device=x.device)
            matches = matches * mask_matrix
            counts = matches.sum(dim=-1)  # shape: (..., L)
            return 0, counts
            # C = counts.sum(dim=-1)  
            
            # return C / (L * (L - 1)), counts
        
     
        c1, count = _phi_optimized(m)
        c2, count1 = _phi_optimized(m + 1)
        

        counts_tensor = torch.zeros(size=(*x.shape, 2)).to(x.device)
        counts_tensor[..., :count.shape[-1], 0] = count
        counts_tensor[..., :count1.shape[-1], 1] = count1

        
        return 0, counts_tensor.to(self.device)

   
    
    def cubic_spline_coefficients(self, x, y):
        """计算三次样条系数"""
        n = len(x)
        h = x[1:] - x[:-1]
        
        # 构建三对角矩阵
        A = torch.zeros(n, n)
        b = torch.zeros(n, )
        
        # 内部点条件
        for i in range(1, n-1):
            A[i, i-1] = h[i-1]
            A[i, i] = 2 * (h[i-1] + h[i])
            A[i, i+1] = h[i]
            b[i] = 3 * ((y[i+1] - y[i]) / h[i] - (y[i] - y[i-1]) / h[i-1])
        
        # 边界条件（自然样条）
        A[0, 0] = 1
        A[-1, -1] = 1
        
        # 解方程求二阶导数
        c = torch.linalg.solve(A, b)
        
        # 计算其他系数
        a = y[:-1]
        b_coeff = (y[1:] - y[:-1]) / h - h * (2 * c[:-1] + c[1:]) / 3
        d = (c[1:] - c[:-1]) / (3 * h)
        
        return a, b_coeff, c[:-1], d, x[:-1], h

    def linear_interpolation(self, x, x_known, y_known):
        """线性插值"""
        idx = torch.searchsorted(x_known, x)
        idx = torch.clamp(idx, 1, len(x_known) - 1)  # 防止越界
        x0, x1 = x_known[idx - 1], x_known[idx]
        y0, y1 = y_known[idx - 1], y_known[idx]
        return y0 + (y1 - y0) * (x - x0) / (x1 - x0)

    def interpolate_batch_vectorized_linear(self, data, mask: torch.Tensor):
        """向量化的线性插值"""
        device = data.device
        batch_size, seq_len = data.shape
        result = data.clone()
        
        x_all = torch.arange(seq_len, device=device, dtype=torch.float32)
        
        for i in range(batch_size):
            sample_mask = mask[i].bool()
            if sample_mask.sum() < 2:
                continue
            
            known_indices = torch.where(sample_mask)[0]
            unknown_indices = torch.where(~sample_mask)[0]
            
            if len(unknown_indices) == 0:
                continue
            
            # 对已知点排序
            sorted_known_indices, sort_idx = torch.sort(known_indices)
            sorted_known_values = data[i, sorted_known_indices]
            
            # 为每个未知点找到最近的左右已知点
            # 使用searchsorted找到插入位置
            right_positions = torch.searchsorted(sorted_known_indices, unknown_indices, right=True)
            left_positions = right_positions - 1
            
            # 处理边界情况
            left_positions = torch.clamp(left_positions, 0, len(sorted_known_indices)-1)
            right_positions = torch.clamp(right_positions, 0, len(sorted_known_indices)-1)
            
            # 获取左右索引和值
            left_indices = sorted_known_indices[left_positions]
            right_indices = sorted_known_indices[right_positions]
            left_values = sorted_known_values[left_positions]
            right_values = sorted_known_values[right_positions]
            
            # 计算插值权重
            # 避免除以零（当左右索引相同时）
            mask_diff = right_indices != left_indices
            weights = torch.zeros_like(unknown_indices, dtype=torch.float32)
            
            if mask_diff.any():
                weights[mask_diff] = (unknown_indices[mask_diff] - left_indices[mask_diff]) / \
                                    (right_indices[mask_diff] - left_indices[mask_diff])
            
            # 线性插值
            interpolated_values = left_values + weights * (right_values - left_values)
            
            # 如果左右索引相同，直接使用该值（最近邻）
            interpolated_values[~mask_diff] = left_values[~mask_diff]
            
            # 更新结果
            result[i, unknown_indices] = interpolated_values
        
        return result



class Encoder(nn.Module):
    def __init__(self, layers, norm_layer=None, projection=None):
        super(Encoder, self).__init__()
        self.layers = nn.ModuleList(layers)
        self.norm = norm_layer
        self.projection = projection

    def forward(self, x, x_ex, x_mask=None, cross_mask=None, tau=None, delta=None, mode="forecast"):
        for layer in self.layers:
            x = layer(x, x_ex, x_ex, x_mask=x_mask, cross_mask=cross_mask, tau=tau, delta=delta, mode=mode)

        if self.norm is not None:
            x = self.norm(x)

        if self.projection is not None:
            x = self.projection(x)
        return x

    

class EncoderLayer(nn.Module):
    def __init__(self, self_attention, cross_attention, d_model, d_ff=None,
                 dropout=0.1, activation="relu"):
        super(EncoderLayer, self).__init__()
        d_ff = d_ff or 4 * d_model
        self.self_attention = self_attention
        self.cross_attention = cross_attention
        self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1)
        self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)
        self.activation = F.relu if activation == "relu" else F.gelu

    def forward(self, x, cross, pseudo, x_mask=None, cross_mask=None, tau=None, delta=None, mode="forecast", num_en=1):
        if mode == "forecast" or num_en == 1:
            B, L, D = cross.shape
            x = x + self.dropout(self.self_attention(
                x, x, x,
                attn_mask=x_mask,
                tau=tau, delta=None
            )[0])
            x = self.norm1(x)

            x_glb_ori = x[:, -1, :].unsqueeze(1)
            x_glb = torch.reshape(x_glb_ori, (B, -1, D))
            x_glb_attn = self.dropout(self.cross_attention(
                x_glb, cross, pseudo,
                attn_mask=cross_mask,
                tau=tau, delta=delta
            )[0])
            x_glb_attn = torch.reshape(x_glb_attn,
                                    (x_glb_attn.shape[0] * x_glb_attn.shape[1], x_glb_attn.shape[2])).unsqueeze(1)
            x_glb = x_glb_ori + x_glb_attn
            x_glb = self.norm2(x_glb)

            y = x = torch.cat([x[:, :-1, :], x_glb], dim=1)

            y = self.dropout(self.activation(self.conv1(y.transpose(-1, 1))))
            y = self.dropout(self.conv2(y).transpose(-1, 1))

            return self.norm3(x + y)
        elif mode == "multi" and num_en > 1:
            B, num_ex, D = cross.shape
            bn, L, D = x.shape
            x = x.reshape(-1, num_en, L, D)
            for i in num_en:
                x_en = x[:, i, :, :] # [B, L, D]
                
                x_glb_ori_en = x[:, :, -1, :] #[B, num_en, D]




class Model(nn.Module):

    def __init__(self, configs):
        super(Model, self).__init__()
        self.task_name = 'long_term_forecast'
        self.features = "MS"
        self.seq_len = configs.seq_len
        self.pred_len = configs.pred_len
        self.use_norm = configs.use_norm
        self.patch_len = configs.patch_len
        self.target_len = configs.seq_len - configs.pseudo_len
        self.pseudo = configs.pseudo_len
        self.patch_num = int(self.target_len // configs.patch_len)
        print("number of patch", self.patch_num)
        self.n_vars =  1 if configs.features == 'MS' else configs.enc_in
        self.device = torch.device('cuda:0')
        # Embedding
        self.en_embedding_x = EnEmbedding(self.n_vars, configs.d_model, self.patch_len, configs.dropout)

        self.en_embedding = EntropyAwareEmbedding(self.n_vars, configs.d_model, self.patch_len, self.patch_num, configs.dropout, self.target_len)

        self.ex_embedding = DataEmbedding_inverted(self.target_len+self.pseudo, configs.d_model, configs.embed, configs.freq,
                                                   configs.dropout)

        
        # self.ex_embedding_pseudo = DataEmbedding_inverted(self.target_len+self.pseudo, configs.d_model, configs.embed, configs.freq,
        #                                            configs.dropout)
        # # Encoder-only architecture
        self.encoder = Encoder(
            [
                EncoderLayer(
                    AttentionLayer(
                        FullAttention(False, configs.factor, attention_dropout=configs.dropout,
                                      output_attention=False),
                        configs.d_model, configs.n_heads),
                    AttentionLayer(
                        FullAttention(False, configs.factor, attention_dropout=configs.dropout,
                                      output_attention=False),
                        configs.d_model, configs.n_heads),
                    configs.d_model,
                    configs.d_ff,
                    dropout=configs.dropout,
                    activation=configs.activation,
                )
                for l in range(configs.e_layers)
            ],
            norm_layer=torch.nn.LayerNorm(configs.d_model)
        )
        self.head_nf = configs.d_model * (self.patch_num + 1)
        self.head = FlattenHead(configs.enc_in, self.head_nf, configs.pred_len,
                                head_dropout=configs.dropout)
        

    def masked_normalize(self, x, mask, eps=1e-5):
        """
        Normalize x using only the positions where mask == 1
        
        Args:
            x: Tensor of shape [batch_size, seq_len, feature_dim]
            mask: Tensor of shape [batch_size, seq_len], 1 for valid, 0 for invalid
            eps: Small constant for numerical stability
        
        Returns:
            normalized_x: Normalized tensor
            means: Computed means
            stdevs: Computed standard deviations
        """
 
        valid_count = mask.sum(dim=1, keepdim=True)  # [batch_size, 1, 1]
        
        # 计算加权均值
        means = (x).sum(dim=1, keepdim=True) / torch.clamp(valid_count, min=1)
        means = means.detach()
        
        # 去均值
        x_centered = x - means
        
        # 计算加权方差
        squared_centered = (x_centered ** 2) * mask
        variance = squared_centered.sum(dim=1, keepdim=True) / torch.clamp(valid_count, min=1)
        stdevs = torch.sqrt(variance + eps)
        
        normalized_x = x_centered / stdevs
        normalized_x = normalized_x * mask
        return normalized_x, means, stdevs

    

    def forecast(self, x_enc, x_mark_enc, x_dec, x_mark_dec, mask=None):
        if mask is None:
            mask = torch.ones_like(x_enc).to(self.device)
        # if self.use_norm:
        #     # Normalization from Non-stationary Transformer
        #     means = x_enc.mean(1, keepdim=True).detach()
        #     x_enc = x_enc - means
        #     stdev = torch.sqrt(torch.var(x_enc, dim=1, keepdim=True, unbiased=False) + 1e-5)
        #     x_enc /= stdev

        if self.use_norm:
            x_enc, means, stdev = self.masked_normalize(x_enc, mask)

        _, _, N = x_enc.shape
        en_embed, n_vars = self.en_embedding(x_enc[:, :self.target_len, -1].unsqueeze(-1).permute(0, 2, 1), 
                                             mask[:, :self.target_len, -1].unsqueeze(-1).permute(0, 2, 1),
                                             x_mark_enc[:, :self.target_len] )
        en_enbed_linear, n_vars = self.en_embedding_x(x_enc[:, :self.target_len, -1].unsqueeze(-1).permute(0, 2, 1), 
                                             mask[:, :self.target_len, -1].unsqueeze(-1).permute(0, 2, 1),
                                             x_mark_enc[:, :self.target_len])
        
        en_embed = en_embed + en_enbed_linear
        ex_embed = self.ex_embedding(x_enc[:, :, :-1], x_mark_enc[:, :, :])
        #ex_embed_pseudo = self.ex_embedding(x_enc[:, :, :-1], x_mark_enc[:, :, :] )

        enc_out = self.encoder(en_embed, ex_embed, ex_embed)
        enc_out = torch.reshape(
            enc_out, (-1, n_vars, enc_out.shape[-2], enc_out.shape[-1]))
        # z: [bs x nvars x d_model x patch_num]
        enc_out = enc_out.permute(0, 1, 3, 2)

        dec_out = self.head(enc_out)  # z: [bs x nvars x target_window]
        dec_out = dec_out.permute(0, 2, 1)

        if self.use_norm:
            # De-Normalization from Non-stationary Transformer
            dec_out = dec_out * (stdev[:, 0, -1:].unsqueeze(1).repeat(1, self.pred_len, 1))
            dec_out = dec_out + (means[:, 0, -1:].unsqueeze(1).repeat(1, self.pred_len, 1))

        return dec_out


    

    def forward(self, x_enc, x_mark_enc, x_dec, x_mark_dec, mask=None, mode='forecast', target:int=1):
        if mode == 'forecast':
            dec_out = self.forecast(x_enc, x_mark_enc, x_dec, x_mark_dec, mask)
        elif mode == 'mutilforecast':
            dec_out = self.mutilforecast(x_enc, x_mark_enc, x_dec, x_mark_dec, mask, target)
        else:
            raise ValueError("mode error!")
        return dec_out[:, -self.pred_len:, :]  # [B, L, D]
    

    def mutilforecast(self, x_enc, x_mark_enc, x_dec, x_mark_dec, mask=None, num_en:int=1):
        if mask is None:
            mask = torch.ones_like(x_enc).to(self.device)  # [B, n_vars, L]
        if self.use_norm:
            # Normalization from Non-stationary Transformer
            means = x_enc.mean(1, keepdim=True).detach()
            x_enc = x_enc - means
            stdev = torch.sqrt(torch.var(x_enc, dim=1, keepdim=True, unbiased=False) + 1e-5)
            x_enc /= stdev
        _, _, N = x_enc.shape

        x_en = x_enc[:, :self.target_len, -num_en:].unsqueeze(-1).permute(0, 2, 1)
        x_ex = x_enc[:, :, :-num_en]
        ex_embed = self.ex_embedding(x_ex, x_mark_enc[:, :, :])
        en_embed, n_vars = self.en_embedding(x_en, mask[:, :self.target_len, -num_en:].unsqueeze(-1))
        
            
