#!/usr/bin/env python3
"""
GeneLinguaLM v5 - 간단한 사용 예시

Usage:
    python inference_example.py
    python inference_example.py --sequence "MKTAYIAKQRQISFVKSH..."
"""
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
os.environ['TRANSFORMERS_VERBOSITY'] = 'error'

import sys
import torch
from pathlib import Path

# Project path
sys.path.insert(0, str(Path(__file__).parent))


class GeneLinguaLM:
    """GeneLinguaLM v5 모델 래퍼"""

    def __init__(self, checkpoint_path: str = None, device: str = "cuda"):
        self.device = device

        if checkpoint_path is None:
            checkpoint_path = "checkpoints/instruct_lora_v5/checkpoint_step15732.pt"

        print("🧬 GeneLinguaLM v5 로딩 중...")
        self._load_model(checkpoint_path)
        print("✅ 모델 로딩 완료!")

    def _load_model(self, checkpoint_path: str):
        from transformers import AutoModel, AutoTokenizer, AutoModelForCausalLM
        from peft import PeftModel
        from src.models.qformer import QFormer
        from src.training.train_instruct import SequenceProjector

        # Load checkpoint
        checkpoint = torch.load(checkpoint_path, map_location=self.device, weights_only=False)

        # ProtBERT (서열 인코더)
        print("  - ProtBERT 로딩...")
        self.seq_tokenizer = AutoTokenizer.from_pretrained("Rostlab/prot_bert_bfd", do_lower_case=False)
        self.seq_encoder = AutoModel.from_pretrained("Rostlab/prot_bert_bfd")
        self.seq_encoder.to(self.device).eval()

        # Q-Former
        print("  - Q-Former 로딩...")
        qformer_ckpt = torch.load('checkpoints/qformer/checkpoint_step7000.pt',
                                   map_location=self.device, weights_only=False)
        ckpt_config = qformer_ckpt.get("config", {})
        self.qformer = QFormer(
            num_queries=ckpt_config.get("num_queries", 32),
            embed_dim=ckpt_config.get("embed_dim", 768),
            num_layers=ckpt_config.get("qformer_layers", 6),
            num_heads=ckpt_config.get("num_heads", 12),
            dna_embed_dim=1024,
        )
        self.qformer.to(self.device)
        self.qformer.load_state_dict(checkpoint["qformer_state_dict"])
        self.qformer.eval()

        # Mistral LLM
        print("  - Mistral-7B 로딩...")
        self.llm_tokenizer = AutoTokenizer.from_pretrained(
            "mistralai/Mistral-7B-Instruct-v0.1", use_fast=False
        )
        if self.llm_tokenizer.pad_token is None:
            self.llm_tokenizer.pad_token = self.llm_tokenizer.eos_token

        import tempfile
        offload_dir = tempfile.mkdtemp()

        self.llm = AutoModelForCausalLM.from_pretrained(
            "mistralai/Mistral-7B-Instruct-v0.1",
            torch_dtype=torch.float16,
            device_map="auto",
            offload_folder=offload_dir,
        )
        lora_path = checkpoint.get("lora_path")
        self.llm = PeftModel.from_pretrained(self.llm, lora_path, offload_folder=offload_dir)
        self.llm.eval()

        # Projector
        print("  - Projector 로딩...")
        self.projector = SequenceProjector(qformer_dim=768, llm_dim=4096, num_queries=32)
        self.projector.load_state_dict(checkpoint["projector_state_dict"])
        self.projector.to(self.device).eval()

    def describe(self, sequence: str, max_length: int = 200) -> str:
        """
        단백질 서열을 자연어로 설명

        Args:
            sequence: 단백질 서열 (예: "MKTAYIAKQRQISFVKSH...")
            max_length: 최대 생성 토큰 수

        Returns:
            단백질 기능 설명 (자연어)
        """
        # 서열 전처리 (공백으로 분리)
        seq_spaced = " ".join(list(sequence.upper().replace(" ", "")))

        with torch.no_grad():
            # 1. 서열 인코딩
            seq_inputs = self.seq_tokenizer(
                seq_spaced,
                return_tensors="pt",
                truncation=True,
                max_length=512,
                padding=True,
            ).to(self.device)

            seq_outputs = self.seq_encoder(**seq_inputs)
            seq_embeds = seq_outputs.last_hidden_state

            # 2. Q-Former
            query_outputs = self.qformer(seq_embeds)

            # 3. Projector
            projected = self.projector(query_outputs)

            # 4. LLM 생성 (beam search + repetition penalty)
            prompt = "<s>[INST] Describe this protein sequence. [/INST]"
            prompt_ids = self.llm_tokenizer(prompt, return_tensors="pt").input_ids.to(self.llm.device)

            prompt_embeds = self.llm.get_input_embeddings()(prompt_ids)
            projected_fp16 = projected.to(device=prompt_embeds.device, dtype=prompt_embeds.dtype)
            combined_embeds = torch.cat([projected_fp16, prompt_embeds], dim=1)

            # 최적 설정: beam_search
            outputs = self.llm.generate(
                inputs_embeds=combined_embeds,
                max_new_tokens=max_length,
                num_beams=4,
                repetition_penalty=1.2,
                no_repeat_ngram_size=3,
                pad_token_id=self.llm_tokenizer.pad_token_id,
                eos_token_id=self.llm_tokenizer.eos_token_id,
            )

            generated = self.llm_tokenizer.decode(outputs[0], skip_special_tokens=True)

            # [/INST] 이후 텍스트만 추출
            if "[/INST]" in generated:
                generated = generated.split("[/INST]")[-1].strip()

            return generated


def main():
    import argparse

    parser = argparse.ArgumentParser(description="GeneLinguaLM v5 - 단백질 서열 설명 생성")
    parser.add_argument("--sequence", type=str, default=None, help="단백질 서열")
    parser.add_argument("--gpu", type=int, default=0, help="GPU 번호")
    args = parser.parse_args()

    # GPU 설정
    os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu)

    # 모델 로드
    model = GeneLinguaLM()

    # 예시 서열 (없으면 기본값 사용)
    if args.sequence:
        sequence = args.sequence
    else:
        # 예시: Human Insulin
        sequence = "MALWMRLLPLLALLALWGPDPAAAFVNQHLCGSHLVEALYLVCGERGFFYTPKTRREAEDLQVGQVELGGGPGAGSLQPLALEGSLQKRGIVEQCCTSICSLYQLENYCN"
        print(f"\n📌 예시 서열 (Human Insulin):")
        print(f"   {sequence[:50]}...")

    # 설명 생성
    print(f"\n🔬 분석 중...")
    description = model.describe(sequence)

    print(f"\n📝 결과:")
    print("-" * 60)
    print(description)
    print("-" * 60)


if __name__ == "__main__":
    main()