
import siconos.numerics as SN
import siconos.kernel as SK
import siconos.xml_parser as xml_parser
import numpy as np
import sys
import os

xsdfile = '@SICONOS_XML_SCHEMA@'


if not os.path.exists(xsdfile):
    tried_folder = 0
    python_path = os.path.dirname(SK.__file__)
    while not os.path.exists(os.path.join(python_path, 'share')) and tried_folder < 10:
        python_path = os.path.join(python_path, os.path.pardir)
        tried_folder += 1

    if tried_folder == 10:
        sys.stderr.write('No XML Schema found, stopping!')
        sys.exit(1)

nsSiconos = {"xsd": "http://www.w3.org/2001/XMLSchema", "cArg": "X-Constructor_Arguments"}

def _createVector(vec, dataType=np.float64):
    '''
    Create vector from xml data.
    The data we get from the xml is horrible.

    :param vec: the xml representation
    :type vec: too ugly to describe
    :param dataType: the type of the matrix (could be integer, double, ...)
    :type dataType: numpy dtype
     '''
    if vec.vectorFile is not None:
        return SK.SimpleVector(vec.vectorFile, True)
    else:
        return np.asarray(vec.valueOf_.split(), dtype=dataType)

def _createMatrix(mat, dataType=np.float64):
    '''
    Create matrix from xml data.
    The data we get from the xml is horrible

    :param mat: the xml representation
    :type mat: too ugly to describe
    :param dataType: the type of the matrix (could be integer, double, ...)
    :type dataType: numpy dtype
    '''
    if mat.matrixFile is not None:
        return SK.SimpleMatrix(mat.matrixFile, True)
    else:
        return np.asarray(mat.row, dtype=dataType)

def _createDS(dsType, dsxml, cArgsDS, dsOpt):
    '''
    Construct a DynamicalSystem based on the xml representation.
    For now we support only one constructor which is specified in the schema
    with the custom cArgs:order attribute

    :param dsType: the type of the DynamicalSystem to create
    :type dsType: string
    :param dsxml: the xml representation of the DynamicalSystem
    :type dsxml: xml representation from the parser
    :param cArgsDS: list of arguments to pass to the constructor
    :type cArgsDS: dict of items to pass as argument
    :param dsOpt: list members that could be set
    :type dsOpt: dict
    '''
    argsCtor = []
    argsCtorOrder = []
    dsD = dsxml.__dict__
    for (k, (eltType, order)) in cArgsDS.items():
        argsCtorOrder.append(order)
        argCtor = dsD[k]
        argCtorT = eltType + 'Plugin'
        argCtorStr = getattr(argCtor, argCtorT)
        if argCtorStr is not None:
            argsCtor.append(argCtorStr)
        else:
            mat = globals()['_create' + eltType.title()](argCtor)
            argsCtor.append(mat)

    # sort arg for constructor
    argsCtorS = [x for (y,x) in sorted(zip(argsCtorOrder, argsCtor))]

    # finger crossed
    newDS = getattr(SK, dsType)(*argsCtorS)

    _additionalInputsDS(dsxml, newDS, dsOpt)
    return newDS

def _create_relation(relType, relxml, cArgsR, relOpt):
    '''
    Construct a Relation based on the xml representation.
    For now we support only one constructor which is specified in the schema
    with the custom cArgs:order attribute

    :param relType: the type of the Relation to create
    :type relType: string
    :param relxml: the xml representation of the Relation
    :type relxml: xml representation from the parser
    :param cArgsR: list of arguments to pass to the constructor
    :type cArgsR: dict of items to pass as argument
    :param relOpt: list members that could be set
    :type relOpt: dict
    '''

    argsCtor = []
    argsCtorOrder = []
    relD = relxml.__dict__
    for (k, (eltType, order)) in cArgsR.items():
        argCtor = relD[k]
        # some arguments are optionals
        if argCtor is not None:
            argsCtor.append(argCtor.plugin)
            argsCtorOrder.append(order)

    # sort arg for constructor
    argsCtorS = [x for (y,x) in sorted(zip(argsCtorOrder, argsCtor))]

    relation = getattr(SK, relType)(*argsCtorS)
    _additionalInputsRelation(relxml, relation, relOpt)
    return relation

def _additionalInputsDS(dsxml, newDS, dsOpt):
    '''
    Set additional member(s) of a DynamicalSystem

    :param dsxml: the xml representation of the DynamicalSystem
    :type dsxml: xml representation from the parser
    :param newDS: the newly created DynamicalSystem to act on
    :type newDS: DynamicalSystem
    :param dsOpt: list members that could be set
    :type dsOpt: dict
    '''

    dsD = dsxml.__dict__
    for (plug, eltType) in dsOpt.items():
        if dsD[plug] is not None:
            plugin = dsD[plug]
            pluginT = eltType + 'Plugin'
            pluginStr = getattr(plugin, pluginT)
            if pluginStr is not None:
                setter = 'setCompute' + plug + 'Function'
                getattr(newDS, setter)(*pluginStr.split(':'))
            else:
                mat = globals()['_create' + eltType.title()](plugin)
                setter = 'set' + plug.title() + 'Ptr'
                getattr(newDS, setter)(mat)

def _additionalInputsRelation(relxml, rel, relOpt):
    '''
    Set additional member(s) of a Relation

    :param relxml: the xml representation of the Relation
    :type relxml: xml representation from the parser
    :param rel: the newly created Relation to act on
    :type rel: Relation
    :param relOpt: list members that could be set
    :type relOpt: dict
    '''

    relD = relxml.__dict__
    for (plug, eltType) in relOpt.items():
        if relD[plug] is not None:
            plugin = relD[plug]
            pluginT = eltType + 'Plugin'
            pluginStr = getattr(plugin, pluginT)
            if pluginStr is not None:
                setter = 'setCompute' + plug.title() + 'Function'
                getattr(rel, setter)(pluginStr.split(':'))
            else:
                mat = globals()['_create' + eltType.title()](relD[plug])
                setter = 'set' + plug.title() + 'Ptr'
                getattr(rel, setter)(mat)

def _additionalInputsNSLaw(nslxml, nslaw, nslOpt):
    '''
    Set additional member(s) of a NSLaw

    :param nslxml: the xml representation of the NSLaw
    :type nslxml: xml representation from the parser
    :param nslaw: the newly created NSLaw to act on
    :type nslaw: NSLaw
    :param nslOpt: list members that could be set
    :type nslOpt: dict
    '''

    nslD = nslxml.__dict__
    for (plug, eltType) in nslOpt.items():
        if nslD[plug] is not None:
            setter = 'set' + plug.title()
            getattr(nslaw, setter)(nslD[plug])

def _additionalInputsOSI(osixml, osi, osiOpt):
    '''
    Set additional member(s) of a OneStepIntegrator

    :param osixml: the xml representation of the OneStepIntegrator
    :type osixml: xml representation from the parser
    :param osi: the newly created OneStepIntegrator to act on
    :type osi: OneStepIntegrator
    :param osiOpt: list members that could be set
    :type osiOpt: dict
    '''

    osiD = osixml.__dict__
    for (plug, eltType) in osiOpt.items():
        if osiD[plug] is not None:
            if plug == 'Theta':
                osi.setTheta(float(osiD[plug].all))
            else:
                print('Not implemented yet')

def _additionalInputsOSNSP(osnspxml, osnsp, osnspOpt):
    '''
    Set additional member(s) of a OneStepNonSmoothProblem

    :param osnspxml: the xml representation of the OneStepNonSmoothProblem
    :type osnspxml: xml representation from the parser
    :param osnsp: the newly created OneStepNonSmoothProblem to act on
    :type osnsp: OneStepNonSmoothProblem
    :param osnspOpt: list members that could be set
    :type osnspOpt: dict
    '''

    osnspD = osnspxml.__dict__
    solverxml = osnspxml.NonSmoothSolver
    if solverxml is not None:
        solverOptions = osnsp.numericsSolverOptions()
        solverId = SN.solver_options_name_to_id(solverxml.Name)
        if solverId == 0:
            print('Error: the solver named {:} has no corresponding id in siconos.numerics'.format(solverxml.Name))
            exit(1)
        solverOptions.solverId = solverId
        if solverxml.dparam is not None:
            dparam = _createVector(solverxml.dparam)
            for i in range(dparam.size):
                solverOptions.dparam[i] = dparam[i]
        if solverxml.iparam is not None:
            iparam = _createVector(solverxml.iparam, np.int)
            for i in range(iparam.size):
                solverOptions.iparam[i] = iparam[i]

    for (plug, eltType) in osnspOpt:
        if osnspD[plug] is not None:
            setter = 'set' + plug.title() + 'Ptr'
            mat = globals()['_create' + eltType.title()](osnspD[plug])
            getattr(osnsp, setter)(mat)

def _parseReferenceXMLSchema(xsdFile):
    '''
    Parse the xsd schema to get information on how to construct object like
    DynamicalSystem amd also the additionla inputs to Relation, NSLaw, OSI, OSNSP

    :param xsdFile: the location of the xsd file
    :type xsdFile: string
    '''

    etree = xml_parser.etree_
    xmlschema_doc = etree.parse(open(xsdFile))

    xsd = etree.XMLSchema(xmlschema_doc)

    root = xmlschema_doc.getroot()

    siconosModel = root.xpath('.//xsd:element[@name = $n]', namespaces={"xsd": "http://www.w3.org/2001/XMLSchema"}, n='SiconosModel')[0]

    nsds = siconosModel.xpath('.//xsd:element[@name = $n]', namespaces={"xsd": "http://www.w3.org/2001/XMLSchema"}, n='NSDS')[0]

    allDS = nsds.xpath('.//xsd:element[@name = $n]//xsd:element/*/..', namespaces=nsSiconos, n='DS_Definition')
    allR = nsds.xpath('.//xsd:element[@name = "Interaction"]//xsd:choice[@id = $n]//xsd:element/*/..', namespaces=nsSiconos, n='Relation')
    allNSL = nsds.xpath('.//xsd:element[@name = "Interaction"]//xsd:choice[@id = $n]//xsd:element/*/..', namespaces=nsSiconos, n='NSLaw')
    allOSI = root.xpath('.//xsd:element[@name = $n]//xsd:element/*/..', namespaces=nsSiconos, n='OneStepIntegrator_Definition')
    allOSNSP = root.xpath('.//xsd:element[@name = $n]//xsd:element/*/..', namespaces=nsSiconos, n='OneStepNSProblems_List')

    constructorArgsDS = {}
    constructorArgsR = {}
    pluginsDS = {}
    pluginsR = {}
    dataNSL = {}
    dataOSI = {}
    dataOSNSP = {}

    bigTuple = ((allNSL, dataNSL), (allOSI, dataOSI), (allOSNSP, dataOSNSP))

    for elt in allDS:
        key = elt.attrib['name']
        cArgs = elt.xpath('.//xsd:element[@cArg:order]', namespaces=nsSiconos)
        constructorArgsDS[key] = dict([(cArg.attrib['name'], (cArg.attrib['type'], cArg.attrib['{X-Constructor_Arguments}order'])) for cArg in cArgs if 'name' in cArg.keys()])
        plugins = elt.xpath('.//xsd:element[not(@cArg:order)]', namespaces=nsSiconos)
        pluginsDS[key] = dict([(plug.attrib['name'], plug.attrib['type']) for plug in plugins if 'name' in plug.keys()])

    for elt in allR:
        key = elt.attrib['name']
        cArgs = elt.xpath('.//xsd:element[@cArg:order]', namespaces=nsSiconos)
        constructorArgsR[key] = dict([(cArg.attrib['name'], (cArg.attrib['type'], cArg.attrib['{X-Constructor_Arguments}order'])) for cArg in cArgs if 'name' in cArg.keys()])
        plugins = elt.xpath('.//xsd:element[not(@cArg:order)]', namespaces=nsSiconos)
        pluginsR[key] = dict([(plug.attrib['name'], plug.attrib['type']) for plug in plugins if 'name' in plug.keys()])

    for (typeElt, dico) in bigTuple:
        for elt in typeElt:
            key = elt.attrib['name']
            plugins = elt.xpath('.//xsd:element', namespaces={"xsd": "http://www.w3.org/2001/XMLSchema"})
            dico[key] = dict([(plug.attrib['name'], plug.attrib['type']) for plug in plugins if 'name' in plug.keys()])

    return (constructorArgsDS, constructorArgsR, pluginsDS, pluginsR, dataNSL, dataOSI, dataOSNSP)


# do the validation
def _validateXML(xmlFile, strict=True):
    '''
    Validate the input XML file against the schema

    :param xmlFile: the location of the XML input file
    :type xmlFile: string
    :param strict: if set to True, the XML input file has to be valid and the
    program exists if this is not the case
    :type strict: bool
    '''
    etree = xml_parser.etree_
    xsd_doc = etree.parse(xsdfile)
    xsd = etree.XMLSchema(xsd_doc)
    xml = etree.parse(xmlFile)
    xsd.validate(xml)
    print(xsd.error_log)


def buildModelXML(xmlFile):
    '''
    Build model from XML file given as argument

    :param xmlFile: the location of the XML input file
    :type xmlFile: string
    '''

    # hahah lots of error
    _validateXML(xmlFile)

    # get
    (constructorArgsDS, constructorArgsR, pluginsDS, pluginsR, dataNSL, dataOSI, dataOSNSP) = _parseReferenceXMLSchema(xsdfile)

    #  Built DOMtree
    modelxml = xml_parser.parse(xmlFile, silence=True)

    T = modelxml.Time.T
    t0 = modelxml.Time.t0
    title = modelxml.Title
    author = modelxml.Author
    description = modelxml.Description
    date = modelxml.Date
    xmlSchema = modelxml.SchemaXML


    # create the Model
    nsds = SK.NonSmoothDynamicalSystem(t0, T)
    nsdsxml = modelxml.NSDS
    #nsds = model.nonSmoothDynamicalSystem()

    # NSDS
    nsds.setBVP(nsdsxml.bvp)
    allDS = {}
    for k, v in nsdsxml.DS_Definition.__dict__.items():
        for dsxml in v:
            newDS = _createDS(k, dsxml, constructorArgsDS[k], pluginsDS[k])
            nsds.insertDynamicalSystem(newDS)
            nsds.setName(newDS, dsxml.Name)
            allDS[dsxml.Name] = newDS
    if len(allDS) == 0:
        print('No DynamicalSystem found !')
        print(nsdsxml.DS_Definition.__dict__.items)
        exit(1)

    # now the Interactions
    for interxml in nsdsxml.Interaction_Definition.Interaction:
        inter_content = interxml.Interaction_Content.__dict__
        for relName in pluginsR.keys():
            if inter_content[relName] is not None:
                relxml = inter_content[relName]
                relation =_create_relation(relName, relxml,
                                           constructorArgsR[relName],
                                           pluginsR[relName])
        for nslawName in dataNSL.keys():
            if inter_content[nslawName] is not None:
                nslxml = inter_content[nslawName]
                nslaw = getattr(SK, nslawName)(nslxml.size)
                _additionalInputsNSLaw(nslxml, nslaw, dataNSL[nslawName])

        # Finally, create the Interaction and insert it in NSDS
        inter = SK.Interaction(nslaw, relation)
        dsL = interxml.DS_Concerned
        if len(dsL) == 1:
            nsds.link(inter, allDS[dsL[0]])
        else:
            nsds.link(inter, allDS[dsL[0]], allDS[dsL[1]])


    # create the Simulation
    simxml = modelxml.Simulation
    if simxml.TimeDiscretisation.tk is not None:
        td = SK.TimeDiscretisation(simxml.TimeDiscretisation.tk)
    elif simxml.TimeDiscretisation.N is not None:
        td = SK.TimeDiscretisation(simxml.TimeDiscretisation.N, t0, T)
    elif simxml.TimeDiscretisation.h is not None:
        td = SK.TimeDiscretisation(t0, simxml.TimeDiscretisation.h)
    else:
        print('Could not define the TimeDiscretisation')
        exit(1)

    try:
        sim = getattr(SK, simxml.type_)(nsds,td)
    except:
        print("wrong type of simulation" + simxml.Type_)
        exit(1)

    # OSI
    allOSI = []
    for osiType in dataOSI.keys():
        for osixml in simxml.OneStepIntegrator_Definition.__dict__[osiType]:
            osi = getattr(SK, osiType)()
            _additionalInputsOSI(osixml, osi, dataOSI[osiType])
            print(osixml.DS_Concerned)
            for dsN in osixml.DS_Concerned:
                sim.associate(osi, allDS[dsN])
            allOSI.append(osi)

    # OSNSPB
    allOSNSP = []
    for osnspType in dataOSNSP.keys():
        if simxml.OneStepNSProblems_List.__dict__[osnspType] is not None:
            for osnspxml in simxml.OneStepNSProblems_List.__dict__[osnspType]:
                if osnspType is 'FrictionContact':
                    typeFC = int(osnspxml.Type)
                    osnsp = SK.FrictionContact(typeFC)
                else:
                    osnsp = getattr(SK, osnspType)()
                if osnspxml.StorageType is not None:
                    osnsp.setMStorageType(osnspxml.StorageType)
                _additionalInputsOSNSP(osnspxml, osnsp, dataOSNSP[osnspType])
                numId = SK.SICONOS_OSNSP_DEFAULT
                if osnspxml.Id is not None:
                    if simxml.type_ == 'EventDriven':
                        if osnspxml.Id == 'impact':
                            numId = SK.SICONOS_OSNSP_ED_IMPACT
                        elif osnspxml.Id == 'acceleration':
                            numId = SK.SICONOS_OSNSP_ED_SMOOTH_ACC
                        else:
                            numId = SK.SICONOS_OSNSP_ED_SMOOTH_POS
                # XXX fixme for TimeStepping
                sim.insertNonSmoothProblem(osnsp, numId)
                allOSNSP.append(osnsp)

    nsds.display()
    return nsds, sim
