#include <webots/robot.h>
#include <webots/distance_sensor.h>
#include <webots/camera.h>
#include <webots/motor.h>
#include <stdio.h>

// time in [ms] of a simulation step
#define TIME_STEP 64

// entree point of the controller
int main(int argc, char **argv)
{
    int  r,g,b, countt;
    const unsigned char *image;
    // initialise the Webots API
    wb_robot_init();
    int drive_counter = 0;

    // internal variables
    int i;
    // initialise distance sensors
    WbDeviceTag ds[8];
    char ds_names[8][16] = {"ds_left_front", "ds_right_front", "ds_left_corner", "ds_right_corner", "ds_left_rear", "ds_right_rear", "ds_left_back", "ds_right_back" };
    for (i=0; i<8 ; i++) {
        ds[i] = wb_robot_get_device(ds_names[i]);
        wb_distance_sensor_enable(ds[i], TIME_STEP);
    }

    // initialise motors
    WbDeviceTag wheels[4];
    char wheels_names[4][8] = {
        "wheel1", "wheel2", "wheel3", "wheel4"
    };
    for (i=0; i<4 ; i++) {
        wheels[i] = wb_robot_get_device(wheels_names[i]);
        wb_motor_set_position(wheels[i], INFINITY);
    }

    // initialise camera
    WbDeviceTag camera;
    camera = wb_robot_get_device("color_camera");
    wb_camera_enable(camera, TIME_STEP);
    int width = wb_camera_get_width(camera);
    int height = wb_camera_get_height(camera);
    int center = 0;
    double left_speed  = 1.0;
    double right_speed = 1.0;
    // feedback loop
    while (wb_robot_step(TIME_STEP) != -1) {
        int i,j;
        image = wb_camera_get_image(camera);
        r = 0;
        g = 0;
        b = 0;
        countt = 0 ;
        center = 0;

        if (drive_counter == 0){
            drive_counter = 10;
            for (i = 3*width/7; i < 4*width/7; i++) {
                for (j = height/2; j < 3*height/4; j++) {
                    r = wb_camera_image_get_red(image, width, i, j);
                    b = wb_camera_image_get_blue(image, width, i, j);
                    g = wb_camera_image_get_green(image, width, i, j);

                    if (b>225 && r >225 && g>225) {
                        countt++;
                        center+=i;
                    }
                }
            }

            double ds_values[8];

            for (i=0; i<8 ; i++)
            ds_values[i] = wb_distance_sensor_get_value(ds[i]);

            bool sheepAtRight = ds_values[3] < 5000.0 || ds_values[5] < 5000.0;


            left_speed  = 1.0;
            right_speed = 1.0;

            if(countt == 0){
                left_speed  =  2.0;
                right_speed = 0.2;
            }

            else
            {
                center /= countt;
                center -= width/2;
                if (center > width/24) {
                    left_speed  =  2.5;
                    right_speed = 0.1;
                }
                else if (center < -width/24) {
                    left_speed  = 0.1;
                    right_speed = 2.5;
                }
            }
            if(sheepAtRight)
            {
                left_speed = 2;
                right_speed = 2;
            }
        }
        else
        {
            drive_counter--;
        }

        // write actuators inputs
        wb_motor_set_velocity(wheels[0], left_speed);
        wb_motor_set_velocity(wheels[1], right_speed);
        wb_motor_set_velocity(wheels[2], left_speed);
        wb_motor_set_velocity(wheels[3], right_speed);
    }

    // cleanup the Webots API
    wb_robot_cleanup();
    return 0; //EXIT_SUCCESS
}