Module neuralops.utils

def compute_metrics_for_task(
    predictions: np.ndarray,
    labels: np.ndarray,
    task_name: str,
    is_regression: bool = False,
) -> dict[str, float]:
    """Compute and return metrics for task."""
    metrics = {}
    if is_regression or task_name == "stsb":
        pred_flat = np.squeeze(predictions)
        labels_flat = np.squeeze(labels)
        metrics["pearson"] = compute_pearson(pred_flat, labels_flat)
        metrics["spearman"] = compute_spearman(pred_flat, labels_flat)
        metrics["pearson_spearman"] = (metrics["pearson"] + metrics["spearman"]) / 2
    else:
        preds = np.argmax(predictions, axis=-1) if predictions.ndim > 1 else predictions
        metrics["accuracy"] = compute_accuracy(preds, labels)
        metrics["f1_macro"] = compute_f1(preds, labels, average="macro")
        metrics["f1_weighted"] = compute_f1(preds, labels, average="weighted")
        if task_name == "cola":
            metrics["matthews_correlation"] = compute_matthews_corrcoef(preds, labels)
        if task_name in ("mrpc", "qqp"):
            metrics["f1"] = compute_f1(preds, labels, average="binary")
    return metrics

def count_parameters(model: torch.nn.Module, trainable_only: bool = True) -> int:
    """Count parameters."""
    if trainable_only:
        return sum(p.numel() for p in model.parameters() if p.requires_grad)
    return sum(p.numel() for p in model.parameters())

def format_time(seconds: float) -> str:
    """Format time."""
    hours = int(seconds // 3600)
    minutes = int((seconds % 3600) // 60)
    secs = int(seconds % 60)
    parts = []
    if hours:
        parts.append(f"{hours}h")
    if minutes:
        parts.append(f"{minutes}m")
    parts.append(f"{secs}s")
    return " ".join(parts)

def get_compression_targets(model_name: str, include_ffn: bool = True) -> list[str]:
    """Return the compression targets."""
    arch = detect_architecture(model_name)
    patterns = get_layer_patterns(arch)
    if not include_ffn:
        ffn_patterns = ["ffn", "intermediate", "mlp", "ff", "feedforward"]
        patterns = [p for p in patterns if not any(fp in p for fp in ffn_patterns)]
    return patterns

def get_hidden_size(model) -> int:
    """Return the hidden size."""
    if hasattr(model, "config") and hasattr(model.config, "hidden_size"):
        return model.config.hidden_size
    return 768

def get_layer_count(model) -> int:
    """Return the layer count."""
    if hasattr(model, "config") and hasattr(model.config, "num_hidden_layers"):
        return model.config.num_hidden_layers
    if hasattr(model, "encoder") and hasattr(model.encoder, "layer"):
        return len(model.encoder.layer)
    if hasattr(model, "transformer") and hasattr(model.transformer, "layer"):
        return len(model.transformer.layer)
    return 12

def get_logger(name: str = "neuralops") -> logging.Logger:
    """Return the logger."""
    return logging.getLogger(name)

def get_model_info(model_name: str) -> dict[str, Any]:
    """Return the model info."""
    arch = detect_architecture(model_name)
    default_configs = {
        "bert": {"hidden_size": 768, "num_layers": 12, "num_heads": 12},
        "distilbert": {"hidden_size": 768, "num_layers": 6, "num_heads": 12},
        "roberta": {"hidden_size": 768, "num_layers": 12, "num_heads": 12},
        "albert": {"hidden_size": 768, "num_layers": 12, "num_heads": 12},
        "electra": {"hidden_size": 768, "num_layers": 12, "num_heads": 12},
        "deberta": {"hidden_size": 768, "num_layers": 12, "num_heads": 12},
        "gpt2": {"hidden_size": 768, "num_layers": 12, "num_heads": 12},
        "xlnet": {"hidden_size": 768, "num_layers": 12, "num_heads": 12},
    }
    if "large" in model_name.lower():
        default_configs["bert"]["hidden_size"] = 1024
        default_configs["roberta"]["hidden_size"] = 1024
    return {
        "architecture": arch,
        "patterns": get_layer_patterns(arch),
        "estimated_config": default_configs.get(arch, default_configs["bert"]),
    }

def init_output_scale(module: nn.Module, name: str = "output_scale") -> None:
    """Register output_scale learnable parameter if not already present.

    Args:
        module: The target PyTorch module.
        name: Parameter attribute name to register.
    """
    if not hasattr(module, name):
        module.register_parameter(name, nn.Parameter(torch.ones(1)))

def is_architecture_error(exc: BaseException) -> bool:
    text = str(exc).lower()
    return any(
        token in text
        for token in (
            "shape",
            "size mismatch",
            "unexpected keyword",
            "unexpected argument",
            "has no attribute",
            "not enough values to unpack",
            "too many values to unpack",
            "values to unpack",
            "attention mask",
            "position_embeddings",
            "mat1 and mat2",
            "cannot be multiplied",
            "unsupported operand type",
            "tensor' and 'tuple'",
        )
    )

def is_qwen_attention(module: nn.Module) -> bool:
    """Detect a native Qwen attention block without importing transformers."""
    cfg = getattr(module, "config", None)
    if cfg is not None and hasattr(cfg, "text_config"):
        cfg = cfg.text_config
    return str(getattr(cfg, "model_type", "")).lower().startswith("qwen") or (
        hasattr(module, "q_proj")
        and hasattr(module, "o_proj")
        and ("qwen" in module.__class__.__name__.lower() or hasattr(module, "rotary_emb"))
    )

def load_json(path: str | Path) -> Any:
    """Load json from disk or config."""
    with open(path) as f:
        return json.load(f)

def measure_inference_latency(
    model, dataloader, device, num_samples: int = 100, warmup: int = 10
) -> dict[str, float]:
    """Measure inference latency."""
    model.eval()
    model.to(device)
    latencies = []
    with torch.no_grad():
        for i, batch in enumerate(dataloader):
            if i >= warmup + num_samples:
                break
            batch = {k: v.to(device) for k, v in batch.items() if isinstance(v, torch.Tensor)}
            if torch.cuda.is_available():
                torch.cuda.synchronize()
            start = time.perf_counter()
            _ = model(**batch)
            if torch.cuda.is_available():
                torch.cuda.synchronize()
            end = time.perf_counter()
            if i >= warmup:
                latencies.append((end - start) * 1000)
    if not latencies:
        return {}
    return {
        "mean_latency_ms": np.mean(latencies),
        "std_latency_ms": np.std(latencies),
        "min_latency_ms": np.min(latencies),
        "max_latency_ms": np.max(latencies),
        "p50_latency_ms": np.percentile(latencies, 50),
        "p95_latency_ms": np.percentile(latencies, 95),
        "p99_latency_ms": np.percentile(latencies, 99),
        "samples_per_second": (1000 / np.mean(latencies) if np.mean(latencies) > 0 else 0),
    }

def measure_memory_usage(model, dataloader, device, num_samples: int = 10) -> dict[str, float]:
    """Measure memory usage."""
    model.eval()
    model.to(device)
    peak_memory = []
    with torch.no_grad():
        for i, batch in enumerate(dataloader):
            if i >= num_samples:
                break
            if torch.cuda.is_available():
                torch.cuda.reset_peak_memory_stats()
            batch = {k: v.to(device) for k, v in batch.items() if isinstance(v, torch.Tensor)}
            _ = model(**batch)
            if torch.cuda.is_available():
                peak_memory.append(torch.cuda.max_memory_allocated() / 1024 / 1024)
    return {
        "peak_memory_mb": np.max(peak_memory) if peak_memory else 0,
        "avg_peak_memory_mb": np.mean(peak_memory) if peak_memory else 0,
    }

def merge_configs(base: dict[str, Any], override: dict[str, Any]) -> dict[str, Any]:
    """Merge configs."""
    result = base.copy()
    return update_config(result, override)

@contextmanager
def profile_context(model, profiler: ModelProfiler | None = None) -> None:
    """Profile context."""
    if profiler is None:
        profiler = ModelProfiler(model)
    profiler.start_profiling()
    yield profiler
    profiler.stop_profiling()

def qwen_config_proxy(model_or_config: Any) -> QwenConfigProxy:
    config = getattr(model_or_config, "config", model_or_config)
    if hasattr(config, "text_config"):
        config = config.text_config
    hidden = int(config.hidden_size)
    heads = int(config.num_attention_heads)
    return QwenConfigProxy(
        source=config,
        hidden_size=hidden,
        num_attention_heads=heads,
        head_dim=int(getattr(config, "head_dim", hidden // heads)),
        num_hidden_layers=int(getattr(config, "num_hidden_layers", 0)),
    )

def qwen_module_policy(module_name: str) -> str:
    """Classify modules whose semantics cannot be promised for causal RoPE Qwen."""
    name = module_name.lower()
    explicit_incompatible = {
        "aegis_v1",
        "aegis_v3",
        "aegis_v4",
        "aegis_v5",
        "aegis_v6",
        "bam_umoe",
        "cosformer",
        "cur",
        "cur_moe",
        "cur_moe_compressor",
        "decay_linear_attn",
        "dobi",
        "dynamic_rank",
        # These legacy replacements currently implement full-MHA, padding
        # masks, or absolute-position assumptions.  Until their forward pass
        # explicitly handles Qwen GQA + RoPE + causal masks, skipping is safer
        # than reporting a numerically valid but semantically invalid result.
        "a3",
        "aasvd",
        "adaptive_svd",
        "asvd",
        "cluster",
        "cluster_attention",
        "goldfinch_ahn",
        "griffin_ssd",
        "hgrn2_layer",
        "hgrn2_outer",
        "fusion_attn",
        "aegis_v2",
        "adaptive_svd_attention",
        "huffman",
        "block_transformer_hiba",
        "chela_slay",
        "fasa",
        "gated_deltanet_block",
        "granite_hybrid",
        "hgrn2_gated_deltanet",
    }
    if name in explicit_incompatible:
        return "incompatible"
    incompatible = (
        "linear",
        "rnn",
        "ssm",
        "mamba",
        "performer",
        "cosformer",
        "random",
        "recurrent",
        "delta_net",
        "gated_linear",
        "span",
    )
    return "incompatible" if any(token in name for token in incompatible) else "adapted"

def qwen_projection(module: nn.Module, legacy: str, native: str) -> nn.Module | None:
    """Return a projection under either its legacy or native spelling."""
    value = getattr(module, native, None)
    if value is None:
        value = getattr(module, legacy, None)
    return value

def resolve_path(relative_path: str) -> str:
    """Resolves a path relative to the project root and returns as string."""
    return str((get_project_root() / relative_path).resolve())

def save_json(data: Any, path: str | Path, indent: int = 2) -> None:
    """Persist json to disk."""
    ensure_dir(os.path.dirname(path))
    with open(path, "w") as f:
        json.dump(data, f, indent=indent)

def set_log_level(logger_name: str, level: int) -> None:
    """Set the log level."""
    logging.getLogger(logger_name).setLevel(level)

def set_seed(seed: int, deterministic: bool = False) -> None:
    """Set the seed."""
    set_all_seeds(seed)
    if deterministic and torch.cuda.is_available():
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False

def unflatten_dict(d: dict, sep: str = ".") -> dict:
    """Unflatten dict."""
    result = {}
    for key, value in d.items():
        parts = key.split(sep)
        current = result
        for part in parts[:-1]:
            if part not in current:
                current[part] = {}
            current = current[part]
        current[parts[-1]] = value
    return result

class CheckpointManager:
    def __init__(self, checkpoint_dir: str = "./checkpoints") -> None:
        """Execute __init__ operation."""
        self.checkpoint_dir = Path(checkpoint_dir)
        self.checkpoint_dir.mkdir(parents=True, exist_ok=True)

    def save_checkpoint(
        self,
        experiment_id: int,
        results: list[dict],
        checkpoint_name: str = "checkpoint.json",
        metadata: dict | None = None,
    ) -> Path:
        """Persist checkpoint to disk."""
        checkpoint_path = self.checkpoint_dir / checkpoint_name
        data = {
            "experiment_id": experiment_id,
            "results": results,
            "timestamp": time.time(),
            "datetime": time.strftime("%Y-%m-%d %H:%M:%S"),
        }
        if metadata:
            data["metadata"] = metadata
        with open(checkpoint_path, "w") as f:
            json.dump(data, f, indent=2)
        logging.info(f"Checkpoint saved: {checkpoint_path}")
        return checkpoint_path

    def load_checkpoint(
        self, checkpoint_name: str = "checkpoint.json"
    ) -> tuple[int, list[dict], dict | None]:
        """Load checkpoint from disk or config."""
        checkpoint_path = self.checkpoint_dir / checkpoint_name
        if not checkpoint_path.exists():
            logging.info(f"No checkpoint found at {checkpoint_path}")
            return 0, [], None
        with open(checkpoint_path) as f:
            data = json.load(f)
        logging.info(f"Checkpoint loaded: {checkpoint_path}")
        return (
            data.get("experiment_id", 0),
            data.get("results", []),
            data.get("metadata"),
        )

    def delete_checkpoint(self, checkpoint_name: str = "checkpoint.json") -> None:
        """Delete checkpoint."""
        checkpoint_path = self.checkpoint_dir / checkpoint_name
        if checkpoint_path.exists():
            checkpoint_path.unlink()
            logging.info(f"Deleted checkpoint: {checkpoint_path}")

    def list_checkpoints(self) -> list[Path]:
        """List checkpoints."""
        return list(self.checkpoint_dir.glob("*.json"))

    def get_latest_checkpoint(self) -> Path | None:
        """Return the latest checkpoint."""
        checkpoints = self.list_checkpoints()
        return max(checkpoints, key=lambda p: p.stat().st_mtime) if checkpoints else None

class ResultsSaver(PrimaryResultsSaver):
    """
    Enhanced ResultsSaver compatible with Django API, experiment runners, and CLI scripts.
    Supports structured report generation and Hugging Face Hub uploads.
    """

    def __init__(self, run_id: str, base_dir: Path, config: dict[str, Any] | None = None) -> None:
        super().__init__(run_id=run_id, base_dir=base_dir, config=config)
        self.config = config or {}
        self.legacy_base_dir = Path(base_dir) / str(run_id)
        self.legacy_base_dir.mkdir(parents=True, exist_ok=True)

    def save_config(self, config: dict[str, Any]) -> None:
        """Persist config to disk in structured format."""
        self.config = config
        self._save_config(config)
        self._write_json("../config.json", config)
        logger.info(f"Saved structured config to {self.run_dir}")

    def save_history(self, history: dict[str, Any]) -> None:
        """Persist training history to disk."""
        self._save_history(history)
        self._write_json("../history.json", history)
        logger.info(f"Saved training history to {self.run_dir}")

    def save_final_metrics(self, metrics: dict[str, Any]) -> None:
        """Persist final metrics and generate run report."""
        self._save_final_metrics(metrics)
        self._write_json("../final_metrics.json", metrics)
        self._generate_run_report(metrics)
        logger.info(f"Saved final metrics and generated run report to {self.run_dir}")

    def generate_plots(self, history: dict[str, Any]) -> None:
        """Generate plots."""
        self._generate_plots(history, {})

    def should_save_model(self, metrics: dict[str, Any]) -> bool:
        """Evaluate whether heavy model artifacts should be saved."""
        return self.should_save_heavy(metrics)

    def handle_storage_backend(self) -> None:
        """Execute configured storage backends including HuggingFace Hub."""
        self.push_all()

    def push_to_hf(self, repo_id: str | None = None) -> None:
        """Upload complete run artifacts to HuggingFace Hub for resource management."""
        cfg = {"hf_options": {"repo_id": repo_id}} if repo_id else self.cfg
        hf_backend.push(self.run_id, self.run_dir, cfg)

    def _generate_run_report(self, metrics: dict[str, Any]) -> None:
        """Generate structured run summary report."""
        report = {
            "run_id": self.run_id,
            "status": "COMPLETED",
            "metrics": metrics,
            "config": self.config
            or (self.cfg.to_dict() if hasattr(self.cfg, "to_dict") else str(self.cfg)),
        }
        self._write_json("artifacts/run_report.json", report)
        self._save_model_card(self.config or {}, metrics)