from __future__ import annotations

__all__ = [
    "ProposalPopulationResult",
    "VirtualPopulationResult",
    "concat_proposals",
    "proposal_population",
    "virtual_population",
]

import ast
import math
import warnings
from dataclasses import dataclass
from typing import Any, Literal

import pandas as pd
import tabeline as tl
from ordered_set import OrderedSet

from abm._helper.create_map import (
    create_distributions_map,
    create_doses_map,
    create_fitting_parameters_map,
    create_models_map,
)
from abm._helper.get_dose_schedules import get_dose_schedules
from abm._helper.helper import (
    DataFrameLike,
    PathLike,
    find_duplicate_keys,
    maybe_add_parentheses,
)
from abm._helper.parameter import Parameter
from abm._helper.validate_columns import (
    necessary_dose_columns,
    necessary_fitting_parameter_columns,
    necessary_model_columns,
    optional_dose_columns,
    optional_fitting_parameter_columns,
    validate_and_process_distributions_table,
    validate_and_process_dose_table,
    validate_and_process_fitting_parameter_tables,
    validate_and_process_model_table,
)
from abm._labelset import LabelSet
from abm._sdk import client
from abm._sdk.contract import Contract
from abm._sdk.data_pipe import (
    DataFrame,
    DataPipe,
    OdeProposalPopulationSampleElement,
    OdeVirtualPopulationSampleElement,
)
from abm._sdk.exceptions import JobRuntimeError
from abm._sdk.job import Job
from abm._sdk.likelihood import DistributionsLikelihoodFunction
from abm._sdk.likelihood.distributions import HistogramDistribution, MultivariateNormalDistribution
from abm._sdk.ode_model_reference import OdeModelReference
from abm._sdk.ode_optimization_configuration import OdeOptimizationConfiguration
from abm._sdk.ode_proposal_population_sample import (
    ClassifierSurrogate,
    OdeProposalPopulationSample,
    TruncatedSSE,
)
from abm._sdk.ode_residual_batch import OdeResidualBatch
from abm._sdk.ode_virtual_population_sample import (
    ExpectedPopulation,
    FixedBeta,
    OdeVirtualPopulationSample,
)
from abm._sdk.scenario import (
    LaplacePrior,
    LogLaplacePrior,
    LogNormalPrior,
    LogUniformPrior,
    NormalPrior,
    Scenario,
    ScenarioParameter,
    UniformPrior,
)
from abm._sdk.solver_configuration import (
    BdfSolver,
    KluSolver,
    SolverConfiguration,
    SpgmrSolver,
)
from abm._sdk.status import Status
from abm._units.comparison import is_equivalent
from abm._units.parser import un

from ._helper.get_model import get_model_types_i, get_remote_model
from ._vpop_canonical import (
    canonicalize_proposal_dataframe,
    rebase_group_ids,
    wide_to_canonical_dataframe,
)


def get_solver_configuration(
    gradient_method: Literal["forward", "adjoint"] = "forward",
    linear_solver: Literal["KLU", "SPGMR"] = "KLU",
    reltol: float = 1e-6,
    abstol: float = 1e-9,
    maxstep: float = math.inf,
    maxord: int = 5,
    nonnegative: bool | dict[str, bool] = False,
    max_steps: int | None = None,
) -> SolverConfiguration:
    match linear_solver:
        case "KLU":
            linear_solver_obj = KluSolver()
        case "SPGMR":
            linear_solver_obj = SpgmrSolver()
        case unsupported:
            raise ValueError(f"Unsupported linear solver {unsupported}")

    return SolverConfiguration(
        ode_solver=BdfSolver(
            linear_solver=linear_solver_obj,
            relative_tolerance=reltol,
            absolute_tolerance=abstol,
            max_step=maxstep,
            max_order=maxord,
            max_steps=max_steps,
        ),
        nonnegative=nonnegative,
        gradient_method=gradient_method,
    )


def _build_scenario(
    distributions,
    models,
    parameters,
    doses,
    solver_configuration,
) -> tuple:
    distributions = validate_and_process_distributions_table(distributions)
    processed_models = validate_and_process_model_table(models)
    parameters = validate_and_process_fitting_parameter_tables(parameters)
    doses = validate_and_process_dose_table(doses)

    distributions_column_names = OrderedSet(distributions.column_names)

    # find label columns:
    maybe_label_columns_from_models = set(processed_models.column_names) - necessary_model_columns
    maybe_label_columns_from_parameters = set(parameters.column_names) - (
        necessary_fitting_parameter_columns | optional_fitting_parameter_columns
    )
    maybe_label_columns_from_doses = set(doses.column_names) - (necessary_dose_columns | optional_dose_columns)

    # label columns are the intersection of the columns from measurements and unnecessary columns the models,
    # parameters, and doses
    label_columns = OrderedSet(
        distributions_column_names.intersection(
            maybe_label_columns_from_models | maybe_label_columns_from_parameters | maybe_label_columns_from_doses
        )
    )

    # Assign a unique parameter name for each row for defining the global parameters
    parameters = parameters.mutate(global_parameter_name="parameter + '_fit_' + to_string(row_index0())")

    not_label_columns = distributions_column_names - label_columns

    # make a simulation table out of label columns
    simulation_table = distributions.deselect(*not_label_columns)
    unique_simulations = simulation_table.unique()

    # Add id column for each simulation for residuals
    unique_simulations = unique_simulations.mutate(id="row_index0()")
    simulation_labels = LabelSet.from_df(unique_simulations)

    distributions_labels = label_columns.intersection(distributions_column_names).items
    distributions_map = create_distributions_map(distributions, distributions_labels)

    # build the parameter map
    parameters_labels = label_columns.intersection(OrderedSet(parameters.column_names)).items
    parameters_map = create_fitting_parameters_map(parameters, parameters_labels)

    # build the dose map
    doses_labels = label_columns.intersection(OrderedSet(doses.column_names)).items
    doses_map = create_doses_map(doses, doses_labels)

    # Parse and store the models
    model_paths = set(processed_models[:, "model"])
    remote_model_map = {
        model_path: get_remote_model(model_path).refresh(include_types=True) for model_path in model_paths
    }
    models_labels = label_columns.intersection(OrderedSet(processed_models.column_names)).items
    models_map = create_models_map(processed_models, models_labels)

    likelihood_functions = []
    global_parameters = {}

    # Optim configuration here? method configuration?

    for i_lbs, lbs in enumerate(simulation_labels):
        # Step 1) get model types
        model_list_i = models_map.get_values([lbs.labels.get(k, "*") for k in models_labels])
        model_id, types = get_model_types_i(model_list_i, remote_model_map, lbs.labels)

        # build the reference time unit from the first model:
        if i_lbs == 0:
            reference_time_unit = types.time_unit

        # check if the time unit of the model is the same as the reference time unit
        if not is_equivalent(types.time_unit, reference_time_unit):
            raise ValueError(
                f"all models must have the same time unit but model {model_id} has time unit"
                f" {types.time_unit} which is different from {reference_time_unit}"
            )

        # Step 2) Data matching this set of labels
        data_i = distributions_map.get_values([lbs.labels.get(k, "*") for k in distributions_labels])

        # This is part of step 3 but we need it to update the error expressions with global parameter name
        parameter_list_i = parameters_map.get_values([lbs.labels.get(k, "*") for k in parameters_labels])
        duplicate_parameters = find_duplicate_keys(parameter_list_i)
        if len(duplicate_parameters) > 0:
            raise ValueError(
                "Duplicate parameters found for label(s)"
                f" {', '.join([f'{key}={value}' for key, value in lbs.original_labels.items()])}:"
                f" {', '.join(duplicate_parameters)}"
            )

        # Build the Distributions object
        distributions_i = []
        for d_i in data_i:
            match d_i["time"]:
                case int() | float():
                    time_measurement_j = f"{d_i['time']}:{maybe_add_parentheses(d_i['time_unit'])}"
                case str() if not d_i["time"].startswith("["):
                    time_measurement_j = f"{d_i['time']}:{maybe_add_parentheses(d_i['time_unit'])}"
                case str() if d_i["time"].startswith("["):
                    time_i = ast.literal_eval(d_i["time"])
                    time_unit_i = ast.literal_eval(d_i["time_unit"])
                    time_measurement_j = [
                        f"{t}:{maybe_add_parentheses(u)}" for t, u in zip(time_i, time_unit_i, strict=True)
                    ]

            match d_i["target"]:
                case str() if not d_i["target"].startswith("["):
                    target_j = d_i["target"]
                case str() if d_i["target"].startswith("["):
                    target_j = ast.literal_eval(d_i["target"])
                case _:
                    raise NotImplementedError(f"{d_i['target']} not correctly formatted.")

            match d_i["distribution"]:
                case "HistogramDistribution":
                    bin_edges_bare_j = ast.literal_eval(d_i["bin_edges"])
                    bin_edges_j = [
                        f"{edge}:{maybe_add_parentheses(d_i['distribution_unit'])}" for edge in bin_edges_bare_j
                    ]
                    bin_probabilities_j = ast.literal_eval(d_i["bin_probabilities"])
                    smoothness_j = f"{d_i['smoothness']}:{maybe_add_parentheses(d_i['distribution_unit'])}"

                    distributions_i.append(
                        HistogramDistribution(
                            time=time_measurement_j,
                            target=target_j,
                            bin_edges=bin_edges_j,
                            bin_probabilities=bin_probabilities_j,
                            smoothness=smoothness_j,
                        )
                    )

                case "MultivariateNormalDistribution":
                    measurements_bare_j = ast.literal_eval(d_i["measurements"])
                    measurements_unit_bare_j = ast.literal_eval(d_i["measurements_unit"])
                    measurements_j = [
                        f"{m}:{maybe_add_parentheses(u)}"
                        for m, u in zip(measurements_bare_j, measurements_unit_bare_j, strict=True)
                    ]
                    cov_mat = ast.literal_eval(d_i["variance"])

                    distributions_i.append(
                        MultivariateNormalDistribution(
                            times=time_measurement_j,
                            measurements=measurements_j,
                            variance=cov_mat,
                            targets=target_j,
                        )
                    )

                case _:
                    raise NotImplementedError(f"{d_i['distribution']} not implemented")

        # Add a boolean column for model parameters
        for p in parameter_list_i:
            for v in p.values():
                v["is_model_parameter"] = v["parameter"] in types.parameters

        # Get the subset of parameters that is fit
        parameter_list_i_is_fit = [p for p in parameter_list_i for value in p.values() if value["is_fit"]]

        # Get the subset of parameters that are model parameters
        parameter_list_i_in_model = [p for p in parameter_list_i for value in p.values() if value["is_model_parameter"]]

        # Get the subset of parameters that is fit and are model parameters
        parameter_list_is_fit_is_model = [
            p for p in parameter_list_i for value in p.values() if value["is_fit"] and value["is_model_parameter"]
        ]

        mapping_i = {
            value["parameter"]: value["global_parameter_name"]
            for dic in parameter_list_is_fit_is_model
            for value in dic.values()
        }

        # Add this simulation's global parameters to the master dict that will be used as the starting values for the
        # fit, along with bounds and log-scaling information.
        global_params_i = {}
        for p in parameter_list_i_is_fit:
            for v in p.values():
                value = v["value"]
                unit = v["unit"]

                if "prior_distribution" in parameters.column_names:
                    prior_distribution = v["prior_distribution"]
                else:
                    prior_distribution = "loguniform"
                match prior_distribution:
                    case "uniform":
                        prior = UniformPrior(
                            lower=f"{v['lower_bound']}:{maybe_add_parentheses(unit)}",
                            upper=f"{v['upper_bound']}:{maybe_add_parentheses(unit)}",
                        )
                    case "loguniform":
                        prior = LogUniformPrior(
                            lower=f"{v['lower_bound']}:{maybe_add_parentheses(unit)}",
                            upper=f"{v['upper_bound']}:{maybe_add_parentheses(unit)}",
                        )
                    case "normal":
                        prior = NormalPrior(
                            mean=f"{v['location']}:{maybe_add_parentheses(unit)}",
                            standard_deviation=f"{v['scale']}:{maybe_add_parentheses(unit)}",
                            lower=f"{v['lower_bound']}:{maybe_add_parentheses(unit)}",
                            upper=f"{v['upper_bound']}:{maybe_add_parentheses(unit)}",
                        )
                    case "lognormal":
                        prior = LogNormalPrior(
                            geometric_mean=f"{v['location']}:{maybe_add_parentheses(unit)}",
                            logspace_standard_deviation=v["scale"],  # dimensionless
                            lower=f"{v['lower_bound']}:{maybe_add_parentheses(unit)}",
                            upper=f"{v['upper_bound']}:{maybe_add_parentheses(unit)}",
                        )
                    case "laplace":
                        prior = LaplacePrior(
                            mean=f"{v['location']}:{maybe_add_parentheses(unit)}",
                            mean_absolute_deviation=f"{v['scale']}:{maybe_add_parentheses(unit)}",
                            lower=f"{v['lower_bound']}:{maybe_add_parentheses(unit)}",
                            upper=f"{v['upper_bound']}:{maybe_add_parentheses(unit)}",
                        )
                    case "loglaplace":
                        prior = LogLaplacePrior(
                            geometric_mean=f"{v['location']}:{maybe_add_parentheses(unit)}",
                            logspace_mean_absolute_deviation=v["scale"],  # dimensionless
                            lower=f"{v['lower_bound']}:{maybe_add_parentheses(unit)}",
                            upper=f"{v['upper_bound']}:{maybe_add_parentheses(unit)}",
                        )
                    case _:
                        raise ValueError(
                            f"Unsupported prior distribution {prior_distribution} provided"
                            f" for fitted parameter {v['parameter']}."
                        )
                global_params_i[v["global_parameter_name"]] = ScenarioParameter(
                    value=f"{value}:{maybe_add_parentheses(unit)}",
                    unit=unit,
                    prior=prior,
                )

        # Could just update global_parameters in the loop above, but we follow the same pattern as the other dicts.
        global_parameters.update(global_params_i)

        # dictionary of parameters to update the OdeModelReference
        parameter_dict_i = {}
        for p in parameter_list_i_in_model:
            for v in p.values():
                parameter_dict_i[v["parameter"]] = Parameter(v["value"], v["unit"])

        # Step 4: get the doses for the simulation and do some validations.
        dose_i = doses_map.get_values([lbs.labels.get(k, "*") for k in doses_labels])
        route_schedules_i = get_dose_schedules(dose_i, lbs.original_labels)

        parameter_dict_i_values = {
            key: f"{param.value}:{maybe_add_parentheses(param.unit)}" for key, param in parameter_dict_i.items()
        }

        # Step 5: create the OdeModelReference
        model_reference = OdeModelReference(model_id, parameter_dict_i_values, route_schedules_i)

        # Step 6: Create a MeasurementsLikelihoodFunction and add it to the list:
        likelihood_functions.append(
            DistributionsLikelihoodFunction(
                mapping=mapping_i,
                model=model_reference,
                distributions=distributions_i,
                configuration=solver_configuration,
            )
        )

    scenario = Scenario(global_parameters=global_parameters, likelihood_functions=likelihood_functions)

    return scenario, label_columns, unique_simulations, parameters


def proposal_population(
    *,
    distributions: DataFrameLike | PathLike,
    models: DataFrameLike | PathLike,
    parameters: DataFrameLike | PathLike,
    doses: DataFrameLike | PathLike | None = None,
    proposal_method: Literal["TruncatedSSE", "ClassifierSurrogate"],
    proposal_options: dict = {},
    gradient_method: Literal["forward", "adjoint"] = "forward",
    linear_solver: Literal["KLU", "SPGMR"] = "KLU",
    reltol: float = 1e-6,
    abstol: float = 1e-9,
    maxstep: float = math.inf,
    maxord: int = 5,
    nonnegative: bool | dict[str, bool] = False,
    max_steps: int | None = None,
    seed: int = 0,
) -> ProposalPopulationResult:
    solver_configuration = get_solver_configuration(
        gradient_method=gradient_method,
        linear_solver=linear_solver,
        reltol=reltol,
        abstol=abstol,
        maxstep=maxstep,
        maxord=maxord,
        nonnegative=nonnegative,
        max_steps=max_steps,
    )

    (
        proposal_scenario,
        proposal_label_columns,
        proposal_unique_simulations,
        _,
    ) = _build_scenario(
        distributions=distributions,
        models=models,
        parameters=parameters,
        doses=doses,
        solver_configuration=solver_configuration,
    )

    proposal_method_obj = _build_proposal_method(proposal_method, proposal_options)

    proposal = OdeProposalPopulationSample(seed=seed, method=proposal_method_obj, scenario=proposal_scenario)
    proposal_jobs = client.create_jobs([proposal], deduplicate=True)
    proposal_contract = client.create_contract(proposal_jobs, progress=True)
    data_pipe = OdeProposalPopulationSampleElement(proposal_population_sample_id=proposal_jobs[0].id)

    result = ProposalPopulationResult(
        label_columns=proposal_label_columns,
        simulations=proposal_unique_simulations,
        models=models,
        parameters=parameters,
        distributions=distributions,
        doses=doses,
        solver_configuration=solver_configuration,
        scenario=proposal_scenario,
        seed=seed,
        _contract=proposal_contract,
        _job=proposal_jobs[0],
        _data_pipe=data_pipe,
        _population=proposal,
        _feasibility_checked=proposal_method == "ClassifierSurrogate",
    )

    if proposal_method == "TruncatedSSE":
        result._ensure_feasible_data_pipe()

    return result


@dataclass
class ProposalPopulationResult:
    label_columns: OrderedSet[str]
    simulations: tl.DataFrame

    models: DataFrameLike | PathLike
    parameters: DataFrameLike | PathLike
    distributions: DataFrameLike | PathLike
    doses: DataFrameLike | PathLike | None

    solver_configuration: SolverConfiguration
    scenario: Scenario
    seed: int

    _contract: Contract | None
    _job: Job[OdeProposalPopulationSample, None] | None
    _data_pipe: DataPipe

    _population: OdeProposalPopulationSample | None
    _feasibility_checked: bool = False

    @property
    def n_candidates(self) -> int:
        return int(self.to_pandas()["id"].nunique())

    @property
    def parameter_names(self) -> list[str]:
        return sorted(self.to_pandas()["parameter_name"].unique())

    def to_pandas(self) -> pd.DataFrame:
        return self._data_pipe.run().to_pandas()

    def to_data_pipe(self) -> DataPipe:
        return self._data_pipe

    @classmethod
    def from_dataframe(
        cls,
        dataframe: pd.DataFrame,
        *,
        distributions: DataFrameLike | PathLike,
        models: DataFrameLike | PathLike,
        parameters: DataFrameLike | PathLike,
        doses: DataFrameLike | PathLike | None = None,
        gradient_method: Literal["forward", "adjoint"] = "forward",
        linear_solver: Literal["KLU", "SPGMR"] = "KLU",
        reltol: float = 1e-6,
        abstol: float = 1e-9,
        maxstep: float = math.inf,
        maxord: int = 5,
        nonnegative: bool | dict[str, bool] = False,
        max_steps: int | None = None,
        seed: int = 0,
        rebase_ids: bool = False,
        filter_feasible: bool = True,
        source_name: str = "proposal dataframe",
    ) -> ProposalPopulationResult:
        solver_configuration = get_solver_configuration(
            gradient_method=gradient_method,
            linear_solver=linear_solver,
            reltol=reltol,
            abstol=abstol,
            maxstep=maxstep,
            maxord=maxord,
            nonnegative=nonnegative,
            max_steps=max_steps,
        )
        scenario, label_columns, simulations, _ = _build_scenario(
            distributions=distributions,
            models=models,
            parameters=parameters,
            doses=doses,
            solver_configuration=solver_configuration,
        )
        canonical = canonicalize_proposal_dataframe(
            dataframe,
            rebase_ids=rebase_ids,
            source_name=source_name,
        )
        canonical = _normalize_proposal_parameter_names(
            canonical,
            parameters=parameters,
            expected_global_parameters=set(scenario.global_parameters),
        )
        _validate_proposal_against_scenario(canonical, scenario)

        result = cls(
            label_columns=label_columns,
            simulations=simulations,
            models=models,
            parameters=parameters,
            distributions=distributions,
            doses=doses,
            solver_configuration=solver_configuration,
            scenario=scenario,
            seed=seed,
            _contract=None,
            _job=None,
            _data_pipe=DataFrame.from_pandas(canonical).to_data_pipe(),
            _population=None,
            _feasibility_checked=False,
        )

        if filter_feasible:
            result._set_feasible_dataframe(canonical)

        return result

    @classmethod
    def from_wide_dataframe(
        cls,
        wide_dataframe: pd.DataFrame,
        *,
        parameter_units: dict[str, str] | pd.DataFrame,
        distributions: DataFrameLike | PathLike,
        models: DataFrameLike | PathLike,
        parameters: DataFrameLike | PathLike,
        doses: DataFrameLike | PathLike | None = None,
        id_column: str | None = None,
        parameter_columns: list[str] | None = None,
        drop_missing_values: bool = False,
        gradient_method: Literal["forward", "adjoint"] = "forward",
        linear_solver: Literal["KLU", "SPGMR"] = "KLU",
        reltol: float = 1e-6,
        abstol: float = 1e-9,
        maxstep: float = math.inf,
        maxord: int = 5,
        nonnegative: bool | dict[str, bool] = False,
        max_steps: int | None = None,
        seed: int = 0,
        filter_feasible: bool = True,
        source_name: str = "proposal wide dataframe",
    ) -> ProposalPopulationResult:
        canonical = wide_to_canonical_dataframe(
            wide_dataframe,
            parameter_units=parameter_units,
            id_column=id_column,
            parameter_columns=parameter_columns,
            drop_missing_values=drop_missing_values,
            source_name=source_name,
        )
        return cls.from_dataframe(
            canonical,
            distributions=distributions,
            models=models,
            parameters=parameters,
            doses=doses,
            gradient_method=gradient_method,
            linear_solver=linear_solver,
            reltol=reltol,
            abstol=abstol,
            maxstep=maxstep,
            maxord=maxord,
            nonnegative=nonnegative,
            max_steps=max_steps,
            seed=seed,
            filter_feasible=filter_feasible,
            source_name=source_name,
        )

    def _set_feasible_dataframe(self, dataframe: pd.DataFrame) -> None:
        feasible = self.filter_to_feasible_samples(dataframe)
        self._data_pipe = DataFrame.from_pandas(feasible).to_data_pipe()
        self._feasibility_checked = True

    def _ensure_feasible_data_pipe(self) -> None:
        if self._feasibility_checked:
            return

        self._set_feasible_dataframe(self.to_pandas())

    def filter_to_feasible_samples(self, dataframe: pd.DataFrame) -> pd.DataFrame:
        canonical = canonicalize_proposal_dataframe(
            dataframe,
            source_name="proposal dataframe",
        )
        if len(canonical) == 0:
            raise ValueError("Proposal dataframe cannot be empty")

        unique_ids = canonical["id"].drop_duplicates().tolist()
        definition = OdeResidualBatch(
            data_pipe=DataFrame.from_pandas(canonical).to_data_pipe(),
            scenario=self.scenario,
        )

        jobs = client.create_jobs([definition], deduplicate=True)
        client.create_contract(jobs, progress=True)

        residual_job_ids = jobs[0].output_or_raise()
        successful_ids = _successful_group_ids(unique_ids, residual_job_ids)
        if not successful_ids:
            raise ValueError("None of the proposal samples were feasible.")

        feasible = canonical.loc[canonical["id"].isin(successful_ids)].copy()
        return feasible.sort_values(["id", "parameter_name"], kind="stable").reset_index(drop=True)

    def sample_virtual_population(
        self,
        distributions: DataFrameLike | PathLike,
        method: Literal["ExpectedPopulation", "FixedBeta"] = "ExpectedPopulation",
        seed: int | None = None,
        **kwargs,
    ) -> VirtualPopulationResult:
        self._ensure_feasible_data_pipe()

        if seed is None:
            seed = self.seed

        vpop_method = _build_virtual_method(method, kwargs)

        (vpop_scenario, virtual_label_columns, virtual_unique_simulations, virtual_parameters) = _build_scenario(
            distributions=distributions,
            models=self.models,
            parameters=self.parameters,
            doses=self.doses,
            solver_configuration=self.solver_configuration,
        )

        virtual_population = OdeVirtualPopulationSample(
            seed=seed,
            data_pipe=self._data_pipe,
            method=vpop_method,
            scenario=vpop_scenario,
        )

        vpop_sample_job = client.create_jobs([virtual_population], deduplicate=True)
        virtual_population_contract = client.create_contract(vpop_sample_job, progress=True)

        return VirtualPopulationResult(
            proposal_result=self,
            virtual_label_columns=virtual_label_columns,
            virtual_parameters=virtual_parameters,
            virtual_simulations=virtual_unique_simulations,
            distributions=distributions,
            virtual_scenario=vpop_scenario,
            _virtual_job=vpop_sample_job[0],
            _virtual_contract=virtual_population_contract,
            _virtual_population=virtual_population,
        )


def concat_proposals(
    proposals: list[ProposalPopulationResult],
) -> ProposalPopulationResult:
    if len(proposals) == 0:
        raise ValueError("concat_proposals requires at least one proposal")

    first = proposals[0]
    combined_frames = []
    next_id = 0
    all_feasible = True
    for proposal in proposals:
        canonical = canonicalize_proposal_dataframe(
            proposal.to_pandas(),
            source_name="proposal dataframe",
        )
        rebased, next_id = rebase_group_ids(canonical, start_id=next_id)
        combined_frames.append(rebased)
        all_feasible = all_feasible and proposal._feasibility_checked

    combined = pd.concat(combined_frames, ignore_index=True)
    result = ProposalPopulationResult(
        label_columns=first.label_columns,
        simulations=first.simulations,
        models=first.models,
        parameters=first.parameters,
        distributions=first.distributions,
        doses=first.doses,
        solver_configuration=first.solver_configuration,
        scenario=first.scenario,
        seed=first.seed,
        _contract=None,
        _job=None,
        _data_pipe=DataFrame.from_pandas(combined).to_data_pipe(),
        _population=None,
        _feasibility_checked=all_feasible,
    )
    if not all_feasible:
        result._set_feasible_dataframe(combined)
    return result


@dataclass
class VirtualPopulationResult:
    proposal_result: ProposalPopulationResult
    virtual_label_columns: OrderedSet[str]
    virtual_parameters: tl.DataFrame
    virtual_simulations: tl.DataFrame
    distributions: tl.DataFrame

    virtual_scenario: Scenario

    _virtual_contract: Contract
    _virtual_job: Job[OdeVirtualPopulationSample, None]
    _virtual_population: OdeVirtualPopulationSample

    # residual job fields
    _residual_batch_definition: OdeResidualBatch | None = None
    _residual_batch_job: Job[OdeResidualBatch, None] | None = None
    _residual_batch_contract: Contract | None = None

    @property
    def _initial_global_parameters_proposal(self):
        return self.proposal_result.scenario.global_parameters

    @property
    def _initial_global_parameters_virtual(self):
        return self.virtual_scenario.global_parameters

    def _residual_batch(self):
        vpop_sample_pipe = OdeVirtualPopulationSampleElement(virtual_population_sample_id=self._virtual_job.id)

        self._residual_batch_definition = OdeResidualBatch(
            data_pipe=vpop_sample_pipe,
            scenario=self.virtual_scenario,
        )

        self._residual_batch_job = client.create_jobs([self._residual_batch_definition], deduplicate=True)[0]
        self._residual_batch_contract = client.create_contract([self._residual_batch_job], progress=True)

        return self._residual_batch_job.output_or_raise()

    def _targets(self):
        output = []
        for i, lf in enumerate(self.virtual_scenario.likelihood_functions):
            labels = ""
            if len(self.virtual_label_columns) > 0:
                simulation_records = self.virtual_simulations[i, :]
                labels = "_".join(str(v) for _, v in simulation_records.items())
            for dist in lf.distributions:
                if hasattr(dist, "target"):
                    targets = [f"{dist.target}_{dist.time}" if not labels else f"{dist.target}_{labels}_{dist.time}"]
                elif hasattr(dist, "targets"):
                    targets = [
                        f"{d}_{t}" if not labels else f"{d}_{labels}_{t}"
                        for (d, t) in zip(dist.targets, dist.times, strict=False)
                    ]
                else:
                    raise ValueError("Distribution does not have target or targets attribute")
                output.extend(targets)
        return output

    def virtual_population_predictions(self):
        residuals = self._residual_batch()

        likelihood_function_residuals = []
        for residual_job_id in residuals:
            residual_job = Job.from_id(residual_job_id)
            if residual_job.status == Status.succeeded:
                likelihood_function_residuals.append(residual_job.output_or_raise().likelihood_function_residuals)
                continue

            if _is_histogram_job_failure(residual_job):
                warnings.warn(
                    "Skipped one or more virtual population residual jobs because histogram support "
                    "errors made those samples infeasible.",
                    RuntimeWarning,
                    stacklevel=2,
                )
                continue

            residual_job.raise_if_error()
            raise RuntimeError("Unexpected residual job state")

        output = []
        for lf_residual in likelihood_function_residuals:
            aggregate = []
            for residuals in lf_residual:
                for residual in residuals.residuals:
                    match residual:
                        case _ if hasattr(residual, "prediction"):
                            predictions = [residual.prediction]
                        case _ if hasattr(residual, "predictions"):
                            predictions = residual.predictions

                    aggregate.extend(predictions)

            output.append(aggregate)

        targets = self._targets()

        return pd.DataFrame(output, columns=targets)

    def canonical_virtual_population_table(self) -> pd.DataFrame:
        vpop_sample_pipe = OdeVirtualPopulationSampleElement(virtual_population_sample_id=self._virtual_job.id)

        job = client.create_jobs([vpop_sample_pipe], deduplicate=True)
        client.create_contract(job, progress=True)

        try:
            return job[0].output_or_raise().to_pandas()
        except JobRuntimeError as error:
            if not _is_histogram_runtime_error(error):
                raise

            warnings.warn(
                "Suppressed a histogram support error while materializing the virtual population "
                "table. Falling back to the proposal population table.",
                RuntimeWarning,
                stacklevel=2,
            )
            return self.proposal_result.to_pandas()

    def virtual_population_tables(self) -> tuple[pd.DataFrame, pd.DataFrame]:
        vpop_parameters = self.canonical_virtual_population_table()
        transformed_params = self.virtual_parameters.to_pandas()
        simulations = self.virtual_simulations.to_pandas()
        label_columns = list(self.virtual_label_columns)

        return self._rebuild_tables(transformed_params, vpop_parameters, simulations, label_columns)

    @staticmethod
    def _rebuild_tables(
        transformed_params: pd.DataFrame,
        results: pd.DataFrame,
        simulations: pd.DataFrame,
        label_columns=None,
    ):
        """
        Reconstructs parameter and simulation tables from VPop results.

        Merges original parameters with fitted results and expands simulation
        tables for each fitted sample, preserving metadata.
        """
        if label_columns is None:
            label_columns = []

        # Working copies
        params = transformed_params.copy()
        sims = simulations.copy()

        # Normalize label columns to use "*" for missing values
        for col in label_columns:
            for df in (params, sims):
                if col not in df.columns:
                    raise ValueError(f"Label column '{col}' missing")
                df[col] = df[col].fillna("*").replace({None: "*"})

        # Identify fit vs unfit parameters
        fit_names = set(results["parameter_name"].unique())
        is_fit = params["global_parameter_name"].isin(fit_names)

        # Columns handling
        core_cols = {"parameter", "value", "unit", "global_parameter_name"}
        meta_cols = [c for c in params.columns if c not in core_cols and c not in label_columns]
        final_cols = ["fit_id", *label_columns, "parameter", "value", "unit", *meta_cols]

        # 1. Process Unfit Parameters (fit_id = "*")
        df_unfit = params[~is_fit].copy()
        df_unfit["fit_id"] = "*"

        # 2. Process Fit Parameters (fit_id = id from results)
        df_fit = (
            params[is_fit]
            .merge(
                results,
                left_on="global_parameter_name",
                right_on="parameter_name",
                how="left",
            )
            .dropna(subset=["id"])
        )

        df_fit["fit_id"] = df_fit["id"].astype(int)
        df_fit["value"] = df_fit["parameter_value"]

        # 3. Combine and Sort
        # Sort by: Unfit first, then Parameter name, then Fit ID
        # Note: We avoid comparing string "*" with int IDs by using a helper key
        combined = pd.concat([df_unfit, df_fit], ignore_index=True)
        combined["_sort_group"] = (combined["fit_id"] != "*").astype(int)

        param_out = combined.sort_values(["_sort_group", "parameter", "fit_id"])[final_cols]

        # 4. Build Simulation Table (Cross join: fit_ids x simulations)
        unique_fit_ids = pd.DataFrame({"fit_id": sorted(df_fit["fit_id"].unique())})

        # Perform cross join using a temporary key for compatibility
        unique_fit_ids["_key"] = 1
        sims_subset = sims[label_columns].copy()
        sims_subset["_key"] = 1

        sims_out = pd.merge(sims_subset, unique_fit_ids, on="_key").drop(columns=["_key"])
        sims_out = sims_out[["fit_id", *label_columns]]

        return param_out, sims_out


def _build_proposal_method(
    proposal_method: Literal["TruncatedSSE", "ClassifierSurrogate"],
    proposal_options: dict[str, Any],
):
    options = dict(proposal_options)
    if "optimization_configuration" in options and isinstance(options["optimization_configuration"], dict):
        options["optimization_configuration"] = OdeOptimizationConfiguration(**options["optimization_configuration"])

    match proposal_method:
        case "ClassifierSurrogate":
            return ClassifierSurrogate(**options)
        case "TruncatedSSE":
            return TruncatedSSE(**options)
        case unsupported:
            raise ValueError(f"Unsupported proposal method {unsupported}")


def _build_virtual_method(
    virtual_method: Literal["ExpectedPopulation", "FixedBeta"],
    virtual_options: dict[str, Any],
):
    match virtual_method:
        case "ExpectedPopulation":
            return ExpectedPopulation(**virtual_options)
        case "FixedBeta":
            return FixedBeta(**virtual_options)
        case unsupported:
            raise ValueError(f"Unsupported virtual method {unsupported}")


def _validate_proposal_against_scenario(
    proposal: pd.DataFrame,
    scenario: Scenario,
) -> None:
    expected_parameters = scenario.global_parameters
    actual_parameter_names = set(proposal["parameter_name"])
    unknown_parameters = sorted(actual_parameter_names - set(expected_parameters))
    if unknown_parameters:
        unknown = ", ".join(unknown_parameters)
        raise ValueError(
            f"Proposal population contains parameters that are not present in the filtering scenario: {unknown}"
        )

    unit_table = (
        proposal.loc[:, ["parameter_name", "parameter_unit"]]
        .drop_duplicates()
        .sort_values(["parameter_name", "parameter_unit"], kind="stable")
    )
    duplicated_units = unit_table.duplicated(subset=["parameter_name"], keep=False)
    if duplicated_units.any():
        duplicates = ", ".join(sorted(unit_table.loc[duplicated_units, "parameter_name"].unique()))
        raise ValueError(f"Proposal population contains conflicting units for one or more parameters: {duplicates}")

    for row in unit_table.itertuples(index=False):
        scenario_unit = expected_parameters[row.parameter_name].unit
        if scenario_unit is None:
            continue
        if not is_equivalent(un(str(row.parameter_unit)), un(str(scenario_unit))):
            raise ValueError(
                "Proposal population parameter units must match the filtering scenario. "
                f"Parameter '{row.parameter_name}' used '{row.parameter_unit}' but expected '{scenario_unit}'."
            )


def _normalize_proposal_parameter_names(
    proposal: pd.DataFrame,
    *,
    parameters,
    expected_global_parameters: set[str],
) -> pd.DataFrame:
    actual_parameter_names = set(proposal["parameter_name"])
    if actual_parameter_names.issubset(expected_global_parameters):
        return proposal

    transformed_parameters = validate_and_process_fitting_parameter_tables(parameters).mutate(
        global_parameter_name="parameter + '_fit_' + to_string(row_index0())"
    )
    transformed_parameters_pd = transformed_parameters.to_pandas()
    fit_parameters = transformed_parameters_pd.loc[transformed_parameters_pd["is_fit"]].copy()

    duplicate_parameter_names = fit_parameters["parameter"].duplicated(keep=False)
    if duplicate_parameter_names.any():
        duplicates = ", ".join(sorted(fit_parameters.loc[duplicate_parameter_names, "parameter"].unique()))
        raise ValueError(
            "Could not normalize proposal parameter names because the fitting parameter "
            f"table contains multiple fitted rows for the same parameter name: {duplicates}"
        )

    parameter_name_map = dict(
        zip(
            fit_parameters["parameter"],
            fit_parameters["global_parameter_name"],
            strict=True,
        )
    )

    unknown_parameters = sorted(actual_parameter_names - set(parameter_name_map))
    if unknown_parameters:
        unknown = ", ".join(unknown_parameters)
        raise ValueError(
            "Proposal population contains parameters that are not present in the filtering scenario "
            f"or the fitted parameter table: {unknown}"
        )

    normalized = proposal.copy()
    normalized["parameter_name"] = normalized["parameter_name"].map(parameter_name_map)
    return normalized


def _successful_group_ids(unique_ids: list[int], residual_job_ids: list[str]) -> list[int]:
    if len(unique_ids) != len(residual_job_ids):
        raise ValueError("Residual batch output did not align with proposal ids")

    successful_ids = []
    for proposal_id, residual_job_id in zip(unique_ids, residual_job_ids, strict=True):
        residual_job = Job.from_id(residual_job_id)
        if residual_job.status == Status.succeeded:
            successful_ids.append(int(proposal_id))

    return successful_ids


def _is_histogram_runtime_error(error: JobRuntimeError) -> bool:
    text = f"{error._type}: {error.message}"
    return _is_histogram_error_text(text)


def _is_histogram_job_failure(job: Job) -> bool:
    if job.exception is None:
        return False
    text = f"{job.exception.get('_type', '')}: {job.exception.get('message', '')}"
    return _is_histogram_error_text(text)


def _is_histogram_error_text(text: str) -> bool:
    return "HistogramException" in text or ("Bin Edges:" in text and "Bin Probabilities:" in text)


def virtual_population(
    *,
    distributions: DataFrameLike | PathLike,
    models: DataFrameLike | PathLike,
    parameters: DataFrameLike | PathLike,
    doses: DataFrameLike | PathLike | None = None,
    proposal_method: Literal["TruncatedSSE", "ClassifierSurrogate"],
    virtual_method: Literal["ExpectedPopulation", "FixedBeta"] = "ExpectedPopulation",
    proposal_distributions: DataFrameLike | PathLike | None = None,
    proposal_options: dict = {},
    virtual_options: dict = {},
    gradient_method: Literal["forward", "adjoint"] = "forward",
    linear_solver: Literal["KLU", "SPGMR"] = "KLU",
    reltol: float = 1e-6,
    abstol: float = 1e-9,
    maxstep: float = math.inf,
    maxord: int = 5,
    nonnegative: bool | dict[str, bool] = False,
    max_steps: int | None = None,
    seed: int = 0,
) -> VirtualPopulationResult:
    proposal_result = proposal_population(
        distributions=distributions if proposal_distributions is None else proposal_distributions,
        models=models,
        parameters=parameters,
        doses=doses,
        proposal_method=proposal_method,
        proposal_options=proposal_options,
        gradient_method=gradient_method,
        linear_solver=linear_solver,
        reltol=reltol,
        abstol=abstol,
        maxstep=maxstep,
        maxord=maxord,
        nonnegative=nonnegative,
        max_steps=max_steps,
        seed=seed,
    )

    return proposal_result.sample_virtual_population(
        distributions=distributions, method=virtual_method, seed=seed, **virtual_options
    )
