"""
Convert CPU cycle count to energy.

Contains classes for some embedded CPUs/MCUs. Given a configuration, each
class can convert a cycle count to an energy consumption.
"""

def get_class(cpu_name):
    """Return model class for cpu_name."""
    if cpu_name == 'MSP430':
        return MSP430
    if cpu_name == 'ATMega168':
        return ATMega168
    if cpu_name == 'ATMega328':
        return ATMega328
    if cpu_name == 'ATTiny88':
        return ATTiny88
    if cpu_name == 'esp8266':
        return ESP8266

def _param_list_to_dict(device, param_list):
    param_dict = dict()
    for i, parameter in enumerate(sorted(device.parameters.keys())):
        param_dict[parameter] = param_list[i]
    return param_dict

class MSP430:
    name = 'MSP430'
    parameters = {
        'cpu_freq': [1e6, 4e6, 8e6, 12e6, 16e6],
        'memory' : ['unified', 'fram0', 'fram50', 'fram66', 'fram75', 'fram100', 'ram'],
        'voltage': [2.2, 3.0],
    }
    default_params = {
        'cpu_freq': 4e6,
        'memory' : 'unified',
        'voltage': 3
    }

    current_by_mem = {
        'unified' : [210,  640, 1220, 1475, 1845],
        'fram0'   : [370, 1280, 2510, 2080, 2650],
        'fram50'  : [240,  745, 1440, 1575, 1990],
        'fram66'  : [200,  560, 1070, 1300, 1620],
        'fram75'  : [170,  480,  890, 1155, 1420],
        'fram100' : [110,  235,  420,  640,  730],
        'ram'     : [130,  320,  585,  890, 1070],
    }

    def get_current(params):
        if type(params) != dict:
            return MSP430.get_current(_param_list_to_dict(MSP430, params))
        cpu_freq_index = MSP430.parameters['cpu_freq'].index(params['cpu_freq'])

        return MSP430.current_by_mem[params['memory']][cpu_freq_index] * 1e-6

    def get_power(params):
        if type(params) != dict:
            return MSP430.get_energy(_param_list_to_dict(MSP430, params))

        return MSP430.get_current(params) * params['voltage']

    def get_energy_per_cycle(params):
        if type(params) != dict:
            return MSP430.get_energy_per_cycle(_param_list_to_dict(MSP430, params))

        return MSP430.get_power(params) / params['cpu_freq']

class ATMega168:
    name = 'ATMega168'
    parameters = {
        'cpu_freq': [1e6, 4e6, 8e6],
        'voltage': [2, 3, 5]
    }
    default_params = {
        'cpu_freq': 4e6,
        'voltage': 3
    }

    def get_current(params):
        if type(params) != dict:
            return ATMega168.get_current(_param_list_to_dict(ATMega168, params))
        if params['cpu_freq'] == 1e6 and params['voltage'] <= 2:
            return 0.5e-3
        if params['cpu_freq'] == 4e6 and params['voltage'] <= 3:
            return 3.5e-3
        if params['cpu_freq'] == 8e6 and params['voltage'] <= 5:
            return 12e-3
        return None

    def get_power(params):
        if type(params) != dict:
            return ATMega168.get_energy(_param_list_to_dict(ATMega168, params))

        return ATMega168.get_current(params) * params['voltage']

    def get_energy_per_cycle(params):
        if type(params) != dict:
            return ATMega168.get_energy_per_cycle(_param_list_to_dict(ATMega168, params))

        return ATMega168.get_power(params) / params['cpu_freq']

class ATMega328:
    name = 'ATMega328'
    parameters = {
        'cpu_freq': [1e6, 4e6, 8e6],
        'voltage': [2, 3, 5]
    }
    default_params = {
        'cpu_freq': 4e6,
        'voltage': 3
    }

    # Source: ATMega328P Datasheet p.316 / table 28.2.4
    def get_current(params):
        if type(params) != dict:
            return ATMega328.get_current(_param_list_to_dict(ATMega328, params))
        # specified for 2V
        if params['cpu_freq'] == 1e6 and params['voltage'] >= 1.8:
            return 0.3e-3
        # specified for 3V
        if params['cpu_freq'] == 4e6 and params['voltage'] >= 1.8:
            return 1.7e-3
        # specified for 5V
        if params['cpu_freq'] == 8e6 and params['voltage'] >= 2.5:
            return 5.2e-3
        return None

    def get_power(params):
        if type(params) != dict:
            return ATMega328.get_energy(_param_list_to_dict(ATMega328, params))

        return ATMega328.get_current(params) * params['voltage']

    def get_energy_per_cycle(params):
        if type(params) != dict:
            return ATMega328.get_energy_per_cycle(_param_list_to_dict(ATMega328, params))

        return ATMega328.get_power(params) / params['cpu_freq']

class ESP8266:
    # Source: ESP8266EX Datasheet, table 5-2 (v2017.11) / table 3-4 (v2018.11)
    # Taken at 3.0V
    name = 'ESP8266'
    parameters = {
        'cpu_freq': [80e6],
        'voltage': [2.5, 3.0, 3.3, 3.6] # min / ... / typ / max
    }
    default_params = {
        'cpu_freq': 80e6,
        'voltage': 3.0
    }

    def get_current(params):
        if type(params) != dict:
            return ESP8266.get_current(_param_list_to_dict(ESP8266, params))

        return 15e-3

    def get_power(params):
        if type(params) != dict:
            return ESP8266.get_power(_param_list_to_dict(ESP8266, params))

        return ESP8266.get_current(params) * params['voltage']

    def get_energy_per_cycle(params):
        if type(params) != dict:
            return ESP8266.get_energy_per_cycle(_param_list_to_dict(ESP8266, params))

        return ESP8266.get_power(params) / params['cpu_freq']

class ATTiny88:
    name = 'ATTiny88'
    parameters = {
        'cpu_freq': [1e6, 4e6, 8e6],
        'voltage': [2, 3, 5]
    }
    default_params = {
        'cpu_freq' : 4e6,
        'voltage' : 3
    }

    def get_current(params):
        if type(params) != dict:
            return ATTiny88.get_current(_param_list_to_dict(ATTiny88, params))
        if params['cpu_freq'] == 1e6 and params['voltage'] <= 2:
            return 0.2e-3
        if params['cpu_freq'] == 4e6 and params['voltage'] <= 3:
            return 1.4e-3
        if params['cpu_freq'] == 8e6 and params['voltage'] <= 5:
            return 4.5e-3
        return None