import sys
import os

sys.path.append(os.path.join(os.path.dirname(__file__)))

# import xgboost as xgb
import torch
import numpy as np
from transformers import AutoModelForMaskedLM
import warnings
import numpy as np
import esm
import torch.nn as nn
from rdkit import Chem, rdBase
import io
import contextlib

rdBase.DisableLog('rdApp.error')
warnings.filterwarnings("ignore", category=DeprecationWarning)
warnings.filterwarnings("ignore", category=UserWarning)
warnings.filterwarnings("ignore", category=FutureWarning)


class Conjugation:
    """
    Conjugation classifier for PDC generation
    
    Takes generated peptide SMILES, conjugates them using AoaDau builder,
    and validates using PDCAnalyzer.
    
    A sequence is valid (score=1) if:
    1. It is a valid peptide (parseable, has peptide bonds)
    2. It has a Lys ε-amine for conjugation (when exclude_n_terminus=True)
    3. At least one of its PDC conjugates passes the analyzer checks
    """
    
    # Builder types that conjugate via Cys thiol (disulfide/thioether)
    # rather than Lys ε-amine (amide bond)
    CYS_THIOL_BUILDERS = {'VcMMAE', 'DSPASN38', 'DSEBSN38', 'DSEPSN38', 'SN38'}

    def __init__(self, tokenizer, device='cpu', builder_type='AoaDau'):
        """
        Initialize conjugation classifier
        
        Args:
            tokenizer: SMILES tokenizer for decoding token IDs
            device: Device for PDCAnalyzer
            builder_type: Type of PDC builder. Options:
                - 'AoaDau': AOA-Daunorubicin
                - 'ValAlaDau': Val-Ala-PABC-Daunorubicin
                - 'ValCitDau': Val-Cit-PABC-Daunorubicin
                - 'GlutarylDau': Glutaryl-Daunorubicin
                - 'Dau': Drug-only Daunorubicin (linker sampled by
                  GNRHR+Dau PDCdb frequency: AoaDau 21/27, others 2/27)
                - 'SSCpt': Disulfide-Camptothecin (ester + carbonate)
                - 'VcMMAE': Val-Cit-MMAE
                - 'DSPACpt': Disulfanylpropanoic acid + Camptothecin
                - 'MalHexCpt': Maleimide-hexanoate + Camptothecin
                - 'DSECCpt': Disulfanylethyl carbonate + Camptothecin
                - 'Cpt': Drug-only Camptothecin (linker sampled by
                  ERBB2+Cpt PDCdb frequency: SSCpt 7/22, others 5/22)
                - 'DSPASN38': Disulfanylpropanoic acid + SN-38 (NRP1)
                - 'DSEBSN38': 4-(2-Disulfanylethoxy)-4-oxobutanoic acid + SN-38
                - 'DSEPSN38': 5-(2-Disulfanylethoxy)-5-oxopentanoic acid + SN-38
                - 'SN38': Drug-only SN-38 (linker sampled by
                  NRP1+SN38 PDCdb frequency: DSPA 9/17, DSEB 6/17, DSEP 2/17)
        """
        self.tokenizer = tokenizer
        self.device = device
        self.builder_type = builder_type
        
        # Import and initialize builder and analyzer
        from smiles.evaluation.pdc_analyzer import PDCAnalyzer
        self.pdc_analyzer = PDCAnalyzer(device, calculate_properties=False)
        
        # Initialize the appropriate builder
        self._init_builder(builder_type)
        
        # Tracking for detailed results
        self.last_detailed_results = None
    
    def _init_builder(self, builder_type):
        """Initialize the PDC builder based on type"""
        if builder_type == 'AoaDau':
            from smiles.conjugation.AoaDau import AoaDaunorubicin_PDC_Builder
            with contextlib.redirect_stdout(io.StringIO()):
                self.builder = AoaDaunorubicin_PDC_Builder()
        elif builder_type == 'SSCpt':
            from smiles.conjugation.SSCpt import DisulfideCamptothecin_PDC_Builder
            with contextlib.redirect_stdout(io.StringIO()):
                self.builder = DisulfideCamptothecin_PDC_Builder()
        elif builder_type == 'VcMMAE':
            from smiles.conjugation.VcMMAE import VcMMAE_PDC_Builder
            with contextlib.redirect_stdout(io.StringIO()):
                self.builder = VcMMAE_PDC_Builder()
        elif builder_type == 'DSPACpt':
            from smiles.conjugation.DSPACpt import DisulfanylpropanoicAcidCpt_PDC_Builder
            with contextlib.redirect_stdout(io.StringIO()):
                self.builder = DisulfanylpropanoicAcidCpt_PDC_Builder()
        elif builder_type == 'MalHexCpt':
            from smiles.conjugation.MalHexCpt import MaleimideHexanoateCpt_PDC_Builder
            with contextlib.redirect_stdout(io.StringIO()):
                self.builder = MaleimideHexanoateCpt_PDC_Builder()
        elif builder_type == 'DSECCpt':
            from smiles.conjugation.DSECCpt import DisulfanylethylCarbonateCpt_PDC_Builder
            with contextlib.redirect_stdout(io.StringIO()):
                self.builder = DisulfanylethylCarbonateCpt_PDC_Builder()
        elif builder_type == 'Cpt':
            from smiles.conjugation.CptFactory import CamptothecinPDC_Factory
            # Drug-only mode: linker sampled per peptide by PDCdb frequency
            self.builder = CamptothecinPDC_Factory(linker='sample')
        elif builder_type == 'DSPASN38':
            from smiles.conjugation.DSPASN38 import DisulfanylpropanoicAcidSN38_PDC_Builder
            with contextlib.redirect_stdout(io.StringIO()):
                self.builder = DisulfanylpropanoicAcidSN38_PDC_Builder()
        elif builder_type == 'DSEBSN38':
            from smiles.conjugation.DSEBSN38 import DisulfanylethOxyButanoicAcidSN38_PDC_Builder
            with contextlib.redirect_stdout(io.StringIO()):
                self.builder = DisulfanylethOxyButanoicAcidSN38_PDC_Builder()
        elif builder_type == 'DSEPSN38':
            from smiles.conjugation.DSEPSN38 import DisulfanylethOxyPentanoicAcidSN38_PDC_Builder
            with contextlib.redirect_stdout(io.StringIO()):
                self.builder = DisulfanylethOxyPentanoicAcidSN38_PDC_Builder()
        elif builder_type == 'SN38':
            from smiles.conjugation.SN38Factory import SN38PDC_Factory
            self.builder = SN38PDC_Factory(linker='sample')
        else:
            raise ValueError(f"Unknown builder type: {builder_type}")
    
    def is_valid_peptide(self, smiles):
        """Check if SMILES represents a valid peptide structure"""
        mol = Chem.MolFromSmiles(smiles)
        if mol is None:
            return False
        
        # Look for peptide bonds: NC(=O) pattern
        peptide_bond_pattern = Chem.MolFromSmarts('[NH][C](=O)')
        n_methyl_pattern = Chem.MolFromSmarts('[N;H0;$(NC)](C)[C](=O)')
        
        return mol.HasSubstructMatch(peptide_bond_pattern) or mol.HasSubstructMatch(n_methyl_pattern)
    
    def has_lys_epsilon_amine(self, smiles):
        """Check if peptide has Lys ε-amine for conjugation"""
        mol = Chem.MolFromSmiles(smiles)
        if mol is None:
            return False
        
        # Lys ε-amine: 4 CH2 chain ending in NH2
        lys_pattern = Chem.MolFromSmarts('[CH2][CH2][CH2][CH2][NH2]')
        return mol.HasSubstructMatch(lys_pattern)
    
    def has_cys_thiol(self, smiles):
        """Check if peptide has free Cys thiol (-SH) for disulfide conjugation"""
        mol = Chem.MolFromSmiles(smiles)
        if mol is None:
            return False
        
        # Iterate over atoms to find S with free H (thiol, not disulfide)
        for atom in mol.GetAtoms():
            if atom.GetAtomicNum() == 16:  # Sulfur
                if atom.GetTotalNumHs() >= 1 and atom.GetDegree() == 1:
                    return True
        return False
    
    def has_conjugation_site(self, smiles):
        """Check if peptide has the required conjugation site for this builder type"""
        if self.builder_type in self.CYS_THIOL_BUILDERS:
            return self.has_cys_thiol(smiles)
        return self.has_lys_epsilon_amine(smiles)
    
    def conjugate_peptide(self, peptide_smiles, exclude_n_terminus=True):
        """
        Attempt to conjugate a peptide SMILES
        
        Args:
            peptide_smiles: SMILES string of peptide
            exclude_n_terminus: If True, only use Lys ε-amines (default)
        
        Returns:
            Tuple of (success: bool, conjugates: list or None, error: str or None)
        """
        try:
            # Suppress builder output
            with contextlib.redirect_stdout(io.StringIO()):
                conjugates = self.builder.conjugate_to_peptide(
                    peptide_smiles, 
                    exclude_n_terminus=exclude_n_terminus
                )
            # Capture which linker was used (meaningful for Factory sampling)
            linker_used = getattr(self.builder, 'last_sampled_linker', self.builder_type)
            return True, conjugates, None, linker_used
        except Exception as e:
            return False, None, str(e), None
    
    def validate_conjugates(self, conjugates):
        """
        Validate a list of conjugates using PDCAnalyzer
        
        Args:
            conjugates: List of conjugate SMILES
        
        Returns:
            Tuple of (has_valid: bool, detailed_results: list)
        """
        if not conjugates:
            return False, []
        
        # PDCAnalyzer expects list of lists (candidates per peptide)
        detailed = self.pdc_analyzer.get_detailed_scores([conjugates])[0]
        
        # Check if any conjugate is valid
        has_valid = any(r.get('valid', 0) == 1 for r in detailed)
        
        return has_valid, detailed
    
    def get_scores(self, x):
        """
        Get validity scores for a batch of token sequences
        
        A sequence scores 1 if:
        1. Valid peptide structure
        2. Has Lys ε-amine (for exclude_n_terminus=True)
        3. At least one PDC conjugate passes analyzer
        
        Args:
            x: Tensor of token IDs (batch_size, seq_len)
        
        Returns:
            Tensor of 0/1 scores
        """
        smiles_list = self.tokenizer.batch_decode(x)
        scores = []
        detailed_results = []
        
        for smiles in smiles_list:
            result = {
                'peptide_smiles': smiles,
                'is_valid_peptide': False,
                'has_conjugation_site': False,
                'conjugation_success': False,
                'conjugates': [],
                'pdc_valid': False,
                'pdc_details': [],
                'error': None,
                'final_score': 0
            }
            
            # Step 1: Check if valid peptide
            if not self.is_valid_peptide(smiles):
                result['error'] = 'Not a valid peptide structure'
                scores.append(0)
                detailed_results.append(result)
                continue
            result['is_valid_peptide'] = True
            
            # Step 2: Check for conjugation site (Cys thiol or Lys ε-amine)
            if not self.has_conjugation_site(smiles):
                if self.builder_type in self.CYS_THIOL_BUILDERS:
                    result['error'] = 'No Cys thiol found for disulfide conjugation'
                else:
                    result['error'] = 'No Lys ε-amine found for conjugation'
                scores.append(0)
                detailed_results.append(result)
                continue
            result['has_conjugation_site'] = True  # "has conjugation site" on peptide level, not builder-specific at this point
            
            # Step 3: Attempt conjugation
            success, conjugates, error, linker_used = self.conjugate_peptide(smiles, exclude_n_terminus=True)
            if not success:
                result['error'] = f'Conjugation failed: {error}'
                scores.append(0)
                detailed_results.append(result)
                continue
            result['conjugation_success'] = True
            result['conjugates'] = conjugates
            result['linker_used'] = linker_used
            
            # Step 4: Validate conjugates with PDCAnalyzer
            has_valid, pdc_details = self.validate_conjugates(conjugates)
            result['pdc_details'] = pdc_details
            result['pdc_valid'] = has_valid
            
            if has_valid:
                result['final_score'] = 1
                scores.append(1)
            else:
                result['error'] = 'No valid PDC conjugates'
                scores.append(0)
            
            detailed_results.append(result)
        
        # Store for later retrieval
        self.last_detailed_results = detailed_results
        
        return torch.tensor(scores, dtype=torch.float)
    
    def get_last_detailed_results(self):
        """Get detailed results from the last get_scores call"""
        return self.last_detailed_results
    
    def get_detailed_scores_with_properties(self, x):
        """
        Get detailed scores with molecular properties for final sequences
        
        This enables property calculation and returns full details
        """
        # Enable property calculation
        self.pdc_analyzer.calculate_properties = True
        self.pdc_analyzer.enable_all_properties()
        
        smiles_list = self.tokenizer.batch_decode(x)
        detailed_results = []
        
        for smiles in smiles_list:
            result = {
                'peptide_smiles': smiles,
                'is_valid_peptide': self.is_valid_peptide(smiles),
                'has_conjugation_site': self.has_conjugation_site(smiles),
                'conjugates': [],
                'linker_used': None,
                'valid_pdcs': [],
                'best_pdc': None,
                'properties': {}
            }
            
            if not result['is_valid_peptide'] or not result['has_conjugation_site']:
                detailed_results.append(result)
                continue
            
            # Conjugate
            success, conjugates, _, linker_used = self.conjugate_peptide(smiles, exclude_n_terminus=True)
            if not success or not conjugates:
                detailed_results.append(result)
                continue
            
            result['conjugates'] = conjugates
            
            # Get detailed validation with properties
            pdc_details = self.pdc_analyzer.get_detailed_scores([conjugates])[0]
            
            # Find valid PDCs
            for detail in pdc_details:
                if detail.get('valid', 0) == 1:
                    result['valid_pdcs'].append(detail)
            
            # Get best PDC (first valid one)
            if result['valid_pdcs']:
                result['best_pdc'] = result['valid_pdcs'][0]
                result['properties'] = result['best_pdc'].get('properties', {})
            
            detailed_results.append(result)
        
        # Reset property calculation to default
        self.pdc_analyzer.calculate_properties = False
        
        return detailed_results
    
    def __call__(self, x):
        """Callable interface for use as objective function"""
        return self.get_scores(x)


class Analyzer:

    def __init__(self, tokenizer):
        self.tokenizer = tokenizer
    
    def get_scores(self, x):
        """Check if the SMILES represents a peptide structure"""
        results = []

        smiles_list = self.tokenizer.batch_decode(x)
        for smiles in smiles_list:
            mol = Chem.MolFromSmiles(smiles)
            if mol is None:
                results.append(0)
                continue
                
            # Look for peptide bonds: NC(=O) pattern
            peptide_bond_pattern = Chem.MolFromSmarts('[NH][C](=O)')

            # Look for N-methylated peptide bonds: N(C)C(=O) pattern
            n_methyl_pattern = Chem.MolFromSmarts('[N;H0;$(NC)](C)[C](=O)')

            if mol.HasSubstructMatch(peptide_bond_pattern) or mol.HasSubstructMatch(n_methyl_pattern):
                results.append(1)
            else:
                results.append(0)

        return torch.tensor(results)
    
    def __call__(self, x):
        scores = self.get_scores(x)
        return torch.tensor(scores)
    
    
# class Hemolysis:
    
#     def __init__(self, device):
#         self.predictor = xgb.Booster(model_file='./scoring/checkpoints/hemolysis-xgboost.json')
#         self.emb_model = AutoModelForMaskedLM.from_pretrained('aaronfeller/PeptideCLM-23M-all').roformer.to(device)
    
#     def get_scores(self, x):
#         scores = np.ones(len(x))
#         features = np.array(self.emb_model(input_ids=x).last_hidden_state.mean(dim=1).detach().cpu())
        
#         if len(features) == 0:
#             return scores
        
#         features = np.nan_to_num(features, nan=0.)
#         features = np.clip(features, np.finfo(np.float32).min, np.finfo(np.float32).max)
        
#         features = xgb.DMatrix(features)
        
#         probs = self.predictor.predict(features)
#         # return the probability of it being not hemolytic
#         return scores - probs
    
#     def __call__(self, x):
#         scores = self.get_scores(x)
#         return torch.tensor(scores)


# class Nonfouling:
    
#     def __init__(self, device):
#         self.predictor = xgb.Booster(model_file='./scoring/checkpoints/nonfouling-xgboost.json')
#         self.emb_model = AutoModelForMaskedLM.from_pretrained('aaronfeller/PeptideCLM-23M-all').roformer.to(device)    

#     def get_scores(self, x):
#         scores = np.zeros(len(x))
#         features = np.array(self.emb_model(input_ids=x).last_hidden_state.mean(dim=1).detach().cpu())
        
#         if len(features) == 0:
#             return scores
        
#         features = np.nan_to_num(features, nan=0.)
#         features = np.clip(features, np.finfo(np.float32).min, np.finfo(np.float32).max)
        
#         features = xgb.DMatrix(features)
        
#         scores = self.predictor.predict(features)
#         # return the probability of it being not hemolytic
#         return scores
    
#     def __call__(self, x):
#         scores = self.get_scores(x)
#         return torch.tensor(scores)


# class Solubility:
#     def __init__(self, device):
#         self.predictor = xgb.Booster(model_file='./scoring/checkpoints/solubility-xgboost.json')
#         self.emb_model = AutoModelForMaskedLM.from_pretrained('aaronfeller/PeptideCLM-23M-all').roformer.to(device)

#     def get_scores(self, x):
#         scores = np.zeros(len(x))
#         features = np.array(self.emb_model(input_ids=x).last_hidden_state.mean(dim=1).detach().cpu())
        
#         if len(features) == 0:
#             return scores
        
#         features = np.nan_to_num(features, nan=0.)
#         features = np.clip(features, np.finfo(np.float32).min, np.finfo(np.float32).max)
        
#         features = xgb.DMatrix(features)
        
#         scores = self.predictor.predict(features)
#         return scores
    
#     def __call__(self, x):
#         scores = self.get_scores(x)
#         return torch.tensor(scores)
    

class ImprovedBindingPredictor(nn.Module):
    def __init__(self, 
                 esm_dim=1280,
                 smiles_dim=768,
                 hidden_dim=512,
                 n_heads=8,
                 n_layers=3,
                 dropout=0.1):
        super().__init__()
        
        # Define binding thresholds
        self.tight_threshold = 7.5    # Kd/Ki/IC50 ≤ ~30nM
        self.weak_threshold = 6.0     # Kd/Ki/IC50 > 1μM
        
        # Project to same dimension
        self.smiles_projection = nn.Linear(smiles_dim, hidden_dim)
        self.protein_projection = nn.Linear(esm_dim, hidden_dim)
        self.protein_norm = nn.LayerNorm(hidden_dim)
        self.smiles_norm = nn.LayerNorm(hidden_dim)
        
        # Cross attention blocks with layer norm
        self.cross_attention_layers = nn.ModuleList([
            nn.ModuleDict({
                'attention': nn.MultiheadAttention(hidden_dim, n_heads, dropout=dropout),
                'norm1': nn.LayerNorm(hidden_dim),
                'ffn': nn.Sequential(
                    nn.Linear(hidden_dim, hidden_dim * 4),
                    nn.ReLU(),
                    nn.Dropout(dropout),
                    nn.Linear(hidden_dim * 4, hidden_dim)
                ),
                'norm2': nn.LayerNorm(hidden_dim)
            }) for _ in range(n_layers)
        ])
        
        # Prediction heads
        self.shared_head = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
        )
        
        # Regression head
        self.regression_head = nn.Linear(hidden_dim, 1)
        
        # Classification head (3 classes: tight, medium, loose binding)
        self.classification_head = nn.Linear(hidden_dim, 3)
        
        
    def get_binding_class(self, affinity):
        """Convert affinity values to class indices
        0: tight binding (>= 7.5)
        1: medium binding (6.0-7.5)
        2: weak binding (< 6.0)
        """
        if isinstance(affinity, torch.Tensor):
            tight_mask = affinity >= self.tight_threshold
            weak_mask = affinity < self.weak_threshold
            medium_mask = ~(tight_mask | weak_mask)
            
            classes = torch.zeros_like(affinity, dtype=torch.long)
            classes[medium_mask] = 1
            classes[weak_mask] = 2
            return classes
        else:
            if affinity >= self.tight_threshold:
                return 0  # tight binding
            elif affinity < self.weak_threshold:
                return 2  # weak binding
            else:
                return 1  # medium binding
        
    def forward(self, protein_emb, smiles_emb):
        protein = self.protein_norm(self.protein_projection(protein_emb))
        smiles = self.smiles_norm(self.smiles_projection(smiles_emb))
        
        #protein = protein.transpose(0, 1)
        #smiles = smiles.transpose(0, 1)
        
        # Cross attention layers
        for layer in self.cross_attention_layers:
            # Protein attending to SMILES
            attended_protein = layer['attention'](
                protein, smiles, smiles
            )[0]
            protein = layer['norm1'](protein + attended_protein)
            protein = layer['norm2'](protein + layer['ffn'](protein))
            
            # SMILES attending to protein
            attended_smiles = layer['attention'](
                smiles, protein, protein
            )[0]
            smiles = layer['norm1'](smiles + attended_smiles)
            smiles = layer['norm2'](smiles + layer['ffn'](smiles))
        
        # Get sequence-level representations
        protein_pool = torch.mean(protein, dim=0)
        smiles_pool = torch.mean(smiles, dim=0)
        
        # Concatenate both representations
        combined = torch.cat([protein_pool, smiles_pool], dim=-1)
        
        # Shared features
        shared_features = self.shared_head(combined)
        
        regression_output = self.regression_head(shared_features)
        classification_logits = self.classification_head(shared_features)
        
        return regression_output, classification_logits
    
class BindingAffinity:
    def __init__(self, prot_seq, device, model_type='PeptideCLM'):
        super().__init__()

        self.pep_model = AutoModelForMaskedLM.from_pretrained('aaronfeller/PeptideCLM-23M-all').roformer.to(device)
        
        self.model = ImprovedBindingPredictor(smiles_dim=768).to(device)
        checkpoint = torch.load('./smiles/scoring/checkpoints/binding-affinity.pt', weights_only=False)
        self.model.load_state_dict(checkpoint['model_state_dict'])
        
        self.model.eval()
        
        self.esm_model, alphabet = esm.pretrained.esm2_t33_650M_UR50D()  # load ESM-2 model
        self.esm_model.to(device)
        self.prot_tokenizer = alphabet.get_batch_converter() # load esm tokenizer

        data = [("target", prot_seq)]  
        # get tokenized protein
        _, _, prot_tokens = self.prot_tokenizer(data)
        prot_tokens = prot_tokens.to(device)
        with torch.no_grad():
            results = self.esm_model.forward(prot_tokens, repr_layers=[33])  # Example with ESM-2
            prot_emb = results["representations"][33]
            
        self.prot_emb = prot_emb[0]
        self.prot_emb = torch.mean(self.prot_emb, dim=0, keepdim=True).to(device)
                
    def forward(self, x, attention_mask=None):
        with torch.no_grad():
            hs = self.pep_model(
                input_ids=x,
                attention_mask=attention_mask,
                output_hidden_states=True,
            ).last_hidden_state
            if attention_mask is not None:
                m = attention_mask.unsqueeze(-1).to(hs.dtype)
                pep_emb = (hs * m).sum(dim=1, keepdim=True) / m.sum(dim=1, keepdim=True).clamp(min=1)
            else:
                pep_emb = hs.mean(dim=1, keepdim=True)
            scores = []
            for pep in pep_emb:
                score, logits = self.model.forward(self.prot_emb, pep)
                scores.append(score.item() / 10)

        return torch.tensor(scores)

    def __call__(self, x, attention_mask=None):
        return self.forward(x, attention_mask=attention_mask)