#!/usr/bin/env python3

from automata import PTA
from codegen import get_simulated_accountingmethod
import unittest

example_json_1 = {
    'parameters': ['datarate', 'txbytes', 'txpower'],
    'initial_param_values': [None, None, None],
    'state': {
        'IDLE': {
            'power': {
                'static': 5,
            }
        },
        'TX': {
            'power': {
                'static': 100,
                'function': {
                    'raw': 'regression_arg(0) + regression_arg(1)'
                    ' * parameter(txpower)',
                    'regression_args': [100, 2]
                },
            }
        },
    },
    'transitions': [
        {
            'name': 'init',
            'origin': ['UNINITIALIZED', 'IDLE'],
            'destination': 'IDLE',
            'duration': {
                'static': 50000,
            },
            'set_param': {
                'txpower': 10
            },
        },
        {
            'name': 'setTxPower',
            'origin': 'IDLE',
            'destination': 'IDLE',
            'duration': {'static': 120},
            'energy ': {'static': 10000},
            'arg_to_param_map': {0: 'txpower'},
            'argument_values': [[10, 20, 30]],
        },
        {
            'name': 'send',
            'origin': 'IDLE',
            'destination': 'TX',
            'duration': {
                'static': 10,
                'function': {
                    'raw': 'regression_arg(0) + regression_arg(1)'
                    ' * function_arg(1)',
                    'regression_args': [48, 8],
                },
            },
            'energy': {
                'static': 3,
                'function': {
                    'raw': 'regression_arg(0) + regression_arg(1)'
                    ' * function_arg(1)',
                    'regression_args': [3, 5],
                },
            },
            'arg_to_param_map': {1: 'txbytes'},
            'argument_values': [['"foo"', '"hodor"'], [3, 5]],
            'argument_combination': 'zip',
        },
        {
            'name': 'txComplete',
            'origin': 'TX',
            'destination': 'IDLE',
            'is_interrupt': 1,
            'timeout': {
                'static': 2000,
                'function': {
                    'raw': 'regression_arg(0) + regression_arg(1)'
                    ' * parameter(txbytes)',
                    'regression_args': [500, 16],
                },
            },
        }
    ],
}


class TestCG(unittest.TestCase):
    def test_statetransition_immediate(self):
        pta = PTA.from_json(example_json_1)
        pta.set_random_energy_model()
        pta.state['IDLE'].power = 9
        cg = get_simulated_accountingmethod('static_statetransition_immediate')(pta, 1000000, 'uint8_t', 'uint8_t', 'uint8_t', 'uint8_t')
        cg.current_state = pta.state['IDLE']
        cg.sleep(7)
        self.assertEqual(cg.get_energy(), 9 * 7)
        pta.transitions[1].energy = 123
        cg.pass_transition(pta.transitions[1])
        self.assertEqual(cg.get_energy(), 9 * 7 + 123)
        cg.pass_transition(pta.transitions[1])
        self.assertEqual(cg.get_energy(), (9 * 7 + 123 + 123) % 256)

        cg = get_simulated_accountingmethod('static_statetransition_immediate')(pta, 100000, 'uint8_t', 'uint8_t', 'uint8_t', 'uint8_t')
        cg.current_state = pta.state['IDLE']
        cg.sleep(7)
        self.assertEqual(cg.get_energy(), 0)
        cg.sleep(15)
        self.assertEqual(cg.get_energy(), 90)
        cg.sleep(90)
        self.assertEqual(cg.get_energy(), 900 % 256)

        cg = get_simulated_accountingmethod('static_statetransition_immediate')(pta, 100000, 'uint8_t', 'uint8_t', 'uint8_t', 'uint16_t')
        cg.current_state = pta.state['IDLE']
        cg.sleep(7)
        self.assertEqual(cg.get_energy(), 0)
        cg.sleep(15)
        self.assertEqual(cg.get_energy(), 90)
        cg.sleep(90)
        self.assertEqual(cg.get_energy(), 900)

        pta.state['IDLE'].power = 9  # -> 90 uW
        pta.transitions[1].energy = 1  # -> 100 pJ
        cg = get_simulated_accountingmethod('static_statetransition_immediate')(pta, 1000000, 'uint8_t', 'uint8_t', 'uint8_t', 'uint8_t', 1e-5, 1e-5, 1e-10)
        cg.current_state = pta.state['IDLE']
        cg.sleep(10)  # 10 us
        self.assertEqual(cg.get_energy(), 90 * 10)
        cg.pass_transition(pta.transitions[1])
        self.assertAlmostEqual(cg.get_energy(), 90 * 10 + 100, places=0)
        cg.pass_transition(pta.transitions[1])
        self.assertAlmostEqual(cg.get_energy(), 90 * 10 + 100 + 100, places=0)

    def test_statetransition(self):
        pta = PTA.from_json(example_json_1)
        pta.set_random_energy_model()
        pta.state['IDLE'].power = 9
        cg = get_simulated_accountingmethod('static_statetransition')(pta, 1000000, 'uint8_t', 'uint8_t', 'uint8_t', 'uint8_t')
        cg.current_state = pta.state['IDLE']
        cg.sleep(7)
        self.assertEqual(cg.get_energy(), 9 * 7)
        pta.transitions[1].energy = 123
        cg.pass_transition(pta.transitions[1])
        self.assertEqual(cg.get_energy(), 9 * 7 + 123)
        cg.pass_transition(pta.transitions[1])
        self.assertEqual(cg.get_energy(), (9 * 7 + 123 + 123) % 256)

    def test_state_immediate(self):
        pta = PTA.from_json(example_json_1)
        pta.set_random_energy_model()
        pta.state['IDLE'].power = 9
        cg = get_simulated_accountingmethod('static_state_immediate')(pta, 1000000, 'uint8_t', 'uint8_t', 'uint8_t', 'uint8_t')
        cg.current_state = pta.state['IDLE']
        cg.sleep(7)
        self.assertEqual(cg.get_energy(), 9 * 7)
        pta.transitions[1].energy = 123
        cg.pass_transition(pta.transitions[1])
        self.assertEqual(cg.get_energy(), 9 * 7)
        cg.pass_transition(pta.transitions[1])
        self.assertEqual(cg.get_energy(), 9 * 7)

    def test_state(self):
        pta = PTA.from_json(example_json_1)
        pta.set_random_energy_model()
        pta.state['IDLE'].power = 9
        cg = get_simulated_accountingmethod('static_state')(pta, 1000000, 'uint8_t', 'uint8_t', 'uint8_t', 'uint8_t')
        cg.current_state = pta.state['IDLE']
        cg.sleep(7)
        self.assertEqual(cg.get_energy(), 9 * 7)
        pta.transitions[1].energy = 123
        cg.pass_transition(pta.transitions[1])
        self.assertEqual(cg.get_energy(), 9 * 7)
        cg.pass_transition(pta.transitions[1])
        self.assertEqual(cg.get_energy(), 9 * 7)

        cg = get_simulated_accountingmethod('static_state')(pta, 1000000, 'uint8_t', 'uint16_t', 'uint16_t', 'uint16_t')


if __name__ == '__main__':
    unittest.main()