"""Tests related to petab.measurements"""
import tempfile
from pathlib import Path

import numpy as np
import pandas as pd
import petab
from petab.C import *


def test_get_measurement_df():
    """Test measurements.get_measurement_df."""
    measurement_df = pd.DataFrame(data={
        OBSERVABLE_ID: ['obs1', 'obs2'],
        OBSERVABLE_PARAMETERS: ['', 'p1;p2'],
        NOISE_PARAMETERS: ['p3;p4', 'p5']
    })

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

    df = petab.get_measurement_df(file_name).replace(np.nan, '')
    assert (df == measurement_df).all().all()

    # test other arguments
    assert (petab.get_measurement_df(measurement_df) == measurement_df) \
        .all().all()
    assert petab.get_measurement_df(None) is None


def test_write_measurement_df():
    """Test measurements.get_measurement_df."""
    measurement_df = pd.DataFrame(data={
        OBSERVABLE_ID: ['obs1', 'obs2'],
        OBSERVABLE_PARAMETERS: ['', 'p1;p2'],
        NOISE_PARAMETERS: ['p3;p4', 'p5']
    })

    with tempfile.TemporaryDirectory() as temp_dir:
        file_name = Path(temp_dir) / "parameters.tsv"
        petab.write_measurement_df(measurement_df, file_name)
        re_df = petab.get_measurement_df(file_name).replace(np.nan, '')
        assert (measurement_df == re_df).all().all()


def test_create_measurement_df():
    """Test measurements.create_measurement_df."""
    df = petab.create_measurement_df()
    assert set(df.columns.values) == set(MEASUREMENT_DF_COLS)


def test_measurements_have_replicates():
    """Test measurements.measurements_have_replicates."""
    measurement_df = pd.DataFrame(data={
        OBSERVABLE_ID: ['obs1', 'obs1'],
        OBSERVABLE_PARAMETERS: ['', 'p1;p2'],
        NOISE_PARAMETERS: ['p3;p4', 'p5'],
        TIME: [0, 1],
        MEASUREMENT: [42, 137.01]
    })
    assert not petab.measurements_have_replicates(measurement_df)

    measurement_df[TIME] = [1, 1]
    assert petab.measurements_have_replicates(measurement_df)


def test_get_simulation_conditions():
    """Test get_simulation_conditions"""

    # only simulation condition
    measurement_df = pd.DataFrame(data={
        SIMULATION_CONDITION_ID: ['c0', 'c1', 'c0', 'c1'],
    })
    expected = pd.DataFrame(data={
        SIMULATION_CONDITION_ID: ['c0', 'c1'],
    })
    actual = petab.get_simulation_conditions(measurement_df)
    assert actual.equals(expected)

    # simulation and preequilibration condition
    measurement_df = pd.DataFrame(data={
        SIMULATION_CONDITION_ID: ['c0', 'c1', 'c0', 'c1'],
        PREEQUILIBRATION_CONDITION_ID: ['c1', 'c0', 'c1', 'c0'],
    })
    expected = pd.DataFrame(data={
        SIMULATION_CONDITION_ID: ['c0', 'c1'],
        PREEQUILIBRATION_CONDITION_ID: ['c1', 'c0'],
    })
    actual = petab.get_simulation_conditions(measurement_df)
    assert actual.equals(expected)

    # simulation with and without preequilibration
    measurement_df = pd.DataFrame(data={
        SIMULATION_CONDITION_ID: ['c0', 'c1', 'c0', 'c1'],
        PREEQUILIBRATION_CONDITION_ID: ['', '', 'c1', 'c0'],
    })
    expected = pd.DataFrame(data={
        SIMULATION_CONDITION_ID: ['c0', 'c1', 'c0', 'c1'],
        PREEQUILIBRATION_CONDITION_ID: ['', '', 'c1', 'c0'],
    }).sort_values([SIMULATION_CONDITION_ID, PREEQUILIBRATION_CONDITION_ID],
                   ignore_index=True)
    actual = petab.get_simulation_conditions(measurement_df)
    assert actual.equals(expected)

    # simulation with and without preequilibration; NaNs
    measurement_df = pd.DataFrame(data={
        SIMULATION_CONDITION_ID: ['c0', 'c1', 'c0', 'c1'],
        PREEQUILIBRATION_CONDITION_ID: [np.nan, np.nan, 'c1', 'c0'],
    })
    expected = pd.DataFrame(data={
        SIMULATION_CONDITION_ID: ['c0', 'c1', 'c0', 'c1'],
        PREEQUILIBRATION_CONDITION_ID: ['', '', 'c1', 'c0'],
    }).sort_values([SIMULATION_CONDITION_ID, PREEQUILIBRATION_CONDITION_ID],
                   ignore_index=True)
    actual = petab.get_simulation_conditions(measurement_df)
    assert actual.equals(expected)
