import os
import torch
from safetensors.torch import safe_open
import yaml

# --- CONFIGURATION ---
YAML_PATH = "B:/24B/C82/mergekit_config.yml"
FINAL_MERGE_DIR = "B:/24B/C82"
LAYERS_TO_SCAN = [
    "model.layers.30.mlp.down_proj.weight"
]
# ---------------------

def load_tensor(model_dir, tensor_name):
    """Finds and loads a tensor from a directory of safetensors."""
    for file in os.listdir(model_dir):
        if file.endswith(".safetensors"):
            filepath = os.path.join(model_dir, file)
            with safe_open(filepath, framework="pt", device="cpu") as f:
                if tensor_name in f.keys():
                    return f.get_tensor(tensor_name).float()
    raise ValueError(f"Tensor {tensor_name} not found in {model_dir}")

def get_latent_metrics(tv):
    """Calculates structural metrics of the task vector."""
    # 1. Global Energy (Frobenius Norm)
    f_norm = torch.linalg.norm(tv).item()
    
    # 2. Spectral Norm (Intensity of dominant feature)
    # Measures the largest singular value (L2 operator norm)
    s_norm = torch.linalg.norm(tv, ord=2).item()
    
    # 3. Stable Rank (Knowledge Breadth / Complexity)
    # Calculation: ||A||_F^2 / ||A||_2^2
    # Higher = more diverse features preserved; Lower = collapsed/overfit
    s_rank = (f_norm**2) / (s_norm**2) if s_norm > 0 else 0
    
    return f_norm, s_norm, s_rank

def main():
    print("Loading YAML config...")
    with open(YAML_PATH, 'r') as f:
        config = yaml.safe_load(f)
    
    base_path = config['base_model']
    donor_paths = [m['model'] for m in config['models']]
    
    print(f"\nScanning {len(LAYERS_TO_SCAN)} MLP layers for structural influence...\n")
    
    for layer in LAYERS_TO_SCAN:
        print(f"--- Layer: {layer} ---")
        try:
            # Load Base and Final
            base_w = load_tensor(base_path, layer)
            final_w = load_tensor(FINAL_MERGE_DIR, layer)
            final_tv = final_w - base_w
            
            # Calculate metrics for the Final Merged Model
            final_f, final_s, final_rk = get_latent_metrics(final_tv)
            
            print(f"Final Model Latent Metrics:")
            print(f"  Energy (Frobenius):   {final_f:.4f}")
            print(f"  Intensity (Spectral): {final_s:.4f}")
            print(f"  Complexity (S-Rank):  {final_rk:.2f}")
            print("")

            results = []
            for donor in donor_paths:
                donor_w = load_tensor(donor, layer)
                donor_tv = donor_w - base_w
                
                # Get structural metrics for donor
                d_f, d_s, d_rk = get_latent_metrics(donor_tv)
                
                # Calculate Cosine Similarity (Alignment)
                cos_sim = torch.nn.functional.cosine_similarity(
                    final_tv.flatten(), donor_tv.flatten(), dim=0
                ).item()
                
                name = donor.split("/")[-1][:40]
                results.append({
                    'name': name,
                    'sim': cos_sim,
                    'energy': d_f,
                    'intens': d_s,
                    'cmplx': d_rk
                })
            
            # Sort by highest alignment
            results.sort(key=lambda x: x['sim'], reverse=True)
            
            print(f"{'Donor Model':<42} | {'Align':<8} | {'Energy':<8} | {'Intens':<8} | {'Cmplx'}")
            print("-" * 90)
            for r in results:
                # Energy is relative to the final model (e.g. 0.90x)
                rel_energy = r['energy'] / final_f if final_f > 0 else 0
                print(f"{r['name']:<42} | {r['sim']:>7.3f} | {rel_energy:>7.2f}x | {r['intens']:>8.3f} | {r['cmplx']:>6.2f}")
            print("\n")
            
        except Exception as e:
            print(f"Skipping layer due to error: {e}")

if __name__ == "__main__":
    main()