# File-IO helpers shared across the codebase: pickle + numpy-aware JSON.

import json
import os
import pickle

import numpy as np


def read_pickle(filename):
    '''Load a pickle file. Appends .pkl extension if missing.'''
    if not filename.endswith('.pkl'):
        filename += '.pkl'
    with open(filename, 'rb') as f:
        return pickle.load(f)


class JsonNumpyEncoder(json.JSONEncoder):
    '''JSON encoder that handles numpy scalar/array types.'''
    def default(self, obj):
        if isinstance(obj, np.integer):
            return int(obj)
        if isinstance(obj, np.floating):
            return float(obj)
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        return super().default(obj)


def save_json(data, fp, indent=4):
    '''Save to JSON file with numpy support. Auto-creates parent dir.'''
    parent = os.path.dirname(fp)
    if parent:
        os.makedirs(parent, exist_ok=True)
    with open(fp, 'w') as f:
        json.dump(data, f, indent=indent, cls=JsonNumpyEncoder)


def load_json(fp):
    '''Load data from JSON file.'''
    with open(fp, 'r') as f:
        return json.load(f)