import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import multivariate_normal
from math import log, sqrt

# 設置隨機種子以確保結果可重現
torch.manual_seed(42)
np.random.seed(42)

# 定義任務分佈類
class TaskDistribution:
    def __init__(self, input_dim=2, output_dim=1, shift_magnitude=0.5):
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.shift_magnitude = shift_magnitude
        
    def sample_source_task(self):
        # 源任務：簡單的線性關係加上高斯噪聲
        W = torch.randn(self.output_dim, self.input_dim) * 0.5
        b = torch.randn(self.output_dim) * 0.1
        noise_std = 0.1
        return W, b, noise_std
    
    def sample_target_task(self):
        # 目標任務：與源任務相關但有分佈偏移
        W_source, b_source, noise_std_source = self.sample_source_task()
        
        # 引入分佈偏移
        W_shift = torch.randn_like(W_source) * self.shift_magnitude
        b_shift = torch.randn_like(b_source) * self.shift_magnitude * 0.5
        
        W_target = W_source + W_shift
        b_target = b_source + b_shift
        noise_std_target = noise_std_source * (1 + self.shift_magnitude * 0.5)
        
        return W_target, b_target, noise_std_target
    
    def generate_task_data(self, W, b, noise_std, num_samples):
        X = torch.randn(num_samples, self.input_dim)
        y = X @ W.t() + b + torch.randn(num_samples, self.output_dim) * noise_std
        return X, y

# 創建元數據集
class MetaDataset(Dataset):
    def __init__(self, task_dist, num_tasks=100, samples_per_task=20, source=True):
        self.task_dist = task_dist
        self.num_tasks = num_tasks
        self.samples_per_task = samples_per_task
        self.source = source
        
        self.tasks = []
        for _ in range(num_tasks):
            if source:
                W, b, noise_std = task_dist.sample_source_task()
            else:
                W, b, noise_std = task_dist.sample_target_task()
            X, y = task_dist.generate_task_data(W, b, noise_std, samples_per_task)
            self.tasks.append((X, y, W, b))
    
    def __len__(self):
        return self.num_tasks
    
    def __getitem__(self, idx):
        return self.tasks[idx]

# 定義基學習器模型
class BaseLearner(nn.Module):
    def __init__(self, input_dim=2, hidden_dim=20, output_dim=1):
        super(BaseLearner, self).__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, output_dim)
        )
        
    def forward(self, x):
        return self.net(x)

# 實現PAC-Bayesian MAML
class PBMAML:
    def __init__(self, model, prior_std=1.0, alpha=0.01, beta=0.001, lambda_reg=0.1):
        self.model = model
        self.prior_std = prior_std
        self.alpha = alpha  # 內循環學習率
        self.beta = beta    # 外循環學習率
        self.lambda_reg = lambda_reg  # KL正則化係數
        
        # 初始化先驗分佈（零均值高斯）
        self.prior_mean = self._get_flat_params()
        self.prior_log_std = torch.log(torch.ones_like(self.prior_mean) * prior_std)
        
        # 優化器
        self.optimizer = optim.Adam(self.model.parameters(), lr=beta)
        
    def _get_flat_params(self):
        """將模型參數展平為一維張量"""
        params = []
        for param in self.model.parameters():
            params.append(param.data.view(-1))
        return torch.cat(params)
    
    def _set_flat_params(self, flat_params):
        """從一維張量設置模型參數"""
        offset = 0
        for param in self.model.parameters():
            numel = param.numel()
            param.data.copy_(flat_params[offset:offset+numel].view_as(param))
            offset += numel
    
    def _compute_kl_divergence(self, mean, log_std):
        """計算後驗分佈與先驗分佈之間的KL散度"""
        # 後驗分佈：對角高斯 N(mean, exp(log_std)^2)
        # 先驗分佈：N(prior_mean, prior_std^2)
        posterior_var = torch.exp(2 * log_std)
        prior_var = self.prior_std ** 2
        
        kl = 0.5 * (torch.log(prior_var / posterior_var) - 1 + 
                    posterior_var / prior_var + 
                    (mean - self.prior_mean) ** 2 / prior_var)
        return kl.sum()
    
    def inner_update(self, task, num_steps=1):
        """在單個任務上進行內循環適應"""
        X, y = task[0], task[1]
        
        # 保存初始參數
        initial_params = self._get_flat_params().detach().clone()
        
        # 創建臨時模型進行內循環更新
        temp_model = BaseLearner()
        self._set_flat_params(initial_params.clone())
        temp_model.load_state_dict(self.model.state_dict())
        
        # 內循環優化器
        inner_optimizer = optim.SGD(temp_model.parameters(), lr=self.alpha)
        
        for step in range(num_steps):
            inner_optimizer.zero_grad()
            y_pred = temp_model(X)
            loss = nn.MSELoss()(y_pred, y)
            loss.backward()
            inner_optimizer.step()
        
        # 計算適應後的參數
        adapted_params = torch.cat([param.data.view(-1) for param in temp_model.parameters()])
        
        return initial_params, adapted_params
    
    def compute_pac_bayes_bound(self, empirical_risk, kl_divergence, n_tasks, delta=0.05):
        """計算PAC-Bayesian泛化上界"""
        # 使用經典的PAC-Bayes邊界
        bound = empirical_risk + sqrt((kl_divergence + log(n_tasks / delta)) / (2 * n_tasks))
        return bound
    
    def meta_train(self, meta_dataloader, num_epochs=100):
        """元訓練過程"""
        bounds_history = []
        empirical_risk_history = []
        kl_history = []
        
        for epoch in range(num_epochs):
            total_meta_loss = 0
            total_empirical_risk = 0
            total_kl = 0
            num_tasks = 0
            
            for task_batch in meta_dataloader:
                batch_meta_loss = 0
                batch_empirical_risk = 0
                batch_kl = 0
                
                for task in task_batch:
                    X, y = task[0], task[1]
                    
                    # 內循環適應
                    initial_params, adapted_params = self.inner_update(task)
                    
                    # 計算後驗分佈的均值和標準差
                    # 這裏我們使用適應後的參數作為後驗均值
                    posterior_mean = adapted_params
                    posterior_log_std = torch.log(torch.ones_like(posterior_mean) * 0.1)
                    
                    # 計算KL散度
                    kl_divergence = self._compute_kl_divergence(posterior_mean, posterior_log_std)
                    
                    # 計算經驗風險
                    adapted_model = BaseLearner()
                    self._set_flat_params(adapted_params.clone())
                    adapted_model.load_state_dict(self.model.state_dict())
                    
                    with torch.no_grad():
                        y_pred = adapted_model(X)
                        empirical_risk = nn.MSELoss()(y_pred, y).item()
                    
                    # 計算元損失（經驗風險 + KL正則化）
                    meta_loss = empirical_risk + self.lambda_reg * kl_divergence
                    
                    batch_meta_loss += meta_loss
                    batch_empirical_risk += empirical_risk
                    batch_kl += kl_divergence.item()
                
                # 平均批次損失
                batch_size = len(task_batch)
                batch_meta_loss /= batch_size
                batch_empirical_risk /= batch_size
                batch_kl /= batch_size
                
                # 元優化步驟
                self.optimizer.zero_grad()
                
                # 為了反向傳播，我們需要重新計算一個任務的損失
                # 這裏我們使用第一個任務作為代表
                task = task_batch[0]
                X, y = task[0], task[1]
                initial_params, adapted_params = self.inner_update(task)
                
                # 重新計算損失用於梯度
                adapted_model = BaseLearner()
                self._set_flat_params(adapted_params.clone())
                adapted_model.load_state_dict(self.model.state_dict())
                
                y_pred = adapted_model(X)
                empirical_risk = nn.MSELoss()(y_pred, y)
                
                posterior_mean = adapted_params
                posterior_log_std = torch.log(torch.ones_like(posterior_mean) * 0.1)
                kl_divergence = self._compute_kl_divergence(posterior_mean, posterior_log_std)
                
                meta_loss = empirical_risk + self.lambda_reg * kl_divergence
                meta_loss.backward()
                self.optimizer.step()
                
                total_meta_loss += batch_meta_loss
                total_empirical_risk += batch_empirical_risk
                total_kl += batch_kl
                num_tasks += batch_size
            
            # 計算平均損失和邊界
            avg_meta_loss = total_meta_loss / len(meta_dataloader)
            avg_empirical_risk = total_empirical_risk / len(meta_dataloader)
            avg_kl = total_kl / len(meta_dataloader)
            
            # 計算PAC-Bayesian邊界
            pac_bound = self.compute_pac_bayes_bound(
                avg_empirical_risk, avg_kl, len(meta_dataloader.dataset)
            )
            
            bounds_history.append(pac_bound)
            empirical_risk_history.append(avg_empirical_risk)
            kl_history.append(avg_kl)
            
            if epoch % 10 == 0:
                print(f"Epoch {epoch}: Meta Loss = {avg_meta_loss:.4f}, "
                      f"Empirical Risk = {avg_empirical_risk:.4f}, "
                      f"KL = {avg_kl:.4f}, PAC-Bound = {pac_bound:.4f}")
        
        return bounds_history, empirical_risk_history, kl_history

# 創建任務分佈和數據集
task_dist = TaskDistribution(shift_magnitude=0.8)
source_dataset = MetaDataset(task_dist, num_tasks=50, samples_per_task=20, source=True)
target_dataset = MetaDataset(task_dist, num_tasks=20, samples_per_task=20, source=False)

source_dataloader = DataLoader(source_dataset, batch_size=5, shuffle=True)
target_dataloader = DataLoader(target_dataset, batch_size=5, shuffle=True)

# 初始化模型和PBMAML
model = BaseLearner()
pbmaml = PBMAML(model, lambda_reg=0.01)

# 在源任務上進行元訓練
print("在源任務上訓練PBMAML...")
bounds_history, empirical_history, kl_history = pbmaml.meta_train(source_dataloader, num_epochs=100)

# 評估在目標任務上的性能
def evaluate_on_target(model, target_dataloader):
    total_loss = 0
    total_tasks = 0
    
    for task_batch in target_dataloader:
        for task in task_batch:
            X, y, W, b = task
            with torch.no_grad():
                y_pred = model(X)
                loss = nn.MSELoss()(y_pred, y).item()
                total_loss += loss
                total_tasks += 1
    
    return total_loss / total_tasks

# 在目標任務上評估
target_loss = evaluate_on_target(model, target_dataloader)
print(f"在目標任務上的平均損失: {target_loss:.4f}")

# 可視化結果
plt.figure(figsize=(12, 4))

plt.subplot(1, 3, 1)
plt.plot(empirical_history, label='經驗風險')
plt.plot(bounds_history, label='PAC-Bayes邊界')
plt.xlabel('訓練輪次')
plt.ylabel('風險')
plt.legend()
plt.title('經驗風險與泛化邊界')

plt.subplot(1, 3, 2)
plt.plot(kl_history)
plt.xlabel('訓練輪次')
plt.ylabel('KL散度')
plt.title('KL散度變化')

# 比較不同分佈偏移程度下的性能
shift_magnitudes = [0.1, 0.3, 0.5, 0.8, 1.0]
performance = []

for shift in shift_magnitudes:
    task_dist = TaskDistribution(shift_magnitude=shift)
    source_dataset = MetaDataset(task_dist, num_tasks=50, samples_per_task=20, source=True)
    target_dataset = MetaDataset(task_dist, num_tasks=20, samples_per_task=20, source=False)
    
    source_dataloader = DataLoader(source_dataset, batch_size=5, shuffle=True)
    target_dataloader = DataLoader(target_dataset, batch_size=5, shuffle=True)
    
    model = BaseLearner()
    pbmaml = PBMAML(model, lambda_reg=0.01)
    
    # 訓練
    bounds_history, _, _ = pbmaml.meta_train(source_dataloader, num_epochs=50)
    
    # 評估
    target_loss = evaluate_on_target(model, target_dataloader)
    performance.append((shift, target_loss, bounds_history[-1]))

plt.subplot(1, 3, 3)
shifts, losses, bounds = zip(*performance)
plt.plot(shifts, losses, 'o-', label='實際風險')
plt.plot(shifts, bounds, 's-', label='PAC-Bayes邊界')
plt.xlabel('分佈偏移程度')
plt.ylabel('風險')
plt.legend()
plt.title('分佈偏移對性能的影響')

plt.tight_layout()
plt.show()

# 輸出分析結果
print("\n=== 分佈偏移分析 ===")
for shift, loss, bound in performance:
    generalization_gap = bound - loss
    print(f"偏移程度 {shift}: 實際風險 = {loss:.4f}, 邊界 = {bound:.4f}, 泛化間隙 = {generalization_gap:.4f}")