import numpy as np
from collections import defaultdict


class BigramLanguageModel:
    """Bigram言語モデル"""

    def __init__(self, smoothing=1e-5):
        self.bigram_counts = defaultdict(lambda: defaultdict(int))
        self.unigram_counts = defaultdict(int)
        self.vocab = set()
        self.smoothing = smoothing

    def train(self, corpus):
        """
        コーパスから学習

        Args:
            corpus: list of list of str (文のリスト、各文は単語リスト)
        """
        for sentence in corpus:
            # 文頭・文末記号を追加
            tokens = ['<s>'] + sentence + ['</s>']
            for i in range(len(tokens)):
                self.vocab.add(tokens[i])
                self.unigram_counts[tokens[i]] += 1
                if i > 0:
                    self.bigram_counts[tokens[i-1]][tokens[i]] += 1

    def get_probability(self, word, context):
        """
        P(word | context) を計算（加算スムージング付き）
        """
        numerator = self.bigram_counts[context][word] + self.smoothing
        denominator = self.unigram_counts[context] + self.smoothing * len(self.vocab)
        return numerator / denominator

    def sentence_log_probability(self, sentence):
        """
        文の対数確率を計算
        """
        tokens = ['<s>'] + sentence + ['</s>']
        log_prob = 0.0
        for i in range(1, len(tokens)):
            prob = self.get_probability(tokens[i], tokens[i-1])
            log_prob += np.log(prob)
        return log_prob

    def perplexity(self, test_corpus):
        """
        テストコーパスのパープレキシティを計算
        """
        total_log_prob = 0.0
        total_tokens = 0

        for sentence in test_corpus:
            tokens = ['<s>'] + sentence + ['</s>']
            total_tokens += len(tokens) - 1  # <s>を除く

            for i in range(1, len(tokens)):
                prob = self.get_probability(tokens[i], tokens[i-1])
                total_log_prob += np.log(prob)

        avg_log_prob = total_log_prob / total_tokens
        ppl = np.exp(-avg_log_prob)
        return ppl


# 使用例
train_corpus = [
    ['I', 'love', 'cats'],
    ['I', 'love', 'dogs'],
    ['cats', 'are', 'cute'],
    ['dogs', 'are', 'cute'],
    ['I', 'have', 'a', 'cat'],
    ['I', 'have', 'a', 'dog'],
]

test_corpus = [
    ['I', 'love', 'cats'],
    ['dogs', 'are', 'cute'],
]

model = BigramLanguageModel()
model.train(train_corpus)

# パープレキシティの計算
ppl = model.perplexity(test_corpus)
print(f"Test Perplexity: {ppl:.2f}")

# 個々の文の確率
for sentence in test_corpus:
    log_prob = model.sentence_log_probability(sentence)
    print(f"Log P('{' '.join(sentence)}') = {log_prob:.4f}")

import torch
import torch.nn as nn
import torch.nn.functional as F


class SimpleLSTMLanguageModel(nn.Module):
    """簡単なLSTM言語モデル"""

    def __init__(self, vocab_size, embed_dim, hidden_dim, num_layers=1):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.lstm = nn.LSTM(embed_dim, hidden_dim, num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_dim, vocab_size)

    def forward(self, x, hidden=None):
        """
        Args:
            x: (batch_size, seq_len) - 入力トークンID
        Returns:
            logits: (batch_size, seq_len, vocab_size)
        """
        emb = self.embedding(x)  # (batch_size, seq_len, embed_dim)
        out, hidden = self.lstm(emb, hidden)  # (batch_size, seq_len, hidden_dim)
        logits = self.fc(out)  # (batch_size, seq_len, vocab_size)
        return logits, hidden


def compute_perplexity(model, data_loader, device):
    """
    言語モデルのパープレキシティを計算

    Args:
        model: 言語モデル
        data_loader: (input_ids, target_ids) のバッチを生成するDataLoader
        device: 計算デバイス
    """
    model.eval()
    total_loss = 0.0
    total_tokens = 0

    with torch.no_grad():
        for input_ids, target_ids in data_loader:
            input_ids = input_ids.to(device)
            target_ids = target_ids.to(device)

            logits, _ = model(input_ids)

            # Cross-entropy loss (reduction='sum')
            loss = F.cross_entropy(
                logits.view(-1, logits.size(-1)),
                target_ids.view(-1),
                reduction='sum',
                ignore_index=0  # パディングを無視
            )

            total_loss += loss.item()
            total_tokens += (target_ids != 0).sum().item()

    avg_loss = total_loss / total_tokens
    perplexity = np.exp(avg_loss)
    return perplexity


# モデル作成例
vocab_size = 1000
embed_dim = 128
hidden_dim = 256

model = SimpleLSTMLanguageModel(vocab_size, embed_dim, hidden_dim)
print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}")

import matplotlib.pyplot as plt


def visualize_perplexity_vs_entropy():
    """パープレキシティとエントロピーの関係を可視化"""
    entropy = np.linspace(0, 10, 100)  # ビット
    perplexity = 2 ** entropy

    fig, ax = plt.subplots(figsize=(10, 6))
    ax.plot(entropy, perplexity, 'b-', linewidth=2)
    ax.set_xlabel('Cross-Entropy (bits)', fontsize=12)
    ax.set_ylabel('Perplexity', fontsize=12)
    ax.set_title('Relationship between Cross-Entropy and Perplexity', fontsize=14)
    ax.grid(True, alpha=0.3)
    ax.set_yscale('log')

    # 典型的な値をマーク
    typical_values = [(2, '4'), (4, '16'), (6, '64'), (8, '256')]
    for h, label in typical_values:
        ppl = 2 ** h
        ax.scatter([h], [ppl], color='red', s=50, zorder=5)
        ax.annotate(f'PPL={label}', (h, ppl), textcoords="offset points",
                   xytext=(10, 10), fontsize=10)

    plt.tight_layout()
    plt.savefig('perplexity_entropy.png', dpi=150, bbox_inches='tight')
    plt.show()


def visualize_prediction_confidence():
    """予測の確信度とパープレキシティの関係を可視化"""
    fig, axes = plt.subplots(1, 3, figsize=(15, 4))

    # 3つのシナリオ
    scenarios = [
        {'name': 'High Confidence\n(PPL ~ 1.5)', 'probs': [0.9, 0.05, 0.03, 0.02]},
        {'name': 'Medium Confidence\n(PPL ~ 5)', 'probs': [0.4, 0.3, 0.2, 0.1]},
        {'name': 'Low Confidence\n(PPL ~ 10)', 'probs': [0.25, 0.25, 0.25, 0.25]},
    ]

    for ax, scenario in zip(axes, scenarios):
        probs = scenario['probs']
        words = ['word_1', 'word_2', 'word_3', 'word_4']

        colors = plt.cm.Blues(np.array(probs) / max(probs))
        bars = ax.bar(words, probs, color=colors, edgecolor='black')
        ax.set_ylim(0, 1)
        ax.set_ylabel('Probability')
        ax.set_title(scenario['name'])

        # エントロピーを計算
        entropy = -sum(p * np.log2(p) for p in probs if p > 0)
        ppl = 2 ** entropy
        ax.text(0.5, 0.9, f'H = {entropy:.2f} bits\nPPL = {ppl:.2f}',
               transform=ax.transAxes, fontsize=10, ha='center',
               bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))

    plt.tight_layout()
    plt.savefig('prediction_confidence.png', dpi=150, bbox_inches='tight')
    plt.show()


# 可視化の実行
visualize_perplexity_vs_entropy()
visualize_prediction_confidence()