#!/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' : { 'txpower' : 0 },
            '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' : { 'txbytes' : 1 },
            '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()