import torch
import torch.nn as nn
from flash_attn import flash_attention  # FlashAttention 패키지 가정

# 간단한 SSM 구현
class SSM(nn.Module):
    def __init__(self, input_dim, state_dim):
        super(SSM, self).__init__()
        self.A = nn.Parameter(torch.randn(state_dim, state_dim))
        self.B = nn.Parameter(torch.randn(state_dim, input_dim))
        self.C = nn.Parameter(torch.randn(input_dim, state_dim))
        self.D = nn.Parameter(torch.randn(input_dim, input_dim))

    def forward(self, x):
        # x: (batch, seq_len, input_dim)
        batch, seq_len, input_dim = x.shape
        state_dim = self.A.shape[0]
        h = torch.zeros(batch, state_dim, device=x.device)
        outputs = []
        for t in range(seq_len):
            h = h @ self.A + x[:, t] @ self.B
            y = h @ self.C + x[:, t] @ self.D
            outputs.append(y)
        return torch.stack(outputs, dim=1)  # (batch, seq_len, input_dim)

# FlashAttention 기반 Transformer 블록
class TransformerBlock(nn.Module):
    def __init__(self, embed_size, heads):
        super(TransformerBlock, self).__init__()
        self.norm1 = nn.LayerNorm(embed_size)
        self.ffn = nn.Sequential(
            nn.Linear(embed_size, 4 * embed_size),
            nn.ReLU(),
            nn.Linear(4 * embed_size, embed_size)
        )
        self.norm2 = nn.LayerNorm(embed_size)
        self.head_dim = embed_size // heads
        self.heads = heads

    def forward(self, x):
        # x: (batch, seq_len, embed_size)
        batch, seq_len, _ = x.shape
        # FlashAttention 적용
        qkv = x.view(batch, seq_len, self.heads, 3 * self.head_dim)
        q, k, v = qkv.chunk(3, dim=-1)
        attn_output = flash_attention(
            q, k, v, causal=True, dropout=0.1
        ).view(batch, seq_len, -1)
        x = self.norm1(x + attn_output)
        x = self.norm2(x + self.ffn(x))
        return x

# 예시 실행
batch, seq_len, embed_size, heads = 32, 128, 512, 8
x = torch.randn(batch, seq_len, embed_size)
ssm = SSM(embed_size, 64)
transformer = TransformerBlock(embed_size, heads)
ssm_output = ssm(x)
trans_output = transformer(x)
print("SSM 출력 크기:", ssm_output.shape)
print("Transformer 출력 크기:", trans_output.shape)

from transformers import AutoModelForCausalLM, AutoTokenizer

# Phi-3 모델 로드
model = AutoModelForCausalLM.from_pretrained("microsoft/Phi-3-mini-4k-instruct")
tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3-mini-4k-instruct")

# 텍스트 생성
text = "AI의 미래는..."
inputs = tokenizer(text, return_tensors="pt")
outputs = model.generate(**inputs, max_length=50)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

from torch.utils.data import DataLoader, TensorDataset
import torch.optim as optim

# 데이터 준비
seq_len, n_samples, embed_size = 128, 1000, 512
X = torch.randn(n_samples, seq_len, embed_size)
y = torch.randint(0, 2, (n_samples,))
dataset = TensorDataset(X, y)
loader = DataLoader(dataset, batch_size=32)

# 모델 학습
def train_model(model, epochs=5):
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    criterion = nn.BCELoss()
    model.train()
    for epoch in range(epochs):
        total_loss = 0
        for data, target in loader:
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output.squeeze(), target.float())
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
        print(f"Epoch {epoch+1}, 평균 손실: {total_loss / len(loader):.4f}")

# SSM 기반 모델
class SSMClassifier(nn.Module):
    def __init__(self):
        super(SSMClassifier, self).__init__()
        self.ssm = SSM(embed_size, 64)
        self.fc = nn.Linear(embed_size, 1)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        x = self.ssm(x)
        return self.sigmoid(self.fc(x[:, -1, :]))

# FlashAttention 기반 Transformer 모델
class TransformerClassifier(nn.Module):
    def __init__(self):
        super(TransformerClassifier, self).__init__()
        self.transformer = TransformerBlock(embed_size, 8)
        self.fc = nn.Linear(embed_size, 1)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        x = self.transformer(x)
        return self.sigmoid(self.fc(x[:, -1, :]))

# 학습 실행
print("SSM 기반 모델 학습:")
train_model(SSMClassifier())
print("\nFlashAttention 기반 Transformer 학습:")
train_model(TransformerClassifier())

def predict_context(model, tokenizer, text):
    inputs = tokenizer(text, return_tensors='pt', padding=True, truncation=True)
    with torch.no_grad():
        outputs = model.generate(**inputs, max_length=100)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# 예시: Phi-3로 문맥 예측
text = "회의는 내일 오전 10시에 시작되는데..."
prediction = predict_context(model, tokenizer, text)
print("예측된 문맥:", prediction)