import java.util.LinkedHashMap;

import com.microsoft.bonsai.client.*;
import com.microsoft.bonsai.generated.Sessions;
import com.microsoft.bonsai.generated.models.*;

import com.fasterxml.jackson.databind.*;

BonsaiClientConfig bcConfig = new BonsaiClientConfig(getWorkspace(),getAccessKey());
BonsaiClient client = new BonsaiClient(bcConfig);

Sessions sessions = client.sessions();
                
SimulatorInterface sim_interface = new SimulatorInterface();
sim_interface.withName("Cartpole-Java");
sim_interface.withTimeout(60.0);
sim_interface.withCapabilities(null);

// minimum required
sim_interface.withSimulatorContext(bcConfig.simulatorContext);

//create only returns an object, so we need to check what type of object
Object registrationResponse = sessions.create(workspaceName, sim_interface);

// if we get an error during registration
if(registrationResponse.getClass() == ProblemDetails.class)
{
    ProblemDetails details = (ProblemDetails)registrationResponse;
    System.out.println(java.time.LocalDateTime.now() + " - ProblemDetails - " +  details.title());
}
// successful registration
else if(registrationResponse.getClass() == SimulatorSessionResponse.class)
{
    registered = registrationResponse;
    SimulatorSessionResponse sessionResponse = (SimulatorSessionResponse)registrationResponse;

    //this is required
    sessionId = sessionResponse.sessionId();
}

System.out.println(java.time.LocalDateTime.now() + " - registered session " + sessionId);

// build the SimulatorState object
SimulatorState simState = new SimulatorState();
simState.withSequenceId(sequenceId); // required
simState.withState(model.getState()); // required
simState.withHalted(model.halted()); // required

// advance only returns an object, so we need to check what type of object
Object response = client.sessions().advance(workspaceName, sessionId, simState);

// if we get an error during advance
if (response.getClass() == ProblemDetails.class) {
    ProblemDetails details = (ProblemDetails) response;

    System.out.println(java.time.LocalDateTime.now() + " - ProblemDetails - " + details.title());
}
// succesful advance
else if (response.getClass() == Event.class) {

    Event event = (Event) response;
    System.out.println(java.time.LocalDateTime.now() + " - received event: " + event.type());
    sequenceId = event.sequenceId(); // get the sequence from the result

    // now check the type of event and handle accordingly

    if (event.type() == EventType.EPISODE_START) {
        // ignored event in Cartpole
        Config config = new Config();

        // If you have config values, pass them to the start event on your model:
        //event.episodeStart().config()
        model.start(config);

    } else if (event.type() == EventType.EPISODE_STEP) {
        Action action = new Action();

        // action() returns an Object with a class value of LinkedHashMap
        LinkedHashMap map = (LinkedHashMap) event.episodeStep().action();

        // get the value of the action
        Object theCommand = map.get("command");

        // sometimes that value is an integer -- somtimes its a double. in either case,
        // Action.command is a double
        if (theCommand.getClass() == Integer.class)
            action.command = ((Integer) theCommand).doubleValue();
        else if (theCommand.getClass() == Double.class)
            action.command = ((Double) theCommand).doubleValue();

        // move the model forward
        model.step(action);
    } else if (event.type() == EventType.EPISODE_FINISH) {
        System.out.println("Episode Finish");
    } else if (event.type() == EventType.IDLE) {
        Thread.sleep(event.idle().callbackTime().longValue() * 1000);
    } else if (event.type() == EventType.UNREGISTER) {
        client.sessions().delete(workspaceName, sessionId);
    }
}