"""Generates a simple one target scenario

Inspiration for how to generate ground truth and measurements are taken from
https://stonesoup.readthedocs.io/en/latest/auto_tutorials/01_KalmanFilterTutorial.html#sphx-glr-auto-tutorials-01-kalmanfiltertutorial-py
"""

from datetime import datetime
from datetime import timedelta
import numpy as np

from stonesoup.models.transition.linear import CombinedLinearGaussianTransitionModel, \
    ConstantVelocity
from stonesoup.types.groundtruth import GroundTruthPath, GroundTruthState
from stonesoup.types.detection import Detection
from stonesoup.models.measurement.linear import LinearGaussian

from utils import store_object


def generate_scenario_1(seed=1996, permanent_save=True, sigma_process=0.01, sigma_meas_radar=3, sigma_meas_ais=1):
    """
    Generates scenario 1. Todo define scenario 1
    :param seed:
    :param permanent_save:
    :param sigma_process:
    :param sigma_meas_radar:
    :param sigma_meas_ais:
    :return:
    """
    # specify seed to be able repeat example
    start_time = datetime.now()

    np.random.seed(seed)

    # combine two 1-D CV models to create a 2-D CV model
    transition_model = CombinedLinearGaussianTransitionModel([ConstantVelocity(sigma_process),
                                                              ConstantVelocity(sigma_process)])

    # starting at 0,0 and moving NE
    truth = GroundTruthPath([GroundTruthState([0, 1, 0, 1], timestamp=start_time)])

    # generate truth using transition_model and noise
    for k in range(1, 21):
        truth.append(GroundTruthState(
            transition_model.function(truth[k - 1], noise=True, time_interval=timedelta(seconds=1)),
            timestamp=start_time + timedelta(seconds=k)))

    # Simulate measurements
    # Specify measurement model for radar
    measurement_model_radar = LinearGaussian(
        ndim_state=4,  # number of state dimensions
        mapping=(0, 2),  # mapping measurement vector index to state index
        noise_covar=np.array([[sigma_meas_radar, 0],  # covariance matrix for Gaussian PDF
                              [0, sigma_meas_radar]])
    )

    # Specify measurement model for AIS
    measurement_model_ais = LinearGaussian(
        ndim_state=4,
        mapping=(0, 2),
        noise_covar=np.array([[sigma_meas_ais, 0],
                              [0, sigma_meas_ais]])
    )

    # generate "radar" measurements
    measurements_radar = []
    for state in truth:
        measurement = measurement_model_radar.function(state, noise=True)
        measurements_radar.append(Detection(measurement, timestamp=state.timestamp))

    # generate "AIS" measurements
    measurements_ais = []
    state_num = 0
    for state in truth:
        state_num += 1
        if not state_num % 2:  # measurement every second time step
            measurement = measurement_model_ais.function(state, noise=True)
            measurements_ais.append(Detection(measurement, timestamp=state.timestamp))

    if permanent_save:
        save_folder_name = seed.__str__()
    else:
        save_folder_name = "temp"

    save_folder = "../scenarios/scenario1/" + save_folder_name + "/"

    # save the ground truth and the measurements for the radar and the AIS
    store_object.store_object(truth, save_folder, "ground_truth.pk1")
    store_object.store_object(measurements_radar, save_folder, "measurements_radar.pk1")
    store_object.store_object(measurements_ais, save_folder, "measurements_ais.pk1")
    store_object.store_object(start_time, save_folder, "start_time.pk1")
    store_object.store_object(measurement_model_radar, save_folder, "measurement_model_radar.pk1")
    store_object.store_object(measurement_model_ais, save_folder, "measurement_model_ais.pk1")
    store_object.store_object(transition_model, save_folder, "transition_model.pk1")


def generate_scenario_2(seed=1996, permanent_save=True, sigma_process=0.01, sigma_meas_radar=3, sigma_meas_ais=1,
                        timesteps=20):
    """
    Generates scenario 2. Scenario 2 consists of radar and ais measurements with the same sampling rate.
    :param timesteps:
    :param seed:
    :param permanent_save:
    :param sigma_process:
    :param sigma_meas_radar:
    :param sigma_meas_ais:
    :return: Nothing. Saves the scenario to a specified folder
    """
    start_time = datetime.now()

    # specify seed to be able repeat example
    np.random.seed(seed)

    # combine two 1-D CV models to create a 2-D CV model
    transition_model = CombinedLinearGaussianTransitionModel(
        [ConstantVelocity(sigma_process), ConstantVelocity(sigma_process)])

    # starting at 0,0 and moving NE
    truth = GroundTruthPath([GroundTruthState([0, 1, 0, 1], timestamp=start_time)])

    # generate truth using transition_model and noise
    for k in range(1, timesteps):
        truth.append(GroundTruthState(
            transition_model.function(truth[k - 1], noise=True, time_interval=timedelta(seconds=1)),
            timestamp=start_time + timedelta(seconds=k)))

    # Simulate measurements
    # Specify measurement model for radar
    measurement_model_radar = LinearGaussian(
        ndim_state=4,  # number of state dimensions
        mapping=(0, 2),  # mapping measurement vector index to state index
        noise_covar=np.array([[sigma_meas_radar, 0],  # covariance matrix for Gaussian PDF
                              [0, sigma_meas_radar]])
    )

    # Specify measurement model for AIS (Same as for radar)
    measurement_model_ais = LinearGaussian(
        ndim_state=4,
        mapping=(0, 2),
        noise_covar=np.array([[sigma_meas_ais, 0],
                              [0, sigma_meas_ais]])
    )

    # generate "radar" measurements
    measurements_radar = []
    for state in truth:
        measurement = measurement_model_radar.function(state, noise=True)
        measurements_radar.append(Detection(measurement, timestamp=state.timestamp))

    # generate "AIS" measurements
    measurements_ais = []
    for state in truth:
        measurement = measurement_model_ais.function(state, noise=True)
        measurements_ais.append(Detection(measurement, timestamp=state.timestamp))

    if permanent_save:
        save_folder_name = seed.__str__()
    else:
        save_folder_name = "temp"

    save_folder = "../scenarios/scenario2/" + save_folder_name + "/"

    # save the ground truth and the measurements for the radar and the AIS
    store_object.store_object(truth, save_folder, "ground_truth.pk1")
    store_object.store_object(measurements_radar, save_folder, "measurements_radar.pk1")
    store_object.store_object(measurements_ais, save_folder, "measurements_ais.pk1")
    store_object.store_object(start_time, save_folder, "start_time.pk1")
    store_object.store_object(measurement_model_radar, save_folder, "measurement_model_radar.pk1")
    store_object.store_object(measurement_model_ais, save_folder, "measurement_model_ais.pk1")
    store_object.store_object(transition_model, save_folder, "transition_model.pk1")


def generate_scenario_3(seed=1996, permanent_save=True, radar_meas_rate=1, ais_meas_rate=5, sigma_process=0.01,
                        sigma_meas_radar=3, sigma_meas_ais=1, timesteps=20):
    """
    Generates scenario 3. Scenario 3 consists of radar and ais measurements with different sampling rate. The sampling
    rate is specified in the input params. A groundtruth is generated for each second.
    :param seed:
    :param permanent_save:
    :param radar_meas_rate:
    :param ais_meas_rate:
    :param sigma_process:
    :param sigma_meas_radar:
    :param sigma_meas_ais:
    :param timesteps: The amount of measurements from the slowest sensor
    :return: Nothing. Saves the scenario to a specified folder
    """
    start_time = datetime.now()

    # specify seed to be able to repeat the example
    np.random.seed(seed)

    # combine two 1-D CV models to create a 2-D CV model
    transition_model = CombinedLinearGaussianTransitionModel(
        [ConstantVelocity(sigma_process), ConstantVelocity(sigma_process)])

    # starting at 0,0 and moving NE
    truth = GroundTruthPath([GroundTruthState([0, 1, 0, 1], timestamp=start_time)])

    # generate truth using transition_model and noise
    end_time = start_time + timedelta(seconds=timesteps * max(radar_meas_rate, ais_meas_rate))
    time = start_time + timedelta(seconds=1)
    while time < end_time:
        truth.append(GroundTruthState(
            transition_model.function(truth[-1], noise=True, time_interval=timedelta(seconds=1)),
            timestamp=time))
        time += timedelta(seconds=1)

    # Simulate measurements
    # Specify measurement model for radar
    measurement_model_radar = LinearGaussian(
        ndim_state=4,  # number of state dimensions
        mapping=(0, 2),  # mapping measurement vector index to state index
        noise_covar=np.array([[sigma_meas_radar, 0],  # covariance matrix for Gaussian PDF
                              [0, sigma_meas_radar]])
    )

    # Specify measurement model for AIS (Same as for radar)
    measurement_model_ais = LinearGaussian(
        ndim_state=4,
        mapping=(0, 2),
        noise_covar=np.array([[sigma_meas_ais, 0],
                              [0, sigma_meas_ais]])
    )

    # generate "radar" measurements
    measurements_radar = []
    measurements_ais = []
    next_radar_meas_time = start_time
    next_ais_meas_time = start_time
    for state in truth:
        # check whether we want to generate a measurement from this gt
        if state.timestamp == next_radar_meas_time:
            measurement = measurement_model_radar.function(state, noise=True)
            measurements_radar.append(Detection(measurement, timestamp=state.timestamp))
            next_radar_meas_time += timedelta(seconds=radar_meas_rate)

        if state.timestamp == next_ais_meas_time:
            measurement = measurement_model_ais.function(state, noise=True)
            measurements_ais.append(Detection(measurement, timestamp=state.timestamp))
            next_ais_meas_time += timedelta(seconds=ais_meas_rate)

    if permanent_save:
        save_folder_name = seed.__str__()
    else:
        save_folder_name = "temp"

    save_folder = "../scenarios/scenario3/" + save_folder_name + "/"

    # save the ground truth and the measurements for the radar and the AIS
    store_object.store_object(truth, save_folder, "ground_truth.pk1")
    store_object.store_object(measurements_radar, save_folder, "measurements_radar.pk1")
    store_object.store_object(measurements_ais, save_folder, "measurements_ais.pk1")
    store_object.store_object(start_time, save_folder, "start_time.pk1")
    store_object.store_object(measurement_model_radar, save_folder, "measurement_model_radar.pk1")
    store_object.store_object(measurement_model_ais, save_folder, "measurement_model_ais.pk1")
    store_object.store_object(transition_model, save_folder, "transition_model.pk1")

b4734781d55876ce0953debb195ad4d3d45e95b2
<mask token>


def generate_scenario_3(seed=1996, permanent_save=True, radar_meas_rate=1,
    ais_meas_rate=5, sigma_process=0.01, sigma_meas_radar=3, sigma_meas_ais
    =1, timesteps=20):
    """
    Generates scenario 3. Scenario 3 consists of radar and ais measurements with different sampling rate. The sampling
    rate is specified in the input params. A groundtruth is generated for each second.
    :param seed:
    :param permanent_save:
    :param radar_meas_rate:
    :param ais_meas_rate:
    :param sigma_process:
    :param sigma_meas_radar:
    :param sigma_meas_ais:
    :param timesteps: The amount of measurements from the slowest sensor
    :return: Nothing. Saves the scenario to a specified folder
    """
    start_time = datetime.now()
    np.random.seed(seed)
    transition_model = CombinedLinearGaussianTransitionModel([
        ConstantVelocity(sigma_process), ConstantVelocity(sigma_process)])
    truth = GroundTruthPath([GroundTruthState([0, 1, 0, 1], timestamp=
        start_time)])
    end_time = start_time + timedelta(seconds=timesteps * max(
        radar_meas_rate, ais_meas_rate))
    time = start_time + timedelta(seconds=1)
    while time < end_time:
        truth.append(GroundTruthState(transition_model.function(truth[-1],
            noise=True, time_interval=timedelta(seconds=1)), timestamp=time))
        time += timedelta(seconds=1)
    measurement_model_radar = LinearGaussian(ndim_state=4, mapping=(0, 2),
        noise_covar=np.array([[sigma_meas_radar, 0], [0, sigma_meas_radar]]))
    measurement_model_ais = LinearGaussian(ndim_state=4, mapping=(0, 2),
        noise_covar=np.array([[sigma_meas_ais, 0], [0, sigma_meas_ais]]))
    measurements_radar = []
    measurements_ais = []
    next_radar_meas_time = start_time
    next_ais_meas_time = start_time
    for state in truth:
        if state.timestamp == next_radar_meas_time:
            measurement = measurement_model_radar.function(state, noise=True)
            measurements_radar.append(Detection(measurement, timestamp=
                state.timestamp))
            next_radar_meas_time += timedelta(seconds=radar_meas_rate)
        if state.timestamp == next_ais_meas_time:
            measurement = measurement_model_ais.function(state, noise=True)
            measurements_ais.append(Detection(measurement, timestamp=state.
                timestamp))
            next_ais_meas_time += timedelta(seconds=ais_meas_rate)
    if permanent_save:
        save_folder_name = seed.__str__()
    else:
        save_folder_name = 'temp'
    save_folder = '../scenarios/scenario3/' + save_folder_name + '/'
    store_object.store_object(truth, save_folder, 'ground_truth.pk1')
    store_object.store_object(measurements_radar, save_folder,
        'measurements_radar.pk1')
    store_object.store_object(measurements_ais, save_folder,
        'measurements_ais.pk1')
    store_object.store_object(start_time, save_folder, 'start_time.pk1')
    store_object.store_object(measurement_model_radar, save_folder,
        'measurement_model_radar.pk1')
    store_object.store_object(measurement_model_ais, save_folder,
        'measurement_model_ais.pk1')
    store_object.store_object(transition_model, save_folder,
        'transition_model.pk1')

<mask token>


def generate_scenario_1(seed=1996, permanent_save=True, sigma_process=0.01,
    sigma_meas_radar=3, sigma_meas_ais=1):
    """
    Generates scenario 1. Todo define scenario 1
    :param seed:
    :param permanent_save:
    :param sigma_process:
    :param sigma_meas_radar:
    :param sigma_meas_ais:
    :return:
    """
    start_time = datetime.now()
    np.random.seed(seed)
    transition_model = CombinedLinearGaussianTransitionModel([
        ConstantVelocity(sigma_process), ConstantVelocity(sigma_process)])
    truth = GroundTruthPath([GroundTruthState([0, 1, 0, 1], timestamp=
        start_time)])
    for k in range(1, 21):
        truth.append(GroundTruthState(transition_model.function(truth[k - 1
            ], noise=True, time_interval=timedelta(seconds=1)), timestamp=
            start_time + timedelta(seconds=k)))
    measurement_model_radar = LinearGaussian(ndim_state=4, mapping=(0, 2),
        noise_covar=np.array([[sigma_meas_radar, 0], [0, sigma_meas_radar]]))
    measurement_model_ais = LinearGaussian(ndim_state=4, mapping=(0, 2),
        noise_covar=np.array([[sigma_meas_ais, 0], [0, sigma_meas_ais]]))
    measurements_radar = []
    for state in truth:
        measurement = measurement_model_radar.function(state, noise=True)
        measurements_radar.append(Detection(measurement, timestamp=state.
            timestamp))
    measurements_ais = []
    state_num = 0
    for state in truth:
        state_num += 1
        if not state_num % 2:
            measurement = measurement_model_ais.function(state, noise=True)
            measurements_ais.append(Detection(measurement, timestamp=state.
                timestamp))
    if permanent_save:
        save_folder_name = seed.__str__()
    else:
        save_folder_name = 'temp'
    save_folder = '../scenarios/scenario1/' + save_folder_name + '/'
    store_object.store_object(truth, save_folder, 'ground_truth.pk1')
    store_object.store_object(measurements_radar, save_folder,
        'measurements_radar.pk1')
    store_object.store_object(measurements_ais, save_folder,
        'measurements_ais.pk1')
    store_object.store_object(start_time, save_folder, 'start_time.pk1')
    store_object.store_object(measurement_model_radar, save_folder,
        'measurement_model_radar.pk1')
    store_object.store_object(measurement_model_ais, save_folder,
        'measurement_model_ais.pk1')
    store_object.store_object(transition_model, save_folder,
        'transition_model.pk1')


<mask token>


def generate_scenario_3(seed=1996, permanent_save=True, radar_meas_rate=1,
    ais_meas_rate=5, sigma_process=0.01, sigma_meas_radar=3, sigma_meas_ais
    =1, timesteps=20):
    """
    Generates scenario 3. Scenario 3 consists of radar and ais measurements with different sampling rate. The sampling
    rate is specified in the input params. A groundtruth is generated for each second.
    :param seed:
    :param permanent_save:
    :param radar_meas_rate:
    :param ais_meas_rate:
    :param sigma_process:
    :param sigma_meas_radar:
    :param sigma_meas_ais:
    :param timesteps: The amount of measurements from the slowest sensor
    :return: Nothing. Saves the scenario to a specified folder
    """
    start_time = datetime.now()
    np.random.seed(seed)
    transition_model = CombinedLinearGaussianTransitionModel([
        ConstantVelocity(sigma_process), ConstantVelocity(sigma_process)])
    truth = GroundTruthPath([GroundTruthState([0, 1, 0, 1], timestamp=
        start_time)])
    end_time = start_time + timedelta(seconds=timesteps * max(
        radar_meas_rate, ais_meas_rate))
    time = start_time + timedelta(seconds=1)
    while time < end_time:
        truth.append(GroundTruthState(transition_model.function(truth[-1],
            noise=True, time_interval=timedelta(seconds=1)), timestamp=time))
        time += timedelta(seconds=1)
    measurement_model_radar = LinearGaussian(ndim_state=4, mapping=(0, 2),
        noise_covar=np.array([[sigma_meas_radar, 0], [0, sigma_meas_radar]]))
    measurement_model_ais = LinearGaussian(ndim_state=4, mapping=(0, 2),
        noise_covar=np.array([[sigma_meas_ais, 0], [0, sigma_meas_ais]]))
    measurements_radar = []
    measurements_ais = []
    next_radar_meas_time = start_time
    next_ais_meas_time = start_time
    for state in truth:
        if state.timestamp == next_radar_meas_time:
            measurement = measurement_model_radar.function(state, noise=True)
            measurements_radar.append(Detection(measurement, timestamp=
                state.timestamp))
            next_radar_meas_time += timedelta(seconds=radar_meas_rate)
        if state.timestamp == next_ais_meas_time:
            measurement = measurement_model_ais.function(state, noise=True)
            measurements_ais.append(Detection(measurement, timestamp=state.
                timestamp))
            next_ais_meas_time += timedelta(seconds=ais_meas_rate)
    if permanent_save:
        save_folder_name = seed.__str__()
    else:
        save_folder_name = 'temp'
    save_folder = '../scenarios/scenario3/' + save_folder_name + '/'
    store_object.store_object(truth, save_folder, 'ground_truth.pk1')
    store_object.store_object(measurements_radar, save_folder,
        'measurements_radar.pk1')
    store_object.store_object(measurements_ais, save_folder,
        'measurements_ais.pk1')
    store_object.store_object(start_time, save_folder, 'start_time.pk1')
    store_object.store_object(measurement_model_radar, save_folder,
        'measurement_model_radar.pk1')
    store_object.store_object(measurement_model_ais, save_folder,
        'measurement_model_ais.pk1')
    store_object.store_object(transition_model, save_folder,
        'transition_model.pk1')

<mask token>


def generate_scenario_1(seed=1996, permanent_save=True, sigma_process=0.01,
    sigma_meas_radar=3, sigma_meas_ais=1):
    """
    Generates scenario 1. Todo define scenario 1
    :param seed:
    :param permanent_save:
    :param sigma_process:
    :param sigma_meas_radar:
    :param sigma_meas_ais:
    :return:
    """
    start_time = datetime.now()
    np.random.seed(seed)
    transition_model = CombinedLinearGaussianTransitionModel([
        ConstantVelocity(sigma_process), ConstantVelocity(sigma_process)])
    truth = GroundTruthPath([GroundTruthState([0, 1, 0, 1], timestamp=
        start_time)])
    for k in range(1, 21):
        truth.append(GroundTruthState(transition_model.function(truth[k - 1
            ], noise=True, time_interval=timedelta(seconds=1)), timestamp=
            start_time + timedelta(seconds=k)))
    measurement_model_radar = LinearGaussian(ndim_state=4, mapping=(0, 2),
        noise_covar=np.array([[sigma_meas_radar, 0], [0, sigma_meas_radar]]))
    measurement_model_ais = LinearGaussian(ndim_state=4, mapping=(0, 2),
        noise_covar=np.array([[sigma_meas_ais, 0], [0, sigma_meas_ais]]))
    measurements_radar = []
    for state in truth:
        measurement = measurement_model_radar.function(state, noise=True)
        measurements_radar.append(Detection(measurement, timestamp=state.
            timestamp))
    measurements_ais = []
    state_num = 0
    for state in truth:
        state_num += 1
        if not state_num % 2:
            measurement = measurement_model_ais.function(state, noise=True)
            measurements_ais.append(Detection(measurement, timestamp=state.
                timestamp))
    if permanent_save:
        save_folder_name = seed.__str__()
    else:
        save_folder_name = 'temp'
    save_folder = '../scenarios/scenario1/' + save_folder_name + '/'
    store_object.store_object(truth, save_folder, 'ground_truth.pk1')
    store_object.store_object(measurements_radar, save_folder,
        'measurements_radar.pk1')
    store_object.store_object(measurements_ais, save_folder,
        'measurements_ais.pk1')
    store_object.store_object(start_time, save_folder, 'start_time.pk1')
    store_object.store_object(measurement_model_radar, save_folder,
        'measurement_model_radar.pk1')
    store_object.store_object(measurement_model_ais, save_folder,
        'measurement_model_ais.pk1')
    store_object.store_object(transition_model, save_folder,
        'transition_model.pk1')


def generate_scenario_2(seed=1996, permanent_save=True, sigma_process=0.01,
    sigma_meas_radar=3, sigma_meas_ais=1, timesteps=20):
    """
    Generates scenario 2. Scenario 2 consists of radar and ais measurements with the same sampling rate.
    :param timesteps:
    :param seed:
    :param permanent_save:
    :param sigma_process:
    :param sigma_meas_radar:
    :param sigma_meas_ais:
    :return: Nothing. Saves the scenario to a specified folder
    """
    start_time = datetime.now()
    np.random.seed(seed)
    transition_model = CombinedLinearGaussianTransitionModel([
        ConstantVelocity(sigma_process), ConstantVelocity(sigma_process)])
    truth = GroundTruthPath([GroundTruthState([0, 1, 0, 1], timestamp=
        start_time)])
    for k in range(1, timesteps):
        truth.append(GroundTruthState(transition_model.function(truth[k - 1
            ], noise=True, time_interval=timedelta(seconds=1)), timestamp=
            start_time + timedelta(seconds=k)))
    measurement_model_radar = LinearGaussian(ndim_state=4, mapping=(0, 2),
        noise_covar=np.array([[sigma_meas_radar, 0], [0, sigma_meas_radar]]))
    measurement_model_ais = LinearGaussian(ndim_state=4, mapping=(0, 2),
        noise_covar=np.array([[sigma_meas_ais, 0], [0, sigma_meas_ais]]))
    measurements_radar = []
    for state in truth:
        measurement = measurement_model_radar.function(state, noise=True)
        measurements_radar.append(Detection(measurement, timestamp=state.
            timestamp))
    measurements_ais = []
    for state in truth:
        measurement = measurement_model_ais.function(state, noise=True)
        measurements_ais.append(Detection(measurement, timestamp=state.
            timestamp))
    if permanent_save:
        save_folder_name = seed.__str__()
    else:
        save_folder_name = 'temp'
    save_folder = '../scenarios/scenario2/' + save_folder_name + '/'
    store_object.store_object(truth, save_folder, 'ground_truth.pk1')
    store_object.store_object(measurements_radar, save_folder,
        'measurements_radar.pk1')
    store_object.store_object(measurements_ais, save_folder,
        'measurements_ais.pk1')
    store_object.store_object(start_time, save_folder, 'start_time.pk1')
    store_object.store_object(measurement_model_radar, save_folder,
        'measurement_model_radar.pk1')
    store_object.store_object(measurement_model_ais, save_folder,
        'measurement_model_ais.pk1')
    store_object.store_object(transition_model, save_folder,
        'transition_model.pk1')


def generate_scenario_3(seed=1996, permanent_save=True, radar_meas_rate=1,
    ais_meas_rate=5, sigma_process=0.01, sigma_meas_radar=3, sigma_meas_ais
    =1, timesteps=20):
    """
    Generates scenario 3. Scenario 3 consists of radar and ais measurements with different sampling rate. The sampling
    rate is specified in the input params. A groundtruth is generated for each second.
    :param seed:
    :param permanent_save:
    :param radar_meas_rate:
    :param ais_meas_rate:
    :param sigma_process:
    :param sigma_meas_radar:
    :param sigma_meas_ais:
    :param timesteps: The amount of measurements from the slowest sensor
    :return: Nothing. Saves the scenario to a specified folder
    """
    start_time = datetime.now()
    np.random.seed(seed)
    transition_model = CombinedLinearGaussianTransitionModel([
        ConstantVelocity(sigma_process), ConstantVelocity(sigma_process)])
    truth = GroundTruthPath([GroundTruthState([0, 1, 0, 1], timestamp=
        start_time)])
    end_time = start_time + timedelta(seconds=timesteps * max(
        radar_meas_rate, ais_meas_rate))
    time = start_time + timedelta(seconds=1)
    while time < end_time:
        truth.append(GroundTruthState(transition_model.function(truth[-1],
            noise=True, time_interval=timedelta(seconds=1)), timestamp=time))
        time += timedelta(seconds=1)
    measurement_model_radar = LinearGaussian(ndim_state=4, mapping=(0, 2),
        noise_covar=np.array([[sigma_meas_radar, 0], [0, sigma_meas_radar]]))
    measurement_model_ais = LinearGaussian(ndim_state=4, mapping=(0, 2),
        noise_covar=np.array([[sigma_meas_ais, 0], [0, sigma_meas_ais]]))
    measurements_radar = []
    measurements_ais = []
    next_radar_meas_time = start_time
    next_ais_meas_time = start_time
    for state in truth:
        if state.timestamp == next_radar_meas_time:
            measurement = measurement_model_radar.function(state, noise=True)
            measurements_radar.append(Detection(measurement, timestamp=
                state.timestamp))
            next_radar_meas_time += timedelta(seconds=radar_meas_rate)
        if state.timestamp == next_ais_meas_time:
            measurement = measurement_model_ais.function(state, noise=True)
            measurements_ais.append(Detection(measurement, timestamp=state.
                timestamp))
            next_ais_meas_time += timedelta(seconds=ais_meas_rate)
    if permanent_save:
        save_folder_name = seed.__str__()
    else:
        save_folder_name = 'temp'
    save_folder = '../scenarios/scenario3/' + save_folder_name + '/'
    store_object.store_object(truth, save_folder, 'ground_truth.pk1')
    store_object.store_object(measurements_radar, save_folder,
        'measurements_radar.pk1')
    store_object.store_object(measurements_ais, save_folder,
        'measurements_ais.pk1')
    store_object.store_object(start_time, save_folder, 'start_time.pk1')
    store_object.store_object(measurement_model_radar, save_folder,
        'measurement_model_radar.pk1')
    store_object.store_object(measurement_model_ais, save_folder,
        'measurement_model_ais.pk1')
    store_object.store_object(transition_model, save_folder,
        'transition_model.pk1')

<mask token>
from datetime import datetime
from datetime import timedelta
import numpy as np
from stonesoup.models.transition.linear import CombinedLinearGaussianTransitionModel, ConstantVelocity
from stonesoup.types.groundtruth import GroundTruthPath, GroundTruthState
from stonesoup.types.detection import Detection
from stonesoup.models.measurement.linear import LinearGaussian
from utils import store_object


def generate_scenario_1(seed=1996, permanent_save=True, sigma_process=0.01,
    sigma_meas_radar=3, sigma_meas_ais=1):
    """
    Generates scenario 1. Todo define scenario 1
    :param seed:
    :param permanent_save:
    :param sigma_process:
    :param sigma_meas_radar:
    :param sigma_meas_ais:
    :return:
    """
    start_time = datetime.now()
    np.random.seed(seed)
    transition_model = CombinedLinearGaussianTransitionModel([
        ConstantVelocity(sigma_process), ConstantVelocity(sigma_process)])
    truth = GroundTruthPath([GroundTruthState([0, 1, 0, 1], timestamp=
        start_time)])
    for k in range(1, 21):
        truth.append(GroundTruthState(transition_model.function(truth[k - 1
            ], noise=True, time_interval=timedelta(seconds=1)), timestamp=
            start_time + timedelta(seconds=k)))
    measurement_model_radar = LinearGaussian(ndim_state=4, mapping=(0, 2),
        noise_covar=np.array([[sigma_meas_radar, 0], [0, sigma_meas_radar]]))
    measurement_model_ais = LinearGaussian(ndim_state=4, mapping=(0, 2),
        noise_covar=np.array([[sigma_meas_ais, 0], [0, sigma_meas_ais]]))
    measurements_radar = []
    for state in truth:
        measurement = measurement_model_radar.function(state, noise=True)
        measurements_radar.append(Detection(measurement, timestamp=state.
            timestamp))
    measurements_ais = []
    state_num = 0
    for state in truth:
        state_num += 1
        if not state_num % 2:
            measurement = measurement_model_ais.function(state, noise=True)
            measurements_ais.append(Detection(measurement, timestamp=state.
                timestamp))
    if permanent_save:
        save_folder_name = seed.__str__()
    else:
        save_folder_name = 'temp'
    save_folder = '../scenarios/scenario1/' + save_folder_name + '/'
    store_object.store_object(truth, save_folder, 'ground_truth.pk1')
    store_object.store_object(measurements_radar, save_folder,
        'measurements_radar.pk1')
    store_object.store_object(measurements_ais, save_folder,
        'measurements_ais.pk1')
    store_object.store_object(start_time, save_folder, 'start_time.pk1')
    store_object.store_object(measurement_model_radar, save_folder,
        'measurement_model_radar.pk1')
    store_object.store_object(measurement_model_ais, save_folder,
        'measurement_model_ais.pk1')
    store_object.store_object(transition_model, save_folder,
        'transition_model.pk1')


def generate_scenario_2(seed=1996, permanent_save=True, sigma_process=0.01,
    sigma_meas_radar=3, sigma_meas_ais=1, timesteps=20):
    """
    Generates scenario 2. Scenario 2 consists of radar and ais measurements with the same sampling rate.
    :param timesteps:
    :param seed:
    :param permanent_save:
    :param sigma_process:
    :param sigma_meas_radar:
    :param sigma_meas_ais:
    :return: Nothing. Saves the scenario to a specified folder
    """
    start_time = datetime.now()
    np.random.seed(seed)
    transition_model = CombinedLinearGaussianTransitionModel([
        ConstantVelocity(sigma_process), ConstantVelocity(sigma_process)])
    truth = GroundTruthPath([GroundTruthState([0, 1, 0, 1], timestamp=
        start_time)])
    for k in range(1, timesteps):
        truth.append(GroundTruthState(transition_model.function(truth[k - 1
            ], noise=True, time_interval=timedelta(seconds=1)), timestamp=
            start_time + timedelta(seconds=k)))
    measurement_model_radar = LinearGaussian(ndim_state=4, mapping=(0, 2),
        noise_covar=np.array([[sigma_meas_radar, 0], [0, sigma_meas_radar]]))
    measurement_model_ais = LinearGaussian(ndim_state=4, mapping=(0, 2),
        noise_covar=np.array([[sigma_meas_ais, 0], [0, sigma_meas_ais]]))
    measurements_radar = []
    for state in truth:
        measurement = measurement_model_radar.function(state, noise=True)
        measurements_radar.append(Detection(measurement, timestamp=state.
            timestamp))
    measurements_ais = []
    for state in truth:
        measurement = measurement_model_ais.function(state, noise=True)
        measurements_ais.append(Detection(measurement, timestamp=state.
            timestamp))
    if permanent_save:
        save_folder_name = seed.__str__()
    else:
        save_folder_name = 'temp'
    save_folder = '../scenarios/scenario2/' + save_folder_name + '/'
    store_object.store_object(truth, save_folder, 'ground_truth.pk1')
    store_object.store_object(measurements_radar, save_folder,
        'measurements_radar.pk1')
    store_object.store_object(measurements_ais, save_folder,
        'measurements_ais.pk1')
    store_object.store_object(start_time, save_folder, 'start_time.pk1')
    store_object.store_object(measurement_model_radar, save_folder,
        'measurement_model_radar.pk1')
    store_object.store_object(measurement_model_ais, save_folder,
        'measurement_model_ais.pk1')
    store_object.store_object(transition_model, save_folder,
        'transition_model.pk1')


def generate_scenario_3(seed=1996, permanent_save=True, radar_meas_rate=1,
    ais_meas_rate=5, sigma_process=0.01, sigma_meas_radar=3, sigma_meas_ais
    =1, timesteps=20):
    """
    Generates scenario 3. Scenario 3 consists of radar and ais measurements with different sampling rate. The sampling
    rate is specified in the input params. A groundtruth is generated for each second.
    :param seed:
    :param permanent_save:
    :param radar_meas_rate:
    :param ais_meas_rate:
    :param sigma_process:
    :param sigma_meas_radar:
    :param sigma_meas_ais:
    :param timesteps: The amount of measurements from the slowest sensor
    :return: Nothing. Saves the scenario to a specified folder
    """
    start_time = datetime.now()
    np.random.seed(seed)
    transition_model = CombinedLinearGaussianTransitionModel([
        ConstantVelocity(sigma_process), ConstantVelocity(sigma_process)])
    truth = GroundTruthPath([GroundTruthState([0, 1, 0, 1], timestamp=
        start_time)])
    end_time = start_time + timedelta(seconds=timesteps * max(
        radar_meas_rate, ais_meas_rate))
    time = start_time + timedelta(seconds=1)
    while time < end_time:
        truth.append(GroundTruthState(transition_model.function(truth[-1],
            noise=True, time_interval=timedelta(seconds=1)), timestamp=time))
        time += timedelta(seconds=1)
    measurement_model_radar = LinearGaussian(ndim_state=4, mapping=(0, 2),
        noise_covar=np.array([[sigma_meas_radar, 0], [0, sigma_meas_radar]]))
    measurement_model_ais = LinearGaussian(ndim_state=4, mapping=(0, 2),
        noise_covar=np.array([[sigma_meas_ais, 0], [0, sigma_meas_ais]]))
    measurements_radar = []
    measurements_ais = []
    next_radar_meas_time = start_time
    next_ais_meas_time = start_time
    for state in truth:
        if state.timestamp == next_radar_meas_time:
            measurement = measurement_model_radar.function(state, noise=True)
            measurements_radar.append(Detection(measurement, timestamp=
                state.timestamp))
            next_radar_meas_time += timedelta(seconds=radar_meas_rate)
        if state.timestamp == next_ais_meas_time:
            measurement = measurement_model_ais.function(state, noise=True)
            measurements_ais.append(Detection(measurement, timestamp=state.
                timestamp))
            next_ais_meas_time += timedelta(seconds=ais_meas_rate)
    if permanent_save:
        save_folder_name = seed.__str__()
    else:
        save_folder_name = 'temp'
    save_folder = '../scenarios/scenario3/' + save_folder_name + '/'
    store_object.store_object(truth, save_folder, 'ground_truth.pk1')
    store_object.store_object(measurements_radar, save_folder,
        'measurements_radar.pk1')
    store_object.store_object(measurements_ais, save_folder,
        'measurements_ais.pk1')
    store_object.store_object(start_time, save_folder, 'start_time.pk1')
    store_object.store_object(measurement_model_radar, save_folder,
        'measurement_model_radar.pk1')
    store_object.store_object(measurement_model_ais, save_folder,
        'measurement_model_ais.pk1')
    store_object.store_object(transition_model, save_folder,
        'transition_model.pk1')

"""Generates a simple one target scenario

Inspiration for how to generate ground truth and measurements are taken from
https://stonesoup.readthedocs.io/en/latest/auto_tutorials/01_KalmanFilterTutorial.html#sphx-glr-auto-tutorials-01-kalmanfiltertutorial-py
"""

from datetime import datetime
from datetime import timedelta
import numpy as np

from stonesoup.models.transition.linear import CombinedLinearGaussianTransitionModel, \
    ConstantVelocity
from stonesoup.types.groundtruth import GroundTruthPath, GroundTruthState
from stonesoup.types.detection import Detection
from stonesoup.models.measurement.linear import LinearGaussian

from utils import store_object


def generate_scenario_1(seed=1996, permanent_save=True, sigma_process=0.01, sigma_meas_radar=3, sigma_meas_ais=1):
    """
    Generates scenario 1. Todo define scenario 1
    :param seed:
    :param permanent_save:
    :param sigma_process:
    :param sigma_meas_radar:
    :param sigma_meas_ais:
    :return:
    """
    # specify seed to be able repeat example
    start_time = datetime.now()

    np.random.seed(seed)

    # combine two 1-D CV models to create a 2-D CV model
    transition_model = CombinedLinearGaussianTransitionModel([ConstantVelocity(sigma_process),
                                                              ConstantVelocity(sigma_process)])

    # starting at 0,0 and moving NE
    truth = GroundTruthPath([GroundTruthState([0, 1, 0, 1], timestamp=start_time)])

    # generate truth using transition_model and noise
    for k in range(1, 21):
        truth.append(GroundTruthState(
            transition_model.function(truth[k - 1], noise=True, time_interval=timedelta(seconds=1)),
            timestamp=start_time + timedelta(seconds=k)))

    # Simulate measurements
    # Specify measurement model for radar
    measurement_model_radar = LinearGaussian(
        ndim_state=4,  # number of state dimensions
        mapping=(0, 2),  # mapping measurement vector index to state index
        noise_covar=np.array([[sigma_meas_radar, 0],  # covariance matrix for Gaussian PDF
                              [0, sigma_meas_radar]])
    )

    # Specify measurement model for AIS
    measurement_model_ais = LinearGaussian(
        ndim_state=4,
        mapping=(0, 2),
        noise_covar=np.array([[sigma_meas_ais, 0],
                              [0, sigma_meas_ais]])
    )

    # generate "radar" measurements
    measurements_radar = []
    for state in truth:
        measurement = measurement_model_radar.function(state, noise=True)
        measurements_radar.append(Detection(measurement, timestamp=state.timestamp))

    # generate "AIS" measurements
    measurements_ais = []
    state_num = 0
    for state in truth:
        state_num += 1
        if not state_num % 2:  # measurement every second time step
            measurement = measurement_model_ais.function(state, noise=True)
            measurements_ais.append(Detection(measurement, timestamp=state.timestamp))

    if permanent_save:
        save_folder_name = seed.__str__()
    else:
        save_folder_name = "temp"

    save_folder = "../scenarios/scenario1/" + save_folder_name + "/"

    # save the ground truth and the measurements for the radar and the AIS
    store_object.store_object(truth, save_folder, "ground_truth.pk1")
    store_object.store_object(measurements_radar, save_folder, "measurements_radar.pk1")
    store_object.store_object(measurements_ais, save_folder, "measurements_ais.pk1")
    store_object.store_object(start_time, save_folder, "start_time.pk1")
    store_object.store_object(measurement_model_radar, save_folder, "measurement_model_radar.pk1")
    store_object.store_object(measurement_model_ais, save_folder, "measurement_model_ais.pk1")
    store_object.store_object(transition_model, save_folder, "transition_model.pk1")


def generate_scenario_2(seed=1996, permanent_save=True, sigma_process=0.01, sigma_meas_radar=3, sigma_meas_ais=1,
                        timesteps=20):
    """
    Generates scenario 2. Scenario 2 consists of radar and ais measurements with the same sampling rate.
    :param timesteps:
    :param seed:
    :param permanent_save:
    :param sigma_process:
    :param sigma_meas_radar:
    :param sigma_meas_ais:
    :return: Nothing. Saves the scenario to a specified folder
    """
    start_time = datetime.now()

    # specify seed to be able repeat example
    np.random.seed(seed)

    # combine two 1-D CV models to create a 2-D CV model
    transition_model = CombinedLinearGaussianTransitionModel(
        [ConstantVelocity(sigma_process), ConstantVelocity(sigma_process)])

    # starting at 0,0 and moving NE
    truth = GroundTruthPath([GroundTruthState([0, 1, 0, 1], timestamp=start_time)])

    # generate truth using transition_model and noise
    for k in range(1, timesteps):
        truth.append(GroundTruthState(
            transition_model.function(truth[k - 1], noise=True, time_interval=timedelta(seconds=1)),
            timestamp=start_time + timedelta(seconds=k)))

    # Simulate measurements
    # Specify measurement model for radar
    measurement_model_radar = LinearGaussian(
        ndim_state=4,  # number of state dimensions
        mapping=(0, 2),  # mapping measurement vector index to state index
        noise_covar=np.array([[sigma_meas_radar, 0],  # covariance matrix for Gaussian PDF
                              [0, sigma_meas_radar]])
    )

    # Specify measurement model for AIS (Same as for radar)
    measurement_model_ais = LinearGaussian(
        ndim_state=4,
        mapping=(0, 2),
        noise_covar=np.array([[sigma_meas_ais, 0],
                              [0, sigma_meas_ais]])
    )

    # generate "radar" measurements
    measurements_radar = []
    for state in truth:
        measurement = measurement_model_radar.function(state, noise=True)
        measurements_radar.append(Detection(measurement, timestamp=state.timestamp))

    # generate "AIS" measurements
    measurements_ais = []
    for state in truth:
        measurement = measurement_model_ais.function(state, noise=True)
        measurements_ais.append(Detection(measurement, timestamp=state.timestamp))

    if permanent_save:
        save_folder_name = seed.__str__()
    else:
        save_folder_name = "temp"

    save_folder = "../scenarios/scenario2/" + save_folder_name + "/"

    # save the ground truth and the measurements for the radar and the AIS
    store_object.store_object(truth, save_folder, "ground_truth.pk1")
    store_object.store_object(measurements_radar, save_folder, "measurements_radar.pk1")
    store_object.store_object(measurements_ais, save_folder, "measurements_ais.pk1")
    store_object.store_object(start_time, save_folder, "start_time.pk1")
    store_object.store_object(measurement_model_radar, save_folder, "measurement_model_radar.pk1")
    store_object.store_object(measurement_model_ais, save_folder, "measurement_model_ais.pk1")
    store_object.store_object(transition_model, save_folder, "transition_model.pk1")


def generate_scenario_3(seed=1996, permanent_save=True, radar_meas_rate=1, ais_meas_rate=5, sigma_process=0.01,
                        sigma_meas_radar=3, sigma_meas_ais=1, timesteps=20):
    """
    Generates scenario 3. Scenario 3 consists of radar and ais measurements with different sampling rate. The sampling
    rate is specified in the input params. A groundtruth is generated for each second.
    :param seed:
    :param permanent_save:
    :param radar_meas_rate:
    :param ais_meas_rate:
    :param sigma_process:
    :param sigma_meas_radar:
    :param sigma_meas_ais:
    :param timesteps: The amount of measurements from the slowest sensor
    :return: Nothing. Saves the scenario to a specified folder
    """
    start_time = datetime.now()

    # specify seed to be able to repeat the example
    np.random.seed(seed)

    # combine two 1-D CV models to create a 2-D CV model
    transition_model = CombinedLinearGaussianTransitionModel(
        [ConstantVelocity(sigma_process), ConstantVelocity(sigma_process)])

    # starting at 0,0 and moving NE
    truth = GroundTruthPath([GroundTruthState([0, 1, 0, 1], timestamp=start_time)])

    # generate truth using transition_model and noise
    end_time = start_time + timedelta(seconds=timesteps * max(radar_meas_rate, ais_meas_rate))
    time = start_time + timedelta(seconds=1)
    while time < end_time:
        truth.append(GroundTruthState(
            transition_model.function(truth[-1], noise=True, time_interval=timedelta(seconds=1)),
            timestamp=time))
        time += timedelta(seconds=1)

    # Simulate measurements
    # Specify measurement model for radar
    measurement_model_radar = LinearGaussian(
        ndim_state=4,  # number of state dimensions
        mapping=(0, 2),  # mapping measurement vector index to state index
        noise_covar=np.array([[sigma_meas_radar, 0],  # covariance matrix for Gaussian PDF
                              [0, sigma_meas_radar]])
    )

    # Specify measurement model for AIS (Same as for radar)
    measurement_model_ais = LinearGaussian(
        ndim_state=4,
        mapping=(0, 2),
        noise_covar=np.array([[sigma_meas_ais, 0],
                              [0, sigma_meas_ais]])
    )

    # generate "radar" measurements
    measurements_radar = []
    measurements_ais = []
    next_radar_meas_time = start_time
    next_ais_meas_time = start_time
    for state in truth:
        # check whether we want to generate a measurement from this gt
        if state.timestamp == next_radar_meas_time:
            measurement = measurement_model_radar.function(state, noise=True)
            measurements_radar.append(Detection(measurement, timestamp=state.timestamp))
            next_radar_meas_time += timedelta(seconds=radar_meas_rate)

        if state.timestamp == next_ais_meas_time:
            measurement = measurement_model_ais.function(state, noise=True)
            measurements_ais.append(Detection(measurement, timestamp=state.timestamp))
            next_ais_meas_time += timedelta(seconds=ais_meas_rate)

    if permanent_save:
        save_folder_name = seed.__str__()
    else:
        save_folder_name = "temp"

    save_folder = "../scenarios/scenario3/" + save_folder_name + "/"

    # save the ground truth and the measurements for the radar and the AIS
    store_object.store_object(truth, save_folder, "ground_truth.pk1")
    store_object.store_object(measurements_radar, save_folder, "measurements_radar.pk1")
    store_object.store_object(measurements_ais, save_folder, "measurements_ais.pk1")
    store_object.store_object(start_time, save_folder, "start_time.pk1")
    store_object.store_object(measurement_model_radar, save_folder, "measurement_model_radar.pk1")
    store_object.store_object(measurement_model_ais, save_folder, "measurement_model_ais.pk1")
    store_object.store_object(transition_model, save_folder, "transition_model.pk1")
