"""
Mutation operators for biological language models.
Adds controlled noise and structural perturbations to model weights.
"""

import os
import torch
import numpy as np
import logging
from typing import Optional
from transformers import AutoModelForMaskedLM
from safetensors.torch import save_file, load_file

logger = logging.getLogger(__name__)


class BioMutator:
    """Base mutation operator."""
    
    def mutate(self, model_path: str, save_path: str) -> str:
        """Apply mutation to a model and save result."""
        raise NotImplementedError


class GaussianNoiseMutator(BioMutator):
    """Add Gaussian noise to model weights."""
    
    def __init__(self, std: float = 0.01, layer_filter: Optional[str] = None):
        self.std = std
        self.layer_filter = layer_filter  # If set, only mutate matching layers
    
    def mutate(self, model_path: str, save_path: str) -> str:
        logger.info(f"Applying Gaussian noise mutation (std={self.std}) to {model_path}")
        
        try:
            state = load_file(os.path.join(model_path, "model.safetensors"))
        except:
            state = torch.load(
                os.path.join(model_path, "pytorch_model.bin"),
                map_location="cpu",
            )
        
        mutated = {}
        for key, tensor in state.items():
            if self.layer_filter is not None and self.layer_filter not in key:
                mutated[key] = tensor.clone()
                continue
            
            # Generate noise matching tensor shape
            noise = torch.randn_like(tensor) * self.std
            mutated[key] = tensor + noise
        
        # Save
        os.makedirs(save_path, exist_ok=True)
        save_file(mutated, os.path.join(save_path, "model.safetensors"))
        
        # Copy auxiliary files
        self._copy_auxiliary_files(model_path, save_path)
        
        logger.info(f"Mutated model saved to {save_path}")
        return save_path
    
    def _copy_auxiliary_files(self, source_path: str, dest_path: str) -> None:
        import shutil
        
        for fname in ["config.json", "tokenizer.json", "tokenizer_config.json", "vocab.txt"]:
            src = os.path.join(source_path, fname)
            if os.path.exists(src):
                shutil.copy2(src, os.path.join(dest_path, fname))


class LayerScaleMutator(BioMutator):
    """Scale specific layers by a random factor."""
    
    def __init__(self, scale_range: tuple = (0.9, 1.1), target_layers: Optional[list] = None):
        self.scale_range = scale_range
        self.target_layers = target_layers  # e.g., ["encoder.layer.0", "encoder.layer.1"]
    
    def mutate(self, model_path: str, save_path: str) -> str:
        logger.info(f"Applying layer scale mutation to {model_path}")
        
        try:
            state = load_file(os.path.join(model_path, "model.safetensors"))
        except:
            state = torch.load(
                os.path.join(model_path, "pytorch_model.bin"),
                map_location="cpu",
            )
        
        mutated = {}
        for key, tensor in state.items():
            # Check if this key matches a target layer
            should_scale = False
            if self.target_layers:
                for layer_name in self.target_layers:
                    if layer_name in key:
                        should_scale = True
                        break
            else:
                # Randomly select ~10% of layers to scale
                should_scale = np.random.random() < 0.1
            
            if should_scale:
                scale = np.random.uniform(*self.scale_range)
                mutated[key] = tensor * scale
                logger.debug(f"Scaled {key} by {scale:.4f}")
            else:
                mutated[key] = tensor.clone()
        
        os.makedirs(save_path, exist_ok=True)
        save_file(mutated, os.path.join(save_path, "model.safetensors"))
        
        self._copy_auxiliary_files(model_path, save_path)
        
        logger.info(f"Layer-scaled model saved to {save_path}")
        return save_path
    
    def _copy_auxiliary_files(self, source_path: str, dest_path: str) -> None:
        import shutil
        
        for fname in ["config.json", "tokenizer.json", "tokenizer_config.json", "vocab.txt"]:
            src = os.path.join(source_path, fname)
            if os.path.exists(src):
                shutil.copy2(src, os.path.join(dest_path, fname))


class DropoutMutator(BioMutator):
    """Randomly zero out a fraction of weights (structural pruning)."""
    
    def __init__(self, dropout_rate: float = 0.01):
        self.dropout_rate = dropout_rate
    
    def mutate(self, model_path: str, save_path: str) -> str:
        logger.info(f"Applying dropout mutation (rate={self.dropout_rate}) to {model_path}")
        
        try:
            state = load_file(os.path.join(model_path, "model.safetensors"))
        except:
            state = torch.load(
                os.path.join(model_path, "pytorch_model.bin"),
                map_location="cpu",
            )
        
        mutated = {}
        for key, tensor in state.items():
            mask = torch.rand_like(tensor) > self.dropout_rate
            mutated[key] = tensor * mask
        
        os.makedirs(save_path, exist_ok=True)
        save_file(mutated, os.path.join(save_path, "model.safetensors"))
        
        self._copy_auxiliary_files(model_path, save_path)
        
        logger.info(f"Dropout-mutated model saved to {save_path}")
        return save_path
    
    def _copy_auxiliary_files(self, source_path: str, dest_path: str) -> None:
        import shutil
        
        for fname in ["config.json", "tokenizer.json", "tokenizer_config.json", "vocab.txt"]:
            src = os.path.join(source_path, fname)
            if os.path.exists(src):
                shutil.copy2(src, os.path.join(dest_path, fname))