This is an active area of research in Morphologically Rich Languages (MRLs), since this problem also appears in other semitic languages like Hebrew, as well as Turkish. There's a nice body of work to learn from, both with and without neural nets. For example, this paper from 2017 (http://aclweb.org/anthology/D17-1073) uses a neural model for morphological disambiguation. You can see a nice comparison of tools in the recent 2018 Universal Dependencies Shared Task results: http://universaldependencies.org/conll18/results-lemmas.html (look for ar_padt).
If you're looking for training data, the Arabic treebanks in http://universaldependencies.org could help. I think some of them contain surface tokens with lemmas. I'm quite sure they also have roots.
Also, you might want to take a look at the SIGMORPHON CONLL shared task (2017 https://sites.google.com/view/conll-sigmorphon2017/ and 2018 https://sigmorphon.github.io/sharedtasks/2018/) on morphological reinflection, which IIRC is a similar task - taking an inflected form and reinflecting it with other morphological properties. They also have a nice data set to train on.