#include "motors.h"
#include "stm32g431xx.h"

void motors_init()
{
  RCC->AHB2ENR |= RCC_AHB2ENR_GPIOAEN;
  RCC->APB2ENR |= RCC_APB2ENR_TIM1EN;
  RCC->APB1ENR1 |= RCC_APB1ENR1_TIM2EN;
  
  motor_gpio_tim1_ch3_init();
  motor_gpio_tim1_ch4_init();
  motors_tim1_ch3_4_init();
  
  motor_gpio_tim2_ch1_init();
  motor_gpio_tim2_ch2_init();
  motors_tim2_ch1_2_init();
}

void motor_gpio_tim1_ch3_init()
{
  // set alt function mode PA10
  GPIOA->MODER &= ~(3 << (10 * 2));
  GPIOA->MODER |= 2 << (10 * 2);
  
  // AF6 for PA10
  GPIOA->AFR[1] &= ~(0xF << 8);
  GPIOA->AFR[1] |= 6 << 8;
  
  // very high speed
  GPIOA->OSPEEDR |= 3 << (10 * 2);
}

void motor_gpio_tim1_ch4_init()
{
  // set alt function mode PA11
  GPIOA->MODER &= ~(3 << (11 * 2));
  GPIOA->MODER |= 2 << (11 * 2);
  
  // AF11 for PA11
  GPIOA->AFR[1] &= ~(0xF << 12);
  GPIOA->AFR[1] |= 11 << 12;
  
  // very high speed
  GPIOA->OSPEEDR |= 3 << (11 * 2);
}

void motors_tim1_ch3_4_init()
{
  // PWM mode 1
  TIM1->CCMR2 &= ~TIM_CCMR2_OC3M;
  TIM1->CCMR2 |= TIM_CCMR2_OC3M_1 | TIM_CCMR2_OC3M_2;
  
  TIM1->CCMR2 &= ~TIM_CCMR2_OC4M;
  TIM1->CCMR2 |= TIM_CCMR2_OC4M_1 | TIM_CCMR2_OC4M_2;
  
  // preload enable
  TIM1->CCMR2 |= TIM_CCMR2_OC3PE;
  TIM1->CCMR2 |= TIM_CCMR2_OC4PE;
  
  // enable capture/compare 3 output
  TIM1->CCER |= TIM_CCER_CC3E;
  TIM1->CCER |= TIM_CCER_CC4E;
  
  TIM1->PSC = 16 - 1;
  TIM1->ARR = 20000 - 1;
  TIM1->CCR3 = 900;
  TIM1->CCR4 = 900;
  
  // TIM1_ARR is buffered
  TIM1->CR1 |= TIM_CR1_ARPE;
  
  // set main output enable
  TIM1->BDTR |= TIM_BDTR_MOE;
  
  // set update generation
  TIM1->EGR |= TIM_EGR_UG;
  
  // set counter enable
  TIM1->CR1 |= TIM_CR1_CEN;
}

void motor_gpio_tim2_ch1_init()
{
  // set alt function mode PA0
  GPIOA->MODER &= ~(3 << (0 * 2));
  GPIOA->MODER |= 2 << (0 * 2);
  
  // AF1 for PA0
  GPIOA->AFR[0] &= ~(0xF << 0);
  GPIOA->AFR[0] |= 1 << 0;
  
  // very high speed
  GPIOA->OSPEEDR |= 3 << (0 * 2);
}

void motor_gpio_tim2_ch2_init()
{
  // set alt function mode PA1
  GPIOA->MODER &= ~(3 << (1 * 2));
  GPIOA->MODER |= 2 << (1 * 2);
  
  // AF1 for PA1
  GPIOA->AFR[0] &= ~(0xF << 4);
  GPIOA->AFR[0] |= 1 << 4;
  
  // very high speed
  GPIOA->OSPEEDR |= 3 << (1 * 2);
}

void motors_tim2_ch1_2_init()
{
  // PWM mode 1
  TIM2->CCMR1 &= ~TIM_CCMR1_OC1M;
  TIM2->CCMR1 |= TIM_CCMR1_OC1M_1 | TIM_CCMR1_OC1M_2;
  
  TIM2->CCMR1 &= ~TIM_CCMR1_OC2M;
  TIM2->CCMR1 |= TIM_CCMR1_OC2M_1 | TIM_CCMR1_OC2M_2;
  
  // preload enable
  TIM2->CCMR1 |= TIM_CCMR1_OC1PE;
  TIM2->CCMR1 |= TIM_CCMR1_OC2PE;
  
  // enable capture/compare 3 output
  TIM2->CCER |= TIM_CCER_CC1E;
  TIM2->CCER |= TIM_CCER_CC2E;
  
  TIM2->PSC = 16 - 1;
  TIM2->ARR = 20000 - 1;
  TIM2->CCR1 = 900;
  TIM2->CCR2 = 900;
  
  // TIM2_ARR is buffered
  TIM2->CR1 |= TIM_CR1_ARPE;
  
  // set main output enable
  TIM2->BDTR |= TIM_BDTR_MOE;
  
  // set update generation
  TIM2->EGR |= TIM_EGR_UG;
  
  // set counter enable
  TIM2->CR1 |= TIM_CR1_CEN;
}

int16_t T;
int16_t P;
int16_t R;
int16_t Y;

int16_t m1;
int16_t m2;
int16_t m3;
int16_t m4;

void motors_set_throttle_mix(const int16_t throttle, const control_channels_t* chs, const int8_t armed)
{
  T = throttle;
  P = (int16_t) chs->pitch;
  R = (int16_t) chs->roll;
  Y = (int16_t) chs->yaw;
  
  m1 = T - P + R - Y;
  m2 = T - P - R + Y;
  m3 = T + P + R + Y;
  m4 = T + P - R - Y;
  
  motor_set_throttle(1, m1, armed);
  motor_set_throttle(2, m2, armed);
  motor_set_throttle(3, m3, armed);
  motor_set_throttle(4, m4, armed);
}

void motor_set_throttle(int8_t motor_number, int16_t us, int8_t armed)
{
  if (armed && us < 1050) us = 1050;
  if (us > 2000) us = 2000;
  
  switch (motor_number)
  {
  case 1:
    TIM1->CCR3 = us;
    break;
  case 2:
    TIM2->CCR2 = us;
    break;
  case 3:
    TIM1->CCR4 = us;
    break;
  case 4:
    TIM2->CCR1 = us;
    break;
  }
}

void motors_turn_off()
{
  motor_set_throttle(1, 900, 0);
  motor_set_throttle(2, 900, 0);
  motor_set_throttle(3, 900, 0);
  motor_set_throttle(4, 900, 0);
}

void motor1_set_throttle(int16_t throttle)
{
  if (throttle < 1000) throttle = 1000;
  if (throttle > 2000) throttle = 2000;
  TIM1->CCR3 = throttle;
}

void motor2_set_throttle(int16_t throttle)
{
  if (throttle < 1000) throttle = 1000;
  if (throttle > 2000) throttle = 2000;
  TIM1->CCR4 = throttle;
}

void motor3_set_throttle(int16_t throttle)
{
  if (throttle < 1000) throttle = 1000;
  if (throttle > 2000) throttle = 2000;
  TIM2->CCR1 = throttle;
}

void motor4_set_throttle(int16_t throttle)
{
  if (throttle < 1000) throttle = 1000;
  if (throttle > 2000) throttle = 2000;
  TIM2->CCR2 = throttle;
}