"""Tests for petab.observables"""
import tempfile
from pathlib import Path

import pandas as pd
import petab
import pytest
from petab.C import *

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


def test_get_observable_df():
    """Test measurements.get_measurement_df."""
    # without id
    observable_df = pd.DataFrame(data={
        OBSERVABLE_NAME: ['observable name 1'],
        OBSERVABLE_FORMULA: ['observable_1'],
        NOISE_FORMULA: [1],
    })

    with tempfile.NamedTemporaryFile(mode='w', delete=False) as fh:
        file_name = fh.name
        observable_df.to_csv(fh, sep='\t', index=False)

    with pytest.raises(KeyError):
        petab.get_observable_df(file_name)

    # with id
    observable_df[OBSERVABLE_ID] = ['observable_1']

    with tempfile.NamedTemporaryFile(mode='w', delete=False) as fh:
        file_name = fh.name
        observable_df.to_csv(fh, sep='\t', index=False)

    df = petab.get_observable_df(file_name)
    assert (df == observable_df.set_index(OBSERVABLE_ID)).all().all()

    # test other arguments
    assert (petab.get_observable_df(observable_df) == observable_df) \
        .all().all()
    assert petab.get_observable_df(None) is None


def test_write_observable_df():
    """Test measurements.get_measurement_df."""
    observable_df = pd.DataFrame(data={
        OBSERVABLE_ID: ['observable_1'],
        OBSERVABLE_NAME: ['observable name 1'],
        OBSERVABLE_FORMULA: ['observable_1'],
        NOISE_FORMULA: [1],
    }).set_index(OBSERVABLE_ID)

    with tempfile.TemporaryDirectory() as temp_dir:
        file_name = Path(temp_dir) / "observables.tsv"
        petab.write_observable_df(observable_df, file_name)
        re_df = petab.get_observable_df(file_name)
        assert (observable_df == re_df).all().all()


def test_get_output_parameters():
    """Test measurements.get_output_parameters."""
    # sbml model
    import simplesbml
    from petab.models.sbml_model import SbmlModel
    ss_model = simplesbml.SbmlModel()
    ss_model.addParameter('fixedParameter1', 1.0)
    ss_model.addParameter('observable_1', 1.0)

    # observable file
    observable_df = pd.DataFrame(data={
        OBSERVABLE_ID: ['observable_1'],
        OBSERVABLE_NAME: ['observable name 1'],
        OBSERVABLE_FORMULA: ['observable_1 * scaling + offset'],
        NOISE_FORMULA: [1],
    }).set_index(OBSERVABLE_ID)

    output_parameters = petab.get_output_parameters(
        observable_df, SbmlModel(sbml_model=ss_model.model))

    assert output_parameters == ['offset', 'scaling']

    # test sympy-special symbols (e.g. N, beta, ...)
    # see https://github.com/ICB-DCM/pyPESTO/issues/1048
    observable_df = pd.DataFrame(data={
        OBSERVABLE_ID: ['observable_1'],
        OBSERVABLE_NAME: ['observable name 1'],
        OBSERVABLE_FORMULA: ['observable_1 * N + beta'],
        NOISE_FORMULA: [1],
    }).set_index(OBSERVABLE_ID)

    output_parameters = petab.get_output_parameters(
        observable_df, SbmlModel(sbml_model=ss_model.model))

    assert output_parameters == ['N', 'beta']


def test_get_formula_placeholders():
    """Test get_formula_placeholders"""

    # no placeholder
    assert petab.get_formula_placeholders('1.0', 'any', 'observable') == []

    # multiple placeholders
    assert petab.get_formula_placeholders(
        'observableParameter1_twoParams * '
        'observableParameter2_twoParams + otherParam',
        'twoParams', 'observable') \
        == ['observableParameter1_twoParams',
            'observableParameter2_twoParams']

    # noise placeholder
    assert petab.get_formula_placeholders(
        '3.0 * noiseParameter1_oneParam', 'oneParam', 'noise') \
        == ['noiseParameter1_oneParam']

    # multiple instances and in 'wrong' order
    assert petab.get_formula_placeholders(
        'observableParameter2_twoParams * '
        'observableParameter1_twoParams + '
        'otherParam / observableParameter2_twoParams',
        'twoParams', 'observable') \
        == ['observableParameter1_twoParams',
            'observableParameter2_twoParams']

    # non-consecutive numbering
    with pytest.raises(AssertionError):
        petab.get_formula_placeholders(
            'observableParameter2_twoParams + observableParameter2_twoParams',
            'twoParams', 'observable')

    # empty
    assert petab.get_formula_placeholders('', 'any', 'observable') == []

    # non-string
    assert petab.get_formula_placeholders(1, 'any', 'observable') == []


def test_create_observable_df():
    """Test observables.create_measurement_df."""
    df = petab.create_observable_df()
    assert set(df.columns.values) == set(OBSERVABLE_DF_COLS)


def test_get_placeholders():
    """Test get_placeholders"""
    observable_df = pd.DataFrame(data={
        OBSERVABLE_ID: ['obs_1', 'obs_2'],
        OBSERVABLE_FORMULA: ['observableParameter1_obs_1 * 2 * foo',
                             '1 + observableParameter1_obs_2'],
    }).set_index(OBSERVABLE_ID)

    # test with missing noiseFormula
    expected = ['observableParameter1_obs_1', 'observableParameter1_obs_2']
    actual = petab.get_placeholders(observable_df)
    assert actual == expected

    # test with noiseFormula
    observable_df[NOISE_FORMULA] = ['noiseParameter1_obs_1', '2.0']
    expected = ['observableParameter1_obs_1', 'noiseParameter1_obs_1',
                'observableParameter1_obs_2']
    actual = petab.get_placeholders(observable_df)
    assert actual == expected
