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

try:
    from safetensors.torch import save_file, load_file
    _HAS_SAFETENSORS = True
except ImportError:
    _HAS_SAFETENSORS = False

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 SwiGLU(nn.Module):
    def __init__(self, dim, hidden_dim):
        super().__init__()
        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
        self.w2 = nn.Linear(dim, hidden_dim, bias=False)
        self.w3 = nn.Linear(hidden_dim, dim, bias=False)

    def forward(self, x):
        return self.w3(F.silu(self.w1(x)) * self.w2(x))

class ConvBlock(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.act = nn.GELU()
        self.drop = nn.Dropout(dropout)

    def forward(self, x):
        r = x
        x = x.transpose(1, 2)
        x = self.conv(x)
        x = self.bn(x)
        x = self.act(x)
        x = self.drop(x)
        x = x.transpose(1, 2)
        if r.shape[-1] != x.shape[-1]:
            return x
        return x + r


class MultiHeadSelfAttention(nn.Module):
    def __init__(self, d_model, n_heads, dropout=0.1):
        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.scale = self.head_dim ** -0.5

        self.qkv = nn.Linear(d_model, d_model * 3, bias=False)
        self.out = nn.Linear(d_model, d_model)
        self.drop = nn.Dropout(dropout)
        self.pos_enc = nn.Parameter(torch.randn(1, 512, d_model) * 0.02)

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

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

        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = F.softmax(attn, dim=-1)
        attn = self.drop(attn)

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


class TransformerLayer(nn.Module):
    def __init__(self, d_model, n_heads, ff_dim, dropout=0.1, sd_prob=0.0):
        super().__init__()
        self.sd_prob = sd_prob
        self.norm1 = nn.LayerNorm(d_model)
        self.attn = MultiHeadSelfAttention(d_model, n_heads, dropout)
        self.norm2 = nn.LayerNorm(d_model)
        self.ff = SwiGLU(d_model, ff_dim)
        self.drop = nn.Dropout(dropout)

    def forward(self, x):
        if self.training and self.sd_prob > 0:
            if torch.rand(1).item() < self.sd_prob:
                return x
        x = x + self.drop(self.attn(self.norm1(x)))
        x = x + self.drop(self.ff(self.norm2(x)))
        return x


class PeptEdgeV2(nn.Module):
    def __init__(self, vocab_size=21, max_len=200, d_model=192, n_heads=6,
                 num_layers=4, ff_dim=384, num_classes=2, dropout=0.15, sd_prob=0.05):
        super().__init__()

        self.config = {
            '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,
            'sd_prob': sd_prob,
            'architectures': ['PeptEdgeV2'],
            'model_type': 'peptedgev2',
        }

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

        self.token_embed = nn.Embedding(vocab_size, d_model - 64, padding_idx=0)
        self.phys_proj = nn.Sequential(
            nn.Linear(12, 64),
            nn.LayerNorm(64),
        )

        self.input_drop = nn.Dropout(dropout)

        self.conv_layers = nn.ModuleList([
            ConvBlock(d_model, d_model, 3, 1, dropout),
            ConvBlock(d_model, d_model, 5, 1, dropout),
            ConvBlock(d_model, d_model, 7, 2, dropout),
            ConvBlock(d_model, d_model, 11, 1, dropout),
        ])

        self.conv_norm = nn.LayerNorm(d_model)

        self.transformer_layers = nn.ModuleList([
            TransformerLayer(d_model, n_heads, ff_dim, dropout, sd_prob if i > 0 else 0.0)
            for i in range(num_layers)
        ])

        self.norm = nn.LayerNorm(d_model)
        self.attn_pool_q = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)

        self.classifier = nn.Sequential(
            nn.Linear(d_model * 3, d_model),
            nn.LayerNorm(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 to_dict(self):
        return dict(self.config)

    @classmethod
    def from_dict(cls, config):
        return cls(
            vocab_size=config['vocab_size'],
            max_len=config['max_len'],
            d_model=config['d_model'],
            n_heads=config['n_heads'],
            num_layers=config['num_layers'],
            ff_dim=config['ff_dim'],
            num_classes=config['num_classes'],
            dropout=config.get('dropout', 0.15),
            sd_prob=config.get('sd_prob', 0.05),
        )

    def save_pretrained(self, save_directory, safe_serialization=True):
        os.makedirs(save_directory, exist_ok=True)
        config = self.to_dict()
        with open(os.path.join(save_directory, 'config.json'), 'w') as f:
            json.dump(config, f, indent=2)
        if safe_serialization and _HAS_SAFETENSORS:
            save_file(self.state_dict(), os.path.join(save_directory, 'model.safetensors'))
        else:
            torch.save({'model_state_dict': self.state_dict()},
                       os.path.join(save_directory, 'pytorch_model.bin'))

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, map_location=None, **kwargs):
        pt_dir = pretrained_model_name_or_path
        if not os.path.isdir(pretrained_model_name_or_path):
            import huggingface_hub
            from huggingface_hub import snapshot_download
            pt_dir = snapshot_download(
                repo_id=pretrained_model_name_or_path,
                allow_patterns=['*.json', '*.safetensors', '*.bin'],
                **kwargs,
            )
        with open(os.path.join(pt_dir, 'config.json')) as f:
            config = json.load(f)
        model = cls.from_dict(config)
        st_dir = os.path.join(pt_dir, 'model.safetensors')
        if os.path.exists(st_dir):
            if not _HAS_SAFETENSORS:
                raise ImportError('safetensors is required to load model.safetensors')
            state = load_file(st_dir)
        else:
            ckpt = torch.load(os.path.join(pt_dir, 'pytorch_model.bin'), map_location=map_location)
            state = ckpt.get('model_state_dict', ckpt)
        model.load_state_dict(state)
        return model

    def forward(self, x, return_embeddings=False):
        b, l = x.shape
        aa_idx = torch.clamp(x, 0, 19).long()
        phys = F.embedding(aa_idx, self.phys_mat)
        phys = self.phys_proj(phys)

        tok = self.token_embed(x)
        h = torch.cat([tok, phys], dim=-1)
        h = self.input_drop(h)

        for conv in self.conv_layers:
            h = conv(h)

        h = self.conv_norm(h)

        for tf in self.transformer_layers:
            h = tf(h)

        h = self.norm(h)
        mean_pool = h.mean(dim=1)
        max_pool = h.max(dim=1)[0]
        attn_scores = torch.matmul(h, self.attn_pool_q.transpose(1, 2))
        attn_weights = F.softmax(attn_scores.squeeze(-1), dim=1).unsqueeze(1)
        attn_pool = torch.matmul(attn_weights, h).squeeze(1)

        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


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


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