import math
import torch
import torch.nn as nn
import torch.utils.checkpoint as checkpoint
import torch.nn.functional as F
from functools import partial
from typing import Optional, Callable
from timm.models.layers import DropPath, to_2tuple, trunc_normal_
from mamba_ssm.ops.selective_scan_interface import selective_scan_fn, selective_scan_ref
from einops import rearrange, repeat
from Mutil_prompt_CLIP.LightTrans import SPE_prompt_injection,SPA_prompt_injection

class Spec_SS1D(nn.Module):
    def __init__(
            self,
            d_model,
            in_chans,
            in_img_size,
            ctx_dim,
            d_state=16,
            d_conv=3,
            expand=1.,
            dt_rank="auto",
            dt_min=0.001,
            dt_max=0.1,
            dt_init="random",
            dt_scale=1.0,
            dt_init_floor=1e-4,
            dropout=0.,
            conv_bias=True,
            bias=False,
            device=None,
            dtype=None,
            **kwargs,
    ):
        factory_kwargs = {"device": device, "dtype": dtype}
        super().__init__()
        self.d_model = d_model
        self.ctx_dim = ctx_dim
        self.d_state = d_state
        self.d_conv = d_conv
        self.expand = expand
        self.d_inner = int(self.expand * self.d_model)
        self.dt_rank = math.ceil(self.d_model / 16) if dt_rank == "auto" else dt_rank

        self.in_proj = nn.Linear(self.d_model, self.d_inner * 2, bias=bias, **factory_kwargs)
        self.proj = nn.Linear(self.d_inner, self.d_inner, bias=bias)
        self.act = nn.SiLU()

        self.x_proj = (  
            nn.Linear(self.d_inner, (self.dt_rank + self.d_state * 2), bias=False, **factory_kwargs),
            nn.Linear(self.d_inner, (self.dt_rank + self.d_state * 2), bias=False, **factory_kwargs),
        )
        self.x_proj_weight = nn.Parameter(torch.stack([t.weight for t in self.x_proj], dim=0))  # (K=4, N, inner)
        del self.x_proj

        self.dt_projs = (
            self.dt_init(self.dt_rank, self.d_inner, dt_scale, dt_init, dt_min, dt_max, dt_init_floor,
                         **factory_kwargs),
            self.dt_init(self.dt_rank, self.d_inner, dt_scale, dt_init, dt_min, dt_max, dt_init_floor,
                         **factory_kwargs),
        )
        self.dt_projs_weight = nn.Parameter(torch.stack([t.weight for t in self.dt_projs], dim=0))  # (K=4, inner, rank)
        self.dt_projs_bias = nn.Parameter(torch.stack([t.bias for t in self.dt_projs], dim=0))  # (K=4, inner)
        del self.dt_projs

        self.A_logs = self.A_log_init(self.d_state, self.d_inner, copies=2, merge=True)  # (K=2, D, N)
        self.Ds = self.D_init(self.d_inner, copies=2, merge=True)  # (K=2, D, N)

        self.selective_scan = selective_scan_fn

        self.out_norm = nn.LayerNorm(self.d_inner)
        self.out_proj = nn.Linear(self.d_inner, self.ctx_dim, bias=bias, **factory_kwargs)
        self.dropout = nn.Dropout(dropout) if dropout > 0. else None
        self.SPE_prompt_injection_text = SPE_prompt_injection(self.d_inner,in_chans)
        self.SPA_prompt_injection_text = SPA_prompt_injection(self.d_inner,in_img_size)

    @staticmethod
    def dt_init(dt_rank, d_inner, dt_scale=1.0, dt_init="random", dt_min=0.001, dt_max=0.1, dt_init_floor=1e-4,
                **factory_kwargs):
        dt_proj = nn.Linear(dt_rank, d_inner, bias=True, **factory_kwargs)

        # Initialize special dt projection to preserve variance at initialization
        dt_init_std = dt_rank ** -0.5 * dt_scale
        if dt_init == "constant":
            nn.init.constant_(dt_proj.weight, dt_init_std)
        elif dt_init == "random":
            nn.init.uniform_(dt_proj.weight, -dt_init_std, dt_init_std)
        else:
            raise NotImplementedError

        dt = torch.exp(
            torch.rand(d_inner, **factory_kwargs) * (math.log(dt_max) - math.log(dt_min))
            + math.log(dt_min)
        ).clamp(min=dt_init_floor)
        inv_dt = dt + torch.log(-torch.expm1(-dt))
        with torch.no_grad():
            dt_proj.bias.copy_(inv_dt)
        dt_proj.bias._no_reinit = True

        return dt_proj

    @staticmethod
    def A_log_init(d_state, d_inner, copies=1, device=None, merge=True):
        # S4D real initialization
        A = repeat(
            torch.arange(1, d_state + 1, dtype=torch.float32, device=device),
            "n -> d n",
            d=d_inner,
        ).contiguous()
        A_log = torch.log(A)  # Keep A_log in fp32
        if copies > 1:
            A_log = repeat(A_log, "d n -> r d n", r=copies)
            if merge:
                A_log = A_log.flatten(0, 1)
        A_log = nn.Parameter(A_log)
        A_log._no_weight_decay = True
        return A_log

    @staticmethod
    def D_init(d_inner, copies=1, device=None, merge=True):
        # D "skip" parameter
        D = torch.ones(d_inner, device=device)
        if copies > 1:
            D = repeat(D, "n1 -> r n1", r=copies)
            if merge:
                D = D.flatten(0, 1)
        D = nn.Parameter(D)  # Keep in fp32
        D._no_weight_decay = True
        return D

    def forward_core(self, x: torch.Tensor):
        B, C, D = x.shape
        L = C
        K = 2
        xs = torch.stack([x.view(B, -1, L), torch.flip(x.view(B, -1, L), dims=[-1])], dim=1)

        x_dbl = torch.einsum("b k d l, k c d -> b k c l", xs.view(B, K, -1, L), self.x_proj_weight)
        dts, Bs, Cs = torch.split(x_dbl, [self.dt_rank, self.d_state, self.d_state], dim=2)
        dts = torch.einsum("b k r l, k d r -> b k d l", dts.view(B, K, -1, L), self.dt_projs_weight)
        xs = xs.float().view(B, -1, L)
        dts = dts.contiguous().float().view(B, -1, L) # (b, k * d, l)
        Bs = Bs.float().view(B, K, -1, L)
        Cs = Cs.float().view(B, K, -1, L) # (b, k, d_state, l)
        Ds = self.Ds.float().view(-1)
        As = -torch.exp(self.A_logs.float()).view(-1, self.d_state)
        dt_projs_bias = self.dt_projs_bias.float().view(-1) # (k * d)
        out_y = self.selective_scan(
            xs, dts,
            As, Bs, Cs, Ds, z=None,
            delta_bias=dt_projs_bias,
            delta_softplus=True,
            return_last_state=False,
        ).view(B, K, -1, L)
        assert out_y.dtype == torch.float
        y2 = torch.flip(out_y[:,1], dims=[-1])
        return out_y[:,0], y2

    def forward(self, x: torch.Tensor, max_spe,max_spa,**kwargs):
        xz = self.in_proj(x)
        x, z = xz.chunk(2, dim=-1)

        x = self.act(self.proj(x))
        SPE = self.SPE_prompt_injection_text(x[:, 0:3, :], max_spe)
        SPA = self.SPA_prompt_injection_text(x[:, 3:6, :], max_spa)
        x = torch.cat((SPE, SPA), 1)
        y1, y2 = self.forward_core(x)

        assert y1.dtype == torch.float32
        y = y1 + y2
        y = torch.transpose(y, dim0=1, dim1=2).contiguous()
        y = self.out_norm(y)
        y = y * F.silu(z) + x
        out = self.out_proj(y)
        if self.dropout is not None:
            out = self.dropout(out)
        return out

class PTLEMambaBlock(nn.Module):
    def __init__(
            self,
            ctx_dim,
            in_chans,
            in_img_size: int = 11,
            hidden_dim: int = 0,
            drop_path: float = 0,
            norm_layer: Callable[..., torch.nn.Module] = partial(nn.LayerNorm, eps=1e-6),
            attn_drop_rate: float = 0,
            d_state: int = 16,
            expand: float = 0.5,
            is_light_sr: bool = False,
            **kwargs,

    ):
        super().__init__()
        self.ln_1 = norm_layer(hidden_dim)
        self.self_attention = Spec_SS1D(d_model=hidden_dim, d_state=d_state,expand=expand,dropout=attn_drop_rate, **kwargs)
        self.drop_path = DropPath(drop_path)
        self.SPE_prompt_injection_text = SPE_prompt_injection(ctx_dim, in_chans)
        self.SPA_prompt_injection_text = SPA_prompt_injection(ctx_dim,in_img_size)

    def forward(self, input):
        # B,C,H,W=input.size()
        # input=input.reshape(B,C,H*W)
        x = self.ln_1(input).unsqueeze(0)
        x = self.self_attention(x)
        x = x + input
        # x = x.reshape(B,C,H,W)
        return x


class PTCOMambaBlock(nn.Module):
    def __init__(
            self,
            ctx_dim,
            in_chans,
            in_img_size,
            hidden_dim: int = 0,
            drop_path: float = 0,
            norm_layer: Callable[..., torch.nn.Module] = partial(nn.LayerNorm, eps=1e-6),
            attn_drop_rate: float = 0,
            d_state: int = 16,
            expand: float = 0.25,
            is_light_sr: bool = False,
            **kwargs,

    ):
        super().__init__()
        self.ln_1 = norm_layer(hidden_dim)
        self.act = nn.SiLU()
        self.self_attention = Spec_SS1D(in_img_size= in_img_size,in_chans=in_chans,ctx_dim=ctx_dim,d_model=hidden_dim, d_state=d_state, expand=expand, dropout=attn_drop_rate,**kwargs)
        self.drop_path = DropPath(drop_path)
        # self.SPE_prompt_injection_text = SPE_prompt_injection(ctx_dim,in_chans)
        # self.SPA_prompt_injection_text = SPA_prompt_injection(ctx_dim)
    def forward(self, input,max_spe,max_spa):
        # B,C,H,W=input.size()
        # input=input.reshape(B,C,H*W)

        x = self.ln_1(input)
        # x1 = x
        # SPE = self.SPE_prompt_injection_text(x1[:, 0:3, :], max_spe)
        # SPA = self.SPA_prompt_injection_text(x1[:, 3:6, :], max_spa)
        #
        # x1 = torch.cat((SPE, SPA), 1)
        # x1 = self.act(x1)
        x = self.self_attention(x,max_spe,max_spa)
        # x = x + input
        # x = x.reshape(B,C,H,W)+x1
        return x
