#pragma once

#include <vector>
#include "api.h"
#include "pros/adi.hpp"
#include "pros/rotation.hpp"

namespace ez{

    enum class ControlPeriod {
        driver,
        auton
    };
    enum class Controller {
        master,
        slave
    };
    enum class Buttons {
        arrow_buttons,
        Y_B_buttons,
        left_shoulder_buttons,
        right_shoulder_buttons
    };

    class ArmMech{
        private:
        std::vector<int> levelOneArmStates;
        std::vector<int> levelTwoArmStates;
        std::vector<int> levelThreeArmStates;
        std::vector<int> levelFourArmStates;
        int numArmStates = 2;
        float armKP;

        public:
        std::vector<int> selectedStates;
        int targetPos;
        int currArmState;
        int armTarget;
        float armError;
        float armVelocity;
        
        ArmMech();

        void arm_driver_control(enum Controller, enum Buttons);
        void arm_auton_control(int pickState);
        void cycle_states();
        void reset_state();
        void arm_control();

        void set_arm_kp(float _armKP);
        void set_state_values(std::vector<int> _levelOneArmStates, std::vector<int> _levelTwoArmStates, std::vector<int> _levelThreeArmStates, std::vector<int> _levelFourArmStates);

        int get_arm_state();
        int get_num_arm_states();
        float get_arm_kp();
    };

    class LiftMech {
        private:
        std::vector<int> liftStates;
        int numLiftStates;
        float liftKP;
        std::vector<int> firstRange;
        std::vector<int> secondRange; 
        std::vector<int> thirdRange; 
        std::vector<int> fourthRange;

        public:
        int currLiftState;
        int liftTarget;
        int liftHeight;
        float liftError;
        float liftVelocity;
    
        LiftMech();
        LiftMech(std::initializer_list<int> _liftStates, int _numLiftStates);

        void lift_driver_control(enum Controller, enum Buttons);
        void lift_auton_control(int pickState);
        void cycle_states();
        void reset_state();
        void lift_control();
        void lift_range_finder();

        void set_lift_kp(float _liftKP);
        void set_ranges(std::vector<int> _firstRange, std::vector<int> _secondRange, std::vector<int> _thirdRange, std::vector<int> _fourthRange);

        int get_lift_state();
        int get_num_lift_states();
        float get_lift_kp();
    };

    class WristMech {
        private:
        std::vector<int> wristStates;
        int numWristStates;
        float wristKP;

        public:
        int currWristState;
        int wristTarget;
        float wristError;
        float wristVelocity;

        WristMech();
        WristMech(std::initializer_list<int> _wristStates, int _numWristStates);

        void wrist_driver_control(enum Controller, enum Buttons);

        void rotate_wrist(int wristThreshold);
        void reset_state();
        void wrist_control();

        int get_wrist_state();
        int get_num_wrist_states();
        float get_wrist_kp();

        void set_wrist_kp(float _wristKP);
    };
};