import torch
import torch.nn as nn
from transformers import BertModel, BertTokenizer

class IntentClassificationModel(nn.Module):
    def __init__(self, bert_path, num_intents, hidden_size=768):
        super().__init__()
        self.bert = BertModel.from_pretrained(bert_path)
        self.lstm = nn.LSTM(768, hidden_size//2, 
                           batch_first=True, bidirectional=True)
        self.classifier = nn.Linear(hidden_size, num_intents)
        self.dropout = nn.Dropout(0.1)
        
    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids, attention_mask=attention_mask)
        sequence_output = outputs.last_hidden_state
        
        lstm_out, _ = self.lstm(sequence_output)
        pooled_output = lstm_out[:, 0, :]  # 取第一个token的输出
        
        pooled_output = self.dropout(pooled_output)
        logits = self.classifier(pooled_output)
        return logits

# 模型初始化
model = IntentClassificationModel(
    bert_path="hfl/chinese-roberta-wwm-ext",
    num_intents=6
)

import torch.nn.functional as F

class MultiIntentModel(IntentClassificationModel):
    def __init__(self, bert_path, num_intents, hidden_size=768, threshold=0.3):
        super().__init__(bert_path, num_intents, hidden_size)
        self.threshold = threshold
        
    def forward(self, input_ids, attention_mask):
        logits = super().forward(input_ids, attention_mask)
        # 使用sigmoid激活函数支持多标签
        probabilities = torch.sigmoid(logits)
        # 基于阈值生成多标签预测
        predictions = (probabilities > self.threshold).long()
        return predictions, probabilities

# 多意图推理示例
def predict_multiple_intents(text, model, tokenizer, device):
    inputs = tokenizer(text, return_tensors="pt", 
                      padding=True, truncation=True, max_length=128)
    inputs = {k: v.to(device) for k, v in inputs.items()}
    
    with torch.no_grad():
        predictions, probabilities = model(**inputs)
    
    intent_labels = ["事业", "财运", "感情", "健康", "学业", "其他"]
    detected_intents = []
    
    for i, prob in enumerate(probabilities[0]):
        if prob > model.threshold:
            detected_intents.append((intent_labels[i], float(prob)))
    
    return detected_intents

class FuzzyExpressionHandler:
    def __init__(self):
        self.fuzzy_patterns = {
            "事业": ["工作怎么样", "职业发展", "前途如何"],
            "财运": ["钱方面", "经济状况", "财务运气"],
            "感情": ["姻缘", "桃花运", "感情生活"]
        }
        self.similarity_threshold = 0.7
    
    def handle_fuzzy_expression(self, text, intent_model, sentence_encoder):
        # 直接意图识别
        direct_intents = predict_multiple_intents(text, intent_model)
        
        if direct_intents:
            return direct_intents
        
        # 模糊表达匹配
        text_embedding = sentence_encoder.encode(text)
        matched_intents = []
        
        for intent, patterns in self.fuzzy_patterns.items():
            pattern_embeddings = sentence_encoder.encode(patterns)
            similarities = cosine_similarity([text_embedding], pattern_embeddings)[0]
            max_similarity = max(similarities)
            
            if max_similarity > self.similarity_threshold:
                matched_intents.append((intent, max_similarity))
        
        return matched_intents

import re
from datetime import datetime

class BirthTimeExtractor:
    def __init__(self):
        # 农历转换器需要额外实现
        self.lunar_converter = LunarConverter()
        
        # 时间表达式模式
        self.patterns = {
            '公历': r'(\d{4})年(\d{1,2})月(\d{1,2})日\s*(\d{1,2})点(\d{1,2})分?',
            '农历': r'农历(\d{1,2})月(\d{1,2})日\s*(\d{1,2})点(\d{1,2})分?',
            '时辰': r'([子丑寅卯辰巳午未申酉戌亥])时'
        }
    
    def extract_birth_time(self, text):
        """从文本中抽取出生时间信息"""
        time_info = {}
        
        # 公历日期匹配
        solar_match = re.search(self.patterns['公历'], text)
        if solar_match:
            year, month, day, hour, minute = solar_match.groups()
            time_info['solar'] = {
                'year': int(year), 'month': int(month), 
                'day': int(day), 'hour': int(hour), 'minute': int(minute or 0)
            }
            time_info['lunar'] = self.lunar_converter.solar_to_lunar(
                time_info['solar']
            )
        
        # 农历日期匹配
        lunar_match = re.search(self.patterns['农历'], text)
        if lunar_match:
            month, day, hour, minute = lunar_match.groups()
            time_info['lunar'] = {
                'month': int(month), 'day': int(day),
                'hour': int(hour), 'minute': int(minute or 0)
            }
        
        return time_info
    
    def calculate_bazi(self, birth_time):
        """根据出生时间计算八字"""
        # 八字计算逻辑（简化版）
        if 'lunar' in birth_time:
            lunar = birth_time['lunar']
            # 实际实现需要完整的八字计算算法
            year_pillar = self.get_year_pillar(lunar['year'])
            month_pillar = self.get_month_pillar(lunar['year'], lunar['month'])
            day_pillar = self.get_day_pillar(lunar['year'], lunar['month'], lunar['day'])
            hour_pillar = self.get_hour_pillar(day_pillar, lunar['hour'])
            
            return {
                'year': year_pillar,
                'month': month_pillar,
                'day': day_pillar,
                'hour': hour_pillar
            }
        return None

class DestinyNERModel(nn.Module):
    def __init__(self, bert_path, num_labels):
        super().__init__()
        self.bert = BertModel.from_pretrained(bert_path)
        self.dropout = nn.Dropout(0.1)
        self.classifier = nn.Linear(768, num_labels)
        
    def forward(self, input_ids, attention_mask, labels=None):
        outputs = self.bert(input_ids, attention_mask=attention_mask)
        sequence_output = outputs.last_hidden_state
        
        sequence_output = self.dropout(sequence_output)
        logits = self.classifier(sequence_output)
        
        if labels is not None:
            loss_fct = nn.CrossEntropyLoss()
            active_loss = attention_mask.view(-1) == 1
            active_logits = logits.view(-1, logits.size(-1))
            active_labels = torch.where(
                active_loss, labels.view(-1), 
                torch.tensor(loss_fct.ignore_index).type_as(labels)
            )
            loss = loss_fct(active_logits, active_labels)
            return loss, logits
        
        return logits

# 实体标签定义
NER_LABELS = {
    'O': 0, 'B-TIME': 1, 'I-TIME': 2,
    'B-STAR': 3, 'I-STAR': 4,          # 星曜
    'B-PATTERN': 5, 'I-PATTERN': 6,    # 格局
    'B-ELEMENT': 7, 'I-ELEMENT': 8,    # 五行
    'B-PALACE': 9, 'I-PALACE': 10      # 宫位
}

class EntityLinker:
    def __init__(self, knowledge_graph):
        self.kg = knowledge_graph
        self.synonym_dict = self.load_synonyms()
    
    def load_synonyms(self):
        """加载同义词词典"""
        return {
            "紫微": ["紫微星", "紫微垣", "北辰"],
            "七杀": ["七杀星", "杀星"],
            "破军": ["破军星", "耗星"],
            "天府": ["天府星", "令星"]
        }
    
    def entity_linking(self, entity_text, context):
        """实体链接到知识图谱"""
        # 同义词扩展
        candidate_entities = self.expand_synonyms(entity_text)
        
        # 基于上下文相似度排序
        ranked_entities = self.rank_by_context_similarity(
            candidate_entities, context
        )
        
        return ranked_entities[0] if ranked_entities else None
    
    def expand_synonyms(self, entity_text):
        """扩展同义词"""
        candidates = [entity_text]
        for canonical, synonyms in self.synonym_dict.items():
            if entity_text in synonyms or entity_text == canonical:
                candidates.extend(synonyms)
                candidates.append(canonical)
        return list(set(candidates))

from enum import Enum
from typing import Dict, Any

class DialogueState(Enum):
    GREETING = "greeting"
    COLLECTING_INFO = "collecting_info"
    ANALYZING = "analyzing"
    PROVIDING_RESULT = "providing_result"
    FOLLOW_UP = "follow_up"
    ENDING = "ending"

class DialogueManager:
    def __init__(self):
        self.current_state = DialogueState.GREETING
        self.user_context = {}
        self.state_handlers = {
            DialogueState.GREETING: self.handle_greeting,
            DialogueState.COLLECTING_INFO: self.handle_info_collection,
            DialogueState.ANALYZING: self.handle_analysis,
            DialogueState.PROVIDING_RESULT: self.handle_result,
            DialogueState.FOLLOW_UP: self.handle_follow_up,
            DialogueState.ENDING: self.handle_ending
        }
    
    def process_user_input(self, user_input: str) -> str:
        """处理用户输入并返回系统响应"""
        # 更新用户上下文
        self.update_context(user_input)
        
        # 状态转移
        self.transition_state()
        
        # 执行当前状态处理
        handler = self.state_handlers[self.current_state]
        response = handler(user_input)
        
        return response
    
    def transition_state(self):
        """状态转移逻辑"""
        if self.current_state == DialogueState.GREETING:
            if self.user_context.get('user_intent'):
                self.current_state = DialogueState.COLLECTING_INFO
        
        elif self.current_state == DialogueState.COLLECTING_INFO:
            if self.has_sufficient_info():
                self.current_state = DialogueState.ANALYZING
        
        elif self.current_state == DialogueState.ANALYZING:
            self.current_state = DialogueState.PROVIDING_RESULT
        
        # 其他状态转移逻辑...
    
    def handle_info_collection(self, user_input):
        """信息收集状态处理"""
        missing_info = self.get_missing_info()
        if missing_info:
            return self.ask_for_missing_info(missing_info)
        return "请告诉我您想咨询的具体问题..."

class ContextManager:
    def __init__(self):
        self.context = {
            'current_intent': None,
            'mentioned_entities': [],
            'collected_info': {},
            'dialogue_history': [],
            'user_profile': {}
        }
        self.max_history_len = 10
    
    def update_context(self, user_input, system_response, nlu_result):
        """更新对话上下文"""
        # 记录对话历史
        self.context['dialogue_history'].append({
            'user': user_input,
            'system': system_response,
            'timestamp': datetime.now()
        })
        
        # 保持历史记录长度
        if len(self.context['dialogue_history']) > self.max_history_len:
            self.context['dialogue_history'] = \
                self.context['dialogue_history'][-self.max_history_len:]
        
        # 更新提及的实体
        if 'entities' in nlu_result:
            self.context['mentioned_entities'].extend(nlu_result['entities'])
        
        # 更新用户信息
        if 'user_info' in nlu_result:
            self.context['collected_info'].update(nlu_result['user_info'])
    
    def get_context_summary(self):
        """获取上下文摘要，用于模型输入"""
        recent_history = self.context['dialogue_history'][-3:]
        history_text = " ".join([
            f"用户: {turn['user']} 系统: {turn['system']}" 
            for turn in recent_history
        ])
        
        return {
            'history': history_text,
            'current_intent': self.context['current_intent'],
            'collected_info': self.context['collected_info']
        }

class DestinyGPTFineTuner:
    def __init__(self, base_model="gpt2-chinese"):
        self.tokenizer = AutoTokenizer.from_pretrained(base_model)
        self.model = AutoModelForCausalLM.from_pretrained(base_model)
        self.tokenizer.pad_token = self.tokenizer.eos_token
    
    def prepare_training_data(self, corpus_path):
        """准备命理领域训练数据"""
        with open(corpus_path, 'r', encoding='utf-8') as f:
            texts = f.readlines()
        
        # 数据预处理
        processed_texts = []
        for text in texts:
            # 清理和标准化文本
            cleaned = self.clean_destiny_text(text.strip())
            if cleaned:
                processed_texts.append(cleaned)
        
        return processed_texts
    
    def clean_destiny_text(self, text):
        """清理命理文本"""
        # 移除特殊符号，保留中文和必要标点
        cleaned = re.sub(r'[^\u4e00-\u9fa5，。！？；：、\s]', '', text)
        return cleaned.strip()
    
    def fine_tune(self, train_texts, output_dir, epochs=3):
        """微调模型"""
        # 编码训练数据
        encodings = self.tokenizer(
            train_texts, 
            truncation=True, 
            padding=True,
            max_length=512,
            return_tensors="pt"
        )
        
        training_args = TrainingArguments(
            output_dir=output_dir,
            num_train_epochs=epochs,
            per_device_train_batch_size=4,
            save_steps=500,
            save_total_limit=2,
            prediction_loss_only=True,
            remove_unused_columns=False
        )
        
        trainer = Trainer(
            model=self.model,
            args=training_args,
            data_collator=DataCollatorForLanguageModeling(
                tokenizer=self.tokenizer, mlm=False
            ),
            train_dataset=encodings
        )
        
        trainer.train()
        trainer.save_model()
        
        return trainer

class HybridResponseGenerator:
    def __init__(self, gpt_model, template_manager):
        self.gpt_model = gpt_model
        self.template_manager = template_manager
        self.generation_config = {
            'max_length': 200,
            'temperature': 0.7,
            'top_p': 0.9,
            'do_sample': True,
            'num_return_sequences': 1
        }
    
    def generate_response(self, context, intent, entities):
        """生成响应"""
        # 1. 尝试模板生成
        template_response = self.template_manager.fill_template(
            intent, entities, context
        )
        
        if template_response and self.is_high_confidence(context):
            return template_response
        
        # 2. GPT生成
        prompt = self.build_generation_prompt(context, intent, entities)
        gpt_response = self.gpt_model.generate(
            prompt, **self.generation_config
        )
        
        # 3. 后处理和质量检查
        processed_response = self.post_process(gpt_response)
        
        return processed_response if self.quality_check(processed_response) else template_response
    
    def build_generation_prompt(self, context, intent, entities):
        """构建生成提示"""
        base_prompt = f"""
        作为专业的命理咨询师，请根据以下信息为用户提供咨询：
        
        用户咨询意图：{intent}
        用户信息：{context['collected_info']}
        相关命理要素：{entities}
        
        请生成专业、准确且易于理解的命理分析：
        """
        return base_prompt

class ClassicalChineseUnderstanding:
    def __init__(self):
        self.modernizer = ClassicalChineseModernizer()
        self.ancient_terms = self.load_ancient_terms()
    
    def modernize_text(self, classical_text):
        """将古文转换为现代文"""
        # 分词和词性标注
        words = jieba.cut(classical_text)
        modernized_words = []
        
        for word in words:
            if word in self.ancient_terms:
                modernized_words.append(
                    self.ancient_terms[word]['modern']
                )
            else:
                modernized_words.append(word)
        
        return ''.join(modernized_words)
    
    def load_ancient_terms(self):
        """加载古今词汇对照表"""
        return {
            "癸水": {"modern": "癸水", "explanation": "天干第十位，属阴水"},
            "伤官": {"modern": "伤官", "explanation": "十神之一，代表才华"},
            "七杀": {"modern": "七杀", "explanation": "十神之一，代表压力"}
        }

class MultilingualSupport:
    def __init__(self):
        self.translator = Translator()
        self.destiny_terms_dict = self.load_terms_dictionary()
    
    def translate_destiny_content(self, text, target_lang='en'):
        """翻译命理内容，保持专业术语准确"""
        # 首先提取和保护专业术语
        protected_terms = self.extract_protected_terms(text)
        protected_text = self.protect_terms(text, protected_terms)
        
        # 翻译文本
        translated_text = self.translator.translate(
            protected_text, dest=target_lang
        ).text
        
        # 恢复专业术语
        final_text = self.restore_terms(translated_text, protected_terms)
        
        return final_text
    
    def extract_protected_terms(self, text):
        """提取需要保护的专业术语"""
        terms = []
        for term in self.destiny_terms_dict.keys():
            if term in text:
                terms.append(term)
        return terms

class DestinyRAGSystem:
    def __init__(self, retriever, generator):
        self.retriever = retriever
        self.generator = generator
        self.reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')
    
    def answer_question(self, question, context):
        """回答用户问题"""
        # 1. 检索相关文档
        retrieved_docs = self.retriever.retrieve(question, k=5)
        
        # 2. 重排序
        reranked_docs = self.rerank_documents(question, retrieved_docs)
        
        # 3. 生成答案
        prompt = self.build_rag_prompt(question, reranked_docs[:3], context)
        answer = self.generator.generate(prompt)
        
        return answer, reranked_docs
    
    def build_rag_prompt(self, question, documents, context):
        """构建RAG提示"""
        context_str = "\n".join([
            f"文档 {i+1}: {doc['content']}" 
            for i, doc in enumerate(documents)
        ])
        
        prompt = f"""
        基于以下命理知识和用户上下文，回答用户问题：
        
        相关知识：
        {context_str}
        
        用户上下文：
        {context}
        
        用户问题：{question}
        
        请提供专业、准确的回答：
        """
        return prompt

class ModelServing:
    def __init__(self):
        self.model_registry = {}
        self.load_balancer = LoadBalancer()
    
    def deploy_model(self, model_name, model_path, replicas=2):
        """部署模型服务"""
        # 创建模型容器
        for i in range(replicas):
            container = ModelContainer(
                model_name=model_name,
                model_path=model_path,
                instance_id=f"{model_name}-{i}",
                resources={"cpu": "2", "memory": "4Gi"}
            )
            container.start()
            
            self.model_registry[f"{model_name}-{i}"] = {
                'container': container,
                'status': 'healthy',
                'load': 0
            }
    
    def predict(self, model_name, input_data):
        """模型预测"""
        # 选择负载最低的实例
        instance_id = self.load_balancer.select_instance(model_name)
        instance = self.model_registry[instance_id]
        
        try:
            result = instance['container'].predict(input_data)
            instance['load'] += 1
            return result
        except Exception as e:
            instance['status'] = 'unhealthy'
            # 重试其他实例
            return self.predict(model_name, input_data)

class IntelligentDestinyAssistant:
    def __init__(self):
        self.nlu_engine = DestinyNLUEngine()
        self.dialogue_manager = DialogueManager()
        self.context_manager = ContextManager()
        self.response_generator = HybridResponseGenerator()
        self.qa_system = DestinyRAGSystem()
    
    def process_query(self, user_input, user_id):
        """处理用户查询"""
        # 1. NLU理解
        nlu_result = self.nlu_engine.parse(user_input)
        
        # 2. 上下文更新
        self.context_manager.update_context(
            user_input, "", nlu_result
        )
        
        # 3. 对话管理
        context_summary = self.context_manager.get_context_summary()
        system_response = self.dialogue_manager.process_user_input(
            user_input, context_summary
        )
        
        # 4. 响应生成
        final_response = self.response_generator.generate_response(
            context_summary, 
            nlu_result['intent'],
            nlu_result['entities']
        )
        
        # 5. 记录交互
        self.log_interaction(user_id, user_input, final_response)
        
        return final_response

class PerformanceOptimizer:
    def __init__(self):
        self.cache = RedisCache()
        self.model_optimizer = ModelOptimizer()
    
    def optimize_response_time(self, query, context):
        """优化响应时间"""
        # 1. 查询缓存
        cache_key = self.generate_cache_key(query, context)
        cached_response = self.cache.get(cache_key)
        
        if cached_response:
            return cached_response
        
        # 2. 模型推理优化
        optimized_model = self.model_optimizer.quantize_model(self.model)
        
        # 3. 异步处理耗时任务
        async_result = self.process_async(query, context)
        
        return async_result.get(timeout=3.0)  # 3秒超时
    
    def generate_cache_key(self, query, context):
        """生成缓存键"""
        key_data = {
            'query': query,
            'intent': context.get('current_intent'),
            'user_info_hash': hash(str(context.get('collected_info')))
        }
        return hashlib.md5(str(key_data).encode()).hexdigest()

class ContinuousLearning:
    def __init__(self, model, feedback_collector):
        self.model = model
        self.feedback_collector = feedback_collector
        self.retraining_threshold = 1000  # 1000个新样本后重训练
    
    def collect_feedback(self, user_input, system_response, user_feedback):
        """收集用户反馈"""
        feedback_data = {
            'input': user_input,
            'response': system_response,
            'feedback': user_feedback,  # 1-5分评分
            'timestamp': datetime.now()
        }
        
        self.feedback_collector.save(feedback_data)
        
        # 检查是否需要重训练
        if self.should_retrain():
            self.retrain_model()
    
    def should_retrain(self):
        """判断是否需要重训练"""
        recent_feedback_count = self.feedback_collector.get_recent_count()
        avg_rating = self.feedback_collector.get_average_rating()
        
        return (recent_feedback_count >= self.retraining_threshold and 
                avg_rating < 4.0)  # 平均评分低于4.0时重训练