bug fixes and added more granularity in the repackaging of reporting

    _prioritize_normalization,

    _rename_params_with_instrument_names,

)

from .lightcurve_plots import (

    plot_models,

    plot_noise_budget,

from .plot_models import plot_models

from .model_components import compute_model_components, write_model_components

from .plot_model_helpers import (

    _compute_planet_model_for_plotting,

    _apply_scoped_values,

)

from .noise_plots import plot_noise_budget

from .periodogram_plots import (

    plot_residuals_histogram,

    plot_periodograms_comparison,

    plot_periodogram,

    plot_correlated_noise,

    plot_residuals_vs_dbasis,

)

from .phase_plots import (

    _time_to_phase,

    plot_phase_folded,

    plot_phase_folded_rv,

    plot_phase_folded_rv_rm_zoom,

)

from .ttv_plots import (

    plot_ttv_oc,

    plot_ttv_per_transit,

    plot_ttv_transit_overlay,

    _compute_planet_model_for_plotting,

    _apply_scoped_values,

    compute_model_components,

    write_model_components,

)

from .pdf_report import (

    generate_report,

# -*- coding: utf-8 -*-

"""

model_components.py

~~~~~~~~~~~~~~~~~~~

Data-layer helpers for assembling and persisting model component arrays.

Functions

---------

compute_model_components

write_model_components

"""

import numpy as np

import jax

import jax.numpy as jnp

from PyAETNA.gp import (

    gp_predict,

    gp_predict_per_visit,

    build_gp_kernel_for_plotting,

    get_gp_x_variable,

)

from PyAETNA.model_utils import build_mean_function

from PyAETNA.body_models import (

    compute_lightcurve as _compute_lightcurve_raw,

    compute_lightcurve_supersampled as _compute_lightcurve_ss,

    compute_radial_velocity as _compute_radial_velocity_raw,

    compute_rossiter_mclaughlin as _compute_rossiter_mclaughlin_raw,

    compute_transit_ttv as _compute_transit_ttv_raw,

    reparametrize_LD as _reparametrize_LD,

)

from . import logger

def _log(msg, level='INFO'):

    logger._log(msg, level=level)

from .plot_model_helpers import (

    _compute_planet_model_for_plotting,

    _apply_scoped_values,

)

def compute_model_components(params, plData, plModel, param_config, skip_gp=False,

                             skip_rm=False):

    """

    Compute individual model components for plotting (scope-aware).

    Args:

        params: Dict of all parameter values (free + fixed)

        plData: PlData object

        plModel: PlModel object

        param_config: Parameter configuration (v3.0 format with scope metadata)

        skip_gp: If True, skip all GP computation and return zeros for GP arrays.

            Use this in sample-band loops where the GP output is superseded by a

            separate grid-based GP computation (avoids expensive per-visit tinygp

            solves on the full data array for each posterior sample).

    Returns:

        dict with keys: 'flat', 'norm', 'planet', 'poly', 'fourier', 'trend', 

                        'mean_function', 'gp_mean', 'gp_std', 'total', 'dy_total',

                        'planet_terms', 'poly_terms', 'fourier_terms', 'gp_terms',

                        'bjd', 'y', 'dy', 'residuals'

        Note: 'norm' contains the per-point normalization values. For LC instruments,

        this is the flux baseline (near 1.0 for multiplicative mode, near 0.0 for

        additive mode). For RV instruments, this is the systemic velocity (gamma) [m/s].

    """

    time = np.array(plData.time)

    y = np.array(plData.y)

    dy = np.array(plData.dy)

    visit = np.array(plData.visit)

    instrument = np.array(plData.instrument)

    n_pts = len(time)

    # Get configs from param_config

    polynomial_configs = param_config.get('polynomial_configs', [])

    fourier_configs = param_config.get('fourier_configs', [])

    gp_configs = param_config.get('gp_configs', {})

    planet_configs = param_config.get('planet_configs', [])

    planet_names = param_config.get('planet_names', ())

    param_scopes = param_config.get('param_scopes', {})

    param_groups = param_config.get('param_groups', {})

    # Compute planet model (scope-aware, LC and RV)

    data_types = param_config.get('data_types', getattr(plModel, 'data_types', None) or {})

    inst_id_to_name = {}

    instrument_names_list = plModel.instrument_names if hasattr(plModel, 'instrument_names') and plModel.instrument_names is not None else None

    for i, inst_id in enumerate(plModel.uniqueInstruments):

        if instrument_names_list is not None and i < len(instrument_names_list):

            inst_id_to_name[inst_id] = instrument_names_list[i]

        else:

            inst_id_to_name[inst_id] = str(inst_id)

    rm_config = param_config.get('rm_config', {'enabled': False})

    planet_model, planet_terms = _compute_planet_model_for_plotting(

        params, time, instrument, plModel.uniqueInstruments, planet_configs, planet_names,

        plData.timeCenter, param_scopes, data_types=data_types,

        inst_id_to_name=inst_id_to_name, exptime=np.array(plData.exptime),

        rm_config=rm_config, skip_rm=skip_rm

    )

    # Flat baseline

    normalization_mode = getattr(plModel, 'normalization_mode', 'multiplicative')

    if normalization_mode == 'additive':

        flat_model = np.zeros(n_pts)

    else:

        flat_model = np.ones(n_pts)

    # Polynomial detrending (per-coefficient scope)

    poly_model = np.zeros(n_pts)

    poly_terms = {}

    for poly_cfg in polynomial_configs:

        poly_name = poly_cfg['name']

        x_variable = poly_cfg['x_variable']

        max_degree = poly_cfg['max_degree']

        coeff_scopes = poly_cfg.get('coeff_scopes', {})

        coeff_hybrid_insts = poly_cfg.get('coeff_hybrid_insts', {})

        # Get x data

        if x_variable == 'BJD':

            x_data = time

        elif x_variable in plData.dBasis:

            x_data = np.array(plData.dBasis[x_variable])

        else:

            continue

        # x-centering is always per-visit for numerical stability

        for uv in plModel.uniqueVisits:

            mask = visit == uv

            x_visit = x_data[mask]

            if len(x_visit) == 0:

                continue

            x_visit_mid = (np.max(x_visit) + np.min(x_visit)) / 2.0

            x_centered = x_data - x_visit_mid

            for d in range(max_degree + 1):

                # Determine parameter name based on the per-coefficient scope

                coeff_scope = coeff_scopes.get(d, 'per_visit')

                if coeff_scope == 'global':

                    param_name = f'{poly_name}_c{d}'

                elif coeff_scope == 'per_instrument':

                    # Find instrument for this visit

                    inst_for_visit = instrument[mask][0] if np.any(mask) else plModel.uniqueInstruments[0]

                    # Hybrid: instrument may use per-visit params within per_instrument

                    if inst_for_visit in coeff_hybrid_insts.get(d, ()):

                        param_name = f'{poly_name}_c{d}_v{uv}'

                    else:

                        param_name = f'{poly_name}_c{d}_{inst_for_visit}'

                else:  # per_visit

                    param_name = f'{poly_name}_c{d}_v{uv}'

                if param_name in params:

                    coef = params[param_name]

                    contrib = np.zeros(n_pts)

                    contrib[mask] = coef * (x_centered[mask] ** d)

                    # Accumulate: for global/per_instrument scopes the same

                    # param_name covers multiple visits, so we must add to the

                    # existing entry rather than overwrite it (otherwise only

                    # the last visit's contribution survives in poly_terms).

                    poly_terms[param_name] = poly_terms.get(

                        param_name, np.zeros(n_pts)) + contrib

                    poly_model[mask] += contrib[mask]

    # Fourier series detrending (per-harmonic scope)

    fourier_model = np.zeros(n_pts)

    fourier_terms = {}

    for fourier_cfg in fourier_configs:

        fourier_name = fourier_cfg['name']

        x_variable = fourier_cfg['x_variable']

        max_harmonic = fourier_cfg['max_harmonic']

        period = float(fourier_cfg['period'])

        coeff_scopes = fourier_cfg.get('coeff_scopes', {})

        coeff_hybrid_insts = fourier_cfg.get('coeff_hybrid_insts', {})

        # Get x data (Fourier)

        if x_variable == 'BJD':

            x_data = time

        elif x_variable in plData.dBasis:

            x_data = np.array(plData.dBasis[x_variable])

        else:

            continue

        x_data = np.array(x_data)

        theta = 2.0 * np.pi * x_data / period

        for uv in plModel.uniqueVisits:

            mask = visit == uv

            for h in range(1, max_harmonic + 1):

                # Determine parameter names based on the per-harmonic scope

                coeff_scope = coeff_scopes.get(h, 'per_visit')

                if coeff_scope == 'global':

                    s_name = f'{fourier_name}_c{h}_s'

                    c_name = f'{fourier_name}_c{h}_c'

                elif coeff_scope == 'per_instrument':

                    inst_for_visit = instrument[mask][0] if np.any(mask) else plModel.uniqueInstruments[0]

                    # Hybrid: instrument may use per-visit params within per_instrument

                    if inst_for_visit in coeff_hybrid_insts.get(h, ()):

                        s_name = f'{fourier_name}_c{h}_s_v{uv}'

                        c_name = f'{fourier_name}_c{h}_c_v{uv}'

                    else:

                        s_name = f'{fourier_name}_c{h}_s_{inst_for_visit}'

                        c_name = f'{fourier_name}_c{h}_c_{inst_for_visit}'

                else:  # per_visit

                    s_name = f'{fourier_name}_c{h}_s_v{uv}'

                    c_name = f'{fourier_name}_c{h}_c_v{uv}'

                if s_name in params:

                    contrib = np.zeros(n_pts)

                    contrib[mask] = params[s_name] * np.sin(h * theta[mask])

                    # Same accumulate pattern as for poly_terms above.

                    fourier_terms[s_name] = fourier_terms.get(

                        s_name, np.zeros(n_pts)) + contrib

                    fourier_model[mask] += contrib[mask]

                if c_name in params:

                    contrib = np.zeros(n_pts)

                    contrib[mask] = params[c_name] * np.cos(h * theta[mask])

                    fourier_terms[c_name] = fourier_terms.get(

                        c_name, np.zeros(n_pts)) + contrib

                    fourier_model[mask] += contrib[mask]

    # Trend = polynomial + fourier

    trend_model = poly_model + fourier_model

    # Compute normalization model (scope-aware, including hybrid per_instrument+per_visit)

    norm_names = param_groups.get('normalization', [])

    norm_mask_specs = param_config.get('param_mask_specs', {}).get('normalization')

    default_norm = 0.0 if normalization_mode == 'additive' else 1.0

    norm_model = _apply_scoped_values(

        params, norm_names, norm_mask_specs,

        visit, instrument, n_pts, default_norm

    )

    # Rescale RV normalization values from km/s to m/s

    # Gamma values from samples are in km/s but internal model needs them in m/s

    data_types = param_config.get('data_types', getattr(plModel, 'data_types', None) or {})

    for inst_id in plModel.uniqueInstruments:

        inst_mask = instrument == inst_id

        if inst_id in data_types and data_types[inst_id] == 'RV':

            norm_model[inst_mask] *= 1000.0

    # Mean function — delegates to the shared helper (model_utils.py) which

    # is also used by sampler.py, ensuring identical assembly logic.

    mean_function = np.asarray(build_mean_function(

        norm_model, planet_model, trend_model,

        normalization_mode, instrument, plModel.uniqueInstruments, data_types

    ))

    # Jitter (scope-aware, including hybrid per_instrument+per_visit)

    jitter_names = param_groups.get('log_jitter', [])

    jitter_mask_specs = param_config.get('param_mask_specs', {}).get('log_jitter')

    log_jitter_arr = _apply_scoped_values(

        params, jitter_names, jitter_mask_specs,

        visit, instrument, n_pts, -np.inf  # default: exp(-inf)=0 → no jitter

    )

    jitter_sq = np.exp(log_jitter_arr) ** 2

    dy_total = np.sqrt(dy ** 2 + jitter_sq)

    # GP prediction using shared GP module

    gp_mean = np.zeros(n_pts)

    gp_std = np.zeros(n_pts)

    gp_terms = {}  # dict of per-kernel GP contribution arrays

    if gp_configs and not skip_gp:

        residuals = y - mean_function

        for inst_id, gp_cfg in gp_configs.items():

            inst_mask = instrument == inst_id

            if np.sum(inst_mask) < 2:

                continue

            # Get GP x-variable using shared helper

            x_inst = get_gp_x_variable(plData, gp_cfg, inst_mask, use_jax=False)

            res_inst = residuals[inst_mask]

            dy_inst = dy_total[inst_mask]

            # Pre-compute per-visit sub-indices if per_visit is enabled

            per_visit_enabled = gp_cfg.get('per_visit', False)

            visit_idx_list = None

            if per_visit_enabled:

                visits_at_inst = np.asarray(plData.visit)[inst_mask]

                unique_visits_for_inst = np.unique(visits_at_inst)

                visit_idx_list = []

                for v in unique_visits_for_inst:

                    local_pos = jnp.array(np.where(visits_at_inst == v)[0])

                    if len(local_pos) >= 2:

                        visit_idx_list.append(local_pos)

            # ── Single composite-kernel GP prediction ────────────────────

            # Build the full (sum-of-terms) kernel once and predict in a

            # single solve.  This is both correct and cheaper than the old

            # approach of predicting each kernel term individually and then

            # re-predicting with the composite kernel just for gp_std.

            # For single-term GP configs (the common case) the cost is

            # halved; for multi-term configs the composite prediction is

            # the proper cross-covariance-aware result anyway.

            inst_label = gp_cfg.get('inst_name', str(inst_id))

            kernel_specs = gp_cfg.get('kernel_specs', [])

            if not kernel_specs:

                continue

            try:

                kernel = build_gp_kernel_for_plotting(params, kernel_specs, inst_label)

            except Exception as e:

                _log(f"WARNING: Failed to build GP kernel for instrument {inst_id}: {str(e)}")

                continue

            x_inst_jax = jnp.array(x_inst)

            res_inst_jax = jnp.array(res_inst)

            dy_sq_jax = jnp.array(dy_inst**2)

            try:

                if per_visit_enabled and visit_idx_list:

                    mu_full, std_full = gp_predict_per_visit(

                        kernel, x_inst_jax, res_inst_jax, dy_sq_jax, visit_idx_list,

                    )

                else:

                    mu_full, std_full = gp_predict(

                        kernel, x_inst_jax, res_inst_jax, dy_sq_jax, x_inst_jax)

            except Exception as e:

                _log(f"WARNING: GP prediction failed for instrument {inst_id}: {str(e)}")

                continue

            mu_full_np  = np.array(mu_full)

            std_full_np = np.array(std_full)

            gp_mean[inst_mask] = mu_full_np

            gp_std[inst_mask]  = std_full_np

            # Record per-kernel-term labels (needed by write_model_components).

            # When there is only one kernel term the composite output IS the

            # single-term output; for multi-term configs we label the

            # composite as the first (dominant) term — per-term decomposition

            # would require separate solves, which we intentionally skip.

            for k_idx, spec in enumerate(kernel_specs):

                ktype = spec.get('type', f'k{k_idx}')

                section = spec.get('section_name', gp_cfg.get('section_name', 'GP'))

                label = f'gp_{section}_{ktype}_inst{inst_id}_idx{k_idx}'

                term_array = np.zeros(n_pts)

                if k_idx == 0:

                    # First (or only) term gets the composite GP mean

                    term_array[inst_mask] = mu_full_np

                # Additional terms remain zero (no per-term decomposition)

                gp_terms[label] = term_array

    # Total model

    total_model = mean_function + gp_mean

    return {

        'flat': flat_model,

        'norm': norm_model,  # Normalization values (for LC: baseline, for RV: gamma)

        'planet': planet_model,

        'planet_terms': planet_terms,

        'poly': poly_model,

        'fourier': fourier_model,

        'trend': trend_model,

        'mean_function': mean_function,

        'gp_mean': gp_mean,

        'gp_std': gp_std,

        'total': total_model,

        'dy_total': dy_total,

        'poly_terms': poly_terms,

        'fourier_terms': fourier_terms,

        'gp_terms': gp_terms,

        'bjd': time + getattr(plData, 'timeCenter', 0.0),

        'y': y,

        'dy': dy,

        'residuals': y - total_model,

    }

def write_model_components(outfile, components, plData=None, mask=None, sort_by_time=True):

    """Write model component columns to a whitespace-separated file.

    The file columns are in this order:

        [planet terms...], [polynomial terms...], [fourier terms...], [gp terms...], RESIDUALS

    The observed data columns (BJD, Y, SIGMA) are NOT repeated here;

    they are already present in the original data file.

    Args:

        outfile: Path to output file

        components: Dict returned by `compute_model_components`

        plData: Optional PlData object (used to compute BJD if not present in components)

        mask: Optional boolean mask or index array to subset rows (e.g., per-visit)

        sort_by_time: Whether to sort rows by BJD (default True)

    """

    # Gather core arrays

    bjd = components.get('bjd', None)

    if bjd is None and plData is not None and hasattr(plData, 'time'):

        bjd = np.array(plData.time) + getattr(plData, 'timeCenter', 0.0)

    y = components.get('y', None)

    dy = components.get('dy', None)

    total = components.get('total', None)

    residuals = components.get('residuals', None)

    n_pts = len(bjd) if bjd is not None else len(y)

    if mask is None:

        sel = np.arange(n_pts)

    else:

        sel = np.where(mask)[0] if isinstance(mask, (np.ndarray, list)) else np.arange(n_pts)[mask]

    # Collect term dicts

    poly_terms = components.get('poly_terms', {})

    fourier_terms = components.get('fourier_terms', {})

    gp_terms = components.get('gp_terms', {})

    # Determine column order (place planet columns first)

    planet_terms = components.get('planet_terms', {})

    planet_keys = list(planet_terms.keys())

    poly_keys = sorted(list(poly_terms.keys()))

    fourier_keys = sorted(list(fourier_terms.keys()))

    gp_keys = sorted(list(gp_terms.keys()))

    # Do not include BJD/Y/SIGMA in the output file (they are present in original data)

    header = planet_keys + poly_keys + fourier_keys + gp_keys + ['RESIDUALS']

    # Build columns (planet columns first, then components, then residuals)

    cols = []

    for k in planet_keys:

        cols.append(np.array(planet_terms[k])[sel])

    for k in poly_keys:

        cols.append(np.array(poly_terms[k])[sel])

    for k in fourier_keys:

        cols.append(np.array(fourier_terms[k])[sel])

    for k in gp_keys:

        cols.append(np.array(gp_terms[k])[sel])

    # Residuals

    cols.append(np.array(residuals)[sel])

    # Stack

    data = np.column_stack(cols)

    if sort_by_time:

        if bjd is None:

            _log('WARNING: BJD not available; skipping time-based sorting for components file')

        else:

            sort_idx = np.argsort(np.array(bjd)[sel])

            data = data[sort_idx, :]

    # Format output with equal-width columns for readability

    # Choose column width based on longest header name and numeric width

    n_cols = len(header)

    longest_name = max(len(h) for h in header)

    value_width = 13  # default value field width

    col_width = max(longest_name, value_width) + 2  # padding

    # Write file with fixed-width columns

    try:

        with open(outfile, 'w') as f:

            # Header as comment line with fixed-width columns

            f.write('# ' + ''.join([h.ljust(col_width) for h in header]) + '\n')

            fmt = f'{{:>{col_width}.6g}}'  # general format, right-aligned into width

            for row in data:

                f.write(''.join([fmt.format(v) for v in row]) + '\n')

        _log(f"Wrote model component file: {outfile}")

    except Exception as e:

        _log(f"ERROR: Failed to write model component file {outfile}: {str(e)}")

from .map_search import find_map_parameters

from .diagnostics import compute_mcmc_diagnostics, print_mcmc_summary_table, print_sampling_summary, compute_goodness_of_fit

from .corner_plot import plot_corner, plot_trace, plot_kde_trace, plot_chain_comparison, plot_visit_kde, _prioritize_normalization

from .lightcurve_plots import (

    plot_models,

    plot_noise_budget,

from .plot_models import plot_models

from .model_components import compute_model_components, write_model_components

from .plot_model_helpers import _precompute_sample_data

from .noise_plots import plot_noise_budget

from .periodogram_plots import (

    plot_residuals_histogram,

    plot_periodograms_comparison,

    plot_correlated_noise,

    plot_residuals_vs_dbasis,

)

from .phase_plots import (

    plot_phase_folded,

    plot_phase_folded_rv,

    plot_phase_folded_rv_rm_zoom,

)

from .ttv_plots import (

    plot_ttv_oc,

    plot_ttv_per_transit,

    plot_ttv_transit_overlay,

    compute_model_components,

    write_model_components,

    _precompute_sample_data,