"""
Shared helpers for model arithmetic (used by both arithmetic.py and arithmetic_torch.py).
"""
import json
from pathlib import Path

import numpy as np
from tqdm import tqdm


def mix_params(params_list, weights):
    """Weighted average of param dicts. Each param is dict[str, np.ndarray]. Returns dict[str, np.ndarray]."""
    weights = np.asarray(weights, dtype=np.float64)
    weights /= weights.sum()
    mixed = {}
    for key in tqdm(params_list[0].keys(), desc="Mixing parameters"):
        stacked = np.stack([np.asarray(p[key], dtype=np.float64) for p in params_list], axis=0)
        mixed[key] = np.average(stacked, axis=0, weights=weights).astype(np.float32)
    return mixed


def load_norm_stats(norm_stats_path: str) -> dict:
    """Load normalization statistics from JSON."""
    with open(norm_stats_path, "r") as f:
        data = json.load(f)
    if "norm_stats" not in data:
        raise ValueError(f"Invalid norm_stats format in {norm_stats_path}")
    return data["norm_stats"]


def mix_norm_stats(norm_stats_list: list, weights: list = None) -> dict:
    """Mix normalization statistics with optional weighting."""
    if len(norm_stats_list) == 1:
        return norm_stats_list[0]
    if weights is None:
        weights = [1.0 / len(norm_stats_list)] * len(norm_stats_list)
    else:
        weight_sum = sum(weights)
        weights = [w / weight_sum for w in weights]
    result = {}
    for key in norm_stats_list[0].keys():
        values = [ns[key] for ns in norm_stats_list]
        if isinstance(values[0], dict):
            result[key] = {}
            for stat_key in values[0].keys():
                arrays = [np.array(v[stat_key]) for v in values]
                stacked = np.stack(arrays, axis=0)
                weighted_avg = np.average(stacked, axis=0, weights=weights)
                result[key][stat_key] = weighted_avg.tolist()
        else:
            result[key] = values[0]
    return result


def save_norm_stats(norm_stats: dict, output_path: str) -> None:
    """Save normalization statistics to JSON."""
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    with open(output_path, "w") as f:
        json.dump({"norm_stats": norm_stats}, f, indent=2)


def compute_optimal_weights(losses):
    """Compute optimal weights based on inverse loss."""
    losses = np.array(losses)
    inv_losses = 1.0 / (losses + 1e-8)
    inv_losses = inv_losses ** 2
    weights = inv_losses / inv_losses.sum()
    return weights.tolist()