"""Method to train the surrogate models"""
import os
from xmlrpc.client import Boolean
import numpy as np
import ast
import h5py
import tqdm
from multiprocessing import Pool, Value

from fiesta.constants import days_to_seconds

try:
    import afterglowpy as grb
except ImportError:
    pass


class AfterglowData:
    def __init__(self,
                 outfile: str,
                 n_training: int, 
                 n_val: int,
                 n_test: int,
                 parameter_distributions: dict = None,
                 jet_type: int = -1,
                 tmin: float = 1.,
                 tmax: float = 1000.,
                 n_times: int = 100,
                 use_log_spacing: bool = True,
                 numin: float = 1e9,
                 numax: float = 2.5e18,
                 n_nu: int = 256,
                 fixed_parameters: dict = None) -> None:
        if fixed_parameters is None:
            fixed_parameters = {}

        outdir = os.path.dirname(outfile)
        if not os.path.exists(outdir):
            os.makedirs(outdir)
        self.outfile = outfile

        self.n_training = n_training
        self.n_val = n_val
        self.n_test = n_test

        if os.path.exists(self.outfile):
            self._read_file()
        else:
            self.jet_type = jet_type
            if self.jet_type not in [-1, 0, 2]:
                raise ValueError(f"Jet type {jet_type} is not supported. Supported jet types are: [-1, 0, 2]")
            self.initialize_times(tmin, tmax, n_times, use_log_spacing) # create time array
            self.initialize_nus(numin, numax, n_nu) # create frequency array
            self.parameter_names = list(parameter_distributions.keys())
            self.parameter_distributions = parameter_distributions
            self._initialize_file() # initialize the h5 file the data is later written to
            self.n_training_exists, self.n_val_exists, self.n_test_exists = 0, 0, 0
        
        print(f"Initialized fiesta.train.AfterglowData \nJet type: {self.jet_type} \nParameters: {self.parameter_names} \nTimes: {self.times[0]} {self.times[-1]} {len(self.times)} \nNus: {self.nus[0]:.3e} {self.nus[-1]:.3e} {len(self.nus)} \nparameter_distributions: {self.parameter_distributions}\nExisting train, val, test: {self.n_training_exists}, {self.n_val_exists}, {self.n_test_exists} \n \n \n")      
        self.fixed_parameters = fixed_parameters

        self.get_raw_data(self.n_training, "train") # create new data and save it to file
        self.get_raw_data(self.n_val, "val")
        self.get_raw_data(self.n_test, "test")

    def initialize_times(self, tmin, tmax, n_times, use_log_spacing: bool = True):
        if use_log_spacing:
            times = np.logspace(np.log10(tmin), np.log10(tmax), num=n_times)
        else:
            times = np.linspace(tmin, tmax, num=n_times)
        self.times = times
    
    def initialize_nus(self, numin: float, numax: float, n_nu: int):
        self.nus = np.logspace(np.log10(numin), np.log10(numax), n_nu)

    def _initialize_file(self,):
        with h5py.File(self.outfile, "w") as f:
            f.create_dataset("times", data = self.times)
            f.create_dataset("nus", data = self.nus)
            f.create_dataset("parameter_names", data = self.parameter_names)
            f.create_dataset("parameter_distributions", data = str(self.parameter_distributions))
            f.create_dataset("jet_type", data = self.jet_type)
            f.create_group("train"); f.create_group("val"); f.create_group("test"); f.create_group("special_train")

    def get_raw_data(self, n: int, group: str):
        if group == "train":
            training = True
        else:
            training = False

        nchunks, rest = divmod(n, self.chunk_size) # create raw data in chunks of chunk_size
        for chunk in tqdm.tqdm([*(nchunks*[self.chunk_size]), rest], desc = f"Calculating {nchunks+1} chunks of {group} data...", leave = True):
            if chunk ==0:
                continue
            X, y = self.create_raw_data(chunk, training)
            X, y = self.fix_nans(X, y)
            self._save_to_file(X, y, group)

    def fix_nans(self,X,y):
        # fixes any nans that remain from create_raw_data
        problematic = np.unique(np.where(np.isnan(y))[0])
        n = len(problematic)
        while n>0:
            if n> 0.1*len(X):
                print(f"Warning: Found many nans for the parameter samples, in total {n} out of {len(X)} samples.")
            X_replacement, y_replacement = self.create_raw_data(n)
            X[problematic] = X_replacement
            y[problematic] = y_replacement
            problematic = np.unique(np.where(np.isnan(y))[0])
            n = len(problematic)
        return X, y
    
    def _read_file(self,):
        with h5py.File(self.outfile, "r") as f:
            self.jet_type = f["jet_type"][()]
            self.times = f["times"][:]
            self.nus = f["nus"][:]
            self.parameter_names = f["parameter_names"][:].astype(str).tolist()
            self.parameter_distributions = ast.literal_eval(f["parameter_distributions"][()].decode('utf-8'))
            try:
                self.n_training_exists = (f["train"]["X"].shape)[0]
            except KeyError:
                self.n_training_exists = 0
            try:
                self.n_val_exists = (f["val"]["X"].shape)[0]
            except KeyError:
                self.n_val_exists = 0
            try:
                self.n_test_exists = (f["test"]["X"].shape)[0]
            except KeyError:
                self.n_test_exists = 0

    def create_raw_data(self, n: int, training: bool = True):
        """
        Create draws X in the parameter space and run the afterglow model on it.
        """
        # Create training data
        X_raw = np.empty((n, len(self.parameter_names)))
       
        if training:
            for j, key in enumerate(self.parameter_names):
                a, b, distribution = self.parameter_distributions[key] # FIXME
                if distribution == "uniform":
                    X_raw[:,j] = np.random.uniform(a, b, size = n)
                elif distribution == "loguniform":
                    X_raw[:,j] = np.exp(np.random.uniform(np.log(a), np.log(b), size = n))
        else:
            for j, key in enumerate(self.parameter_distributions.keys()):
                a, b, _ = self.parameter_distributions[key]
                X_raw[:, j] = np.random.uniform(a, b, size = n)

        # Ensure that epsilon_e +  epsilon_B < 1
        epsilon_e_ind = self.parameter_names.index("log10_epsilon_e")
        epsilon_B_ind = self.parameter_names.index("log10_epsilon_B")
        epsilon_tot = (10**(X_raw[:, epsilon_e_ind]) + 10**(X_raw[:, epsilon_B_ind]))
        mask = epsilon_tot>=1 
        X_raw[mask, epsilon_B_ind] += np.log10(0.99/epsilon_tot[mask])
        X_raw[mask, epsilon_e_ind] += np.log10(0.99/epsilon_tot[mask])
        
        # Ensure that thetaWing is smaller than pi/2
        if self.jet_type !=-1:
            alphaw_ind = self.parameter_names.index("alphaWing")
            thetac_ind = self.parameter_names.index("thetaCore")
            mask = X_raw[:, alphaw_ind]*X_raw[:, thetac_ind] >= np.pi/2
            X_raw[mask, alphaw_ind] = np.pi/2 * 1/X_raw[mask, thetac_ind]

        X, y = self.run_afterglow_model(X_raw)
        return X, y
    
    def create_special_data(self, X_raw, label:str, comment: str = None):
        """Create special training data with pre-specified parameters X. These will be stored in the 'special_train' hdf5 group."""

        # Ensure that epsilon_e +  epsilon_B < 1
        epsilon_e_ind = self.parameter_names.index("log10_epsilon_e")
        epsilon_B_ind = self.parameter_names.index("log10_epsilon_B")
        epsilon_tot = (10**(X_raw[:, epsilon_e_ind]) + 10**(X_raw[:, epsilon_B_ind]))
        mask = epsilon_tot>=1 
        X_raw[mask, epsilon_B_ind] += np.log10(0.99/epsilon_tot[mask])
        X_raw[mask, epsilon_e_ind] += np.log10(0.99/epsilon_tot[mask])
        
        # Ensure that thetaWing is smaller than pi/2
        if self.jet_type !=-1:
            alphaw_ind = self.parameter_names.index("alphaWing")
            thetac_ind = self.parameter_names.index("thetaCore")
            mask = X_raw[:, alphaw_ind]*X_raw[:, thetac_ind] >= np.pi/2
            X_raw[mask, alphaw_ind] = np.pi/2 * 1/X_raw[mask, thetac_ind]
        
        X, y = self.run_afterglow_model(X_raw)
        X, y = self.fix_nans(X,y)
        self._save_to_file(X, y, "special_train", label = label, comment= comment)

    def run_afterglow_model(self, X):
        raise NotImplementedError

    def _save_to_file(self, X, y, group: str, label: str = None, comment: str = None):
        with h5py.File(self.outfile, "a") as f:
            if "y" in f[group]: # checks if the dataset already exists
                Xset = f[group]["X"]
                Xset.resize(Xset.shape[0]+X.shape[0], axis = 0)
                Xset[-X.shape[0]:] = X

                yset = f[group]["y"]
                yset.resize(yset.shape[0]+y.shape[0], axis = 0)
                yset[-y.shape[0]:] = y
            
            elif label is not None: # or if we have special training data
                if label in f["special_train"]:
                    Xset = f["special_train"][label]["X"]
                    Xset.resize(Xset.shape[0]+X.shape[0], axis = 0)
                    Xset[-X.shape[0]:] = X

                    yset = f[group][label]["y"]
                    yset.resize(yset.shape[0]+y.shape[0], axis = 0)
                    yset[-y.shape[0]:] = y

                else: 
                    f["special_train"].create_group(label)
                    if comment is not None:
                        f["special_train"][label].attrs["comment"] = comment
                    f["special_train"][label].create_dataset("X", data = X, maxshape=(None, len(self.parameter_names)), chunks = (self.chunk_size, len(self.parameter_names)))
                    f["special_train"][label].create_dataset("y", data = y, maxshape=(None, len(self.nus), len(self.times)), chunks = (self.chunk_size, len(self.times), len(self.nus)))

            else: # or if we need to create a new data set
                f[group].create_dataset("X", data = X, maxshape=(None, len(self.parameter_names)), chunks = (self.chunk_size, len(self.parameter_names)))
                f[group].create_dataset("y", data = y, maxshape=(None, len(self.nus), len(self.times)), chunks = (self.chunk_size, len(self.nus), len(self.times)))

class AfterglowpyData(AfterglowData):

    def __init__(self, n_pool: int, *args, **kwargs):
        self.n_pool = n_pool
        self.chunk_size = 1000
        super().__init__(*args, **kwargs)


    def run_afterglow_model(self, X):
        """Uses multiprocessing to run afterglowpy on the supplied parameters in X."""
        y = np.empty((len(X), len(self.nus), len(self.times)))
        afgpy = RunAfterglowpy(self.jet_type, self.times, self.nus, X, self.parameter_names, self.fixed_parameters)
        pool = Pool(processes=self.n_pool)
        jobs = [pool.apply_async(func=afgpy, args=(argument,)) for argument in range(len(X))]
        pool.close()
        for Idx, job in enumerate(tqdm.tqdm(jobs, desc = f"Computing {len(X)} afterglowpy calculations.", leave = False)):
            try:
                idx, out = job.get()
                y[idx] = out
            except Exception:
                y[Idx] = np.full(len(self.nus), len(self.times), np.nan)
        return X, y


class PyblastafterglowData(AfterglowData):

    def __init__(self, path_to_exec: str, pbag_kwargs: dict = None, rank: int = 0, *args, **kwargs):
        self.chunk_size = 10
        self.rank = rank
        self.path_to_exec = path_to_exec
        self.pbag_kwargs = pbag_kwargs
        super().__init__(*args, **kwargs)


    def run_afterglow_model(self, X):
        """Should be run in parallel with different mpi processes to run pyblastafterglow on the parameters in the array X."""
        y = np.empty((len(X), len(self.nus), len(self.times)))

        pbag = RunPyblastafterglow(self.jet_type,
                                   self.times, 
                                   self.nus, 
                                   X, 
                                   self.parameter_names, 
                                   self.fixed_parameters, 
                                   rank=self.rank, 
                                   path_to_exec=self.path_to_exec,
                                   **self.pbag_kwargs)
        
        for j in range(len(X)):
            try:
                idx, out = pbag(j)
                y[idx] = out
            except Exception:
                try:
                    # increase blast wave evolution time grid if there is an error
                    old_ntb = pbag.ntb
                    pbag.ntb = 3000
                    idx, out = pbag(j)
                    y[idx] = out
                    pbag.ntb = old_ntb
                except Exception:
                    y[j] = np.full(len(self.nus), len(self.times), np.nan)           
        return X, y
    
    def supplement_time(self,t_supp):
        """ WARNING: NOT READY TO BE USED"""
        self.times = t_supp

        for group in ["train", "val", "test"]:
            with h5py.File(self.outfile) as f:
                if "y" not in f[group].keys():
                    continue
                if f[group]["y"].shape[1]>f["times"].shape[0] * f["nus"].shape[0]:
                    continue
                X = f[group]["X"][:]

            _, y_new = self.run_afterglow_model(X)
            y_new = y_new.reshape(-1, len(self.nus), len(self.times))

            with h5py.File(self.outfile, "r+") as f:
                y_old = f[group]["y"][:]
                y_old = y_old.reshape(-1, f["nus"].shape[0], f["times"].shape[0])
                y = np.concatenate((y_new, y_old), axis=-1)

                new_time_shape = len(self.times) + f["times"].shape[0]
                y = y.reshape(-1, len(self.nus), new_time_shape)
                del f[group]["y"]
                f[group].create_dataset("y", data=y, maxshape=(None, len(self.nus), new_time_shape), chunks = (self.chunk_size, len(self.nus), new_time_shape) )
        

        with h5py.File(self.outfile,"r+") as f:
            t_old = f["times"][:]
            del f["times"]
            time = np.concatenate((t_supp, t_old))
            f.create_dataset("times", data=time)

class RunAfterglowpy:
    def __init__(self, jet_type, times, nus, X, parameter_names, fixed_parameters=None):
        if fixed_parameters is None:
            fixed_parameters = {}
        self.jet_type = jet_type
        self.times = times
        self._times_afterglowpy = self.times * days_to_seconds # afterglowpy takes seconds as input
        self.nus = nus
        self.X = X
        self.parameter_names = parameter_names
        self.fixed_parameters = fixed_parameters

    def _call_afterglowpy(self,
                         params_dict: dict[str, float]):
        """
        Call afterglowpy to generate a single flux density output, for a given set of parameters. Note that the parameters_dict should contain all the parameters that the model requires, as well as the nu value.
        The output will be a set of mJys.

        Args:
            Float[Array, "n_times"]: The flux density in mJys at the given times.
        """
        
        # Preprocess the params_dict into the format that afterglowpy expects, which is usually called Z
        Z = {}
        
        Z["jetType"]  = params_dict.get("jetType", self.jet_type)
        Z["specType"] = params_dict.get("specType", 0)
        Z["z"] = params_dict.get("redshift", 0.0)
        Z["xi_N"] = params_dict.get("xi_N", 1.0)
        Z["counterjet"] = True
            
        Z["E0"]        = 10 ** params_dict["log10_E0"]
        Z["n0"]        = 10 ** params_dict["log10_n0"]
        Z["p"]         = params_dict["p"]
        Z["epsilon_e"] = 10 ** params_dict["log10_epsilon_e"]
        Z["epsilon_B"] = 10 ** params_dict["log10_epsilon_B"]
        Z["d_L"]       = 3.086e19 # fix at 10 pc, so that AB magnitude equals absolute magnitude

        if "inclination_EM" in list(params_dict.keys()):
            Z["thetaObs"]  = params_dict["inclination_EM"]
        else:
            Z["thetaObs"]  = params_dict["thetaObs"]

        if self.jet_type == -1:
             Z["thetaCore"] = params_dict["thetaCore"]
        
        elif self.jet_type == 0:
             Z["thetaCore"] = params_dict["thetaCore"]
             Z["thetaWing"] = params_dict["thetaCore"]*params_dict["alphaWing"]

        elif self.jet_type == 4:
            Z["thetaCore"] = params_dict["thetaCore"]
            Z["thetaWing"] = params_dict["thetaCore"]*params_dict["alphaWing"]
            Z["b"] = params_dict["b"]
        
        else:
            raise ValueError(f"Provided jet type {self.jet_type} invalid.")
        
        # Afterglowpy returns flux in mJys
        tt, nunu = np.meshgrid(self._times_afterglowpy, self.nus)
        mJys = grb.fluxDensity(tt, nunu, **Z)
        return mJys

    def __call__(self, idx):
        param_dict = dict(zip(self.parameter_names, self.X[idx], strict=True))
        param_dict.update(self.fixed_parameters)
        mJys = self._call_afterglowpy(param_dict)
        return  idx, np.log10(mJys)



class RunPyblastafterglow:
    def __init__(self,
                 jet_type: int,
                 times, 
                 nus, 
                 X, 
                 parameter_names, 
                 fixed_parameters=None,
                 rank = 0,
                 path_to_exec: str="./pba.out", 
                 grb_resolution: int=12,
                 ntb: int=1000,
                 tb0: float=1e1,
                 tb1: float=1e11,
                 rtol: float=1e-1,
                 loglevel: str="err",
                 ):
        
        jet_conversion = {"-1": "tophat",
                          "0": "gaussian"}
        self.jet_type = jet_conversion[str(jet_type)]
        times_seconds = times * days_to_seconds # pyblastafterglow takes seconds as input

        # preparing the pyblastafterglow string argument for time array
        is_log_uniform = np.allclose(np.diff(np.log(times_seconds)), np.log(times_seconds[1])-np.log(times_seconds[0]), atol=0.01)
        if is_log_uniform:
            log_dt = np.log(times_seconds[1])-np.log(times_seconds[0])
            self.lc_times = f'array logspace {times_seconds[0]:e} {np.exp(log_dt)*times_seconds[-1]:e} {len(times_seconds)}' # pyblastafterglow only takes this string format
        else:
            raise ValueError("Time array must be loguniform.")

        # preparing the pyblastafterglow string argument for frequency array
        log_dnu = np.log(nus[1]/nus[0])
        self.lc_freqs = f'array logspace {nus[0]:e} {np.exp(log_dnu)*nus[-1]:e} {len(nus)}' # pyblastafterglow only takes this string format

        self.X = X
        self.parameter_names = parameter_names
        if fixed_parameters is None:
            fixed_parameters = {}
        self.fixed_parameters = fixed_parameters
        self.rank = rank
        self.path_to_exec = path_to_exec
        self.grb_resolution = grb_resolution
        self.ntb = ntb
        self.tb0 = tb0
        self.tb1 = tb1
        self.rtol = rtol
        self.loglevel = loglevel

    def _call_pyblastafterglow(self,
                         params_dict: dict[str, float]):
        """
        Run pyblastafterglow to generate a single flux density output, for a given set of parameters. Note that the parameters_dict should contain all the parameters that the model requires.
        The output will be a set of mJys.

        Args:
            Float[Array, "n_times"]: The flux density in mJys at the given times.
        """
        
        try:
            from PyBlastAfterglowMag.wrappers import run_grb
        except ImportError:
            raise ImportError("PyBlastAfterglowMag is not installed. Please install it from source")
        
        # Define jet structure (analytic; gaussian) -- 3 free parameters 
        struct = dict(
            struct= self.jet_type, # type of the structure tophat or gaussian
            Eiso_c=np.power(10, params_dict["log10_E0"]),  # isotropic equivalent energy of the burst 
            Gamma0c=params_dict["Gamma0"],    # lorentz factor of the core of the jet 
            M0c=-1.,         # mass of the ejecta (if -1 -- inferr from Eiso_c and Gamma0c)
            n_layers_a=self.grb_resolution    # resolution of the jet (number of individual blastwaves)
        )

        if self.jet_type == "tophat":
            struct["theta_c"] =  params_dict['thetaCore'] # half-opening angle of the winds of the jet
        
        elif self.jet_type == "gaussian":
            struct["theta_c"] =  params_dict['thetaCore'] # half-opening angle of the winds of the jet
            struct["theta_w"] = params_dict["thetaCore"] * params_dict["alphaWing"]
        
        else:
            raise ValueError(f"Provided jet type {self.jet_type} invalid.")

        # set model parameters
        P = dict(
                # main model parameters; Uniform ISM -- 2 free parameters
                main=dict(
                    d_l= 3.086e19, # luminocity distance to the source [cm], fix at 10 pc, so that AB magnitude equals absolute magnitude
                    z = params_dict.get("redshift", 0.0),   # redshift of the source (used in Doppler shifring and EBL table)
                    n_ism=np.power(10, params_dict["log10_n0"]), # ISM density [cm^-3] (assuming uniform)
                    theta_obs= params_dict["inclination_EM"], # observer angle [rad] (from pol to jet axis)  
                    lc_freqs= self.lc_freqs, # frequencies for light curve calculation
                    lc_times= self.lc_times, # times for light curve calculation
                    tb0=self.tb0, tb1=self.tb1, ntb=self.ntb, # burster frame time grid boundary, resolution, for the simulation
                ),

                # ejecta parameters; FS only -- 3 free parameters 
                grb=dict(
                    structure=struct, # structure of the ejecta
                    eps_e_fs=np.power(10, params_dict["log10_epsilon_e"]), # microphysics - FS - frac. energy in electrons
                    eps_b_fs=np.power(10, params_dict["log10_epsilon_B"]), # microphysics - FS - frac. energy in magnetic fields
                    p_fs= params_dict["p"], # microphysics - FS - slope of the injection electron spectrum
                    do_lc='yes',      # task - compute light curves
                    rtol_theta = self.rtol,
                    method_limit_spread=None,
                    # save_spec='yes' # save comoving spectra 
                    # method_synchrotron_fs = 'Joh06',
                    # method_ne_fs = 'usenprime',
                    # method_ele_fs = 'analytic',
                    # method_comp_mode = 'observFlux'
                )
        )
        pba_run = run_grb(working_dir= os.getcwd() + f'/tmp_{self.rank}/', # directory to save/load from simulation data
                              P=P,                     # all parameters 
                              run=True,                # run code itself (if False, it will try to load results)
                              path_to_cpp=self.path_to_exec, # absolute path to the C++ executable of the code
                              loglevel=self.loglevel,         # logging level of the code (info or err)
                              process_skymaps=False    # process unstractured sky maps. Only useed if `do_skymap = yes`
                             )
        
        mJys = pba_run.GRB.get_lc()
        return mJys

    def __call__(self, idx):
        param_dict = dict(zip(self.parameter_names, self.X[idx], strict=True))
        param_dict.update(self.fixed_parameters)
        mJys = self._call_pyblastafterglow(param_dict)
        return  idx, np.log10(mJys)


class BlastwaveData(AfterglowData):

    def __init__(self, *args, n_pool=None, **kwargs):
        if n_pool is not None:
            import warnings
            warnings.warn("n_pool is ignored; blastwave uses rayon for parallelism", stacklevel=2)
        self.chunk_size = 100
        super().__init__(*args, **kwargs)

    def _sample_parameters(self, n, training):
        """Sample parameters and enforce the forward-shock energy constraint."""
        X_raw = np.empty((n, len(self.parameter_names)))

        if training:
            for j, key in enumerate(self.parameter_names):
                a, b, distribution = self.parameter_distributions[key]
                if distribution == "uniform":
                    X_raw[:, j] = np.random.uniform(a, b, size=n)
                elif distribution == "loguniform":
                    X_raw[:, j] = np.exp(np.random.uniform(np.log(a), np.log(b), size=n))
        else:
            for j, key in enumerate(self.parameter_distributions.keys()):
                a, b, _ = self.parameter_distributions[key]
                X_raw[:, j] = np.random.uniform(a, b, size=n)

        # Ensure that epsilon_e + epsilon_B < 1
        epsilon_e_ind = self.parameter_names.index("log10_epsilon_e")
        epsilon_B_ind = self.parameter_names.index("log10_epsilon_B")
        epsilon_tot = (10**(X_raw[:, epsilon_e_ind]) + 10**(X_raw[:, epsilon_B_ind]))
        mask = epsilon_tot >= 1
        X_raw[mask, epsilon_B_ind] += np.log10(0.99 / epsilon_tot[mask])
        X_raw[mask, epsilon_e_ind] += np.log10(0.99 / epsilon_tot[mask])

        return X_raw

    def create_raw_data(self, n: int, training: bool = True):
        """Create draws X in the parameter space and run the blastwave model on it."""
        X_raw = self._sample_parameters(n, training)
        X, y = self.run_afterglow_model(X_raw)
        return X, y

    def run_afterglow_model(self, X):
        """Run blastwave model on the supplied parameters in X.

        No multiprocessing.Pool needed — the Rust extension uses rayon for
        automatic parallelism within each FluxDensity call (releases the GIL
        via py.allow_threads).
        """
        y = np.empty((len(X), len(self.nus), len(self.times)))
        runner = RunBlastwave(self.times, self.nus, X, self.parameter_names, self.fixed_parameters)
        for idx in tqdm.tqdm(range(len(X)), desc=f"Computing {len(X)} blastwave calculations.", leave=False):
            try:
                i, out = runner(idx)
                y[i] = out
            except Exception as e:
                print(f"[WARNING] Blastwave call failed for sample {idx}: {e}")
                y[idx] = np.full((len(self.nus), len(self.times)), np.nan)
        return X, y


class RunBlastwave:
    def __init__(self, times, nus, X, parameter_names, fixed_parameters=None,
                 ncells=33, spread_mode="ode"):
        if fixed_parameters is None:
            fixed_parameters = {}
        self.times = np.array(times)
        self._times_seconds = self.times * days_to_seconds
        self.nus = np.array(nus)
        self.X = np.array(X)
        self.parameter_names = list(parameter_names)
        self.fixed_parameters = dict(fixed_parameters)
        self.ncells = ncells
        if spread_mode not in ("ode", "pde", "none"):
            raise ValueError(f"Invalid spread_mode '{spread_mode}'. Must be 'ode', 'pde', or 'none'.")
        self.spread_mode = spread_mode

        # Pre-compute the meshgrid once
        tt, nunu = np.meshgrid(self._times_seconds, self.nus)
        self._tt_flat = tt.flatten()
        self._nunu_flat = nunu.flatten()
        self._tmax = float(np.max(self._times_seconds)) * 2

        try:
            import blastwave as _bw
            self._bw = _bw
        except ImportError as err:
            raise ImportError(
                "blastwave is not installed. Install it from "
                "https://github.com/nuclear-multimessenger-astronomy/blastwave"
            ) from err

    def _call_blastwave(self, params_dict):
        """
        Call blastwave model to generate a single flux density output.
        Uses the Jet API directly with tuned settings for data generation.
        Returns flux density in mJy as a 2D array (n_nu, n_times).
        """
        bw = self._bw

        theta_c = params_dict["thetaCore"]
        Eiso = 10 ** params_dict["log10_E0"]
        lf = 10 ** params_dict["log10_lf"]
        n0 = 10 ** params_dict["log10_n0"]
        p = params_dict["p"]
        eps_e = 10 ** params_dict["log10_epsilon_e"]
        eps_b = 10 ** params_dict["log10_epsilon_B"]
        theta_v = params_dict["inclination_EM"]

        P = dict(
            Eiso=Eiso, lf=lf, theta_c=theta_c, n0=n0, A=0,
            eps_e=eps_e, eps_b=eps_b, p=p, theta_v=theta_v,
            d=9.99e-6, z=0.0,
        )

        # Solve dynamics with reduced cell count and ODE spreading
        spread_kwargs = {}
        if self.spread_mode == "ode":
            spread_kwargs["spread_mode"] = "ode"
        elif self.spread_mode == "pde":
            spread_kwargs["spread"] = True
        else:
            spread_kwargs["spread"] = False

        jet = bw.Jet(
            bw.Gaussian(theta_c, Eiso, lf0=lf),
            0.0, n0,
            tmin=10.0,
            tmax=self._tmax,
            grid=bw.ForwardJetRes(theta_c, self.ncells),
            tail=True,
            cal_level=1,
            rtol=1e-6,
            eps_e=eps_e,
            eps_b=eps_b,
            p_fwd=p,
            **spread_kwargs,
        )

        mJys = jet.FluxDensity(
            self._tt_flat, self._nunu_flat, P,
            force_return=True,
            flux_method="forward",
        )
        mJys = mJys.reshape(len(self.nus), len(self._times_seconds))
        return mJys

    def __call__(self, idx):
        param_dict = dict(zip(self.parameter_names, self.X[idx], strict=True))
        param_dict.update(self.fixed_parameters)
        mJys = self._call_blastwave(param_dict)
        mJys = np.clip(mJys, 1e-300, None)
        return idx, np.log10(mJys)


class BlastwaveRSData(BlastwaveData):
    """BlastwaveData variant with reverse shock enabled.

    Following Japelj+ 2014 (1402.3701), the RS microphysics are tied to
    the FS values via a magnetization ratio RB::

        eps_e_rs = eps_e_f
        eps_b_rs = RB * eps_b_f
        p_rs     = p_f

    Extra sampled parameter: log10_RB, log10_duration.
    sigma is kept fixed at 0.0 (unmagnetized ejecta).
    """

    def create_raw_data(self, n: int, training: bool = True):
        """Sample parameters, enforce FS and RS energy constraints, then run model."""
        X_raw = self._sample_parameters(n, training)

        # Ensure that eps_b_rs = RB * eps_b_f < 1 and eps_e + eps_b_rs < 1 (reverse shock)
        epsilon_e_ind = self.parameter_names.index("log10_epsilon_e")
        epsilon_B_ind = self.parameter_names.index("log10_epsilon_B")
        RB_ind = self.parameter_names.index("log10_RB")
        eps_e = 10**(X_raw[:, epsilon_e_ind])
        eps_b_f = 10**(X_raw[:, epsilon_B_ind])
        # Cap RB so that RB * eps_b_f < min(1 - eps_e, 0.99)
        RB_max = np.clip((0.99 - eps_e) / eps_b_f, 1.0, None)
        log10_RB_max = np.log10(RB_max)
        mask_rs = X_raw[:, RB_ind] > log10_RB_max
        X_raw[mask_rs, RB_ind] = log10_RB_max[mask_rs]

        X, y = self.run_afterglow_model(X_raw)
        return X, y

    def run_afterglow_model(self, X):
        y = np.empty((len(X), len(self.nus), len(self.times)))
        runner = RunBlastwaveRS(self.times, self.nus, X, self.parameter_names, self.fixed_parameters)
        for idx in tqdm.tqdm(range(len(X)), desc=f"Computing {len(X)} blastwave+RS calculations.", leave=False):
            try:
                i, out = runner(idx)
                y[i] = out
            except Exception as e:
                print(f"[WARNING] Blastwave+RS call failed for sample {idx}: {e}")
                y[idx] = np.full((len(self.nus), len(self.times)), np.nan)
        return X, y


class RunBlastwaveRS:
    """Like RunBlastwave but with reverse shock enabled.

    Following Japelj+ 2014: RS microphysics derived from FS values via
    magnetization ratio RB = eps_b_rs / eps_b_f, with eps_e_rs = eps_e_f
    and p_rs = p_f.
    """
    def __init__(self, times, nus, X, parameter_names, fixed_parameters=None,
                 ncells=33, spread_mode="ode"):
        if fixed_parameters is None:
            fixed_parameters = {}
        self.times = np.array(times)
        self._times_seconds = self.times * days_to_seconds
        self.nus = np.array(nus)
        self.X = np.array(X)
        self.parameter_names = list(parameter_names)
        self.fixed_parameters = dict(fixed_parameters)
        self.ncells = ncells
        if spread_mode not in ("ode", "pde", "none"):
            raise ValueError(f"Invalid spread_mode '{spread_mode}'. Must be 'ode', 'pde', or 'none'.")
        self.spread_mode = spread_mode

        tt, nunu = np.meshgrid(self._times_seconds, self.nus)
        self._tt_flat = tt.flatten()
        self._nunu_flat = nunu.flatten()
        self._tmax = float(np.max(self._times_seconds)) * 2

        try:
            import blastwave as _bw
            self._bw = _bw
        except ImportError as err:
            raise ImportError(
                "blastwave is not installed. Install it from "
                "https://github.com/nuclear-multimessenger-astronomy/blastwave"
            ) from err

    def _call_blastwave(self, params_dict):
        bw = self._bw

        theta_c = params_dict["thetaCore"]
        Eiso = 10 ** params_dict["log10_E0"]
        lf = 10 ** params_dict["log10_lf"]
        n0 = 10 ** params_dict["log10_n0"]
        p = params_dict["p"]
        eps_e = 10 ** params_dict["log10_epsilon_e"]
        eps_b = 10 ** params_dict["log10_epsilon_B"]
        theta_v = params_dict["inclination_EM"]

        # Reverse shock parameters derived from FS (Japelj+ 2014)
        RB = 10 ** params_dict["log10_RB"]
        eps_e_rs = eps_e          # same as forward shock
        eps_b_rs = RB * eps_b     # magnetization ratio
        p_rs = p                  # same as forward shock
        duration = 10 ** params_dict["log10_duration"]
        sigma = params_dict.get("sigma", 0.0)

        P = dict(
            Eiso=Eiso, lf=lf, theta_c=theta_c, n0=n0, A=0,
            eps_e=eps_e, eps_b=eps_b, p=p, theta_v=theta_v,
            d=9.99e-6, z=0.0,
        )

        spread_kwargs = {}
        if self.spread_mode == "ode":
            spread_kwargs["spread_mode"] = "ode"
        elif self.spread_mode == "pde":
            spread_kwargs["spread"] = True
        else:
            spread_kwargs["spread"] = False

        jet = bw.Jet(
            bw.Gaussian(theta_c, Eiso, lf0=lf),
            0.0, n0,
            tmin=10.0,
            tmax=self._tmax,
            grid=bw.ForwardJetRes(theta_c, self.ncells),
            tail=True,
            cal_level=1,
            rtol=1e-6,
            eps_e=eps_e,
            eps_b=eps_b,
            p_fwd=p,
            include_reverse_shock=True,
            sigma=sigma,
            eps_e_rs=eps_e_rs,
            eps_b_rs=eps_b_rs,
            p_rs=p_rs,
            duration=duration,
            **spread_kwargs,
        )

        mJys = jet.FluxDensity(
            self._tt_flat, self._nunu_flat, P,
            force_return=True,
        )
        mJys = mJys.reshape(len(self.nus), len(self._times_seconds))
        return mJys

    def __call__(self, idx):
        param_dict = dict(zip(self.parameter_names, self.X[idx], strict=True))
        param_dict.update(self.fixed_parameters)
        mJys = self._call_blastwave(param_dict)
        mJys = np.clip(mJys, 1e-300, None)
        return idx, np.log10(mJys)


# Backward-compatibility aliases
JetsimpyData = BlastwaveData
RunJetsimpy = RunBlastwave
