import requests
import numpy as np
import tensorflow as tf
from tensorflow.keras import layers  
import asyncio
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
import sentencepiece as spm
from typing import List
import re
app = FastAPI()
dialogue_history = []
from fastapi.middleware.cors import CORSMiddleware

origins = [
    "https://insect5386.github.io",
    "https://insect5386.github.io/insect5386"
]

app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)
      
sp = spm.SentencePieceProcessor()      
sp.load("kolig_unigram.model")      
      
pad_id = sp.piece_to_id("<pad>")      
if pad_id == -1: pad_id = 0      
start_id = sp.piece_to_id("<start>")      
if start_id == -1: start_id = 1      
end_id = sp.piece_to_id("<end>")      
if end_id == -1: end_id = 2      
unk_id = sp.piece_to_id("<unk>")      
if unk_id == -1: unk_id = 3      
      
vocab_size = sp.get_piece_size()      
max_len = 100      
      
def text_to_ids(text):      
    return sp.encode(text, out_type=int)      
      
def ids_to_text(ids):      
    return sp.decode(ids)      
    
class RotaryPositionalEmbedding(layers.Layer):      
    def __init__(self, dim):      
        super().__init__()      
        inv_freq = 1.0 / (10000 ** (np.arange(0, dim, 2) / dim))      
        self.inv_freq = tf.constant(inv_freq, dtype=tf.float32)      
      
    def call(self, x):      
        batch, heads, seq_len, depth = tf.unstack(tf.shape(x))      
        t = tf.range(seq_len, dtype=tf.float32)      
        freqs = tf.einsum('i,j->ij', t, self.inv_freq)      
        emb_sin = tf.sin(freqs)      
        emb_cos = tf.cos(freqs)      
        emb_cos = tf.reshape(emb_cos, [1, 1, seq_len, -1])      
        emb_sin = tf.reshape(emb_sin, [1, 1, seq_len, -1])      
        x1 = x[..., ::2]      
        x2 = x[..., 1::2]      
        x_rotated = tf.stack([      
            x1 * emb_cos - x2 * emb_sin,      
            x1 * emb_sin + x2 * emb_cos      
        ], axis=-1)      
        x_rotated = tf.reshape(x_rotated, tf.shape(x))      
        return x_rotated    
    
class SwiGLU(tf.keras.layers.Layer):    
    def __init__(self, d_model, d_ff):    
        super().__init__()    
        self.proj = tf.keras.layers.Dense(d_ff * 2)    
        self.out = tf.keras.layers.Dense(d_model)    
    
    def call(self, x):    
        x_proj = self.proj(x)    
        x_val, x_gate = tf.split(x_proj, 2, axis=-1)    
        return self.out(x_val * tf.nn.silu(x_gate))    
            
class GPTBlock(tf.keras.layers.Layer):    
    def __init__(self, d_model, d_ff, num_heads=8, dropout_rate=0.1, adapter_dim=64):      
        super().__init__()      
        self.ln1 = tf.keras.layers.LayerNormalization(epsilon=1e-5)      
        self.mha = tf.keras.layers.MultiHeadAttention(num_heads=num_heads, key_dim=d_model // num_heads)      
        self.dropout1 = tf.keras.layers.Dropout(dropout_rate)     
        self.adapter_down = tf.keras.layers.Dense(adapter_dim, activation='gelu')     
        self.adapter_up = tf.keras.layers.Dense(d_model)      
      
        self.ln2 = tf.keras.layers.LayerNormalization(epsilon=1e-5)      
        self.ffn = SwiGLU(d_model, d_ff)      
        self.dropout2 = tf.keras.layers.Dropout(dropout_rate)     
        self.rope = RotaryPositionalEmbedding(d_model // num_heads)      
      
    def call(self, x, training=False):      
        x_norm = self.ln1(x)      
        b, s, _ = tf.shape(x_norm)[0], tf.shape(x_norm)[1], tf.shape(x_norm)[2]      
        h = self.mha.num_heads      
        d = x_norm.shape[-1] // h      
      
        qkv = tf.reshape(x_norm, [b, s, h, d])      
        qkv = tf.transpose(qkv, [0, 2, 1, 3])      
        q = self.rope(qkv)      
        k = self.rope(qkv)      
        q = tf.reshape(tf.transpose(q, [0, 2, 1, 3]), [b, s, h * d])      
        k = tf.reshape(tf.transpose(k, [0, 2, 1, 3]), [b, s, h * d])      
      
        attn_out = self.mha(query=q, value=x_norm, key=k, use_causal_mask=True, training=training)      
        attn_out = self.dropout1(attn_out, training=training)      
    
        adapter_out = self.adapter_up(self.adapter_down(attn_out))    
        attn_out = attn_out + adapter_out      
      
        x = x + attn_out      
        ffn_out = self.ffn(self.ln2(x))      
        x = x + self.dropout2(ffn_out, training=training)      
        return x    
    
class InteractGPT(tf.keras.Model):      
    def __init__(self, vocab_size, seq_len, d_model, d_ff, n_layers, num_heads=8, dropout_rate=0.1):      
        super().__init__()      
        self.token_embedding = tf.keras.layers.Embedding(vocab_size, d_model)      
        self.blocks = [GPTBlock(d_model, d_ff, num_heads, dropout_rate) for _ in range(n_layers)]      
        self.ln_f = tf.keras.layers.LayerNormalization(epsilon=1e-5)      
      
    def call(self, x, training=False):      
        x = self.token_embedding(x)      
        for block in self.blocks:      
            x = block(x, training=training)      
        x = self.ln_f(x)      
        logits = tf.matmul(x, self.token_embedding.embeddings, transpose_b=True)      
        return logits      
    
model = InteractGPT(vocab_size=vocab_size, seq_len=max_len, d_model=256, d_ff=1024, n_layers=6)        
      
dummy_input = tf.zeros((1, max_len), dtype=tf.int32)  # 배치1, 시퀀스길이 max_len      
_ = model(dummy_input)  # 모델이 빌드됨      
model.load_weights("InteractGPT.weights.h5")      
print("모델 가중치 로드 완료!")      

import re
import math
import numpy as np
import requests
import tensorflow as tf
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
from fastapi import Request
from fastapi.responses import PlainTextResponse

def is_greedy_response_acceptable(text):
    text = text.strip()

    # 너무 짧은 문장 거르기
    if len(text) < 5:
        return False

    # 단어 수 너무 적은 것도 거름
    if len(text.split()) < 3:
        return False

    # ㅋㅋㅋ 같은 자모 연속만 있으면 거름 (단, 'ㅋㅋ' 포함되면 허용)
    if re.search(r'[ㄱ-ㅎㅏ-ㅣ]{3,}', text) and 'ㅋㅋ' not in text:
        return False

    # 문장 끝이 어색한 경우 (다/요/죠 등 일반적 형태로 끝나지 않으면 거름)
    if not re.search(r'(다|요|죠|다\.|요\.|죠\.|다!|요!|죠!|\!|\?|\.)$', text):
        return False

    return True

def generate_text_greedy_strong(model, prompt, max_len=100, max_gen=98,  
                                repetition_penalty=1.2, min_len=20):  
    model_input = text_to_ids(f"<start> {prompt} <sep>")  
    model_input = model_input[:max_len]  
    generated = list(model_input)  

    for _ in range(max_gen):  
        pad_len = max(0, max_len - len(generated))  
        input_padded = np.pad(generated, (0, pad_len), constant_values=pad_id)  
        input_tensor = tf.convert_to_tensor([input_padded])  
        logits = model(input_tensor, training=False)  
        next_logits = logits[0, len(generated) - 1].numpy()  

        # Repetition Penalty
        for t in set(generated):  
            count = generated.count(t)  
            next_logits[t] /= (repetition_penalty ** count)

        # Stop token filtering  
        stop_tokens = ["음", "어", "그", "뭐지", "..."]
        for tok in stop_tokens:
            tok_id = sp.piece_to_id(tok)
            next_logits[tok_id] -= 5.0

        next_logits[pad_id] -= 10.0
        next_token = np.argmax(next_logits)
        generated.append(int(next_token))  

        decoded = sp.decode(generated)  
        for t in ["<start>", "<sep>", "<end>"]:  
            decoded = decoded.replace(t, "")  
        decoded = decoded.strip()

        if len(generated) >= min_len and (next_token == end_id or decoded.endswith(('요', '다', '.', '!', '?'))):  
            if is_greedy_response_acceptable(decoded):
                return decoded
            else:
                continue

    return sp.decode(generated)

def mismatch_tone(input_text, output_text):  
    if "ㅋㅋ" in input_text and not re.search(r'ㅋㅋ|ㅎ|재밌|놀|만나|맛집|여행', output_text):  
        return True  
    return False

# 유효한 응답인지 검사
def is_valid_response(response):
    if len(response.strip()) < 2:
        return False
    if re.search(r'[ㄱ-ㅎㅏ-ㅣ]{3,}', response):
        return False
    if len(response.split()) < 2:
        return False
    if response.count(' ') < 2:
        return False
    if any(tok in response.lower() for tok in ['hello', 'this', 'ㅋㅋ']):
        return False
    return True

# 위키 요약 관련
def extract_main_query(text):
    sentences = re.split(r'[.?!]\s*', text)
    sentences = [s.strip() for s in sentences if s.strip()]
    if not sentences:
        return text
    last = sentences[-1]
    last = re.sub(r'[^가-힣a-zA-Z0-9 ]', '', last)
    particles = ['이', '가', '은', '는', '을', '를', '의', '에서', '에게', '한테', '보다']
    for p in particles:
        last = re.sub(rf'\b(\w+){p}\b', r'\1', last)
    return last.strip()

def get_wikipedia_summary(query):
    cleaned_query = extract_main_query(query)
    url = f"https://ko.wikipedia.org/api/rest_v1/page/summary/{cleaned_query}"
    res = requests.get(url)
    if res.status_code == 200:
        return res.json().get("extract", "요약 정보를 찾을 수 없습니다.")
    else:
        return "위키백과에서 정보를 가져올 수 없습니다."

def textrank_summarize(text, top_n=3):
    sentences = re.split(r'(?<=[.!?])\s+', text.strip())
    sentences = [s.strip() for s in sentences if len(s.strip()) > 10]
    if len(sentences) <= top_n:
        return text
    vectorizer = TfidfVectorizer()
    tfidf_matrix = vectorizer.fit_transform(sentences)
    sim_matrix = cosine_similarity(tfidf_matrix)
    np.fill_diagonal(sim_matrix, 0)
    def pagerank(matrix, damping=0.85, max_iter=100, tol=1e-4):
        N = matrix.shape[0]
        ranks = np.ones(N) / N
        row_sums = np.sum(matrix, axis=1)
        row_sums[row_sums == 0] = 1
        for _ in range(max_iter):
            prev_ranks = ranks.copy()
            for i in range(N):
                incoming = matrix[:, i]
                ranks[i] = (1 - damping) / N + damping * np.sum(incoming * prev_ranks / row_sums)
            if np.linalg.norm(ranks - prev_ranks) < tol:
                break
        return ranks
    scores = pagerank(sim_matrix)
    ranked_idx = np.argsort(scores)[::-1]
    selected_idx = sorted(ranked_idx[:top_n])
    summary = ' '.join([sentences[i] for i in selected_idx])
    return summary

def summarize_from_wikipedia(query, top_n=3):
    raw_summary = get_wikipedia_summary(query)
    return textrank_summarize(raw_summary, top_n=top_n)

def build_contexted_prompt(history: List[str], user_input: str):
    # 최근 대화 3개를 합쳐서 요약
    recent = ' '.join(history[-3:])
    summary = textrank_summarize(recent, top_n=2)
    # 요약 + 최신 사용자 입력을 같이 던져줌
    prompt = f"{summary} {user_input}"
    return prompt

# 의도 분류기
def simple_intent_classifier(text):
    text = text.lower()
    greet_keywords = ["안녕", "반가워", "이름", "누구", "소개", "어디서 왔", "정체", "몇 살", "너 뭐야"]
    info_keywords = ["설명", "정보", "무엇", "뭐야", "어디", "누구", "왜", "어떻게", "종류", "개념"]
    math_keywords = ["더하기", "빼기", "곱하기", "나누기", "루트", "제곱", "+", "-", "*", "/", "=", "^", "√", "계산", "몇이야", "얼마야"]
    if any(kw in text for kw in greet_keywords):
        return "인사"
    elif any(kw in text for kw in info_keywords):
        return "정보질문"
    elif any(kw in text for kw in math_keywords):
        return "수학질문"
    else:
        return "일상대화"

def parse_math_question(text):
    text = text.replace("곱하기", "*").replace("더하기", "+").replace("빼기", "-").replace("나누기", "/").replace("제곱", "*2")
    text = re.sub(r'루트\s(\d+)', r'math.sqrt(\1)', text)
    try:
        result = eval(text)
        return f"정답은 {result}입니다."
    except:
        return "계산할 수 없는 수식이에요. 다시 한번 확인해 주세요!"

# respond 함수 수정
def respond(input_text):
    global dialogue_history

    intent = simple_intent_classifier(input_text)

    if "이름" in input_text:
        return "제 이름은 Ector.V입니다."

    if "누구" in input_text:
        return "저는 Ector.V라고 해요."

    if intent == "수학질문":
        dialogue_history.append(f"사용자: {input_text}")
        response = parse_math_question(input_text)
        dialogue_history.append(f"Ector: {response}")
        return response

    if intent == "인사":
        response = "반가워요! 무엇을 도와드릴까요?"
        dialogue_history.append(f"사용자: {input_text}")
        dialogue_history.append(f"Ector: {response}")
        return response

    if intent == "정보질문":
        keyword = re.sub(r"(에 대해|에 대한|에 대해서)?\s*(설명해줘|알려줘|뭐야|개념|정의|정보)?", "", input_text).strip()
        if not keyword:
            return "어떤 주제에 대해 궁금한가요?"
        summary = summarize_from_wikipedia(keyword)
        response = f"{summary}\n다른 궁금한 점 있으신가요?"
        dialogue_history.append(f"사용자: {input_text}")
        dialogue_history.append(f"Ector: {response}")
        return response

    # 일상 대화: 요약 기반 컨텍스트 생성
    contexted_prompt = build_contexted_prompt(dialogue_history, input_text)
    response = generate_text_greedy_strong(model, contexted_prompt)

    # fallback
    if not is_valid_response(response) or mismatch_tone(input_text, response):
        response = generate_text_greedy_strong(model, contexted_prompt)

    # 히스토리 추가
    dialogue_history.append(f"사용자: {input_text}")
    dialogue_history.append(f"Ector: {response}")
    return response

@app.get("/generate", response_class=PlainTextResponse)
async def generate(request: Request):
    prompt = request.query_params.get("prompt", "안녕하세요")
    response_text = respond(prompt)
    return response_text