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

# --- CONFIGURATION ---
YAML_PATH = "B:/24B/qliphoth2e/mergekit_config.yml"
FINAL_MERGE_DIR = "B:/24B/qliphoth2e"
LAYERS_TO_SCAN =[
    "model.layers.10.mlp.up_proj.weight",
    "model.layers.20.mlp.gate_proj.weight",
    "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 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:
            base_w = load_tensor(base_path, layer)
            final_w = load_tensor(FINAL_MERGE_DIR, layer)
            final_tv = final_w - base_w
            
            results =[]
            for donor in donor_paths:
                donor_w = load_tensor(donor, layer)
                donor_tv = donor_w - base_w
                
                # Calculate Cosine Similarity (How much does the final model align with this donor?)
                cos_sim = torch.nn.functional.cosine_similarity(
                    final_tv.flatten(), donor_tv.flatten(), dim=0
                ).item()
                
                # Calculate Relative Magnitude
                rel_mag = (donor_tv.norm() / final_tv.norm()).item()
                
                name = donor.split("/")[-1][:50]
                results.append((name, cos_sim, rel_mag))
            
            # Sort by highest similarity
            results.sort(key=lambda x: x[1], reverse=True)
            
            print(f"{'Donor Model':<55} | {'Alignment (Cos Sim)':<20} | {'Relative Mag'}")
            print("-" * 95)
            for name, sim, mag in results:
                print(f"{name:<55} | {sim:>18.4f} | {mag:>10.2f}x")
            print("\n")
            
        except Exception as e:
            print(f"Skipping layer due to error: {e}")

if __name__ == "__main__":
    main()