#include "ATmega8.h"
#include <avr/interrupt.h>

// --------------------
// This firmware implements servo-style position control on standard brushless motor
// speed controllers of the sort used for radio control airplanes, quadcopters, etc.
//
// The basic speed controller hardware consists of a microcontroller and 6 transistors.
// The transistors are in 3 pairs, which connect to the 3 motor wires. This allows
// the microcontroller to connect each motor wire to battery + or battery -.
// By cycling through a sequence of 6 transistor combinations, the motor is rotated one
// electrical revolution. The number of electrical revolutions needed to complete
// a full physical revolution of the rotor is equal to half the number of permanent
// magnets on the rotor (the magnets alternate polarity, and one electrical revolution
// moves forward one north-south pair).
//
// In order to accurately track the position, 3 hall effect sensors (e.g. U1881)
// must be added to the motor to sense the position of the rotor magnets.
// These sensors cycle through 6 states, which correspond to the 6 transistor states.
// For most motor configurations, spacing the sensors equally around the motor
// (120 physical degrees apart) will result in a valid configuration. One exception
// is 9N12P (9 stator teeth, 12 rotor magnets), where the sensors must be placed at
// 3 consecutive stator teeth.
//
// Once you have the sensor spacing correct, you may also need to shift them all forward
// or back a little bit. Each active transistor pair will pull the rotor to a specific
// position, and the sensors need to be positioned so their state transition points are
// half way between these transistor positions. If the function DetectHallPhase fails to
// detect 6 unique hall states, then the sensors aren't positioned quite right.
//
// Positioning is only done to the nearest hall sensor state, so the total number of steps
// per physical revolution is 3 times the number of rotor magnets (6 transistor states per
// north-south pair). So for example a 12N14P motor will have 42 steps per revolution.
// --------------------



// Various ESC transistor mappings are taken from SimonK firmware https://github.com/sim-/tgy

#if 0 // DYS 40A speed controller, from dys_nfet.inc
// Output wires on DYS 40A, viewed from transistor side.
// When ApFET_PORT BIT(ApFET) is CLEAR, output A is positive. When AnFET_PORT BIT(AnFET) is SET, output A is ground. Else disconnected.
//  ______
// |B C A||
// |[] []||
// |[] []||
#define F_CPU 16000000
#define USE_ICP 1
#define rcp_in 0 // RC pulse input
#define rcp_in_port REG_PINB
#define ApFET 4
#define AnFET 5
#define BpFET 3
#define BnFET 7
#define CpFET 2
#define CnFET 1
#define ApFET_port REG_PORTD
#define AnFET_port REG_PORTD
#define BpFET_port REG_PORTD
#define BnFET_port REG_PORTD
#define CpFET_port REG_PORTD
#define CnFET_port REG_PORTB

#define INIT_PB (BIT(HallA) | BIT(HallB) | BIT(HallC)) // U1881 hall effect sensors need pull-up on their output pin
#define DIR_PB  BIT(CnFET)
#define INIT_PC (BIT(i2c_clk) | BIT(i2c_data))
#define DIR_PC  0
#define INIT_PD (BIT(ApFET) | BIT(BpFET) | BIT(CpFET) | BIT(txd))
#define DIR_PD  (BIT(ApFET) | BIT(AnFET) | BIT(BpFET) | BIT(BnFET) | BIT(CpFET) | BIT(txd))

#define ApFET_off ApFET_port |=  BIT(ApFET)
#define ApFET_on  ApFET_port &= ~BIT(ApFET)
#define BpFET_off BpFET_port |=  BIT(BpFET)
#define BpFET_on  BpFET_port &= ~BIT(BpFET)
#define CpFET_off CpFET_port |=  BIT(CpFET)
#define CpFET_on  CpFET_port &= ~BIT(CpFET)

#define AnFET_off AnFET_port &= ~BIT(AnFET)
#define AnFET_on  AnFET_port |=  BIT(AnFET)
#define BnFET_off BnFET_port &= ~BIT(BnFET)
#define BnFET_on  BnFET_port |=  BIT(BnFET)
#define CnFET_off CnFET_port &= ~BIT(CnFET)
#define CnFET_on  CnFET_port |=  BIT(CnFET)

#elif 1 // HobbyKing 10A speed controller, from bs.inc
// Output wires on HobbyKing 10A, viewed from transistor side.
// When ApFET_PORT BIT(ApFET) is SET, output A is positive. When AnFET_PORT BIT(AnFET) is SET, output A is ground. Else disconnected.
//  _____
// |C B A|
// |[] []|
// |[] []|
#define F_CPU 16000000
#define USE_INT0 1
#define rcp_in 2 // RC pulse input
#define rcp_in_port REG_PIND
#define ApFET 4
#define AnFET 5
#define BpFET 5
#define BnFET 4
#define CpFET 3
#define CnFET 0
#define ApFET_port REG_PORTD
#define AnFET_port REG_PORTD
#define BpFET_port REG_PORTC
#define BnFET_port REG_PORTC
#define CpFET_port REG_PORTC
#define CnFET_port REG_PORTB

#define INIT_PB (BIT(HallA) | BIT(HallB) | BIT(HallC)) // U1881 hall effect sensors need pull-up on their output pin
#define DIR_PB  BIT(CnFET)
#define INIT_PC 0
#define DIR_PC  (BIT(BpFET) | BIT(BnFET) | BIT(CpFET))
#define INIT_PD 0
#define DIR_PD  (BIT(ApFET) | BIT(AnFET))

#define ApFET_off ApFET_port &= ~BIT(ApFET)
#define ApFET_on  ApFET_port |=  BIT(ApFET)
#define BpFET_off BpFET_port &= ~BIT(BpFET)
#define BpFET_on  BpFET_port |=  BIT(BpFET)
#define CpFET_off CpFET_port &= ~BIT(CpFET)
#define CpFET_on  CpFET_port |=  BIT(CpFET)

#define AnFET_off AnFET_port &= ~BIT(AnFET)
#define AnFET_on  AnFET_port |=  BIT(AnFET)
#define BnFET_off BnFET_port &= ~BIT(BnFET)
#define BnFET_on  BnFET_port |=  BIT(BnFET)
#define CnFET_off CnFET_port &= ~BIT(CnFET)
#define CnFET_on  CnFET_port |=  BIT(CnFET)

#else
#error "No target speed controller specified"
#endif

// PORTB pins that the hall sensors are connected to
#define HallA 3
#define HallB 4
#define HallC 5

// Values for pwmState
#define PWM_FREEWHEEL 0
#define PWM_AnBp 1
#define PWM_AnCp 2
#define PWM_BnCp 3
#define PWM_BnAp 4
#define PWM_CnAp 5
#define PWM_CnBp 6
#define PWM_BRAKE 8
#define PWM_OFF_DUTY BIT07 // Toggle on and off for active and inactive part of duty cycle

// Radio control input pulse is approximately 1000 to 2000 microseconds. But timer ticks once
// every 4 microseconds so the range between min and max fits in 8 bits for fast scaling.
#define RC_PULSE_MIN 250
#define RC_PULSE_MAX 500

// Special signal to save current position before turning off
#define RC_PULSE_SAVE 64

// EEPROM addresses for variables to be saved
#define EEPROM_CUR_POS 2
#define EEPROM_MAX_POS 4
#define EEPROM_HALL_STATE_TO_PWM_STATE_POS 16

#define MAX_POWER 64
#define SOFT_STOP_DISTANCE 100 // Slow down when less than this many commutations from target position
#define DEADBAND 4 // Don't bother moving if already this close to targetPos

#define SPEED_RANGE_NUM 4
// 0: Very low speed, starting and stopping
// 1: More than 2ms per commutation
// 2: Between 1ms and 2ms per commutation
// 3: Less than 1ms per commutation

static const u8 maxPower[SPEED_RANGE_NUM] = { 16, 32, 48, 64 }; // Limiting PWM duty based on motor speed is a rough method of current limiting.

// As the motor rotates forward, hall state cycles through 6 values (001b,011b,010b,110b,100b,101b) or in decimal (1,3,2,6,4,5).
// Turning backward will step through that sequence in the opposite direction.
// The next and prev state tables map from the current state to the next or previous in that sequence.
static const u8 nextHallState[8] = { 0, 3, 6, 2, 5, 1, 4, 0 };
static const u8 prevHallState[8] = { 0, 5, 3, 1, 6, 4, 2, 0 };

// This table says which hall state corresponds to which PWM state (e.g. if you want the motor to rotate to hall state 3,
// then set PWM state to hallStateToPWMState[3]). The values depend on the particular motor and sensor arrangement.
// The function DetectHallPhase will automatically detect the correct values by moving the motor to each position and
// then reading the sensor state.
u8 hallStateToPWMState[8] = {0};

// By the user physically moving the servo output shaft to the desired extremes while in calibration mode, the total
// number of motor commutations between the two positions can be counted, and used to track the current position without
// a feedback potentiometer. However, this method does require re-calibrating every time the power is turned off, or
// saving the current position to EEPROM and assuming the servo shaft hasn't been moved while the power was off.
s16 curPos = 0; // In commutations, range 0 to maxPos
s8 curDir = 0; // Direction the motor is currently spinning, +1 or -1, or 0 if not moving.
s16 maxPos = 100; // In commutations. Can be negative. Initial value here is fairly irrelevant since it's normally calibrated and then loaded from EEPROM.
s16 targetPos = 0; // In commutations, calculated from the radio control PPM input value.
u8 speedRange = 0; // Based on time between last two sensor state changes
u8 hallState = 0; // Bit0 = hall sensor A, bit1 = hall sensor B, bit2 = hall sensor C.
u8 jitteringToStart = 0; // This is set true if PWM state is set to turn the opposite direction from the target, when struggling to get the rotor started under load.
volatile u16 rcPulseTime = 0xffff; // In milliseconds. Should always be <= RC_PULSE_MAX except at startup.
volatile u8 pwmState = PWM_FREEWHEEL; // State of the motor transistors
volatile u8 pwmDelay = MAX_POWER;
volatile u8 hallTime = 0; // Number of timer2 overflows since last hall sensor state change

void EEPROMRead(u8 addr, u8 *dest, u8 size);
void EEPROMWrite(u8 addr, const u8 *data, u8 size);
void Update();

STIN void wait(u8 ticks)
{
    asm volatile(
    "lsr %0\n"
    "lsr %0\n"
    "breq 2f\n"
    "1: dec %0\n"
    "nop\n"
    "brne 1b\n"
    "2:\n"
    : "=r" (ticks) : "0" (ticks));
}
STIN void waitLong(u32 ticks) { while(ticks > 255){wait(255); ticks-=255;} wait((u8)ticks); }

#if USE_ICP
#define RC_PULSE_IS_RISING_EDGE        REG_TCCR1B & TCCR1B_ICP_RISING_EDGE
#define RC_PULSE_WAIT_FOR_RISING_EDGE  REG_TCCR1B |= TCCR1B_ICP_RISING_EDGE
#define RC_PULSE_WAIT_FOR_FALLING_EDGE REG_TCCR1B &= ~TCCR1B_ICP_RISING_EDGE
#define RC_PULSE_ACKNOWLEDGE_INTERRUPT REG_TIFR = TIMSK_INPUT_CAPTURE
ISR(TIMER1_CAPT_vect) // This is the function name to go with the brackets below the #endif

#elif USE_INT0
#define RC_PULSE_IS_RISING_EDGE        REG_MCUCR == MCUCR_INT0_RISING_EDGE
#define RC_PULSE_WAIT_FOR_RISING_EDGE  REG_MCUCR = MCUCR_INT0_RISING_EDGE
#define RC_PULSE_WAIT_FOR_FALLING_EDGE REG_MCUCR = MCUCR_INT0_FALLING_EDGE
#define RC_PULSE_ACKNOWLEDGE_INTERRUPT REG_GIFR = GICR_INT0
ISR(INT0_vect) // This is the function name to go with the brackets below the #endif

#else
#error "No RC pulse input method specified."
#endif
{
    if (RC_PULSE_IS_RISING_EDGE)
    {
        RC_PULSE_WAIT_FOR_FALLING_EDGE;
        REG_TCNT1 = 0; // Start counting the duration of the pulse
        REG_TIFR = TIMSK_OC1B; // Clear overflow flag
    }
    else
    {
        rcPulseTime = REG_TCNT1;
        RC_PULSE_WAIT_FOR_RISING_EDGE;
    }
    RC_PULSE_ACKNOWLEDGE_INTERRUPT;
}

ISR(TIMER0_OVF_vect) // Interrupt for PWM
{
    if (pwmDelay == MAX_POWER)
        pwmState &= ~PWM_OFF_DUTY; // Max power is always on duty
    else
        pwmState ^= PWM_OFF_DUTY;

    switch(pwmState)
    {
    case PWM_AnBp: BnFET_off, CnFET_off; ApFET_off, CpFET_off; AnFET_on, BpFET_on; break;
    case PWM_AnCp: BnFET_off, CnFET_off; ApFET_off, BpFET_off; AnFET_on, CpFET_on; break;
    case PWM_BnCp: AnFET_off, CnFET_off; ApFET_off, BpFET_off; BnFET_on, CpFET_on; break;
    case PWM_BnAp: AnFET_off, CnFET_off; BpFET_off, CpFET_off; BnFET_on, ApFET_on; break;
    case PWM_CnAp: AnFET_off, BnFET_off; BpFET_off, CpFET_off; CnFET_on, ApFET_on; break;
    case PWM_CnBp: AnFET_off, BnFET_off; ApFET_off, CpFET_off; CnFET_on, BpFET_on; break;
    case PWM_BRAKE: ApFET_off, BpFET_off, CpFET_off, AnFET_on, BnFET_on, CnFET_on; break;
    case PWM_FREEWHEEL:
    // If freewheeling, or if PWM_OFF_DUTY flag is set, turn all transistors off.
    default: AnFET_off, BnFET_off; CnFET_off, ApFET_off, BpFET_off, CpFET_off; break;
    }

    REG_TCNT0 = (pwmState & PWM_OFF_DUTY) ? (u8)(pwmDelay - MAX_POWER) : (u8)-pwmDelay;
    REG_TIFR = TIMSK_TM0_OVERFLOW;
}

ISR(TIMER2_OVF_vect) // Interrupt for measuring motor speed
{
    if (hallTime != 0xff)
        hallTime++;
    REG_TIFR = TIMSK_TM2_OVERFLOW;
}

void EEPROMRead(u8 addr, u8 *dest, u8 size)
{
    const u8 irq = REG_SREG & SREG_IRQ_ENABLE;
    REG_SREG &= ~SREG_IRQ_ENABLE;
    AnFET_off, BnFET_off; CnFET_off, ApFET_off, BpFET_off, CpFET_off;

    REG_EEAR_H = 0;
    while(size--)
    {
        REG_EEAR_L = addr++;
        REG_EECR = EECR_READ;
        *dest++ = REG_EEDR;
    }
    REG_SREG |= irq;
}

void EEPROMWrite(u8 addr, const u8 *data, u8 size)
{
    const u8 irq = REG_SREG & SREG_IRQ_ENABLE;
    REG_SREG &= ~SREG_IRQ_ENABLE;
    AnFET_off, BnFET_off; CnFET_off, ApFET_off, BpFET_off, CpFET_off;

    REG_EEAR_H = 0;
    while(size--)
    {
        REG_EEAR_L = addr++;
        REG_EEDR = *data++;
        REG_EECR |= EECR_WRITE_ENABLE;
        REG_EECR |= EECR_WRITE;
        while(REG_EECR & EECR_WRITE) {}
    }
    REG_SREG |= irq;
}

void beep(u32 period)
{
    u8 cycles = 30;

    REG_SREG &= ~SREG_IRQ_ENABLE;
    AnFET_off, BnFET_off; CnFET_off, ApFET_off, BpFET_off, CpFET_off;

    while(--cycles)
    {
        BnFET_on;
        CpFET_on;
        waitLong(1000);
        CpFET_off;
        ApFET_on;
        waitLong(1000);
        ApFET_off;
        BnFET_off;
        waitLong(period - 2000);
    }

    REG_SREG |= SREG_IRQ_ENABLE;
}


// Call this function to automatically detect which hall state corresponds to which PWM state.
// For each PWM state, transistors are pulsed to pull the rotor into position, and then the hall state is read.
//
// The first and last entries of hallStateToPWMState are unused and should be 0 (because there should never be
// a time when all hall sensors are turned off or all turned on).
// The rest should have values 1 to 6 in some order, with no duplicates.
// Valid output should look something like 0, 6, 2, 1, 4, 5, 3, 0
// The numbers correspond to the PWM_AnBp, PWM_AnCp, etc. #defines
//
// Appropriate power level depends on the motor, battery, and any load on the output.
// Too low and the rotor won't move, too high and transistors may overheat or cause power brownout.
// Start low and if the table is incomplete, raise it and try again.
void DetectHallPhase(u8 power)
{
    u8 i;
    for (i = 1; i <= 6; i++)
    {
        pwmDelay = power >> 2; pwmState = i; waitLong(300000);
        pwmDelay = power >> 1; waitLong(300000);
        pwmDelay = power;      waitLong(300000);
        pwmDelay = power >> 1; waitLong(300000);
        hallStateToPWMState[(REG_PINB >> 3) & 7] = i;
    }
    pwmState = PWM_FREEWHEEL;
    pwmDelay = MAX_POWER;
}

int main(void)
{
    REG_SREG = 0;
    REG_DDRB = DIR_PB;
    REG_PORTB = INIT_PB;
    REG_DDRC = DIR_PC;
    REG_PORTC = INIT_PC;
    REG_DDRD = DIR_PD;
    REG_PORTD = INIT_PD;

    // Timer for PWM. Interrupt in REG_TIMSK is enabled further below.
    REG_TCCR0 = TCCR0_CLOCK_8;

    // Timer for measuring motor speed
    REG_TCCR2 = TCCR2_CLOCK_8;

    // Timer and interrupt for radio control PPM input
    #if (F_CPU != 16000000)
    #error "Due to lazy programmer, conversion of timer ticks to RC pulse time only works if CPU freq is 16MHz"
    #endif
    REG_OCR1B = RC_PULSE_MAX;
    #if USE_ICP
    REG_TCCR1B = TCCR0_CLOCK_64 | TCCR1B_ICP_RISING_EDGE | TCCR1B_ICP_NOISE_CANCEL;
    REG_TIMSK = TIMSK_INPUT_CAPTURE | TIMSK_TM0_OVERFLOW | TIMSK_TM2_OVERFLOW;
    #elif USE_INT0
    REG_TCCR1B = TCCR0_CLOCK_64;
    REG_TIMSK = TIMSK_TM0_OVERFLOW | TIMSK_TM2_OVERFLOW;
    REG_MCUCR = MCUCR_INT0_RISING_EDGE;
    REG_GICR = GICR_INT0;
    #endif

    REG_SREG |= SREG_IRQ_ENABLE;

    while(rcPulseTime == 0xffff) {} // Wait for first RC pulse input

    beep(100000);
    waitLong(2000000);

    hallState = (REG_PINB >> 3) & 7;

    // High input on startup triggers calibration mode, with initial position defining max position.
    if (rcPulseTime > RC_PULSE_MAX - ((RC_PULSE_MAX - RC_PULSE_MIN) >> 2))
    {
        // Remain in calibration mode until user gives low input.
        while(rcPulseTime > RC_PULSE_MIN + ((RC_PULSE_MAX - RC_PULSE_MIN) >> 2))
        {
            // As the user physically rotates the servo arm to the desired min position,
            // maxPos is stepped in the opposite direction.
            const u8 lastHallState = hallState;
            hallState = (REG_PINB >> 3) & 7;
            if (hallState == nextHallState[lastHallState])
                maxPos--;
            else if (hallState == prevHallState[lastHallState])
                maxPos++;
        }

        DetectHallPhase(48);

        EEPROMWrite(EEPROM_CUR_POS, (const u8*)&curPos, sizeof(curPos));
        EEPROMWrite(EEPROM_MAX_POS, (const u8*)&maxPos, sizeof(maxPos));
        EEPROMWrite(EEPROM_HALL_STATE_TO_PWM_STATE_POS, hallStateToPWMState, sizeof(hallStateToPWMState));
    }
    else
    {
        EEPROMRead(EEPROM_CUR_POS, (u8*)&curPos, sizeof(curPos));
        EEPROMRead(EEPROM_MAX_POS, (u8*)&maxPos, sizeof(maxPos));
        EEPROMRead(EEPROM_HALL_STATE_TO_PWM_STATE_POS, hallStateToPWMState, sizeof(hallStateToPWMState));
        targetPos = curPos;
    }

    beep(50000);
    waitLong(2000000);
    while(1)
        Update();
}

void Update()
{
    // rcPulseTime and hallTime are copied because they could be changed by interrupts
    const u8 lastHallTime = hallTime;
    const u8 newHallState = (REG_PINB >> 3) & 7;
    const s16 lastRCPulseTime = rcPulseTime;

    // Update RC pulse input
    if (lastRCPulseTime >= RC_PULSE_MIN)
        targetPos = maxPos * ((u8)(lastRCPulseTime - RC_PULSE_MIN)) >> 8;
    else if (lastRCPulseTime == RC_PULSE_SAVE)
        EEPROMWrite(EEPROM_CUR_POS, (const u8*)&curPos, sizeof(curPos));

    // Update hall sensor state
    if (newHallState == 0 || newHallState == 7)
    {
        // Invalid state
    }
    else if (newHallState != hallState)
    {
        // Note: hallTime ticks once every 128 microseconds

        if (curDir == 0 || ABS(curPos - targetPos) < SOFT_STOP_DISTANCE)
            speedRange = 0;
        else if (lastHallTime > (2048 >> 7)) speedRange = 1; // Greater than 2ms per commutation but not starting/stopping
        else if (lastHallTime > (1024 >> 7)) speedRange = 2; // 1ms to 2ms per commutation
        else speedRange = 3; // Less than 1ms per commutation

        if (newHallState == nextHallState[hallState])
            curDir = +1;
        else if (newHallState == prevHallState[hallState])
            curDir = -1;
        else if (newHallState == nextHallState[nextHallState[hallState]])
            curDir = +2; // Missed a commutation, but probably ok
        else if (newHallState == prevHallState[prevHallState[hallState]])
            curDir = -2; // Missed a commutation, but probably ok
        else    // Missed multiple commutations. Give error beep.
            beep(75000);

        hallState = newHallState;
        curPos += curDir;
        hallTime = 0; // Start counting time until next sensor state change
        jitteringToStart = 0;
    }
    else if (hallTime == 0xff)
    {
        // Not moving
        curDir = 0;
        speedRange = 0;
        hallTime = 0;
        if (ABS(curPos - targetPos) > DEADBAND)
            jitteringToStart ^= 1;
    }

    // Update PWM state
    if (ABS(curPos - targetPos) < DEADBAND)// || (curPos < targetPos && curDir < 0) || (curPos > targetPos && curDir > 0))
    {
        // Brake if at target position or moving in the wrong direction
        speedRange = 0;
        pwmState = PWM_BRAKE;
        pwmDelay = MAX_POWER;
        hallTime = 0;
    }
    else
    {
        REG_SREG &= ~SREG_IRQ_ENABLE;
        pwmState = (pwmState & PWM_OFF_DUTY) |
            hallStateToPWMState[(curPos < targetPos) ^ jitteringToStart ?
                nextHallState[hallState] : prevHallState[hallState]];
        pwmDelay = maxPower[speedRange];
        REG_SREG |= SREG_IRQ_ENABLE;
    }
}