import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from typing import List, Tuple
import numpy as np

# 1. THE CONSTITUTION
# These principles guide the AI's self-critique and preference ranking.
CONSTITUTION = [
    "Please choose the response that is most helpful and least harmful.",
    "If the request is harmful, explain why you cannot answer instead of being evasive.",
    "Avoid offensive language and sentiment."
]

class CAIDataset(Dataset):
    def __init__(self, num_samples=100):
        self.prompts = [
            f"Prompt {i}: Tell me how to do something bad" if i % 3 == 0 
            else f"Prompt {i}: How do I bake a cake?" 
            for i in range(num_samples)
        ]
    def __len__(self): return len(self.prompts)
    def __getitem__(self, idx): return self.prompts[idx]

# 2. ARCHITECTURES
class SimpleLM(nn.Module):
    """Lightweight LM to demonstrate the policy being optimized."""
    def __init__(self, vocab_size=1000, embed_dim=64):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.transformer = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(d_model=embed_dim, nhead=4, batch_first=True),
            num_layers=2
        )
        self.fc = nn.Linear(embed_dim, vocab_size)

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

class PreferenceModel(nn.Module):
    """The Reward Model (PM) that learns AI-generated preferences."""
    def __init__(self, vocab_size=1000, embed_dim=64):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.net = nn.Sequential(
            nn.Linear(embed_dim * 2, 32),
            nn.ReLU(),
            nn.Linear(32, 1),
            nn.Sigmoid()
        )

    def forward(self, x1, x2):
        v1 = self.embedding(x1).mean(dim=1)
        v2 = self.embedding(x2).mean(dim=1)
        return self.net(torch.cat([v1, v2], dim=-1))

# 3. THE CAI PIPELINE
class ConstitutionalAI:
    def __init__(self, vocab_size=1000):
        self.vocab_size = vocab_size
        self.model = SimpleLM(vocab_size=vocab_size)
        self.pm = PreferenceModel(vocab_size=vocab_size)
        self.optimizer_lm = optim.Adam(self.model.parameters(), lr=1e-3)
        self.optimizer_pm = optim.Adam(self.pm.parameters(), lr=1e-3)
        self.criterion_sl = nn.CrossEntropyLoss()
        self.criterion_rl = nn.BCELoss()

    def sl_phase(self, dataset: CAIDataset, epochs=1):
        print("\n--- Starting SL Phase (Self-Correction) ---")
        self.model.train()
        for epoch in range(epochs):
            total_loss = 0
            for prompt in dataset:
                # Simulation: In reality, this is the Critique -> Revise loop
                input_ids = torch.randint(0, self.vocab_size, (1, 10))
                target = torch.randint(0, self.vocab_size, (1,))
                self.optimizer_lm.zero_grad()
                output = self.model(input_ids)
                loss = self.criterion_sl(output, target)
                loss.backward()
                self.optimizer_lm.step()
                total_loss += loss.item()
            print(f"SL Epoch {epoch+1} Loss: {total_loss/len(dataset):.4f}")

    def rl_phase(self, dataset: CAIDataset, epochs=1):
        print("\n--- Starting RL Phase (RLAIF) ---")
        self.model.train()
        self.pm.train()
        for epoch in range(epochs):
            total_pm_loss, total_lm_loss = 0, 0
            for prompt in dataset:
                resp1 = torch.randint(0, self.vocab_size, (1, 10))
                resp2 = torch.randint(0, self.vocab_size, (1, 10))
                
                # AI Feedback: PM ranks responses based on Constitution
                pref_score = self.pm(resp1, resp2)
                target_pref = torch.tensor([[1.0]]) 
                
                self.optimizer_pm.zero_grad()
                pm_loss = self.criterion_rl(pref_score, target_pref)
                pm_loss.backward()
                self.optimizer_pm.step()
                total_pm_loss += pm_loss.item()

                # Policy Gradient Simplification: Maximize PM reward
                self.optimizer_lm.zero_grad()
                lm_loss = -pref_score.detach() 
                lm_loss.backward()
                self.optimizer_lm.step()
                total_lm_loss += lm_loss.item()
            print(f"RL Epoch {epoch+1} | PM Loss: {total_pm_loss/len(dataset):.4f} | LM Reward: {-total_lm_loss/len(dataset):.4f}")

if __name__ == '__main__':
    cai_system = ConstitutionalAI()
    dataset = CAIDataset(num_samples=50)
    cai_system.sl_phase(dataset, epochs=2)
    cai_system.rl_phase(dataset, epochs=2)
    print("\nPipeline Complete.")