from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import torch

from ..models.model import MRPT, MRPTConfig


def save_model(model: MRPT, path: Path, extra: dict | None = None) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    torch.save({"cfg": model.cfg.__dict__, "model": model.state_dict(), "extra": extra or {}}, path)


def load_model(path: Path, map_location="cpu") -> MRPT:
    ck = torch.load(path, map_location=map_location, weights_only=False)
    m = MRPT(MRPTConfig(**ck["cfg"]))
    m.load_state_dict(ck["model"])
    return m


def save_predictions(res: dict[str, np.ndarray], path: Path, drop_samples: bool = True) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    keep = {k: v for k, v in res.items() if not (drop_samples and ("samples" in k or k == "paths"))}
    np.savez_compressed(path, **keep)


def write_json(obj, path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(obj, indent=1, default=lambda o: float(o) if isinstance(o, (np.floating,)) else
                               int(o) if isinstance(o, (np.integer,)) else o.tolist() if isinstance(o, np.ndarray) else str(o)))