import torch
import torch.nn as nn
import torch.nn.functional as F
import math

PHYSICOCHEMICAL_FEATURES = {
    'A': [1.8, 0.0, 0.0, 89.0, 0.360, 0.76, 0.83, 0.37, 1.0, 0.0, 0.0, 0.0],
    'R': [-4.5, 1.0, 1.0, 174.0, 0.293, 0.76, 0.93, 0.58, 0.0, 0.0, 1.0, 0.0],
    'N': [-3.5, 0.0, 0.0, 132.0, 0.337, 0.79, 0.89, 0.54, 0.0, 0.0, 1.0, 0.0],
    'D': [-3.5, -1.0, 0.0, 133.0, 0.281, 0.74, 0.72, 0.54, 0.0, 0.0, 1.0, 0.0],
    'C': [2.5, 0.0, 0.0, 121.0, 0.160, 0.81, 0.87, 0.39, 1.0, 0.0, 0.0, 0.0],
    'Q': [-3.5, 0.0, 0.0, 146.0, 0.316, 0.80, 0.91, 0.57, 0.0, 0.0, 1.0, 0.0],
    'E': [-3.5, -1.0, 0.0, 147.0, 0.282, 0.77, 0.73, 0.60, 0.0, 0.0, 1.0, 0.0],
    'G': [-0.4, 0.0, 0.0, 75.0, 0.360, 0.69, 0.75, 0.39, 1.0, 0.0, 0.0, 0.0],
    'H': [-3.2, 0.1, 1.0, 155.0, 0.290, 0.80, 0.80, 0.49, 0.0, 0.0, 1.0, 0.0],
    'I': [4.5, 0.0, 0.0, 131.0, 0.321, 0.79, 0.83, 0.24, 1.0, 0.0, 0.0, 0.0],
    'L': [3.8, 0.0, 0.0, 131.0, 0.313, 0.78, 0.85, 0.22, 1.0, 0.0, 0.0, 0.0],
    'K': [-3.9, 1.0, 0.0, 146.0, 0.329, 0.73, 0.85, 0.58, 0.0, 0.0, 1.0, 0.0],
    'M': [1.9, 0.0, 0.0, 149.0, 0.302, 0.80, 0.85, 0.34, 1.0, 0.0, 0.0, 0.0],
    'F': [2.8, 0.0, 0.0, 165.0, 0.287, 0.77, 0.82, 0.22, 1.0, 0.0, 0.0, 0.0],
    'P': [-1.6, 0.0, 0.0, 115.0, 0.314, 0.64, 0.63, 0.42, 1.0, 0.0, 0.0, 0.0],
    'S': [-0.8, 0.0, 0.0, 105.0, 0.384, 0.74, 0.74, 0.47, 0.0, 0.0, 1.0, 0.0],
    'T': [-0.7, 0.0, 0.0, 119.0, 0.360, 0.76, 0.77, 0.45, 0.0, 0.0, 1.0, 0.0],
    'W': [-0.9, 0.0, 0.0, 204.0, 0.277, 0.75, 0.82, 0.27, 1.0, 0.0, 0.0, 0.0],
    'Y': [-1.3, 0.0, 0.0, 181.0, 0.292, 0.77, 0.79, 0.34, 0.0, 0.0, 1.0, 0.0],
    'V': [4.2, 0.0, 0.0, 117.0, 0.329, 0.76, 0.85, 0.24, 1.0, 0.0, 0.0, 0.0],
}
AA_ORDER = 'ARNDCQEGHILKMFPSTWYV'

def get_physicochem_matrix():
    mat = []
    for aa in AA_ORDER:
        mat.append(PHYSICOCHEMICAL_FEATURES[aa])
    return torch.tensor(mat, dtype=torch.float32)

class PhysicochemicalEmbedding(nn.Module):
    def __init__(self, in_dim=12, out_dim=32):
        super().__init__()
        self.proj = nn.Linear(in_dim, out_dim)
        self.norm = nn.LayerNorm(out_dim)

    def forward(self, x):
        return self.norm(F.gelu(self.proj(x)))

class ConvBranch(nn.Module):
    def __init__(self, in_dim, out_dim, kernel_size, dilation=1, dropout=0.1):
        super().__init__()
        pad = dilation * (kernel_size - 1) // 2
        self.conv = nn.Conv1d(in_dim, out_dim, kernel_size, padding=pad, dilation=dilation, bias=False)
        self.bn = nn.BatchNorm1d(out_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        x = x.transpose(1, 2)
        x = self.conv(x)
        x = self.bn(x)
        x = F.gelu(x)
        x = self.dropout(x)
        return x.transpose(1, 2)

class SqueezeExcitation(nn.Module):
    def __init__(self, channels, reduction=16):
        super().__init__()
        self.fc1 = nn.Linear(channels, channels // reduction, bias=False)
        self.fc2 = nn.Linear(channels // reduction, channels, bias=False)

    def forward(self, x):
        b, l, c = x.shape
        y = x.mean(dim=1)
        y = F.gelu(self.fc1(y))
        y = torch.sigmoid(self.fc2(y)).unsqueeze(1)
        return x * y

class LightweightAttention(nn.Module):
    def __init__(self, d_model, n_heads, dropout=0.1, max_len=512):
        super().__init__()
        assert d_model % n_heads == 0
        self.d_model = d_model
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads

        self.qkv = nn.Linear(d_model, d_model * 3, bias=False)
        self.out = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)

        self.pos_encoding = nn.Parameter(torch.randn(1, max_len, d_model) * 0.02)

    def forward(self, x, mask=None):
        b, l, _ = x.shape
        x = x + self.pos_encoding[:, :l, :]

        qkv = self.qkv(x).reshape(b, l, 3, self.n_heads, self.head_dim).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]

        scale = self.head_dim ** -0.5
        attn = torch.matmul(q, k.transpose(-2, -1)) * scale

        if mask is not None:
            attn = attn.masked_fill(mask.unsqueeze(1).unsqueeze(1) == 0, float('-inf'))

        attn = F.softmax(attn, dim=-1)
        attn = self.dropout(attn)

        out = torch.matmul(attn, v).transpose(1, 2).contiguous().reshape(b, l, self.d_model)
        return self.out(out)

class TransformerBlock(nn.Module):
    def __init__(self, d_model, n_heads, ff_dim, dropout=0.1):
        super().__init__()
        self.norm1 = nn.LayerNorm(d_model)
        self.attn = LightweightAttention(d_model, n_heads, dropout)
        self.norm2 = nn.LayerNorm(d_model)
        self.ff = nn.Sequential(
            nn.Linear(d_model, ff_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(ff_dim, d_model),
            nn.Dropout(dropout),
        )

    def forward(self, x):
        x = x + self.attn(self.norm1(x))
        x = x + self.ff(self.norm2(x))
        return x

class AttentionPooling(nn.Module):
    def __init__(self, d_model):
        super().__init__()
        self.query = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)
        self.attn = nn.Linear(d_model, 1)

    def forward(self, x):
        scores = self.attn(torch.tanh(x)).transpose(1, 2)
        weights = F.softmax(scores, dim=-1)
        return torch.matmul(weights, x).squeeze(1)


class PeptEdge(nn.Module):
    def __init__(self, vocab_size=21, max_len=200, d_model=128, n_heads=4,
                 num_layers=3, ff_dim=256, num_classes=2, dropout=0.15):
        super().__init__()

        self.max_len = max_len
        self.vocab_size = vocab_size

        phys_mat = get_physicochem_matrix()
        self.register_buffer('phys_mat', phys_mat)

        self.token_embed = nn.Embedding(vocab_size, d_model - 32, padding_idx=0)
        self.phys_embed = PhysicochemicalEmbedding(12, 32)

        conv_dims = [48, 48, 32, 32, 32]
        conv_ks = [3, 5, 7, 15, 31]
        conv_dils = [1, 1, 2, 1, 1]

        self.conv_branches = nn.ModuleList([
            ConvBranch(d_model, conv_dims[i], conv_ks[i], conv_dils[i], dropout)
            for i in range(len(conv_dims))
        ])

        total_conv_out = sum(conv_dims)
        self.conv_proj = nn.Sequential(
            nn.Linear(total_conv_out, d_model),
            nn.LayerNorm(d_model),
        )
        self.se = SqueezeExcitation(d_model)

        self.transformer_blocks = nn.ModuleList([
            TransformerBlock(d_model, n_heads, ff_dim, dropout)
            for _ in range(num_layers)
        ])

        self.attn_pool = AttentionPooling(d_model)
        self.classifier = nn.Sequential(
            nn.LayerNorm(d_model * 3),
            nn.Linear(d_model * 3, d_model),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(d_model, d_model // 2),
            nn.GELU(),
            nn.Dropout(dropout * 0.5),
            nn.Linear(d_model // 2, num_classes),
        )

        self._init_weights()

    def _init_weights(self):
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.kaiming_normal_(p, mode='fan_out', nonlinearity='relu')

    def forward(self, x, return_embeddings=False):
        b, l = x.shape
        mask = (x != 0).float()

        aa_indices = torch.clamp(x, 0, 19).long()
        phys_feats = F.embedding(aa_indices, self.phys_mat)
        phys_emb = self.phys_embed(phys_feats)

        tok_emb = self.token_embed(x)

        h = torch.cat([tok_emb, phys_emb], dim=-1)

        conv_out = []
        for branch in self.conv_branches:
            conv_out.append(branch(h))

        max_len = max(c.shape[1] for c in conv_out)
        conv_out_padded = []
        for c in conv_out:
            if c.shape[1] < max_len:
                pad = max_len - c.shape[1]
                c = F.pad(c, (0, 0, 0, pad))
            conv_out_padded.append(c)

        h_conv = torch.cat(conv_out_padded, dim=-1)
        h_conv = self.conv_proj(h_conv)
        h_conv = self.se(h_conv)

        for block in self.transformer_blocks:
            h_conv = block(h_conv)

        mean_pool = h_conv.mean(dim=1)
        max_pool = h_conv.max(dim=1)[0]
        attn_pool = self.attn_pool(h_conv)

        h_pool = torch.cat([mean_pool, max_pool, attn_pool], dim=-1)

        if return_embeddings:
            return h_pool

        logits = self.classifier(h_pool)
        return logits


class PeptEdgeMultilabel(nn.Module):
    def __init__(self, vocab_size=27, max_len=200, d_model=128, n_heads=4,
                 num_layers=3, ff_dim=256, num_classes=5, dropout=0.15):
        super().__init__()
        self.backbone = PeptEdge(
            vocab_size=vocab_size, max_len=max_len,
            d_model=d_model, n_heads=n_heads,
            num_layers=num_layers, ff_dim=ff_dim,
            num_classes=num_classes, dropout=dropout
        )
        in_features = d_model * 3
        self.backbone.classifier = nn.Sequential(
            nn.LayerNorm(in_features),
            nn.Linear(in_features, d_model),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(d_model, d_model // 2),
            nn.GELU(),
            nn.Dropout(dropout * 0.5),
            nn.Linear(d_model // 2, num_classes),
        )

    def forward(self, x, return_embeddings=False):
        return self.backbone(x, return_embeddings)


def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


if __name__ == '__main__':
    model = PeptEdge(vocab_size=21, max_len=200, num_classes=2)
    total = count_parameters(model)
    print(f'PeptEdge params: {total:,}')
    x = torch.randint(0, 20, (4, 100))
    out = model(x)
    print(f'Input: {x.shape}, Output: {out.shape}')