import sys
import os
from transformers import AutoModel, AutoTokenizer
from transformers.utils import cached_file

# Load model and tokenizer from Hugging Face Hub
print("Loading model and tokenizer...")
model = AutoModel.from_pretrained("hemantn/ablang2", trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("hemantn/ablang2", trust_remote_code=True)

# Find the cached model directory and import adapter
adapter_path = cached_file("hemantn/ablang2", "adapter.py")
cached_model_dir = os.path.dirname(adapter_path)
sys.path.insert(0, cached_model_dir)

# Import and create the adapter
from adapter import AbLang2PairedHuggingFaceAdapter
ablang = AbLang2PairedHuggingFaceAdapter(model=model, tokenizer=tokenizer)

def restore_sequences(heavy_chain, light_chain, use_align=False):
    try:
        # Prepare input sequences
        if heavy_chain.strip() and light_chain.strip():
            sequences = [[heavy_chain.strip(), light_chain.strip()]]
        elif heavy_chain.strip():
            sequences = [[heavy_chain.strip(), ""]]
        elif light_chain.strip():
            sequences = [["", light_chain.strip()]]
        else:
            return "Please provide at least one antibody chain sequence."
        
        # Perform restoration
        restored = ablang(sequences, mode='restore', align=use_align)
        
        # Format output
        if hasattr(restored, '__len__') and len(restored) > 0:
            result = restored[0]  # Get the first (and only) result
            return result
        else:
            return "Error: No restoration result obtained."
            
    except Exception as e:
        return f"Error during restoration: {str(e)}"

# Test the function
heavy_chain = "EVQ***SGGEVKKPGASVKVSCRASGYTFRNYGLTWVRQAPGQGLEWMGWISAYNGNTNYAQKFQGRVTLTTDTSTSTAYMELRSLRSDDTAVYFCAR**PGHGAAFMDVWGTGTTVTVSS"
light_chain = "DIQLTQSPLSLPVTLGQPASISCRSS*SLEASDTNIYLSWFQQRPGQSPRRLIYKI*NRDSGVPDRFSGSGSGTHFTLRISRVEADDVAVYYCMQGTHWPPAFGQGTKVDIK"

result = restore_sequences(heavy_chain, light_chain, False)
print("="*80)
print("APP OUTPUT TEST:")
print("="*80)
print(result)
print("="*80)
print(f"Result length: {len(result)}")
print(f"Result type: {type(result)}")