import os
from math import nan

import numpy as np
import pandas as pd

import petab
from petab.C import *
from petab.models.sbml_model import SbmlModel
from petab.parameter_mapping import _apply_parameter_table

# import fixtures
pytest_plugins = [
    "tests.test_petab",
]


class TestGetSimulationToOptimizationParameterMapping(object):
    @staticmethod
    def test_no_condition_specific(condition_df_2_conditions):
        # Trivial case - no condition-specific parameters

        condition_df = condition_df_2_conditions

        measurement_df = pd.DataFrame(
            data={
                OBSERVABLE_ID: ["obs1", "obs2"],
                SIMULATION_CONDITION_ID: ["condition1", "condition2"],
                PREEQUILIBRATION_CONDITION_ID: ["", ""],
                OBSERVABLE_PARAMETERS: ["", ""],
                NOISE_PARAMETERS: ["", ""],
            }
        )

        import simplesbml

        ss_model = simplesbml.SbmlModel()
        ss_model.addParameter("dynamicParameter1", 1.0)
        ss_model.addParameter("dynamicParameter2", 2.0)
        ss_model.addParameter("dynamicParameter3", 3.0)

        # add species, which will have initial concentration in condition table
        #  but which should not show up in mapping
        ss_model.addSpecies("[someSpecies]", 1.0)
        condition_df["someSpecies"] = [0.0, 0.0]

        # Test without parameter table
        expected = [
            (
                {},
                {
                    "dynamicParameter1": 1.0,
                    "dynamicParameter2": 2.0,
                    "dynamicParameter3": 3.0,
                    "fixedParameter1": 1.0,
                },
                {},
                {
                    "dynamicParameter1": LIN,
                    "dynamicParameter2": LIN,
                    "dynamicParameter3": LIN,
                    "fixedParameter1": LIN,
                },
            ),
            (
                {},
                {
                    "dynamicParameter1": 1.0,
                    "dynamicParameter2": 2.0,
                    "dynamicParameter3": 3.0,
                    "fixedParameter1": 2.0,
                },
                {},
                {
                    "dynamicParameter1": LIN,
                    "dynamicParameter2": LIN,
                    "dynamicParameter3": LIN,
                    "fixedParameter1": LIN,
                },
            ),
        ]

        model = SbmlModel(sbml_model=ss_model.model)
        actual = petab.get_optimization_to_simulation_parameter_mapping(
            model=model,
            measurement_df=measurement_df,
            condition_df=condition_df,
        )
        assert actual == expected

        # Test with parameter table
        parameter_df = pd.DataFrame(
            data={
                PARAMETER_ID: [
                    "dynamicParameter1",
                    "dynamicParameter2",
                    "dynamicParameter3",
                ],
                ESTIMATE: [0, 1, 1],
                NOMINAL_VALUE: [11.0, 12.0, None],
                PARAMETER_SCALE: [LOG, LOG10, LIN],
            }
        )
        parameter_df.set_index(PARAMETER_ID, inplace=True)

        expected = [
            (
                {},
                {
                    "dynamicParameter1": 11.0,
                    "dynamicParameter2": "dynamicParameter2",
                    "dynamicParameter3": "dynamicParameter3",
                    "fixedParameter1": 1.0,
                },
                {},
                {
                    "dynamicParameter1": LIN,
                    "dynamicParameter2": LOG10,
                    "dynamicParameter3": LIN,
                    "fixedParameter1": LIN,
                },
            ),
            (
                {},
                {
                    "dynamicParameter1": 11.0,
                    "dynamicParameter2": "dynamicParameter2",
                    "dynamicParameter3": "dynamicParameter3",
                    "fixedParameter1": 2.0,
                },
                {},
                {
                    "dynamicParameter1": LIN,
                    "dynamicParameter2": LOG10,
                    "dynamicParameter3": LIN,
                    "fixedParameter1": LIN,
                },
            ),
        ]

        actual = petab.get_optimization_to_simulation_parameter_mapping(
            model=model,
            measurement_df=measurement_df,
            condition_df=condition_df,
            parameter_df=parameter_df,
        )

        assert actual == expected

        # Test with applied scaling

        expected = [
            (
                {},
                {
                    "dynamicParameter1": np.log(11.0),
                    "dynamicParameter2": "dynamicParameter2",
                    "dynamicParameter3": "dynamicParameter3",
                    "fixedParameter1": 1.0,
                },
                {},
                {
                    "dynamicParameter1": LOG,
                    "dynamicParameter2": LOG10,
                    "dynamicParameter3": LIN,
                    "fixedParameter1": LIN,
                },
            ),
            (
                {},
                {
                    "dynamicParameter1": np.log(11.0),
                    "dynamicParameter2": "dynamicParameter2",
                    "dynamicParameter3": "dynamicParameter3",
                    "fixedParameter1": 2.0,
                },
                {},
                {
                    "dynamicParameter1": LOG,
                    "dynamicParameter2": LOG10,
                    "dynamicParameter3": LIN,
                    "fixedParameter1": LIN,
                },
            ),
        ]

        actual = petab.get_optimization_to_simulation_parameter_mapping(
            model=model,
            measurement_df=measurement_df,
            condition_df=condition_df,
            parameter_df=parameter_df,
            scaled_parameters=True,
        )

        assert actual == expected

        # Test without fixed overrides

        expected = [
            (
                {},
                {
                    "dynamicParameter1": "dynamicParameter1",
                    "dynamicParameter2": "dynamicParameter2",
                    "dynamicParameter3": "dynamicParameter3",
                    "fixedParameter1": 1.0,
                },
                {},
                {
                    "dynamicParameter1": LOG,
                    "dynamicParameter2": LOG10,
                    "dynamicParameter3": LIN,
                    "fixedParameter1": LIN,
                },
            ),
            (
                {},
                {
                    "dynamicParameter1": "dynamicParameter1",
                    "dynamicParameter2": "dynamicParameter2",
                    "dynamicParameter3": "dynamicParameter3",
                    "fixedParameter1": 2.0,
                },
                {},
                {
                    "dynamicParameter1": LOG,
                    "dynamicParameter2": LOG10,
                    "dynamicParameter3": LIN,
                    "fixedParameter1": LIN,
                },
            ),
        ]

        actual = petab.get_optimization_to_simulation_parameter_mapping(
            model=model,
            measurement_df=measurement_df,
            condition_df=condition_df,
            parameter_df=parameter_df,
            fill_fixed_parameters=False,
        )

        assert actual == expected

    @staticmethod
    def test_all_override(condition_df_2_conditions):
        # Condition-specific parameters overriding original parameters
        condition_df = condition_df_2_conditions

        import simplesbml

        ss_model = simplesbml.SbmlModel()
        ss_model.addParameter("dynamicParameter1", 0.0)
        ss_model.addParameter("dynamicParameter2", 0.0)
        model = SbmlModel(sbml_model=ss_model.model)

        measurement_df = pd.DataFrame(
            data={
                OBSERVABLE_ID: ["obs1", "obs2", "obs1", "obs2"],
                SIMULATION_CONDITION_ID: [
                    "condition1",
                    "condition1",
                    "condition2",
                    "condition2",
                ],
                PREEQUILIBRATION_CONDITION_ID: ["", "", "", ""],
                OBSERVABLE_PARAMETERS: [
                    "obs1par1override;obs1par2cond1override",
                    "obs2par1cond1override",
                    "obs1par1override;obs1par2cond2override",
                    "obs2par1cond2override",
                ],
                NOISE_PARAMETERS: ["", "", "", ""],
            }
        )

        parameter_df = pd.DataFrame(
            data={
                PARAMETER_ID: [
                    "dynamicParameter1",
                    "dynamicParameter2",
                    "obs1par1override",
                    "obs1par2cond1override",
                    "obs1par2cond2override",
                    "obs2par1cond1override",
                    "obs2par1cond2override",
                ],
                ESTIMATE: [1] * 7,
            }
        )
        parameter_df.set_index(PARAMETER_ID, inplace=True)

        expected = [
            (
                {},
                {
                    "fixedParameter1": 1.0,
                    "dynamicParameter1": "dynamicParameter1",
                    "dynamicParameter2": "dynamicParameter2",
                    "observableParameter1_obs1": "obs1par1override",
                    "observableParameter2_obs1": "obs1par2cond1override",
                    "observableParameter1_obs2": "obs2par1cond1override",
                },
                {},
                {
                    "fixedParameter1": LIN,
                    "dynamicParameter1": LIN,
                    "dynamicParameter2": LIN,
                    "observableParameter1_obs1": LIN,
                    "observableParameter2_obs1": LIN,
                    "observableParameter1_obs2": LIN,
                },
            ),
            (
                {},
                {
                    "fixedParameter1": 2.0,
                    "dynamicParameter1": "dynamicParameter1",
                    "dynamicParameter2": "dynamicParameter2",
                    "observableParameter1_obs1": "obs1par1override",
                    "observableParameter2_obs1": "obs1par2cond2override",
                    "observableParameter1_obs2": "obs2par1cond2override",
                },
                {},
                {
                    "fixedParameter1": LIN,
                    "dynamicParameter1": LIN,
                    "dynamicParameter2": LIN,
                    "observableParameter1_obs1": LIN,
                    "observableParameter2_obs1": LIN,
                    "observableParameter1_obs2": LIN,
                },
            ),
        ]

        actual = petab.get_optimization_to_simulation_parameter_mapping(
            measurement_df=measurement_df,
            condition_df=condition_df,
            model=model,
            parameter_df=parameter_df,
        )
        assert actual == expected

        # For one case we test parallel execution, which must yield the same
        # result
        os.environ[petab.ENV_NUM_THREADS] = "4"
        actual = petab.get_optimization_to_simulation_parameter_mapping(
            measurement_df=measurement_df,
            condition_df=condition_df,
            model=model,
            parameter_df=parameter_df,
        )
        assert actual == expected

    @staticmethod
    def test_partial_override(condition_df_2_conditions):
        # Condition-specific parameters, keeping original parameters
        condition_df = pd.DataFrame(
            data={
                "conditionId": ["condition1", "condition2"],
                "conditionName": ["", "Condition 2"],
                "fixedParameter1": [1.0, 2.0],
                "fixedParameter2": [nan, 2.5],
            }
        )
        condition_df.set_index("conditionId", inplace=True)

        import simplesbml

        ss_model = simplesbml.SbmlModel()
        ss_model.addParameter("fixedParameter1", 0.5)
        ss_model.addParameter("fixedParameter2", 1.0)
        ss_model.addParameter("dynamicParameter1", 0.0)
        ss_model.addParameter("observableParameter1_obs1", 0.0)
        ss_model.addParameter("observableParameter2_obs1", 0.0)
        ss_model.addParameter("observableParameter1_obs2", 0.0)

        measurement_df = pd.DataFrame(
            data={
                OBSERVABLE_ID: ["obs1", "obs2", "obs1", "obs2"],
                SIMULATION_CONDITION_ID: [
                    "condition1",
                    "condition1",
                    "condition2",
                    "condition2",
                ],
                PREEQUILIBRATION_CONDITION_ID: ["", "", "", ""],
                OBSERVABLE_PARAMETERS: [
                    "obs1par1override;obs1par2cond1override",
                    "",
                    "obs1par1override;obs1par2cond2override",
                    "obs2par1cond2override",
                ],
                NOISE_PARAMETERS: ["", "", "", ""],
            }
        )

        parameter_df = pd.DataFrame(
            data={
                PARAMETER_ID: [
                    "dynamicParameter1",
                    "obs1par1override",
                    "obs1par2cond1override",
                    "obs1par2cond2override",
                    "obs2par1cond2override",
                ],
                ESTIMATE: [1, 1, 1, 1, 1],
            }
        )
        parameter_df.set_index(PARAMETER_ID, inplace=True)

        expected = [
            (
                {},
                {
                    "fixedParameter1": 1.0,
                    "fixedParameter2": 1.0,
                    "dynamicParameter1": "dynamicParameter1",
                    "observableParameter1_obs1": "obs1par1override",
                    "observableParameter2_obs1": "obs1par2cond1override",
                    "observableParameter1_obs2": np.nan,
                },
                {},
                {
                    "fixedParameter1": LIN,
                    "fixedParameter2": LIN,
                    "dynamicParameter1": LIN,
                    "observableParameter1_obs1": LIN,
                    "observableParameter2_obs1": LIN,
                    "observableParameter1_obs2": LIN,
                },
            ),
            (
                {},
                {
                    "fixedParameter1": 2.0,
                    "fixedParameter2": 2.5,
                    "dynamicParameter1": "dynamicParameter1",
                    "observableParameter1_obs1": "obs1par1override",
                    "observableParameter2_obs1": "obs1par2cond2override",
                    "observableParameter1_obs2": "obs2par1cond2override",
                },
                {},
                {
                    "fixedParameter1": LIN,
                    "fixedParameter2": LIN,
                    "dynamicParameter1": LIN,
                    "observableParameter1_obs1": LIN,
                    "observableParameter2_obs1": LIN,
                    "observableParameter1_obs2": LIN,
                },
            ),
        ]

        actual = petab.get_optimization_to_simulation_parameter_mapping(
            measurement_df=measurement_df,
            condition_df=condition_df,
            model=petab.models.sbml_model.SbmlModel(ss_model.model),
            parameter_df=parameter_df,
        )

        # Comparison with NaN containing expected results fails after pickling!
        # Need to test first for correct NaNs, then for the rest.
        assert np.isnan(expected[0][1]["observableParameter1_obs2"])
        assert np.isnan(actual[0][1]["observableParameter1_obs2"])
        expected[0][1]["observableParameter1_obs2"] = 0.0
        actual[0][1]["observableParameter1_obs2"] = 0.0

        assert actual == expected

    @staticmethod
    def test_parameterized_condition_table():
        condition_df = pd.DataFrame(
            data={
                CONDITION_ID: ["condition1", "condition2", "condition3"],
                CONDITION_NAME: ["", "Condition 2", ""],
                "dynamicParameter1": [
                    "dynamicOverride1_1",
                    "dynamicOverride1_2",
                    0,
                ],
            }
        )
        condition_df.set_index(CONDITION_ID, inplace=True)

        measurement_df = pd.DataFrame(
            data={
                SIMULATION_CONDITION_ID: [
                    "condition1",
                    "condition2",
                    "condition3",
                ],
                OBSERVABLE_ID: ["obs1", "obs2", "obs1"],
                OBSERVABLE_PARAMETERS: "",
                NOISE_PARAMETERS: "",
            }
        )

        parameter_df = pd.DataFrame(
            data={
                PARAMETER_ID: ["dynamicOverride1_1", "dynamicOverride1_2"],
                PARAMETER_NAME: ["", "..."],
                ESTIMATE: [1, 1],
            }
        )
        parameter_df.set_index(PARAMETER_ID, inplace=True)

        import simplesbml

        ss_model = simplesbml.SbmlModel()
        ss_model.addParameter("dynamicParameter1", 1.0)

        assert petab.get_model_parameters(ss_model.model) == [
            "dynamicParameter1"
        ]

        actual = petab.get_optimization_to_simulation_parameter_mapping(
            measurement_df=measurement_df,
            condition_df=condition_df,
            parameter_df=parameter_df,
            model=petab.models.sbml_model.SbmlModel(ss_model.model),
        )

        expected = [
            (
                {},
                {"dynamicParameter1": "dynamicOverride1_1"},
                {},
                {"dynamicParameter1": LIN},
            ),
            (
                {},
                {"dynamicParameter1": "dynamicOverride1_2"},
                {},
                {"dynamicParameter1": LIN},
            ),
            ({}, {"dynamicParameter1": 0}, {}, {"dynamicParameter1": LIN}),
        ]

        assert actual == expected

    @staticmethod
    def test_parameterized_condition_table_changed_scale():
        """Test overriding a dynamic parameter `overridee` with
        - a log10 parameter to be estimated (condition 1)
        - lin parameter not estimated (condition2)
        - log10 parameter not estimated (condition 3)
        - constant override (condition 4)"""

        # overridden parameter
        overridee_id = "overridee"

        # set up model
        import simplesbml

        ss_model = simplesbml.SbmlModel()
        ss_model.addParameter(overridee_id, 2.0)
        assert petab.get_model_parameters(ss_model.model) == [overridee_id]
        assert petab.get_model_parameters(
            ss_model.model, with_values=True
        ) == {overridee_id: 2.0}

        # set up condition table
        condition_df = pd.DataFrame(
            data={
                CONDITION_ID: [
                    "condition1",
                    "condition2",
                    "condition3",
                    "condition4",
                ],
                overridee_id: [
                    "dynamicOverrideLog10",
                    "fixedOverrideLin",
                    "fixedOverrideLog10",
                    10.0,
                ],
            }
        )
        condition_df.set_index("conditionId", inplace=True)

        # set up measurement table
        measurement_df = pd.DataFrame(
            data={
                SIMULATION_CONDITION_ID: [
                    "condition1",
                    "condition2",
                    "condition3",
                    "condition4",
                ],
                OBSERVABLE_ID: ["obs1", "obs2", "obs1", "obs2"],
                OBSERVABLE_PARAMETERS: "",
                NOISE_PARAMETERS: "",
            }
        )

        # set up parameter table
        parameter_df = pd.DataFrame(
            data={
                PARAMETER_ID: [
                    "dynamicOverrideLog10",
                    "fixedOverrideLin",
                    "fixedOverrideLog10",
                ],
                ESTIMATE: [1, 0, 0],
                NOMINAL_VALUE: [np.nan, -2, 1000],
                PARAMETER_SCALE: [LOG10, LIN, LOG10],
            }
        )
        parameter_df.set_index(PARAMETER_ID, inplace=True)

        # test without preequilibration condition; unscaled known parameters

        actual = petab.get_optimization_to_simulation_parameter_mapping(
            measurement_df=measurement_df,
            condition_df=condition_df,
            parameter_df=parameter_df,
            model=petab.models.sbml_model.SbmlModel(ss_model.model),
        )

        expected = [
            (
                {},
                {overridee_id: "dynamicOverrideLog10"},
                {},
                {overridee_id: LOG10},
            ),
            ({}, {overridee_id: -2.0}, {}, {overridee_id: LIN}),
            # not scaled:
            ({}, {overridee_id: 1000.0}, {}, {overridee_id: LIN}),
            ({}, {overridee_id: 10.0}, {}, {overridee_id: LIN}),
        ]

        assert actual == expected

        # test without preequilibration condition; scaled known parameters

        actual = petab.get_optimization_to_simulation_parameter_mapping(
            measurement_df=measurement_df,
            condition_df=condition_df,
            parameter_df=parameter_df,
            model=petab.models.sbml_model.SbmlModel(ss_model.model),
            scaled_parameters=True,
        )

        expected = [
            (
                {},
                {overridee_id: "dynamicOverrideLog10"},
                {},
                {overridee_id: LOG10},
            ),
            ({}, {overridee_id: -2.0}, {}, {overridee_id: LIN}),
            # scaled fixedOverrideLog10:
            ({}, {overridee_id: 3.0}, {}, {overridee_id: LOG10}),
            ({}, {overridee_id: 10.0}, {}, {overridee_id: LIN}),
        ]

        assert actual == expected

        # Add preeq condition

        measurement_df[PREEQUILIBRATION_CONDITION_ID] = [
            "condition1",
            "condition1",
            "condition3",
            "condition3",
        ]
        actual = petab.get_optimization_to_simulation_parameter_mapping(
            measurement_df=measurement_df,
            condition_df=condition_df,
            parameter_df=parameter_df,
            model=petab.models.sbml_model.SbmlModel(ss_model.model),
        )

        expected = [
            (
                {overridee_id: "dynamicOverrideLog10"},
                {overridee_id: "dynamicOverrideLog10"},
                {overridee_id: LOG10},
                {overridee_id: LOG10},
            ),
            (
                {overridee_id: "dynamicOverrideLog10"},
                {overridee_id: -2.0},
                {overridee_id: LOG10},
                {overridee_id: LIN},
            ),
            # not rescaled:
            (
                {overridee_id: 1000.0},
                {overridee_id: 1000.0},
                {overridee_id: LIN},
                {overridee_id: LIN},
            ),
            (
                {overridee_id: 1000.0},
                {overridee_id: 10.0},
                {overridee_id: LIN},
                {overridee_id: LIN},
            ),
        ]

        assert actual == expected


def test_fill_in_nominal_values():
    parameter_df = pd.DataFrame(
        data={
            PARAMETER_ID: ["estimated", "not_estimated"],
            NOMINAL_VALUE: [0.0, 2.0],
            ESTIMATE: [1, 0],
        }
    )
    parameter_df.set_index([PARAMETER_ID], inplace=True)

    par_mapping = {"estimated": "estimated", "not_estimated": "not_estimated"}
    scale_mapping = {"estimated": LIN, "not_estimated": LIN}
    _apply_parameter_table(par_mapping, scale_mapping, parameter_df)

    expected_par = {"estimated": "estimated", "not_estimated": 2.0}
    expected_scale = {"estimated": LIN, "not_estimated": LIN}

    assert expected_par == par_mapping
    assert expected_scale == scale_mapping
