/* main.c
   Keil - STM32F103C8 (Blue Pill) using SPL (Standard Peripheral Library)
   Auto torque distribution: 2 motors (L298N), 2x ACS712, MPU6050 IMU (I2C)
*/

#include "stm32f10x.h"
#include <math.h>
#include <stdio.h>

/* ---------------- CONFIG ---------------- */
#define PWM_MAX            1000U
#define CONTROL_MS         100U   // control loop period in ms
#define ACS_VOLTAGE_ZERO   2.50f  // measured Vout at 0A (measure & set)
#define ACS_SENSITIVITY    0.185f // V/A typical for 5A variant @ Vcc=5V
#define VREF               3.3f
#define ADC_MAX            4095.0f
#define K_SLOPE            0.6f   // slope compensation factor
#define ALPHA              0.35f  // max torque re-distribution
#define I_MAX              10.0f  // A, safe current max

/* I2C/MPU6050 addresses */
#define MPU6050_ADDR       0x68   // 7-bit address
#define MPU6050_REG_PWRMG  0x6B
#define MPU6050_REG_ACCEL  0x3B

/* ---------- global volatile counters (if encoders used) ---------- */
volatile uint32_t enc1_count = 0;
volatile uint32_t enc2_count = 0;

/* ---------- function prototypes ---------- */
void RCC_Config(void);
void GPIO_Config(void);
void TIM3_PWM_Config(void);
void ADC1_Config(void);
void I2C1_Config(void);
uint16_t ADC_ReadChannel(uint8_t channel);
void Delay_ms(uint32_t ms);

void MPU6050_Init(void);
int MPU6050_ReadAccel(float *ax, float *ay, float *az);

void SetMotorDirFront(int forward);
void SetMotorDirRear(int forward);
void SetPWM_MotorFront(uint16_t pwm);
void SetPWM_MotorRear(uint16_t pwm);

/* ---------- Simple I2C helper (SPL) ---------- */
int I2C_WriteReg(uint8_t devAddr, uint8_t regAddr, uint8_t data) {
  I2C_AcknowledgeConfig(I2C1, ENABLE);
  I2C_GenerateSTART(I2C1, ENABLE);
  /* wait start */
  while(!I2C_GetFlagStatus(I2C1, I2C_FLAG_SB));
  I2C_Send7bitAddress(I2C1, devAddr<<1, I2C_Direction_Transmitter);
  while(!I2C_CheckEvent(I2C1, I2C_EVENT_MASTER_TRANSMITTER_MODE_SELECTED));
  I2C_SendData(I2C1, regAddr);
  while(!I2C_GetFlagStatus(I2C1, I2C_FLAG_TXE));
  I2C_SendData(I2C1, data);
  while(!I2C_GetFlagStatus(I2C1, I2C_FLAG_TXE));
  I2C_GenerateSTOP(I2C1, ENABLE);
  return 0;
}

int I2C_ReadBytes(uint8_t devAddr, uint8_t regAddr, uint8_t *buf, uint8_t len) {
  /* send register */
  I2C_GenerateSTART(I2C1, ENABLE);
  while(!I2C_GetFlagStatus(I2C1, I2C_FLAG_SB));
  I2C_Send7bitAddress(I2C1, devAddr<<1, I2C_Direction_Transmitter);
  while(!I2C_CheckEvent(I2C1, I2C_EVENT_MASTER_TRANSMITTER_MODE_SELECTED));
  I2C_SendData(I2C1, regAddr);
  while(!I2C_GetFlagStatus(I2C1, I2C_FLAG_TXE));
  /* restart for read */
  I2C_GenerateSTART(I2C1, ENABLE);
  while(!I2C_GetFlagStatus(I2C1, I2C_FLAG_SB));
  I2C_Send7bitAddress(I2C1, devAddr<<1, I2C_Direction_Receiver);
  while(!I2C_CheckEvent(I2C1, I2C_EVENT_MASTER_RECEIVER_MODE_SELECTED));
  while(len) {
    if(len == 1) {
      I2C_AcknowledgeConfig(I2C1, DISABLE);
      I2C_GenerateSTOP(I2C1, ENABLE);
    }
    if(I2C_CheckEvent(I2C1, I2C_EVENT_MASTER_BYTE_RECEIVED)) {
      *buf++ = I2C_ReceiveData(I2C1);
      len--;
    }
  }
  I2C_AcknowledgeConfig(I2C1, ENABLE);
  return 0;
}

/* ---------- main ---------- */
int main(void) {
  SystemInit();
  RCC_Config();
  GPIO_Config();
  TIM3_PWM_Config();
  ADC1_Config();
  I2C1_Config();

  MPU6050_Init();

  /* start PWM with zero duty */
  SetPWM_MotorFront(0);
  SetPWM_MotorRear(0);

  /* Main control loop */
  uint32_t lastTick = 0;
  while(1) {
    uint32_t t0 = SysTick->VAL; (void)t0;

    /* 1) Read sensors */
    uint16_t adc_acs_f = ADC_ReadChannel(ADC_Channel_4); // PA4
    uint16_t adc_acs_r = ADC_ReadChannel(ADC_Channel_5); // PA5
    uint16_t adc_pot_t = ADC_ReadChannel(ADC_Channel_6); // PA6
    uint16_t adc_pot_b = ADC_ReadChannel(ADC_Channel_7); // PA7

    float v_acs_f = (adc_acs_f * VREF) / ADC_MAX;
    float v_acs_r = (adc_acs_r * VREF) / ADC_MAX;
    float I_front = (v_acs_f - ACS_VOLTAGE_ZERO) / ACS_SENSITIVITY;
    float I_rear  = (v_acs_r - ACS_VOLTAGE_ZERO) / ACS_SENSITIVITY;
    if(I_front < 0) I_front = 0;
    if(I_rear < 0)  I_rear = 0;

    float throttle = (float)adc_pot_t / ADC_MAX; // 0..1
    float biasRaw  = (float)adc_pot_b / ADC_MAX; // 0..1
    float bias = (biasRaw - 0.5f) * 2.0f;        // -1..1

    /* IMU read */
    float ax, ay, az;
    if(MPU6050_ReadAccel(&ax, &ay, &az) != 0) {
      ax = ay = 0; az = 1;
    }
    /* compute pitch (deg) using accelerometer */
    float pitch_rad = atan2f(-ax, sqrtf(ay*ay + az*az));
    float pitch_deg = pitch_rad * 180.0f / 3.14159265f;

    /* 2) Control logic */
    /* Basic slope compensation: T_total scales with sin(pitch) */
    float T_base = throttle; // normalized 0..1
    float T_total = T_base * (1.0f + K_SLOPE * sinf(pitch_rad));

    /* front/rear ratio from pitch + bias */
    float pitch_bias = (K_SLOPE * pitch_deg) / 30.0f; // scale
    if(pitch_bias > 1.0f) pitch_bias = 1.0f;
    if(pitch_bias < -1.0f) pitch_bias = -1.0f;
    float front_ratio = 0.5f - pitch_bias * ALPHA;
    float rear_ratio  = 0.5f + pitch_bias * ALPHA;
    /* also apply user bias (adds/subtracts small amount) */
    front_ratio += -bias * 0.1f;
    rear_ratio  += bias * 0.1f;

    /* normalize */
    float sum = front_ratio + rear_ratio;
    if(sum <= 0.001f) { front_ratio = 0.5f; rear_ratio = 0.5f; }
    else { front_ratio /= sum; rear_ratio /= sum; }

    /* optional slip detection using currents (simple) */
    const float SLIP_CURR_DIFF = 2.0f; // A
    if(I_rear - I_front > SLIP_CURR_DIFF) {
      rear_ratio *= 0.7f;
      front_ratio = 1.0f - rear_ratio;
    } else if(I_front - I_rear > SLIP_CURR_DIFF) {
      front_ratio *= 0.7f;
      rear_ratio = 1.0f - front_ratio;
    }

    /* compute PWM targets (0..PWM_MAX) */
    uint16_t pwm_front = (uint16_t)(T_total * front_ratio * PWM_MAX);
    uint16_t pwm_rear  = (uint16_t)(T_total * rear_ratio  * PWM_MAX);

    /* apply simple direction policy: throttle > small threshold -> forward */
    if(throttle > 0.02f) {
      SetMotorDirFront(1);
      SetMotorDirRear(1);
    } else { // no throttle -> stop
      SetMotorDirFront(0);
      SetMotorDirRear(0);
      pwm_front = 0;
      pwm_rear  = 0;
    }

    /* set PWM outputs */
    SetPWM_MotorFront(pwm_front);
    SetPWM_MotorRear(pwm_rear);

    /* small delay for control period */
    Delay_ms(CONTROL_MS);
  }
}

/* ========== Hardware init functions ========== */

void RCC_Config(void) {
  /* Enable clocks: GPIOA/B, ADC1, TIM3, AFIO, I2C1 */
  RCC_APB2PeriphClockCmd(RCC_APB2Periph_GPIOA | RCC_APB2Periph_GPIOB |
                        RCC_APB2Periph_ADC1 | RCC_APB2Periph_AFIO, ENABLE);
  RCC_APB1PeriphClockCmd(RCC_APB1Periph_TIM3 | RCC_APB1Periph_I2C1, ENABLE);
  /* ADC clock config */
  RCC_ADCCLKConfig(RCC_PCLK2_Div6); // ADC clock = PCLK2/6
}

void GPIO_Config(void) {
  GPIO_InitTypeDef gpio;

  /* Motor IN pins PA0..PA3 output push-pull */
  gpio.GPIO_Pin = GPIO_Pin_0 | GPIO_Pin_1 | GPIO_Pin_2 | GPIO_Pin_3;
  gpio.GPIO_Mode = GPIO_Mode_Out_PP;
  gpio.GPIO_Speed = GPIO_Speed_50MHz;
  GPIO_Init(GPIOA, &gpio);

  /* PWM pins PB0, PB1 alternate function push-pull */
  gpio.GPIO_Pin = GPIO_Pin_0 | GPIO_Pin_1;
  gpio.GPIO_Mode = GPIO_Mode_AF_PP;
  gpio.GPIO_Speed = GPIO_Speed_50MHz;
  GPIO_Init(GPIOB, &gpio);

  /* I2C pins PB6 (SCL), PB7 (SDA) open-drain */
  gpio.GPIO_Pin = GPIO_Pin_6 | GPIO_Pin_7;
  gpio.GPIO_Mode = GPIO_Mode_AF_OD;
  gpio.GPIO_Speed = GPIO_Speed_50MHz;
  GPIO_Init(GPIOB, &gpio);

  /* ADC pins PA4..PA7 as analog - no init needed (floating input) */
}

void TIM3_PWM_Config(void) {
  TIM_TimeBaseInitTypeDef tb;
  TIM_OCInitTypeDef oc;

  /* TIM3 base */
  TIM_TimeBaseStructInit(&tb);
  tb.TIM_Period = PWM_MAX;
  tb.TIM_Prescaler = 72 - 1; // 72MHz/72 = 1MHz -> period counts
  tb.TIM_ClockDivision = 0;
  tb.TIM_CounterMode = TIM_CounterMode_Up;
  TIM_TimeBaseInit(TIM3, &tb);

  /* PWM mode for CH3 (PB0) and CH4 (PB1) */
  TIM_OCStructInit(&oc);
  oc.TIM_OCMode = TIM_OCMode_PWM1;
  oc.TIM_OutputState = TIM_OutputState_Enable;
  oc.TIM_Pulse = 0;
  oc.TIM_OCPolarity = TIM_OCPolarity_High;

  TIM_OC3Init(TIM3, &oc);
  TIM_OC4Init(TIM3, &oc);

  TIM_Cmd(TIM3, ENABLE);
}

void ADC1_Config(void) {
  ADC_InitTypeDef adc;
  ADC_DeInit(ADC1);
  ADC_StructInit(&adc);
  adc.ADC_Mode = ADC_Mode_Independent;
  adc.ADC_ScanConvMode = DISABLE;
  adc.ADC_ContinuousConvMode = DISABLE;
  adc.ADC_ExternalTrigConv = ADC_ExternalTrigConv_None;
  adc.ADC_DataAlign = ADC_DataAlign_Right;
  adc.ADC_NbrOfChannel = 1;
  ADC_Init(ADC1, &adc);
  ADC_Cmd(ADC1, ENABLE);

  /* calibration */
  ADC_ResetCalibration(ADC1);
  while(ADC_GetResetCalibrationStatus(ADC1));
  ADC_StartCalibration(ADC1);
  while(ADC_GetCalibrationStatus(ADC1));
}

uint16_t ADC_ReadChannel(uint8_t channel) {
  ADC_RegularChannelConfig(ADC1, channel, 1, ADC_SampleTime_55Cycles5);
  ADC_SoftwareStartConvCmd(ADC1, ENABLE);
  while(!ADC_GetFlagStatus(ADC1, ADC_FLAG_EOC));
  return ADC_GetConversionValue(ADC1);
}

void I2C1_Config(void) {
  I2C_InitTypeDef i2c;
  I2C_DeInit(I2C1);
  I2C_StructInit(&i2c);
  i2c.I2C_Mode = I2C_Mode_I2C;
  i2c.I2C_DutyCycle = I2C_DutyCycle_2;
  i2c.I2C_OwnAddress1 = 0x00;
  i2c.I2C_Ack = I2C_Ack_Enable;
  i2c.I2C_AcknowledgedAddress = I2C_AcknowledgedAddress_7bit;
  i2c.I2C_ClockSpeed = 100000;
  I2C_Init(I2C1, &i2c);
  I2C_Cmd(I2C1, ENABLE);
}

/* ---------- MPU6050 functions ---------- */
void MPU6050_Init(void) {
  /* wake up */
  I2C_WriteReg(MPU6050_ADDR, MPU6050_REG_PWRMG, 0x00);
  Delay_ms(10);
}

int MPU6050_ReadAccel(float *ax, float *ay, float *az) {
  uint8_t buf[6];
  if(I2C_ReadBytes(MPU6050_ADDR, MPU6050_REG_ACCEL, buf, 6) != 0) return -1;
  /* raw values */
  int16_t rx = (buf[0] << 8) | buf[1];
  int16_t ry = (buf[2] << 8) | buf[3];
  int16_t rz = (buf[4] << 8) | buf[5];
  /* Convert to g (assuming FS = ±2g, scale = 16384 LSB/g) */
  *ax = (float)rx / 16384.0f;
  *ay = (float)ry / 16384.0f;
  *az = (float)rz / 16384.0f;
  return 0;
}

/* ---------- Motor helpers ---------- */
void SetMotorDirFront(int forward) {
  if(forward) {
    GPIO_SetBits(GPIOA, GPIO_Pin_0);
    GPIO_ResetBits(GPIOA, GPIO_Pin_1);
  } else {
    GPIO_ResetBits(GPIOA, GPIO_Pin_0);
    GPIO_SetBits(GPIOA, GPIO_Pin_1);
  }
}
void SetMotorDirRear(int forward) {
  if(forward) {
    GPIO_SetBits(GPIOA, GPIO_Pin_2);
    GPIO_ResetBits(GPIOA, GPIO_Pin_3);
  } else {
    GPIO_ResetBits(GPIOA, GPIO_Pin_2);
    GPIO_SetBits(GPIOA, GPIO_Pin_3);
  }
}
void SetPWM_MotorFront(uint16_t pwm) {
  if(pwm > PWM_MAX) pwm = PWM_MAX;
  TIM3->CCR3 = pwm;
}
void SetPWM_MotorRear(uint16_t pwm) {
  if(pwm > PWM_MAX) pwm = PWM_MAX;
  TIM3->CCR4 = pwm;
}

/* ---------- small delay (blocking) ---------- */
void Delay_ms(uint32_t ms) {
  /* crude busy loop using SysTick */
  SysTick_Config(SystemCoreClock / 1000);
  uint32_t start = SysTick->VAL;
  for(uint32_t i=0;i<ms;i++) {
    uint32_t t = SysTick->VAL;
    while((SysTick->VAL - t) > 0); // no-op (SysTick counts down)
  }
  SysTick->CTRL = 0; // disable
}