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

class DPOLoss(nn.Module):
    def __init__(self, beta=0.1):
        """DPO損失関数

        Args:
            beta: KL制約の強さ（温度パラメータ）
        """
        super().__init__()
        self.beta = beta

    def forward(self, policy_chosen_logps, policy_rejected_logps,
                reference_chosen_logps, reference_rejected_logps):
        """
        DPO損失を計算

        Args:
            policy_chosen_logps: 学習モデルの選好応答の対数確率
            policy_rejected_logps: 学習モデルの非選好応答の対数確率
            reference_chosen_logps: 参照モデルの選好応答の対数確率
            reference_rejected_logps: 参照モデルの非選好応答の対数確率

        Returns:
            loss: DPO損失
            metrics: 追加メトリクス
        """
        # 対数確率比を計算
        chosen_rewards = self.beta * (policy_chosen_logps - reference_chosen_logps)
        rejected_rewards = self.beta * (policy_rejected_logps - reference_rejected_logps)

        # DPO損失
        logits = chosen_rewards - rejected_rewards
        loss = -F.logsigmoid(logits).mean()

        # メトリクス
        chosen_rewards_mean = chosen_rewards.mean().item()
        rejected_rewards_mean = rejected_rewards.mean().item()
        reward_margin = (chosen_rewards - rejected_rewards).mean().item()
        accuracy = (logits > 0).float().mean().item()

        metrics = {
            'chosen_rewards': chosen_rewards_mean,
            'rejected_rewards': rejected_rewards_mean,
            'reward_margin': reward_margin,
            'accuracy': accuracy
        }

        return loss, metrics

# 使用例
dpo_loss = DPOLoss(beta=0.1)

# サンプルデータ（バッチサイズ4）
batch_size = 4
policy_chosen_logps = torch.randn(batch_size)
policy_rejected_logps = torch.randn(batch_size) - 0.5  # 非選好は低い確率
reference_chosen_logps = torch.randn(batch_size)
reference_rejected_logps = torch.randn(batch_size)

loss, metrics = dpo_loss(
    policy_chosen_logps, policy_rejected_logps,
    reference_chosen_logps, reference_rejected_logps
)

print(f"DPO Loss: {loss.item():.4f}")
print(f"Metrics: {metrics}")

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

def compute_log_probs(model, tokenizer, prompt, response):
    """応答の対数確率を計算"""
    # プロンプトと応答を連結
    full_text = prompt + response

    # トークン化
    inputs = tokenizer(full_text, return_tensors='pt')
    prompt_inputs = tokenizer(prompt, return_tensors='pt')
    prompt_length = prompt_inputs['input_ids'].shape[1]

    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits

    # 応答部分のトークンに対する対数確率を計算
    labels = inputs['input_ids']
    shift_logits = logits[:, :-1, :]
    shift_labels = labels[:, 1:]

    # クロスエントロピーを計算（負の対数確率）
    log_probs = F.log_softmax(shift_logits, dim=-1)
    token_log_probs = torch.gather(
        log_probs, 2, shift_labels.unsqueeze(-1)
    ).squeeze(-1)

    # 応答部分のみを合計
    response_log_probs = token_log_probs[:, prompt_length-1:].sum()

    return response_log_probs

# 使用例（小さなモデルでデモ）
# model_name = "gpt2"
# model = AutoModelForCausalLM.from_pretrained(model_name)
# tokenizer = AutoTokenizer.from_pretrained(model_name)
#
# prompt = "What is machine learning?"
# response = " Machine learning is a subset of AI."
# log_prob = compute_log_probs(model, tokenizer, prompt, response)
# print(f"Log probability: {log_prob.item():.4f}")

import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
from tqdm import tqdm

class DPOTrainer:
    def __init__(self, model, ref_model, tokenizer, beta=0.1, lr=1e-6):
        """DPOトレーナー

        Args:
            model: 学習対象のモデル
            ref_model: 参照モデル（凍結）
            tokenizer: トークナイザー
            beta: KL制約の強さ
            lr: 学習率
        """
        self.model = model
        self.ref_model = ref_model
        self.tokenizer = tokenizer
        self.beta = beta

        # 参照モデルは凍結
        for param in self.ref_model.parameters():
            param.requires_grad = False

        self.optimizer = torch.optim.AdamW(model.parameters(), lr=lr)

    def compute_logps(self, model, input_ids, attention_mask, labels):
        """対数確率を計算"""
        outputs = model(input_ids=input_ids, attention_mask=attention_mask)
        logits = outputs.logits

        # シフトして対応を合わせる
        shift_logits = logits[:, :-1, :]
        shift_labels = labels[:, 1:]
        shift_mask = attention_mask[:, 1:]

        # 対数確率
        log_probs = F.log_softmax(shift_logits, dim=-1)
        token_log_probs = torch.gather(
            log_probs, 2, shift_labels.unsqueeze(-1)
        ).squeeze(-1)

        # マスクを適用して合計
        masked_log_probs = token_log_probs * shift_mask
        sequence_log_probs = masked_log_probs.sum(dim=1)

        return sequence_log_probs

    def dpo_loss(self, policy_chosen_logps, policy_rejected_logps,
                 ref_chosen_logps, ref_rejected_logps):
        """DPO損失を計算"""
        chosen_rewards = self.beta * (policy_chosen_logps - ref_chosen_logps)
        rejected_rewards = self.beta * (policy_rejected_logps - ref_rejected_logps)

        logits = chosen_rewards - rejected_rewards
        loss = -F.logsigmoid(logits).mean()

        return loss

    def train_step(self, batch):
        """1ステップの学習"""
        self.model.train()
        self.optimizer.zero_grad()

        # 選好応答の対数確率
        policy_chosen_logps = self.compute_logps(
            self.model,
            batch['chosen_input_ids'],
            batch['chosen_attention_mask'],
            batch['chosen_labels']
        )

        # 非選好応答の対数確率
        policy_rejected_logps = self.compute_logps(
            self.model,
            batch['rejected_input_ids'],
            batch['rejected_attention_mask'],
            batch['rejected_labels']
        )

        # 参照モデルの対数確率（勾配不要）
        with torch.no_grad():
            ref_chosen_logps = self.compute_logps(
                self.ref_model,
                batch['chosen_input_ids'],
                batch['chosen_attention_mask'],
                batch['chosen_labels']
            )
            ref_rejected_logps = self.compute_logps(
                self.ref_model,
                batch['rejected_input_ids'],
                batch['rejected_attention_mask'],
                batch['rejected_labels']
            )

        # 損失計算
        loss = self.dpo_loss(
            policy_chosen_logps, policy_rejected_logps,
            ref_chosen_logps, ref_rejected_logps
        )

        # 逆伝播
        loss.backward()
        self.optimizer.step()

        return loss.item()

    def train(self, dataloader, epochs=3):
        """学習ループ"""
        for epoch in range(epochs):
            total_loss = 0
            pbar = tqdm(dataloader, desc=f"Epoch {epoch+1}/{epochs}")

            for batch in pbar:
                loss = self.train_step(batch)
                total_loss += loss
                pbar.set_postfix({'loss': f'{loss:.4f}'})

            avg_loss = total_loss / len(dataloader)
            print(f"Epoch {epoch+1} - Average Loss: {avg_loss:.4f}")

# 使用例
# trainer = DPOTrainer(model, ref_model, tokenizer, beta=0.1)
# trainer.train(dataloader, epochs=3)