#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Emotion Summary Model - 推理示例
"""

from transformers import MT5ForConditionalGeneration, MT5Tokenizer
import json
import torch

def load_model(model_path="./emotion_summary"):
    """加载模型和tokenizer"""
    print(f"Loading model from {model_path}...")
    model = MT5ForConditionalGeneration.from_pretrained(model_path)
    tokenizer = MT5Tokenizer.from_pretrained(model_path)
    
    # 如果有GPU就使用GPU
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = model.to(device)
    model.eval()
    
    print(f"Model loaded on {device}")
    return model, tokenizer, device

def summarize_case(model, tokenizer, device, case_data, field="cause"):
    """
    对案例进行总结
    
    Args:
        model: 模型
        tokenizer: tokenizer
        device: 设备
        case_data: 案例数据（字典）
        field: 要生成的字段 (cause/symptoms/treatment_process/illness_characteristics/treatment_effect)
    
    Returns:
        生成的总结文本
    """
    
    # 构建输入
    case_desc = " ".join(case_data.get("case_description", []))
    consultation = " ".join(case_data.get("consultation_process", []))
    reflection = case_data.get("experience_and_reflection", "")
    
    full_text = f"Case: {case_desc}\nConsultation: {consultation}\nReflection: {reflection}"
    
    # 根据字段构建不同的prompt
    prompts = {
        "cause": f"Extract cause from: {full_text}",
        "symptoms": f"Extract symptoms from: {full_text}",
        "treatment_process": f"Extract treatment process from: {full_text}",
        "illness_characteristics": f"Extract illness characteristics from: {full_text}",
        "treatment_effect": f"Extract treatment effect from: {full_text}"
    }
    
    input_text = prompts.get(field, full_text)
    
    # 编码
    input_ids = tokenizer.encode(
        input_text,
        return_tensors="pt",
        max_length=512,
        truncation=True
    ).to(device)
    
    # 生成
    with torch.no_grad():
        output_ids = model.generate(
            input_ids,
            max_length=256,
            num_beams=4,
            early_stopping=True,
            no_repeat_ngram_size=3
        )
    
    # 解码
    output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    return output_text

def process_test_file(input_file, output_file, model_path="./emotion_summary"):
    """处理测试文件"""
    
    # 加载模型
    model, tokenizer, device = load_model(model_path)
    
    # 读取测试数据
    test_data = []
    with open(input_file, 'r', encoding='utf-8') as f:
        for line in f:
            if line.strip():
                test_data.append(json.loads(line))
    
    print(f"\nProcessing {len(test_data)} samples...")
    
    results = []
    for i, sample in enumerate(test_data, 1):
        print(f"  [{i}/{len(test_data)}] Processing ID: {sample['id']}...")
        
        result = {
            "id": sample["id"],
            "predicted_cause": summarize_case(model, tokenizer, device, sample, "cause"),
            "predicted_symptoms": summarize_case(model, tokenizer, device, sample, "symptoms"),
            "predicted_treatment_process": summarize_case(model, tokenizer, device, sample, "treatment_process"),
            "predicted_illness_Characteristics": summarize_case(model, tokenizer, device, sample, "illness_characteristics"),
            "predicted_treatment_effect": summarize_case(model, tokenizer, device, sample, "treatment_effect")
        }
        
        results.append(result)
    
    # 保存结果
    with open(output_file, 'w', encoding='utf-8') as f:
        for result in results:
            json.dump(result, f, ensure_ascii=False)
            f.write('\n')
    
    print(f"\n✓ Results saved to {output_file}")

if __name__ == "__main__":
    # 示例：处理单个案例
    sample_case = {
        "id": 1,
        "case_description": ["A 34-year-old male with health anxiety..."],
        "consultation_process": ["The consultation began with..."],
        "experience_and_reflection": "This case demonstrates..."
    }
    
    model, tokenizer, device = load_model()
    
    print("\nGenerating summaries...")
    cause = summarize_case(model, tokenizer, device, sample_case, "cause")
    print(f"\nCause: {cause}")
    
    # 如果要处理整个测试文件，取消下面的注释：
    # process_test_file("data/test/Emotion_Summary.jsonl", "results/predictions.jsonl")