from transformers import AutoModel
model = AutoModel.from_pretrained(
    "muskwff/amp4multitask_124M",
    task_name="mic_regression",     # choose the task ["amp_classification", "hemolysis_regression", "mic_regression", "half_life_regression"]
    trust_remote_code=True      # required for loading the model
)
      

class AMPLayer(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.attention = AMPAttention(config)
        self.attn_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.intermediate = nn.Linear(config.hidden_size, config.intermediate_size)
        self.output = nn.Linear(config.intermediate_size, config.hidden_size)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.final_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.activation = nn.GELU()
        self.config = config
        
    def forward(self, hidden_states, attention_mask=None):
        residual = hidden_states
        attention_output = self.attention(hidden_states, attention_mask)
        attention_output = self.dropout(attention_output)
        attention_output = residual + attention_output
        attention_output = self.attn_layer_norm(attention_output)
        
        residual = attention_output
        intermediate_output = self.intermediate(attention_output)
        intermediate_output = self.activation(intermediate_output)
        layer_output = self.output(intermediate_output)
        layer_output = self.dropout(layer_output)
        layer_output = residual + layer_output
        layer_output = self.final_layer_norm(layer_output)
        return layer_output
      

class RotaryEmbedding(nn.Module):
    def __init__(self, dim, max_seq_len=512):
        super().__init__()
        self.dim = dim
        self.max_seq_len = max_seq_len
        freqs_cis = precompute_freqs_cis(dim, max_seq_len)
        self.register_buffer("freqs_cis", freqs_cis, persistent=False)
    
    def forward(self, seq_len=None):
        seq_len = seq_len if seq_len is not None else self.max_seq_len
        return self.freqs_cis[:seq_len]
      

class ActivityRegressionHead(nn.Module):
    def __init__(self, embedding_dim, num_organisms):
        super().__init__()        
        self.organism_encoder = nn.Sequential(
            nn.Embedding(num_organisms, embedding_dim),
            nn.Dropout(0.1),
            nn.Linear(embedding_dim, embedding_dim),
            nn.GELU(),
            nn.Dropout(0.1),
        )

        self.interactor = nn.Sequential(
            nn.Linear(embedding_dim * 2, embedding_dim * 2),
            nn.GELU(),
            nn.Dropout(0.2),
            nn.Linear(embedding_dim * 2, embedding_dim),
            nn.GELU(),
            nn.Dropout(0.2)
        )

        self.regressor = nn.Sequential(
            nn.Linear(embedding_dim, embedding_dim // 2),
            nn.GELU(),
            nn.Dropout(0.3),
            nn.Linear(embedding_dim // 2, embedding_dim // 4),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(embedding_dim // 4, 1)
        )

    def forward(self, features, organism_ids):
        org_features = self.organism_encoder(organism_ids)
        combined = torch.cat([features, org_features], dim=1)
        interaction = self.interactor(combined)
        interaction += features  # residual fusion with sequence features
        return self.regressor(interaction)

self.task_heads = nn.ModuleDict()
for task_name, task_config in task_config.items():
    num_orgs = task_config.get("num_organisms")
    self.task_heads[task_name] = TaskAwareHead(
        config,
        str(task_name),
        num_organisms=num_orgs
    )

from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("muskwff/amp4multitask_124M")

{
  "model_max_length": 128,
  "tokenizer_class": "EsmTokenizer",
  "special_tokens_map": {
    "cls_token": "<cls>",
    "pad_token": "<pad>",
    "eos_token": "<eos>",
    "unk_token": "<unk>",
    "mask_token": "<mask>"
  }
}

if 'bioactivity_classification' in label:
    multihot = [0] * len(bioactivity_projection)
    for l in label['bioactivity_classification']:
        if l in bioactivity_projection:
            multihot[bioactivity_projection[l]] = 1
    label['bioactivity_classification'] = torch.tensor(multihot, dtype=torch.float)

["Anti-Gram+"] 
--->
tensor([1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])

def calculate_parameters4half_life(labels):
    values = []
    for label in labels:
        if 'half_life_regression' in label and label['half_life_regression'] is not None:
            values.append(label['half_life_regression'])
    assert len(values) > 0, "No half life values found"
    converted = -np.log10(np.array(values))
    mean = np.mean(converted)
    std = np.std(converted)
    return mean, std

E_clipped = max(1, min(99, E))
EC50 = C * (100 - E_clipped) / E_clipped
Activity = -log10(EC50)
MeanActivity = average(Activity)
Activity_clipped = max(-5, min(-1, MeanActivity))
Score = 0.1 + 0.8 * (Activity_clipped - (-5)) / ((-1) - (-5))

def calculate_parameters4hemolysis(labels):
    values = []
    for label in labels:
        if 'hemolysis_regression' in label and label['hemolysis_regression'] is not None:
            values.append(label['hemolysis_regression'])
    assert len(values) > 0, "No hemolysis values found"
    converted = np.log1p(values)
    mean = np.mean(converted)
    std = np.std(converted)
    return mean, std

def calculate_mic_stats_by_organism(labels, organism_projection):
    idx_to_org = {idx: org for org, idx in organism_projection.items()}
    mic_values = {org_idx: [] for org_idx in idx_to_org.keys()}
    
    for label in labels:
        if 'mic_regression' in label:
            for org, value in label['mic_regression'].items():
                if org in organism_projection and value > 0:
                    org_idx = organism_projection[org]
                    log_val = -math.log10(value)
                    mic_values[org_idx].append(log_val)
    
    # calculate average value and mean standard deviation for each organism
    mic_stats = {}
    for org_idx, values in mic_values.items():
        if values:  # make sure there are values
            mean = np.mean(values)
            std = np.std(values)
            mic_stats[org_idx] = (mean, std)
    
    return mic_stats

if __name__ == '__main__':
    file_path = "merged_data.json"
    train_loader, val_loader, test_loader = build_dataloader(file_path, "bioactivity_classification", 64)

norm_params_path = "normalization_parameters.csv"       # change this to your path
norm_params = pd.read_csv(norm_params_path)
def denormalize(task_name, normalized_value, organism_id=None):
    """
    Denormalize the value for regression tasks
    """
    if isinstance(organism_id, torch.Tensor):
        organism_id = organism_id.cpu().item()

    if task_name == "mic_regression":
        org_name = norm_params.iloc[organism_id]['organism']
        params = norm_params[
            (norm_params['parameter_type'] == 'mic_regression') &
            (norm_params['organism'] == org_name)
        ]
        mean = params['mean'].values[0]
        std = params['std'].values[0]
        log_val = normalized_value * std + mean
        return 10 ** (-log_val)  # 10^(-log10(value)) = value
    elif task_name == "half_life_regression":
        params = norm_params[norm_params['parameter_type'] == 'half_life_regression']
        mean = params['mean'].values[0]
        std = params['std'].values[0]
        log_val = normalized_value * std + mean
        return 10 ** (-log_val)    
    elif task_name == "hemolysis_regression":
        params = norm_params[norm_params['parameter_type'] == 'hemolysis_regression']
        mean = params['mean'].values[0]
        std = params['std'].values[0]
        log_val = normalized_value * std + mean
        return np.expm1(log_val)  # e^(log_val) - 1

    return normalized_value  # the rest of tasks are not normalized