/*
  www.aifes.ai
  Copyright (C) 2020-2021 Fraunhofer Institute for Microelectronic Circuits and Systems. All rights reserved.

  AIfES is free software: you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation, either version 3 of the License, or
  (at your option) any later version.
  
  This program is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.
  
  You should have received a copy of the GNU General Public License
  along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/

#include <limits.h>

aimodel_t model;

// Layer definition
uint16_t input_layer_shape[] = {1, 9};
ailayer_input_q7_t   input_layer    = AILAYER_INPUT_Q7_A(2, input_layer_shape);
ailayer_dense_q7_t   dense_layer_1  = AILAYER_DENSE_Q7_A(30);
ailayer_relu_q7_t    relu_layer_1   = AILAYER_RELU_Q7_A();
ailayer_dense_q7_t   dense_layer_2  = AILAYER_DENSE_Q7_A(20);
ailayer_relu_q7_t    relu_layer_2   = AILAYER_RELU_Q7_A();
ailayer_dense_q7_t   dense_layer_3  = AILAYER_DENSE_Q7_A(9);
ailayer_softmax_q7_t softmax_layer  = AILAYER_SOFTMAX_Q7_A();


#if INT_MAX == 32767
  
// Model parameters (Weights, Biases and quantization parameters) with 2 byte alignment (e.g. for Arduino UNO)
const uint32_t parameter_memory_size = 1332;
const uint16_t model_parameters[ 666 ] PROGMEM = {
    0x0006, 0x0000, 0x0003, 0x00FA, 0x0002, 0x00EC, 0x0001, 0x00F4, 0x0005, 0x0000, 0xD708, 0x3025,
    0x31F8, 0xE0EB, 0x1DD2, 0x1610, 0x0CBD, 0x0C23, 0x20FD, 0x01D9, 0xFC32, 0x11EC, 0x0CDB, 0x0A1C,
    0x1AC4, 0x02F2, 0x23EC, 0x1727, 0xE503, 0x26EA, 0xF337, 0xE902, 0xD1E6, 0x01F6, 0xDEFD, 0x3417,
    0xD609, 0xFC0A, 0x1019, 0x0B1C, 0x072C, 0x12C5, 0x142F, 0x0A10, 0xBE02, 0x071A, 0x12DC, 0x0603,
    0x220E, 0xFF15, 0x140C, 0xDFD6, 0x2E08, 0xFA05, 0x1AD4, 0xE80A, 0x37D3, 0x0BE9, 0x0DED, 0x1C14,
    0xD9E8, 0xFE1C, 0x0D2E, 0x02E1, 0xCCEA, 0xE413, 0x1526, 0xFEEF, 0xF81F, 0xC91E, 0xF8DF, 0x0AE4,
    0xFF19, 0xD631, 0xFEED, 0xFAF1, 0xDAEC, 0xD6E7, 0x05D6, 0xDF14, 0x1C22, 0x2BED, 0xDADC, 0xF7E5,
    0xEF03, 0x201B, 0x382D, 0x01EF, 0xC9FB, 0xF6F1, 0x0FEC, 0x1F27, 0xE126, 0xD7E8, 0xECDD, 0x2EDB,
    0x0706, 0xC817, 0x0BF4, 0xFE02, 0x04E5, 0x1ACA, 0xF025, 0x1E36, 0x03EF, 0xE635, 0xE419, 0xF4F5,
    0xDED7, 0x21F6, 0x0E0A, 0xBB13, 0x0B10, 0xFDEE, 0x37E9, 0xE8EA, 0xFCEF, 0x00F3, 0x36D6, 0xEB09,
    0x1EE9, 0x0FE4, 0xD409, 0xF70F, 0x301B, 0xF518, 0xE41F, 0x0CFB, 0x020E, 0x0233, 0xE5D5, 0x10F2,
    0x040F, 0xCD16, 0x2210, 0x1313, 0x3AF5, 0xC8F3, 0xE303, 0x0CF4, 0xBA19, 0x0412, 0x2217, 0xFD0C,
    0xFF0C, 0x000B, 0x0000, 0x0000, 0x0550, 0x0000, 0x03BB, 0x0000, 0xFF8D, 0xFFFF, 0x044C, 0x0000,
    0x03AC, 0x0000, 0x08C5, 0x0000, 0x03FF, 0x0000, 0x0476, 0x0000, 0x037F, 0x0000, 0x039C, 0x0000,
    0x00A6, 0x0000, 0x0093, 0x0000, 0x03D5, 0x0000, 0x0932, 0x0000, 0xFC9E, 0xFFFF, 0x04C8, 0x0000,
    0x0400, 0x0000, 0xFFCC, 0xFFFF, 0x0AA7, 0x0000, 0x029A, 0x0000, 0x01E9, 0x0000, 0x07A6, 0x0000,
    0x0538, 0x0000, 0xFBAA, 0xFFFF, 0x058C, 0x0000, 0xFF17, 0xFFFF, 0xFED9, 0xFFFF, 0x03E1, 0x0000,
    0x03E1, 0x0000, 0x0478, 0x0000, 0x0006, 0x0000, 0xE6DF, 0x4336, 0x3EE6, 0x231B, 0x04F0, 0x0DBF,
    0x0723, 0xD8F1, 0xFEE7, 0x333A, 0x1237, 0x43F4, 0xABF1, 0xE9CB, 0xF046, 0xACFE, 0x22F7, 0xE7C5,
    0xE029, 0x4F0E, 0x550B, 0x4742, 0xDC25, 0x0504, 0xDAF2, 0xEC9C, 0x2559, 0x163D, 0xE8CA, 0x1A31,
    0x00F6, 0xFAC0, 0x37B9, 0x1F2A, 0x500C, 0x1DD0, 0xD8FD, 0x2EC7, 0x3818, 0x05E6, 0xF754, 0xE102,
    0x5130, 0x011D, 0xD3AF, 0x28CD, 0xD403, 0xEE36, 0xD744, 0x34D6, 0x034D, 0xF004, 0x17B2, 0x1409,
    0x4738, 0x2D0A, 0xE3F1, 0x2836, 0x3406, 0xD426, 0xAB2F, 0xE41E, 0x3857, 0x48DE, 0x3D08, 0xF227,
    0xCCF6, 0x3517, 0xFAC9, 0xC4EA, 0x253A, 0xEDEB, 0x12FD, 0x01CE, 0x41D9, 0xBCB6, 0xEB0F, 0xBA22,
    0xDF51, 0xE614, 0x3124, 0x2A6D, 0x0AEE, 0xDB18, 0x4313, 0xEDDF, 0x18F3, 0x3A0A, 0x081C, 0xFF65,
    0x28FE, 0x29EB, 0x093D, 0x5199, 0x2EE1, 0xF3E4, 0x1806, 0xC6E5, 0xCE09, 0xDBE0, 0x2814, 0x2835,
    0xD307, 0xFA21, 0xD3F4, 0x42BD, 0xE9F8, 0xE528, 0x17F6, 0xEC13, 0xF302, 0x1F58, 0x2CEC, 0xE5DF,
    0xE150, 0xFD74, 0xA51F, 0xFFA8, 0xD7C3, 0xEDA2, 0x51D9, 0xE105, 0x13DD, 0x6D1C, 0x1F34, 0xD6E5,
    0x36DD, 0x3D43, 0x473C, 0xD037, 0xCE2B, 0xDF14, 0xFEDD, 0x06B6, 0x4DFF, 0x08E9, 0x8713, 0xF1E5,
    0x07F4, 0x0D13, 0x40E8, 0xFAC9, 0xCA85, 0xDAE1, 0x2416, 0x3F16, 0xE33E, 0xDB50, 0xD1ED, 0xFF38,
    0x0E49, 0x26D1, 0xF4FF, 0xD109, 0xDCCD, 0x3121, 0xEF11, 0x4BD7, 0xF3E5, 0xE010, 0x5510, 0x604A,
    0x1E23, 0xAFF0, 0xF1D1, 0x042F, 0xEB62, 0x4509, 0x54D5, 0x4806, 0xEDC2, 0xD60A, 0xF5FA, 0xF7EC,
    0xF244, 0x1F18, 0x3E18, 0xCA07, 0x7119, 0xB70D, 0x1707, 0xDF04, 0x3405, 0xF8D9, 0xF623, 0xDE10,
    0xFA1F, 0xEF4E, 0x6DB2, 0xF742, 0x0BD5, 0xC018, 0xCC0D, 0xB7EA, 0x2EE1, 0x3226, 0x17E4, 0x41F4,
    0x13C8, 0x11A9, 0x1A21, 0x1D26, 0xE5DB, 0xC750, 0x2308, 0xFA08, 0xC215, 0xDD54, 0x1907, 0xCBCE,
    0xFDC0, 0x2EC5, 0xFFE0, 0x0CEE, 0x36E1, 0x2928, 0x76CC, 0x3334, 0x02DF, 0x2681, 0x3931, 0xEDE2,
    0x4DCD, 0x2CFB, 0x120A, 0x0907, 0xD9EF, 0xC425, 0x080E, 0x08FD, 0x20D6, 0x0F0D, 0xBB20, 0xE215,
    0xD8F1, 0x3F14, 0xC9F7, 0x19E9, 0x1B70, 0xEB36, 0x2CEC, 0x1521, 0x310F, 0x96E0, 0xD141, 0xF810,
    0x4209, 0xF64D, 0xFFF8, 0x0916, 0x0C00, 0xF446, 0xDB4D, 0x5825, 0x63E0, 0xA84C, 0xCBD9, 0x06D5,
    0x2F29, 0xDF31, 0x2A09, 0xE7F8, 0x2C29, 0x4A5C, 0x3B3F, 0xFFF1, 0xFA3D, 0x24EA, 0x2916, 0xFD05,
    0x0C67, 0x5F10, 0xCB2D, 0x0E62, 0xE5FA, 0xFE17, 0xD235, 0x55D3, 0x370D, 0x04CE, 0xB132, 0x08E3,
    0xCEE2, 0xF649, 0xCB19, 0x5A06, 0xCAFC, 0x46CA, 0xF505, 0xE4F7, 0x461E, 0x2AFF, 0xBEAC, 0xD7EF,
    0x4019, 0xBAF1, 0x27E9, 0x0D1F, 0xDC07, 0xF85F, 0x0009, 0x0000, 0x0000, 0x0236, 0x0000, 0x018A,
    0x0000, 0x00B3, 0x0000, 0xFF26, 0xFFFF, 0x01AA, 0x0000, 0xFFB4, 0xFFFF, 0xFEED, 0xFFFF, 0x010C,
    0x0000, 0x000F, 0x0000, 0x01EF, 0x0000, 0x019A, 0x0000, 0xFFAB, 0xFFFF, 0x012D, 0x0000, 0xFF3D,
    0xFFFF, 0x0185, 0x0000, 0xFFBF, 0xFFFF, 0xFF7B, 0xFFFF, 0xFFD4, 0xFFFF, 0xFF24, 0xFFFF, 0x0076,
    0x0000, 0x0005, 0x0000, 0xEB05, 0x2B30, 0x28E9, 0x2A1E, 0x0ECF, 0x2915, 0x2EB8, 0xE605, 0x13C9,
    0x3D06, 0x1ADA, 0xF605, 0x270B, 0x28F2, 0xB820, 0x3700, 0xDBFC, 0xE515, 0x10FA, 0x26BD, 0x05EB,
    0x2C0D, 0xE907, 0x32FE, 0x332D, 0xDEC9, 0x0C11, 0xF4BB, 0x0E19, 0xDF30, 0x392E, 0xD53B, 0x2A05,
    0xE4CB, 0x1EDD, 0x1BCC, 0xD4FF, 0x0429, 0x16FA, 0xCFFB, 0x2C26, 0xEC05, 0xD128, 0xCD2D, 0x283E,
    0xDC30, 0xDD31, 0xC91D, 0xD4CB, 0x1C01, 0xBC3E, 0x1CC6, 0x1D33, 0xFBD8, 0xFD08, 0xFAD9, 0xE9F7,
    0x121B, 0xCC1B, 0xFA2C, 0x23DB, 0x27D9, 0x1AC1, 0xD1EA, 0xFDD4, 0xCC2C, 0x1934, 0x2200, 0xEA38,
    0x0024, 0xDBCE, 0xC309, 0x0CC0, 0x09FE, 0x1623, 0x0DFF, 0xDDE7, 0x3A1E, 0x3408, 0x1D28, 0xE4F9,
    0xC8D3, 0xBCF4, 0x222A, 0x1310, 0x3BFC, 0x2E0E, 0x1C29, 0xD028, 0xBE03, 0x0007, 0x0000, 0x0000,
    0x0000, 0x0000, 0xFFE7, 0xFFFF, 0xFFEC, 0xFFFF, 0x000F, 0x0000, 0x002C, 0x0000, 0xFFED, 0xFFFF,
    0x000B, 0x0000, 0x0009, 0x0000, 0xFFF7, 0xFFFF
};

#else

// Model parameters (Weights, Biases and quantization parameters) with 4 byte alignment (e.g. for Arduino Nano 33)
const uint16_t parameter_memory_size = 1340;
const uint32_t model_parameters[ 335 ] PROGMEM = {
    0x00000006, 0x00FA0003, 0x00EC0002, 0x00F40001, 0x00000005, 0x3025D708, 0xE0EB31F8, 0x16101DD2,
    0x0C230CBD, 0x01D920FD, 0x11ECFC32, 0x0A1C0CDB, 0x02F21AC4, 0x172723EC, 0x26EAE503, 0xE902F337,
    0x01F6D1E6, 0x3417DEFD, 0xFC0AD609, 0x0B1C1019, 0x12C5072C, 0x0A10142F, 0x071ABE02, 0x060312DC,
    0xFF15220E, 0xDFD6140C, 0xFA052E08, 0xE80A1AD4, 0x0BE937D3, 0x1C140DED, 0xFE1CD9E8, 0x02E10D2E,
    0xE413CCEA, 0xFEEF1526, 0xC91EF81F, 0x0AE4F8DF, 0xD631FF19, 0xFAF1FEED, 0xD6E7DAEC, 0xDF1405D6,
    0x2BED1C22, 0xF7E5DADC, 0x201BEF03, 0x01EF382D, 0xF6F1C9FB, 0x1F270FEC, 0xD7E8E126, 0x2EDBECDD,
    0xC8170706, 0xFE020BF4, 0x1ACA04E5, 0x1E36F025, 0xE63503EF, 0xF4F5E419, 0x21F6DED7, 0xBB130E0A,
    0xFDEE0B10, 0xE8EA37E9, 0x00F3FCEF, 0xEB0936D6, 0x0FE41EE9, 0xF70FD409, 0xF518301B, 0x0CFBE41F,
    0x0233020E, 0x10F2E5D5, 0xCD16040F, 0x13132210, 0xC8F33AF5, 0x0CF4E303, 0x0412BA19, 0xFD0C2217,
    0x0000FF0C, 0x0000000B, 0x00000000, 0x00000550, 0x000003BB, 0xFFFFFF8D, 0x0000044C, 0x000003AC,
    0x000008C5, 0x000003FF, 0x00000476, 0x0000037F, 0x0000039C, 0x000000A6, 0x00000093, 0x000003D5,
    0x00000932, 0xFFFFFC9E, 0x000004C8, 0x00000400, 0xFFFFFFCC, 0x00000AA7, 0x0000029A, 0x000001E9,
    0x000007A6, 0x00000538, 0xFFFFFBAA, 0x0000058C, 0xFFFFFF17, 0xFFFFFED9, 0x000003E1, 0x000003E1,
    0x00000478, 0x00000006, 0x4336E6DF, 0x231B3EE6, 0x0DBF04F0, 0xD8F10723, 0x333AFEE7, 0x43F41237,
    0xE9CBABF1, 0xACFEF046, 0xE7C522F7, 0x4F0EE029, 0x4742550B, 0x0504DC25, 0xEC9CDAF2, 0x163D2559,
    0x1A31E8CA, 0xFAC000F6, 0x1F2A37B9, 0x1DD0500C, 0x2EC7D8FD, 0x05E63818, 0xE102F754, 0x011D5130,
    0x28CDD3AF, 0xEE36D403, 0x34D6D744, 0xF004034D, 0x140917B2, 0x2D0A4738, 0x2836E3F1, 0xD4263406,
    0xE41EAB2F, 0x48DE3857, 0xF2273D08, 0x3517CCF6, 0xC4EAFAC9, 0xEDEB253A, 0x01CE12FD, 0xBCB641D9,
    0xBA22EB0F, 0xE614DF51, 0x2A6D3124, 0xDB180AEE, 0xEDDF4313, 0x3A0A18F3, 0xFF65081C, 0x29EB28FE,
    0x5199093D, 0xF3E42EE1, 0xC6E51806, 0xDBE0CE09, 0x28352814, 0xFA21D307, 0x42BDD3F4, 0xE528E9F8,
    0xEC1317F6, 0x1F58F302, 0xE5DF2CEC, 0xFD74E150, 0xFFA8A51F, 0xEDA2D7C3, 0xE10551D9, 0x6D1C13DD,
    0xD6E51F34, 0x3D4336DD, 0xD037473C, 0xDF14CE2B, 0x06B6FEDD, 0x08E94DFF, 0xF1E58713, 0x0D1307F4,
    0xFAC940E8, 0xDAE1CA85, 0x3F162416, 0xDB50E33E, 0xFF38D1ED, 0x26D10E49, 0xD109F4FF, 0x3121DCCD,
    0x4BD7EF11, 0xE010F3E5, 0x604A5510, 0xAFF01E23, 0x042FF1D1, 0x4509EB62, 0x480654D5, 0xD60AEDC2,
    0xF7ECF5FA, 0x1F18F244, 0xCA073E18, 0xB70D7119, 0xDF041707, 0xF8D93405, 0xDE10F623, 0xEF4EFA1F,
    0xF7426DB2, 0xC0180BD5, 0xB7EACC0D, 0x32262EE1, 0x41F417E4, 0x11A913C8, 0x1D261A21, 0xC750E5DB,
    0xFA082308, 0xDD54C215, 0xCBCE1907, 0x2EC5FDC0, 0x0CEEFFE0, 0x292836E1, 0x333476CC, 0x268102DF,
    0xEDE23931, 0x2CFB4DCD, 0x0907120A, 0xC425D9EF, 0x08FD080E, 0x0F0D20D6, 0xE215BB20, 0x3F14D8F1,
    0x19E9C9F7, 0xEB361B70, 0x15212CEC, 0x96E0310F, 0xF810D141, 0xF64D4209, 0x0916FFF8, 0xF4460C00,
    0x5825DB4D, 0xA84C63E0, 0x06D5CBD9, 0xDF312F29, 0xE7F82A09, 0x4A5C2C29, 0xFFF13B3F, 0x24EAFA3D,
    0xFD052916, 0x5F100C67, 0x0E62CB2D, 0xFE17E5FA, 0x55D3D235, 0x04CE370D, 0x08E3B132, 0xF649CEE2,
    0x5A06CB19, 0x46CACAFC, 0xE4F7F505, 0x2AFF461E, 0xD7EFBEAC, 0xBAF14019, 0x0D1F27E9, 0xF85FDC07,
    0x00000009, 0x00000000, 0x00000236, 0x0000018A, 0x000000B3, 0xFFFFFF26, 0x000001AA, 0xFFFFFFB4,
    0xFFFFFEED, 0x0000010C, 0x0000000F, 0x000001EF, 0x0000019A, 0xFFFFFFAB, 0x0000012D, 0xFFFFFF3D,
    0x00000185, 0xFFFFFFBF, 0xFFFFFF7B, 0xFFFFFFD4, 0xFFFFFF24, 0x00000076, 0x00000005, 0x2B30EB05,
    0x2A1E28E9, 0x29150ECF, 0xE6052EB8, 0x3D0613C9, 0xF6051ADA, 0x28F2270B, 0x3700B820, 0xE515DBFC,
    0x26BD10FA, 0x2C0D05EB, 0x32FEE907, 0xDEC9332D, 0xF4BB0C11, 0xDF300E19, 0xD53B392E, 0xE4CB2A05,
    0x1BCC1EDD, 0x0429D4FF, 0xCFFB16FA, 0xEC052C26, 0xCD2DD128, 0xDC30283E, 0xC91DDD31, 0x1C01D4CB,
    0x1CC6BC3E, 0xFBD81D33, 0xFAD9FD08, 0x121BE9F7, 0xFA2CCC1B, 0x27D923DB, 0xD1EA1AC1, 0xCC2CFDD4,
    0x22001934, 0x0024EA38, 0xC309DBCE, 0x09FE0CC0, 0x0DFF1623, 0x3A1EDDE7, 0x1D283408, 0xC8D3E4F9,
    0x222ABCF4, 0x3BFC1310, 0x1C292E0E, 0xBE03D028, 0x00000007, 0x00000000, 0x00000000, 0xFFFFFFE7,
    0xFFFFFFEC, 0x0000000F, 0x0000002C, 0xFFFFFFED, 0x0000000B, 0x00000009, 0xFFFFFFF7
};

#endif 


void init_ai_agent()
{
  // Layer pointer to perform the connection
  ailayer_t *x;

  #if __AVR__
  // Only on AVR contollers (e.g. Arduino UNO) with weights in PROGMEM

  // Connect the layers to build a model
  model.input_layer = ailayer_input_q7_default(&input_layer);
  x = ailayer_dense_wt_q7_avr_pgm(&dense_layer_1, model.input_layer);   // "wt" implementation, because the weights matrix is transposed
  x = ailayer_relu_q7_avr_pgm(&relu_layer_1, x);
  x = ailayer_dense_wt_q7_avr_pgm(&dense_layer_2, x);
  x = ailayer_relu_q7_avr_pgm(&relu_layer_2, x);
  x = ailayer_dense_wt_q7_avr_pgm(&dense_layer_3, x);
  x = ailayer_softmax_q7_avr_pgm(&softmax_layer, x);
  model.output_layer = x;
  
  #elif __arm__ && defined USE_CMSIS_ACCELERATION_ON_ARM
  // Only on ARM contollers (e.g. Arduino Nano 33 BLE). AIFES_WITH_CMSIS must be activated in aifes_config.h and CMSIS-NN library must be available in the AIfES_for_Arduino library.
  
  // Connect the layers to build a model
  model.input_layer = ailayer_input_q7_default(&input_layer);
  x = ailayer_dense_wt_q7_cmsis(&dense_layer_1, model.input_layer);   // "wt" implementation, because the weights matrix is transposed
  x = ailayer_relu_q7_default(&relu_layer_1, x);
  x = ailayer_dense_wt_q7_cmsis(&dense_layer_2, x);
  x = ailayer_relu_q7_default(&relu_layer_2, x);
  x = ailayer_dense_wt_q7_cmsis(&dense_layer_3, x);
  x = ailayer_softmax_q7_default(&softmax_layer, x);
  model.output_layer = x;

  #else

  // Connect the layers to build a model
  model.input_layer = ailayer_input_q7_default(&input_layer);
  x = ailayer_dense_wt_q7_default(&dense_layer_1, model.input_layer);   // "wt" implementation, because the weights matrix is transposed
  x = ailayer_relu_q7_default(&relu_layer_1, x);
  x = ailayer_dense_wt_q7_default(&dense_layer_2, x);
  x = ailayer_relu_q7_default(&relu_layer_2, x);
  x = ailayer_dense_wt_q7_default(&dense_layer_3, x);
  x = ailayer_softmax_q7_default(&softmax_layer, x);
  model.output_layer = x;

  #endif // __AVR__

  // Finish the model creation by checking the connections and setting some parameters for further processing
  aialgo_compile_model(&model);

  // Distribute the memory to the trainable parameters of the model
  aialgo_distribute_parameter_memory(&model, (void *) model_parameters, parameter_memory_size);

  Serial.print(F("\n-------------- AI agent model structure ---------------\n"));
  aialgo_print_model_structure(&model);
  Serial.print(F("--------------------------------------------------------\n\n"));
  
  Serial.print("The weights need "); Serial.print(parameter_memory_size); Serial.println(" bytes of memory.");

  // Allocate memory for intermediate results of the inference
  uint32_t inference_memory_size = aialgo_sizeof_inference_memory(&model);
  Serial.print("The model needs "); Serial.print(inference_memory_size); Serial.println(" bytes of memory for inference.\n");
  void *inference_memory = malloc(inference_memory_size);

  // Schedule the memory to the model
  aialgo_schedule_inference_memory(&model, inference_memory, inference_memory_size);
}

uint8_t run_ai_agent(int8_t *board)
{
  uint16_t board_shape[2] = {1, 9};
  aimath_q7_params_t input_q_params;
  input_q_params.shift = pgm_read_word(&(((aimath_q7_params_t *) model.input_layer->result.tensor_params)->shift));
  input_q_params.zero_point = pgm_read_word(&(((aimath_q7_params_t *) model.input_layer->result.tensor_params)->zero_point));
  aitensor_t board_tensor = AITENSOR_2D_Q7(board_shape, &input_q_params, board); // Has the same quantization parameters as the input layer of the nn

  aitensor_t *out_tensor = aialgo_forward_model(&model, &board_tensor);
  int8_t *out_data = (int8_t *) out_tensor->data;

  // Calculate the most probable index (argmax)
  int8_t maxi = out_data[0];
  uint8_t argmax = 0;
  for(uint8_t j = 1; j < 9; j++){
    if(out_data[j] > maxi){
        maxi = out_data[j];
        argmax = j;
    }
  }

  return argmax;
}
