from os import path as osp, chdir
import numpy as np
import matplotlib.pyplot as plt
from co2mpas_driver import dsp as driver
from scipy.interpolate import interp1d

my_dir = osp.dirname(osp.abspath(__file__))
chdir(my_dir)


def simple_run():
    """
    This example computes and plots the driving trajectories with the varying desired speed.

    """
    # A sample car id from the database = {
    car_id_db = {
        "fuel_engine": [
            35135,
            39393,
            27748,
            8188,
            40516,
            35452,
            40225,
            7897,
            7972,
            41388,
            5766,
            9645,
            9639,
            5798,
            8280,
            34271,
            34265,
            6378,
            39723,
            34092,
            2592,
            5635,
            5630,
            7661,
            7683,
            8709,
            9769,
            1872,
            10328,
            35476,
            41989,
            26799,
            26851,
            27189,
            23801,
            3079,
            36525,
            47766,
            6386,
            33958,
            33986,
            5725,
            5718,
            36591,
            4350,
            39396,
            40595,
            5909,
            5897,
            5928,
            5915,
            40130,
            42363,
            34760,
            34766,
            1835,
            36101,
            42886,
            1431,
            46547,
            44799,
            41045,
            39820,
            34183,
            34186,
            20612,
            20605,
            1324,
            9882,
            4957,
            5595,
            18831,
            18833,
            9575,
            5380,
            9936,
            7995,
            6331,
            18173,
            34286,
            34279,
            20706,
            34058,
            34057,
            24268,
            19028,
            19058,
            7979,
            22591,
            34202,
            40170,
            44599,
            5358,
            5338,
            34015,
            9872,
            9856,
            6446,
            8866,
            9001,
            9551,
            6222,
        ],
        "electric_engine": [47844],
    }
    car_id = 47844

    # The desired speed
    vdes = 124 / 3.6

    # Current speed
    v_start = 0

    # The simulation step in seconds
    sim_step = 0.1

    # The driving style as described in the TRR paper.
    driver_style = 0.6

    # Duration of the simulation in seconds.
    duration = 300

    # sample time series
    times = np.arange(0, duration + sim_step, sim_step)

    # core model, this will select and execute the proper functions for the given inputs and returns the output
    # You can also pass vehicle database path db_path='path to vehicle db'
    if car_id in car_id_db["electric_engine"]:
        my_veh = driver(
            dict(
                vehicle_id=car_id,
                inputs=dict(
                    inputs=dict(
                        starting_velocity=v_start,
                        driver_style=driver_style,
                        sim_start=0,
                        sim_step=sim_step,
                        duration=duration,
                        use_linear_gs=True,
                        use_cubic=False,
                    )
                ),
            )
        )["outputs"]["driver_simulation_model"]
    else:
        # The gear shifting style as described in the TRR paper.
        gs_style = 0.8
        my_veh = driver(
            dict(
                vehicle_id=car_id,
                inputs=dict(
                    inputs=dict(
                        gear_shifting_style=gs_style,
                        starting_velocity=v_start,
                        driver_style=driver_style,
                        sim_start=0,
                        sim_step=sim_step,
                        duration=duration,
                        degree=4,
                        use_linear_gs=True,
                        use_cubic=False,
                    )
                ),
            )
        )["outputs"]["driver_simulation_model"]
    # Plots workflow of the core model, this will automatically open an internet browser and show the work flow
    # of the core model. you can click all the rectangular boxes to see in detail sub models like load, model,
    # write and plot.
    # you can uncomment to plot the work flow of the model
    # driver.plot(1)  # Note: run your IDE as Administrator if file permission error.

    f_des = interp1d(
        [0, 210, 700, 1050, 1400, 3150, 4550, 5250],  # Position (m)
        [20, 20, 30, 5, 25, 35, 15, 0],  # Desired speed (m/s)
        kind="next",
    )
    res = {}
    for myt in times:
        if myt == times[0]:
            my_veh.reset(v_start)
            res = {
                "accel": [0],
                "speed": [v_start],
                "position": [0],
                "gear": [0],
                "v_des": [f_des(0)],
                "fc": [0],
            }
            continue
        res["v_des"].append(f_des(res["position"][-1]))
        gear, gear_count, next_velocity, acc = my_veh(sim_step, res["v_des"][-1])

        if car_id in car_id_db["electric_engine"]:
            fc = 0
        else:
            fc = my_veh.calculate_fuel_consumption(
                next_velocity, acc, gear, gear_count, sim_step
            )

        res["accel"].append(acc)
        res["speed"].append(next_velocity)
        res["gear"].append(gear)
        res["position"].append(res["position"][-1] + next_velocity * sim_step)
        res["fc"].append(fc)

    fig = plt.figure("Time-Speed")
    plt.plot(times, res["speed"], label=f"Simulation car:{car_id}, DS: {driver_style}")
    plt.plot(times, res["v_des"], label=f"Desired speed (m/s)")
    plt.legend()
    fig.suptitle("Speed over time", x=0.54, y=1, fontsize=12)
    plt.xlabel("Time (s)", fontsize=14)
    plt.ylabel("Speed (m/s)", fontsize=12)
    plt.grid()

    fig = plt.figure("Time-Acceleration")
    plt.plot(times, res["accel"], label=f"Simulation car:{car_id}, DS: {driver_style}")
    plt.legend()
    fig.suptitle("Acceleration over time", x=0.54, y=1, fontsize=12)
    plt.xlabel("Time (s)", fontsize=14)
    plt.ylabel("Acceleration (m/s2)", fontsize=12)
    plt.grid()

    if car_id in car_id_db["electric_engine"]:
        pass
    else:
        fig = plt.figure("Time-Fuel Consumption")
        plt.plot(
            times,
            np.dot(res["fc"], sim_step),
            label=f"Simulation car:{car_id}, DS: {driver_style}",
        )
        plt.legend()
        fig.suptitle("Fuel Consumption over time", x=0.54, y=1, fontsize=12)
        plt.xlabel("Time (s)", fontsize=14)
        plt.ylabel("Fuel consumption rate (g/s)", fontsize=12)
        plt.grid()

    fig = plt.figure("Distance-Speed")
    plt.plot(
        res["position"],
        res["speed"],
        label=f"Simulation car:{car_id}, DS: {driver_style}",
    )
    plt.plot(res["position"], res["v_des"], label=f"Desired speed (m/s)")
    plt.legend()
    fig.suptitle("Speed over distance", x=0.54, y=1, fontsize=12)
    plt.xlabel("Distance (m)", fontsize=14)
    plt.ylabel("Speed (m/s)", fontsize=12)
    plt.grid()

    fig = plt.figure("Distance-Acceleration")
    plt.plot(
        res["position"],
        res["accel"],
        label=f"Simulation car:{car_id}, DS: {driver_style}",
    )
    plt.legend()
    fig.suptitle("Acceleration over distance", x=0.54, y=1, fontsize=12)
    plt.xlabel("Distance (m)", fontsize=14)
    plt.ylabel("Acceleration (m/s2)", fontsize=12)
    plt.grid()
    plt.show()


if __name__ == "__main__":
    simple_run()
