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

# --- CONFIGURATION ---
YAML_PATH = "B:/24B/karcher_stock_24b/mergekit_config.yml"
FINAL_MERGE_DIR = "B:/24B/karcher_stock_24b"
LAYERS_TO_SCAN =[
    "model.layers.10.mlp.down_proj.weight" # "model.language_model.layers.10.mlp.gate_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)
            
            # Use float64 for norm calculations to prevent precision loss in energy ratios
            final_norm = torch.norm(final_w.double()).item()
            final_tv = final_w - base_w
            final_tv_norm = torch.norm(final_tv.double()).item()
            
            results = []
            
            # 1. Collect raw magnitudes of the components
            base_norm = torch.norm(base_w.double()).item()
            donor_tvs = []
            for donor in donor_paths:
                dw = load_tensor(donor, layer)
                donor_tvs.append(dw - base_w)
            
            donor_tv_norms = [torch.norm(dtv.double()).item() for dtv in donor_tvs]
            
            # 2. Calculate Total Component Energy (Base + all Donor Deltas)
            total_component_energy = base_norm + sum(donor_tv_norms)
            
            results = []
            # 3. Assign Share to Base Model
            base_share = (base_norm / total_component_energy) * 100
            results.append(("(Base Model)", -1.0, 0.0, base_share))

            # 4. Assign Share to Donors
            for i, donor in enumerate(donor_paths):
                donor_tv = donor_tvs[i]
                
                cos_sim = torch.nn.functional.cosine_similarity(
                    final_tv.flatten(), donor_tv.flatten(), dim=0
                ).item()
                
                rel_mag = (torch.norm(donor_tv.double()).item() / final_tv_norm)
                
                # Compositional Share: How much of the total energy sum belongs to this donor's delta
                comp_share = (donor_tv_norms[i] / total_component_energy) * 100
                
                name = donor.split("/")[-1][:50]
                results.append((name, cos_sim, rel_mag, comp_share))
            
            # Sort by highest similarity (Donors first, Base at the very bottom)
            results.sort(key=lambda x: x[1], reverse=True)
            
            print(f"{'Model Name':<55} | {'Alignment {Cos}':<12} | {'Rel Mag (TV)':<12} | {'Merge Composition'}")
            print("-" * 105)
            for name, sim, mag, energy in results:
                sim_str = f"{sim:12.4f}" if sim >= 0 else "    N/A     "
                mag_str = f"{mag:11.2f}x" if mag > 0 else "    N/A    "
                print(f"{name:<55} | {sim_str} | {mag_str} | {energy:>13.2f}%")
            
        except Exception as e:
            print(f"Skipping layer due to error: {e}")

if __name__ == "__main__":
    main()