import json
import sys
from string import Template
import os.path
import argparse
import re
from copy import copy
from functools import reduce
# import set in a python 2->3 safe way
try:
    set
except NameError:
    from sets import Set as set
import warnings


def find(key, dictionary):
    """ Recursively search a dictionary for a key """
    for k, v in dictionary.items():
        if k == key:
            yield v
        if isinstance(v, dict):
            for result in find(key, v):
                yield result
        elif isinstance(v, list):
            for d in v:
                for result in find(key, d):
                    yield result


def getDependencies(program):
    """ Traverses the program AST and extracts all dependencies """
    allClasses = set()
    interfaceClasses = set()
    enums = set()
    globalFunctions = set()

    # All classes used
    for objectType in find("ObjectType", program):
        allClasses.add(objectType)

    for shogunSGType in find("ShogunSGType", program):
        allClasses.add(shogunSGType)

    # All classes where the constructor is called
    for initialisation in find("Init", program):
        # Constructor interface is used if an argument list is passed
        if list(initialisation[2].keys())[0] == "ArgumentList":
            typeDict = initialisation[0]
            objectKey = list(typeDict.keys())[0]
            interfaceClasses.add(typeDict[objectKey])

    # All classes on which a static method is called
    for staticCall in find("StaticCall", program):
        typeDict = staticCall[0]
        objectKey = list(typeDict.keys())[0]
        interfaceClasses.add(typeDict[objectKey])

    # All global function calls
    for globalCall in find("GlobalCall", program):
        identifier = list(globalCall[0].values())[0]
        globalFunctions.add(identifier)

    # All enums used
    for enum in find("Enum", program):
        enumType = enum[0]["Identifier"]
        enumValue = enum[1]["Identifier"]
        enums.add((enumType, enumValue))

    return allClasses, interfaceClasses, enums, globalFunctions


def getBasicTypesToStore():
    """ Returns all basic types which will be serialized """
    return ("real", "float", "int")

def getSGTypesToStore():
    """ Returns all SG* types which will be serialized """
    return ("RealVector","RealMatrix","FloatVector","FloatMatrix",
            "StringCharList", "StringWordList")

def getSGTypeToStoreMethodName(sgType):
    """ Translates given SG* type into meta language type """
    assert sgType in getSGTypesToStore()

    if sgType=="RealVector":
        return "real_vector"
    elif sgType=="FloatVector":
        return "float_vector"
    elif sgType=="RealMatrix":
        return "real_matrix"
    elif sgType=="FloatMatrix":
        return "float_matrix"
    elif sgType=="StringCharList":
        return "string_char_list"
    elif sgType=="StringWordList":
        return "string_word_list"
    else:
        raise RuntimeError("Given Shogun type \"%s\" cannot be translated to meta type", sgType)


def getVarsToStore(program):
    """ Extracts all variables in program that should be stored """
    varsToStore = []

    # store only real valued matrices, vectors and scalars
    for init in find("Init", program):
        typeObject = init[0]
        nameString = init[1]["Identifier"]

        appendStore = False
        if "ShogunSGType" in init[0]:

            if init[0]["ShogunSGType"] in getSGTypesToStore():
                appendStore = True

        elif "BasicType" in init[0]:
            if init[0]["BasicType"] in getBasicTypesToStore():
                appendStore = True

        if appendStore:
            varsToStore.append((typeObject, nameString))

    return varsToStore

def findPath(key, dictionary, path=[]):
    """ Recursively search a dictionary for a key """
    for k, v in dictionary.items():
        if k == key:
            yield v, path + [k]
        if isinstance(v, dict):
            for result, p in findPath(key, v, path + [k]):
                yield result, p
        if isinstance(v, list):
            for idx, d in enumerate(v):
                for result, p in findPath(key, d, path + [k, idx]):
                    yield result, p

def validateProgram(program, filePath=None):
    """ Checks that a program is valid. Throws exceptions if not. """

    for d, path in findPath('KeywordArgument', program):
        # Check that any keyword argument appears either in a constructor call
        # or in the first global call of a init by copy statement
        try:
            initByGlobalCall = path[-8:-2] == ['Init', 2, 'Expr', 'GlobalCall', 1, 'ArgumentList']
            initByConstructor = path[-5:-2] == ['Init', 2, 'ArgumentList']
        except:
            initByGlobalCall = False
            initByConstructor = False

        if not (initByGlobalCall or initByConstructor):
            raise TranslationFailure('Keyword arguments must only appear in initialisation of variables')

    for d, path in findPath('ArgumentList', program):
        # Check that normal arguments always come before kwargs
        if not isinstance(d, list):
            raise TranslationFailure('ArgumentList should be a list')

        kwargObserved = False
        for arg in d:
            if 'KeywordArgument' in arg:
                kwargObserved = True
            else:
                if kwargObserved:
                    raise TranslationFailure('Keyword argument must come after normal arguments')

    # Warn against misspelled boolean literals
    for identifier, path in findPath('Identifier', program):
        # Retrieve the line number of the identifier
        lineNo = None
        if "Statement" in path:
            lineNoDict = program
            for p in path:
                if p == "Statement":
                    break
                lineNoDict = lineNoDict[p]
            lineNo = lineNoDict.get("__PARSER_INFO_LINE_NO")

        if identifier.lower() == 'true' or identifier.lower() == 'false':
            boolLiteral = identifier[0].upper() + identifier[1:].lower()