from alabebm import generate, get_params_path
import os
import numpy as np 
import json 

# Get path to default parameters
params_file = get_params_path()

experiment_names = [
    "sn_kjOrdinalDM_xnjNormal",     # Experiment 1: Ordinal kj with Dirichlet-Multinomial, Normal Xnj
    "sn_kjOrdinalDM_xnjNonNormal",  # Experiment 2: Ordinal kj with Dirichlet-Multinomial, Non-Normal Xnj
    "sn_kjOrdinalUniform_xnjNormal", # Experiment 3: Ordinal kj with Uniform distribution, Normal Xnj
    "sn_kjOrdinalUniform_xnjNonNormal", # Experiment 4: Ordinal kj with Uniform distribution, Non-Normal Xnj
    "sn_kjContinuousUniform",       # Experiment 5: Continuous kj with Uniform distribution
    "sn_kjContinuousBeta",          # Experiment 6: Continuous kj with Beta distribution
    "xiNearNormal_kjContinuousUniform", # Experiment 7: Near-normal Xi with Continuous Uniform kj
    "xiNearNormal_kjContinuousBeta", # Experiment 8: Near-normal Xi with Continuous Beta kj
    "xiNearNormalWithNoise_kjContinuousBeta", # Experiment 9: Same as Exp 8 but with noises to xi
]

class OneLineArrayEncoder(json.JSONEncoder):
    def __init__(self, *args, **kwargs):
        kwargs['indent'] = 2  # keep dicts pretty
        super().__init__(*args, **kwargs)

    def iterencode(self, o, _one_shot=False):
        from json.encoder import _make_iterencode
        # Patch the list serializer to skip indentation
        def _my_list_repr(lst, current_indent_level):
            return '[' + ', '.join(self.encode(el) for el in lst) + ']'
        _encoder = _make_iterencode(
            None, self.default, self.indent, self.key_separator,
            self.item_separator, self.sort_keys, self.skipkeys,
            _my_list_repr, None, self.allow_nan
        )
        return _encoder(o, 0)

def convert_np_types(obj):
    """Convert numpy types in a nested dictionary to Python standard types."""
    if isinstance(obj, dict):
        return {k: convert_np_types(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [convert_np_types(item) for item in obj]
    elif isinstance(obj, np.integer):
        return int(obj)
    elif isinstance(obj, np.floating):
        return float(obj)
    elif isinstance(obj, np.ndarray):
        return convert_np_types(obj.tolist())
    else:
        return obj

all_exp_dicts = []

for exp_name in experiment_names:
    bm_et_dict = generate(
            experiment_name = exp_name,
            params_file=params_file,
            js = [200],
            rs = [0.25],
            num_of_datasets_per_combination=1,
            output_dir='my_data',
            seed=42,
            keep_all_cols = False,
        )
    all_exp_dicts.append(bm_et_dict)

combined = {k: v for d in all_exp_dicts for k, v in d.items()}

combined = convert_np_types(combined)

# Dump the JSON
with open("true_order_and_stages.json", "w") as f:
    json.dump(combined, f, indent=2)

    

