import torch
import torch.nn as nn
import numpy as np
from scipy import stats
import matplotlib.pyplot as plt

class SOCAnalyzer:
    def __init__(self, model, layer_names):
        self.model = model
        self.layer_names = layer_names
        self.gradient_history = []
        
    def compute_gradient_covariance(self, dataloader, loss_fn):
        """計算梯度協方差矩陣"""
        gradients = []
        
        # 收集梯度樣本
        for batch_idx, (data, target) in enumerate(dataloader):
            if batch_idx >= 100:  # 限制樣本數量
                break
                
            self.model.zero_grad()
            output = self.model(data)
            loss = loss_fn(output, target)
            loss.backward()
            
            # 拼接所有參數梯度
            layer_grads = []
            for name, param in self.model.named_parameters():
                if any(layer_name in name for layer_name in self.layer_names):
                    if param.grad is not None:
                        layer_grads.append(param.grad.view(-1))
            
            if layer_grads:
                full_grad = torch.cat(layer_grads)
                gradients.append(full_grad.cpu().numpy())
        
        # 計算協方差矩陣的特徵值
        gradients = np.array(gradients)
        cov_matrix = np.cov(gradients, rowvar=False)
        eigenvalues = np.linalg.eigvalsh(cov_matrix)
        eigenvalues = eigenvalues[eigenvalues > 1e-10]  # 過濾極小值
        
        return eigenvalues
    
    def analyze_power_law(self, eigenvalues):
        """分析特徵值的冪律分佈"""
        # 排序特徵值
        sorted_evals = np.sort(eigenvalues)[::-1]
        
        # 計算累積分佈
        ccdf = 1 - np.arange(len(sorted_evals)) / len(sorted_evals)
        
        # 冪律擬合
        mask = sorted_evals > sorted_evals.max() * 0.01
        x = np.log(sorted_evals[mask])
        y = np.log(ccdf[mask])
        
        slope, intercept, r_value, _, _ = stats.linregress(x, y)
        
        return {
            'exponent': -slope,
            'r_squared': r_value**2,
            'eigenvalues': sorted_evals,
            'ccdf': ccdf
        }

# 使用示例
if __name__ == "__main__":
    # 創建簡易Transformer層進行分析
    class MiniTransformer(nn.Module):
        def __init__(self, d_model=512, nhead=8):
            super().__init__()
            self.attention = nn.MultiheadAttention(d_model, nhead)
            self.ffn = nn.Sequential(
                nn.Linear(d_model, d_model*4),
                nn.ReLU(),
                nn.Linear(d_model*4, d_model)
            )
            self.norm1 = nn.LayerNorm(d_model)
            self.norm2 = nn.LayerNorm(d_model)
            
        def forward(self, x):
            attn_out, _ = self.attention(x, x, x)
            x = self.norm1(x + attn_out)
            ffn_out = self.ffn(x)
            x = self.norm2(x + ffn_out)
            return x
    
    model = MiniTransformer()
    analyzer = SOCAnalyzer(model, ['attention', 'ffn'])
    
    # 模擬訓練過程並分析臨界狀態演化
    training_stages = []
    for epoch in range(5):
        # 模擬訓練（此處簡化）
        eigenvalues = analyzer.compute_gradient_covariance(None, None)
        result = analyzer.analyze_power_law(eigenvalues)
        training_stages.append(result)
        
        print(f"Epoch {epoch}: Power law exponent = {result['exponent']:.3f}, "
              f"R² = {result['r_squared']:.3f}")

class ActivationDynamics:
    def __init__(self, model):
        self.model = model
        self.activation_hooks = []
        
    def setup_activation_recording(self):
        """設置鈎子記錄激活值"""
        activations = {}
        
        def get_hook(name):
            def hook(module, input, output):
                activations[name] = output.detach()
            return hook
        
        for name, module in self.model.named_modules():
            if isinstance(module, (nn.Linear, nn.MultiheadAttention)):
                hook = module.register_forward_hook(get_hook(name))
                self.activation_hooks.append(hook)
        
        return activations
    
    def compute_avalanche_statistics(self, activations_sequence):
        """
        計算激活值雪崩統計
        基於閾值方法檢測激活級聯
        """
        avalanches = []
        current_avalanche = 0
        
        for t, act_dict in enumerate(activations_sequence):
            total_activation = 0
            for name, activation in act_dict.items():
                # 計算超過閾值的激活比例
                threshold = activation.abs().mean()
                above_threshold = (activation.abs() > threshold).float().mean()
                total_activation += above_threshold.item()
            
            # 檢測雪崩開始/結束
            if total_activation > 0.1:  # 激活閾值
                current_avalanche += 1
            else:
                if current_avalanche > 0:
                    avalanches.append(current_avalanche)
                    current_avalanche = 0
        
        return avalanches
    
    def analyze_avalanche_distribution(self, avalanches):
        """分析雪崩規模的分佈"""
        from collections import Counter
        
        if len(avalanches) == 0:
            return None
        
        # 計算分佈
        counter = Counter(avalanches)
        sizes = np.array(list(counter.keys()))
        counts = np.array(list(counter.values()))
        frequencies = counts / counts.sum()
        
        # 冪律擬合
        log_sizes = np.log(sizes)
        log_freq = np.log(frequencies)
        
        mask = sizes > 1
        if mask.sum() > 2:
            slope, intercept, r_value, _, _ = stats.linregress(
                log_sizes[mask], log_freq[mask]
            )
            
            return {
                'sizes': sizes,
                'frequencies': frequencies,
                'exponent': -slope,
                'r_squared': r_value**2
            }
        
        return None

import torch
from torch.utils.data import DataLoader, TensorDataset
from tqdm import tqdm

class CriticalityMonitor:
    def __init__(self, model, optimizer):
        self.model = model
        self.optimizer = optimizer
        self.metrics_history = {
            'loss': [],
            'gradient_norm': [],
            'weight_updates': [],
            'correlation_length': []
        }
    
    def compute_correlation_length(self, activations):
        """
        通過激活值的空間相關性估計關聯長度
        使用指數衰減擬合：C(r) ~ exp(-r/ξ)
        """
        if len(activations.shape) < 3:
            return 0.0
        
        # 計算空間相關性
        C_r = []
        distances = range(1, min(10, activations.shape[1]))
        
        for r in distances:
            # 計算距離為r的神經元間的相關性
            corr_sum = 0
            count = 0
            
            for i in range(activations.shape[1] - r):
                corr = torch.corrcoef(torch.stack([
                    activations[:, i].flatten(),
                    activations[:, i+r].flatten()
                ]))[0, 1]
                if not torch.isnan(corr):
                    corr_sum += corr.abs()
                    count += 1
            
            if count > 0:
                C_r.append(corr_sum.item() / count)
        
        # 指數擬合求關聯長度ξ
        if len(C_r) > 2:
            x = np.array(distances[:len(C_r)])
            y = np.log(np.array(C_r) + 1e-10)
            
            try:
                slope, intercept = np.polyfit(x, y, 1)
                correlation_length = -1/slope if slope < 0 else 0.0
                return max(0.0, correlation_length)
            except:
                return 0.0
        
        return 0.0
    
    def monitor_training(self, train_loader, epochs=10):
        """監控訓練過程中的臨界指標"""
        
        for epoch in range(epochs):
            epoch_loss = 0
            epoch_grad_norm = 0
            epoch_updates = 0
            epoch_corr_len = 0
            batch_count = 0
            
            pbar = tqdm(train_loader, desc=f'Epoch {epoch+1}/{epochs}')
            for batch_idx, (data, target) in enumerate(pbar):
                # 前向傳播
                self.optimizer.zero_grad()
                output = self.model(data)
                loss = nn.CrossEntropyLoss()(output, target)
                
                # 反向傳播
                loss.backward()
                
                # 記錄梯度範數
                total_norm = 0
                for p in self.model.parameters():
                    if p.grad is not None:
                        param_norm = p.grad.data.norm(2)
                        total_norm += param_norm.item() ** 2
                grad_norm = total_norm ** 0.5
                
                # 優化步驟
                self.optimizer.step()
                
                # 記錄參數更新量
                update_norm = 0
                for p in self.model.parameters():
                    if hasattr(p, 'old_data'):
                        update_norm += (p.data - p.old_data).norm(2).item() ** 2
                    p.old_data = p.data.clone()
                update_norm = update_norm ** 0.5
                
                # 計算激活值相關性
                with torch.no_grad():
                    # 獲取中間層激活
                    activations = []
                    def hook(module, input, output):
                        activations.append(output.detach())
                    
                    hook_handles = []
                    for name, module in self.model.named_modules():
                        if isinstance(module, nn.Linear) and 'ffn' in name:
                            handle = module.register_forward_hook(hook)
                            hook_handles.append(handle)
                    
                    _ = self.model(data[:1])  # 小樣本計算相關性
                    corr_len = 0
                    if activations:
                        corr_len = self.compute_correlation_length(activations[0])
                    
                    for handle in hook_handles:
                        handle.remove()
                
                # 更新統計
                epoch_loss += loss.item()
                epoch_grad_norm += grad_norm
                epoch_updates += update_norm
                epoch_corr_len += corr_len
                batch_count += 1
                
                # 更新進度條
                pbar.set_postfix({
                    'loss': loss.item(),
                    'grad_norm': grad_norm,
                    'corr_len': corr_len
                })
            
            # 記錄平均指標
            self.metrics_history['loss'].append(epoch_loss / batch_count)
            self.metrics_history['gradient_norm'].append(epoch_grad_norm / batch_count)
            self.metrics_history['weight_updates'].append(epoch_updates / batch_count)
            self.metrics_history['correlation_length'].append(epoch_corr_len / batch_count)
            
            # 檢測相變點
            if len(self.metrics_history['correlation_length']) >= 3:
                recent_corr = self.metrics_history['correlation_length'][-3:]
                if recent_corr[-1] > 2 * np.mean(recent_corr[:-1]):
                    print(f"\n⚠️  檢測到可能的相變點在 Epoch {epoch+1}")
                    print(f"   關聯長度突增: {recent_corr[-2]:.3f} -> {recent_corr[-1]:.3f}")
        
        return self.metrics_history

def finite_size_scaling_analysis(model_sizes, critical_exponents):
    """
    有限尺度標度分析
    model_sizes: 不同規模的模型列表
    critical_exponents: 臨界指數測量結果
    """
    
    plt.figure(figsize=(12, 4))
    
    # 標度函數擬合
    plt.subplot(131)
    for size, exponents in zip(model_sizes, critical_exponents):
        sizes = np.array([size] * len(exponents))
        plt.scatter(sizes, exponents, alpha=0.6, label=f'Size {size}')
    
    plt.xscale('log')
    plt.yscale('log')
    plt.xlabel('Model Size')
    plt.ylabel('Critical Exponent')
    plt.title('Finite Size Scaling')
    plt.legend()
    
    # 數據整理和擬合
    all_sizes = []
    all_exponents = []
    for size, exponents in zip(model_sizes, critical_exponents):
        all_sizes.extend([size] * len(exponents))
        all_exponents.extend(exponents)
    
    all_sizes = np.array(all_sizes)
    all_exponents = np.array(all_exponents)
    
    # 擬合標度關係: exponent ~ size^(-ν)
    log_sizes = np.log(all_sizes)
    log_exponents = np.log(all_exponents)
    
    mask = ~np.isinf(log_exponents) & ~np.isnan(log_exponents)
    if mask.sum() > 2:
        slope, intercept, r_value, _, _ = stats.linregress(
            log_sizes[mask], log_exponents[mask]
        )
        
        plt.subplot(132)
        plt.scatter(log_sizes[mask], log_exponents[mask], alpha=0.6)
        x_fit = np.linspace(log_sizes[mask].min(), log_sizes[mask].max(), 100)
        y_fit = slope * x_fit + intercept
        plt.plot(x_fit, y_fit, 'r--', 
                label=f'ν = {-slope:.3f}\nR² = {r_value**2:.3f}')
        plt.xlabel('log(Model Size)')
        plt.ylabel('log(Critical Exponent)')
        plt.title(f'Scaling Relation Fit')
        plt.legend()
    
    # 數據 collapse 分析
    plt.subplot(133)
    
    # 嘗試數據 collapse
    for size, exponents in zip(model_sizes, critical_exponents):
        scaled_exponents = exponents * (size ** 0.25)  # 假設 ν = 0.25
        hist, bins = np.histogram(scaled_exponents, bins=20, density=True)
        bin_centers = (bins[:-1] + bins[1:]) / 2
        plt.plot(bin_centers, hist, 'o-', label=f'Size {size}', alpha=0.7)
    
    plt.xlabel('Scaled Exponent')
    plt.ylabel('Probability Density')
    plt.title('Data Collapse Test')
    plt.legend()
    
    plt.tight_layout()
    plt.show()
    
    return {
        'scaling_exponent': -slope if 'slope' in locals() else None,
        'r_squared': r_value**2 if 'r_value' in locals() else None
    }

class NeuralSandpile:
    """
    神經網絡參數的沙堆模型模擬
    將參數更新類比為沙粒添加
    """
    
    def __init__(self, num_params, critical_threshold=4):
        self.num_params = num_params
        self.threshold = critical_threshold
        self.state = np.zeros(num_params)
        self.avalanche_sizes = []
        
    def add_grain(self, position=None):
        """添加一粒沙子（模擬梯度更新）"""
        if position is None:
            position = np.random.randint(0, self.num_params)
        
        self.state[position] += 1
        avalanche_size = 0
        
        # 檢查是否觸發崩塌
        topples = np.where(self.state >= self.threshold)[0]
        
        while len(topples) > 0:
            avalanche_size += len(topples)
            
            for pos in topples:
                # 沙粒向鄰居傳遞
                self.state[pos] -= self.threshold
                if pos > 0:
                    self.state[pos-1] += 1
                if pos < self.num_params - 1:
                    self.state[pos+1] += 1
                if pos >= 10:  # 跨層連接
                    self.state[pos-10] += 0.3
                if pos < self.num_params - 10:
                    self.state[pos+10] += 0.3
            
            # 檢查新的崩塌點
            topples = np.where(self.state >= self.threshold)[0]
        
        if avalanche_size > 0:
            self.avalanche_sizes.append(avalanche_size)
        
        return avalanche_size
    
    def train(self, num_steps=10000):
        """模擬訓練過程"""
        print("模擬神經網絡沙堆訓練...")
        
        for step in range(num_steps):
            # 模擬梯度更新：90%小更新，10%大更新
            if np.random.random() < 0.9:
                # 小批量更新：添加少量沙粒
                for _ in range(3):
                    self.add_grain()
            else:
                # 大批量更新：添加較多沙粒
                for _ in range(10):
                    pos = np.random.randint(0, self.num_params)
                    self.add_grain(pos)
            
            # 定期記錄狀態
            if step % 1000 == 0:
                self._analyze_state(step)
    
    def _analyze_state(self, step):
        """分析當前狀態"""
        if len(self.avalanche_sizes) < 10:
            return
        
        sizes = np.array(self.avalanche_sizes[-1000:])
        
        # 計算冪律指數
        unique_sizes, counts = np.unique(sizes, return_counts=True)
        if len(unique_sizes) > 3:
            mask = counts > 0
            x = np.log(unique_sizes[mask])
            y = np.log(counts[mask] / counts[mask].sum())
            
            if len(x) > 2:
                slope, intercept, r_value, _, _ = stats.linregress(x, y)
                
                print(f"Step {step}: "
                      f"Avg avalanche size = {sizes.mean():.2f}, "
                      f"Power law exponent = {-slope:.3f}, "
                      f"R² = {r_value**2:.3f}")
                
                # 檢查是否達到臨界狀態
                if 1.0 < -slope < 3.0 and r_value**2 > 0.8:
                    print(f"  → 已達到自組織臨界狀態!")

def predict_emergence_point(training_metrics, window_size=100):
    """
    基於訓練指標預測能力涌現點
    使用突變檢測算法
    """
    
    # 提取關鍵指標
    losses = training_metrics['loss']
    grad_norms = training_metrics['gradient_norm']
    corr_lengths = training_metrics['correlation_length']
    
    # 計算指標的變化率
    loss_changes = np.abs(np.diff(losses))
    grad_changes = np.abs(np.diff(grad_norms))
    corr_changes = np.abs(np.diff(corr_lengths))
    
    # 綜合突變分數
    combined_score = (
        0.4 * loss_changes / loss_changes.max() +
        0.3 * grad_changes / grad_changes.max() +
        0.3 * corr_changes / corr_changes.max()
    )
    
    # 滑動窗口檢測突變
    emergence_points = []
    for i in range(window_size, len(combined_score)):
        window_mean = np.mean(combined_score[i-window_size:i-10])
        current_mean = np.mean(combined_score[i-10:i])
        
        # 檢測顯著變化
        if current_mean > 2 * window_mean and current_mean > 0.5:
            emergence_points.append(i)
    
    # 分析突變模式
    if len(emergence_points) > 0:
        print(f"檢測到 {len(emergence_points)} 個可能的涌現點")
        
        # 聚類連續的涌現點
        clusters = []
        current_cluster = [emergence_points[0]]
        
        for i in range(1, len(emergence_points)):
            if emergence_points[i] - emergence_points[i-1] <= window_size:
                current_cluster.append(emergence_points[i])
            else:
                clusters.append(current_cluster)
                current_cluster = [emergence_points[i]]
        
        if current_cluster:
            clusters.append(current_cluster)
        
        # 輸出主要涌現區域
        for j, cluster in enumerate(clusters):
            avg_point = int(np.mean(cluster))
            print(f"涌現區域 {j+1}: 步驟 {avg_point} 附近")
    
    return {
        'emergence_points': emergence_points,
        'combined_score': combined_score,
        'clusters': clusters if 'clusters' in locals() else []
    }