# encoding: utf-8

import os, re, clr, Project, Parameters, codecs, json
clr.AddReference("Ans.ProjectSchematic")
import Ansys.ProjectSchematic

def isMechanicalSystem(system):
    props = Project.GetEntityProperties(system)
    attributes = props.get('Attributes', {})
    solverType = attributes.get('SolverType', '')
    anType = attributes.get('AnalysisType', '')
    return solverType in ['ANSYS', 'LSDYNA', 'AUTODYN', 'ABAQUS', 'NASTRAN', 'RBD', 'Samcef'] or anType == 'Mechanical'

def writeScriptToExtractMechanicalData(outputDirForExtractedFiles):
    list = outputDirForExtractedFiles.split('\\')
    outputDirForExtractedFiles = '\\\\'.join(list)
    str = """
    WB.AppletList.ActivateApplet("DSApplet");
    ds = WB.AppletList("DSApplet");
    savedir = "%s";
    pr = ds.Applet.App.Tree.Projects(1);
    if (pr.Models.Count > 0) 
    {
        ds.Applet.App.ExtractEkmData(savedir);
    }
    """ % outputDirForExtractedFiles
    jspath = os.path.join(os.getcwd(), 'extract.js')
    f = open(jspath, 'w');
    f.write(str)
    f.close()
    return jspath

def exportMechanicalData(outputFolder):
    jspath = writeScriptToExtractMechanicalData(outputFolder)
    for system in Project.GetAllSystems():
        if isMechanicalSystem(system):
            cont = system.GetContainer(ComponentName="Model")
            name = cont.Name.ToString()
            Simulation.RunScript(jspath, IsMeshing=False, ModelName=name)
            Simulation.ExitSimulationEditor(name)
    os.remove(jspath)

def IsEntireProjectUpToDate(includeParametricSystems):
    upToDate = True
    if not includeParametricSystems:
        upToDate = Project.IsProjectUpToDate()
    else:
        systems = Project.GetAllSystems()
        components = []
        for system in systems:
            components.extend(system.Components)
        
        for c in components:
            state = Project.GetSystemComponentState(c).State
            if state != Ansys.ProjectSchematic.State.UpToDate:
                upToDate = False
                break
        
        if upToDate:
            for c in components:
                state = Project.GetComponentState(c).State
                if state != Ansys.ProjectSchematic.State.UpToDate:
                    upToDate = False
                    break
    
    return upToDate


def extractProjectMetadata():
    systemNames = []
    systemTypes = []
    for s in Project.GetAllSystems():
        systemNames.append(s.DisplayText)
        systemTypes.append(s.GetEntityProperty(Name='SystemType'))
    designPoints = []
    for dp in Parameters.GetAllDesignPoints():
        designPoints.append(dp.DisplayText)
    parameters = []
    for param in Parameters.GetAllParameters():
        parameters.Add(param.DisplayText)

    def joinWithComma(strings):
        return ', '.join(strings)

    return {                                                                        \
            'persistedProjectVersion': Project.GetPersistedExternalVersionString(), \
            'updateRequired': str(not IsEntireProjectUpToDate(True)),               \
            'systemNames': joinWithComma(systemNames),                              \
            'systemTypeNames': joinWithComma(systemTypes),                          \
            'parameterNames': joinWithComma(parameters),                            \
            'designPointNames': joinWithComma(designPoints),                        \
            'isProjectInAIMMode': str(Project.IsProjectInAIMMode())
           }

def extractAimMetadata(system):
    try:
        return {'name': system.DisplayText, 'spec': 'AimStudyPropertyModel', 'properties': GetStudyMetadataQuery()}
    except:
        return None
    
def extractMetadata():
    metadata = {}
    metadata['properties'] = extractProjectMetadata()
    systemMetadata = []
    for sys in Project.GetAllSystems():
        sysType = sys.GetEntityProperty(Name='SystemType')
        if sysType == 'Study':
            aimMetadata = extractAimMetadata(sys)
            if aimMetadata:
                systemMetadata.append(aimMetadata)
    metadata['children'] = systemMetadata    
    return metadata
    
# Returns 'Modal 2' with 'Modal', 'Modal 3' for 'Modal 2', and so on.
def incrementDisplayTextCounter(dt):
    m = re.match('.*\d+$', dt) # string ends with digit(s)
    if m:
        s = re.search('\d+$', dt) # get digit(s) at end of string 
        if s:
            ctr = int(s.group(0))
            ctr += 1 # Increment last number
            return dt[:s.start()] + str(ctr)
    return dt + ' 2'


class EkmMatadataXmlFormatter:

    def __init__(self, metadata):
        self.__lines = []
        self.metadata = metadata
    
    def __addXmlProlog(self):
        self.__lines.append(r'<?xml version="1.0" encoding="UTF-8" standalone="no" ?>')
    
    def __openMetadataRootNode(self):
        self.__lines.append(r'<ekmMetadata:meta-data xmlns:ekmMetadata="http://www.ansys.com/ekm/metaData">')
        
    def __closeMetadataRootNode(self):
        self.__lines.append('</ekmMetadata:meta-data>')
    
    def __addMetadatumNode(self, name, value):
        # DE91750: Format pre-defined XML entities for correct rendering 
        value = value.replace('&', '&amp;')
        value = value.replace('<', '&lt;')
        value = value.replace('>', '&gt;')
        value = value.replace("'", '&apos;')
        value = value.replace('"', '&quot;')
        str = """    <data name="%s" value="%s"/>""" % (name, value)
        self.__lines.append(str)
    
    def __addChildNode(self, childMetadata):
        str = """    <child name="%s" type="%s">""" % (childMetadata['name'], childMetadata['spec'])
        self.__lines.append(str)
        properties = childMetadata['properties']
        for key in properties:
            self.__addMetadatumNode(key, properties[key])
        self.__lines.append("    </child>")
    
    def getXml(self):
        self.__addXmlProlog()
        self.__openMetadataRootNode()
        properties = self.metadata['properties']
        for key in properties:
            self.__addMetadatumNode(key, properties[key])
        for childMetadata in self.metadata['children']:
            self.__addChildNode(childMetadata)
        self.__closeMetadataRootNode()
        return '\n'.join(self.__lines)


class EkmDataExporter:

    def exportData(self,
                   exportMetadata=True,
                   exportReport=True,
                   metadataFileName='ekm-meta-data.xml', 
                   systemFilesMapName='system-files-association.txt',
                   reportPath=r"report\report.xml"):
        if exportMetadata:
            self.__writeMetadata(metadataFileName)
            self.__writeSystemFilesMap(systemFilesMapName)
        if exportReport and not os.getenv('DONT_EXPORT_REPORT'): # DE112608
            Project.ExportReport(FilePath=reportPath)
        parentPath = os.path.dirname(os.path.abspath(metadataFileName))
        exportMaterialData(parentPath)
        self.__writeRepositoryFilesData(parentPath)

    def __writeMetadata(self, fileName):
        metadata = extractMetadata()
        w = EkmMatadataXmlFormatter(metadata)
        f = codecs.open(fileName, encoding='utf-8', mode='w')
        f.write(w.getXml())
        f.close()

    def __writeSystemFilesMap(self, fileName):
        f2 = open(fileName, 'w')
        f2.write('sysFilesMap = ' + str(self.__getSystemFilesMap()))
        f2.close()

    def __getSystemFilesMap(self):
        sysFilesMap = {}
        projectPath = Project.GetProjectFile()
        projectFilesPath = projectPath.replace('.wbpj', '_files')
        for system in Project.GetAllSystems():
            files = []
            for component in system.Components:
                container = component.DataContainer
                # Loop over all file data references associated with the container for the component
                for compFile in container.GetFiles():
                    files.append(compFile.Location)
            for entry in os.listdir(projectFilesPath):
                entryPath = os.path.join(projectFilesPath, entry)
                if os.path.isdir(entryPath) and re.match(r'dp\d+', entry):
                    sysDirName = Project.GetEntityProperty(Entity=system, Name='UniqueSystemDirectoryName')
                    sysDirPath = os.path.join(entryPath, sysDirName)
                    if os.path.exists(sysDirPath):
                    	self.__addMapdlResultsRecursively(sysDirPath, files)
            dt = system.DisplayText
            if sysFilesMap.has_key(dt): # DE68559: multiple systems can have same display text. Adjust by affixing index. Better to replace map in future.
                while sysFilesMap.has_key(dt):
                    dtNext = incrementDisplayTextCounter(dt)
                    if dtNext == dt:
                        dt = system.Name
                        break
                    else:
                        dt = dtNext
                        continue
                sysFilesMap[dt] = files
            else:
                sysFilesMap[dt] = files
        return self.__updateSystemFilesMapForAnsoftSystems(sysFilesMap)
    
    def __addMapdlResultsRecursively(self, directoryPath, files):
        for dirEntry in os.listdir(directoryPath):
            path = os.path.join(directoryPath, dirEntry)
            if os.path.isdir(path):
                self.__addMapdlResultsRecursively(path, files)
            elif os.path.splitext(path)[1] in ['.rst', '.rth', '.rmg', '.rfl']:
                files.append(path)

    def __updateSystemFilesMapForAnsoftSystems(self, sysFilesMap):
        projectPath = Project.GetProjectFile()
        projectFilesPath = projectPath.replace('.wbpj', '_files')
        ansoftSystemTypes = ['HFSS Design', 'Simplorer Design', 'Maxwell2D Design', 'Maxwell3D Design', 'RMxprt Design', 'DesignerCircuit Design', 'DesignerEM Design', 'NexximNetlist Design']
        ansoftSystems = [s for s in Project.GetAllSystems() if Project.GetEntityProperty(s, 'SystemType') in ansoftSystemTypes]
        for sys in ansoftSystems:
            components = sys.Components
            for comp in components:
                cont = sys.GetContainer(ComponentName=comp.Name)
                entities = [e for e in Project.GetDataEntities(Container=cont) if e.Type=='AnsoftDesktopDesignEntityType']
                if len(entities) > 0:
                    e0 = entities[0]
                    relPath = Project.GetEntityProperty(Entity=e0, Name='DesktopProjectRelPath')
                    relPathOs = relPath.replace('/', os.sep)
                    absPath = os.path.join(projectFilesPath, relPathOs)
                    if sysFilesMap.has_key(sys.DisplayText):
                        sysFilesMap[sys.Name] = [absPath]
                    else:
                        sysFilesMap[sys.DisplayText] = [absPath]
        return sysFilesMap
        
    def __writeRepositoryFilesData(self, parentPath):
        repoFiles = Project.GetRepositoryFiles()
        if len(repoFiles) > 0:
            repoFilesInfo = [rf.GetEntityProperty(Name='SourceFileInfo') for rf in repoFiles]
            repoFilesInfoStr = json.dumps(repoFilesInfo)
            path = os.path.join(parentPath, 'repository-files.json')
            f = open(path, 'w')
            f.write(repoFilesInfoStr)
            f.close()


def disableAddingRecentlyOpenedFiles():
    clr.AddReference("Ans.Core")
    import Ansys.Core.FileManagement
    Ansys.Core.FileManagement.RecentlyUsedFilesManager.SetAdditionsEnabled(False)


class EkmXmlReport:
    
    def __init__(self):
        self.__lines = []
     
    def getContents(self):
        return '\n'.join(self.__lines)
        
    def openSectionNode(self, rootName):
        self.__lines.append(r'<?xml version="1.0" encoding="UTF-8" standalone="no" ?>')
        self.__lines.append('<section name="' + rootName + '">')
        
    def closeSectionNode(self):
        self.__lines.append('</section>')
        
    def openTableNode(self, tableName):
        str = """    <table name="%s" showBorder="1" showHeaders="1">""" % tableName
        self.__lines.append(str)
        
    def closeTableNode(self):
        self.__lines.append('    </table>')
    
    def addColumnNode(self, columnName):
        str = """        <column name="%s"/>""" % columnName
        self.__lines.append(str)
    
    def addRowNode(self, propName, propValue):
        str = """        <row Name="%s" Value="%s"/>""" % (propName, propValue)
        self.__lines.append(str)


# Material meta-data

propertiesIgnored = [
"Coefficient of Thermal Expansion",
"Compressive Ultimate Strength",
"Compressive Yield Strength",
"Density",
"Specific Heat",
"Tensile Yield Strength",
"Tensile Ultimate Strength",
"Thermal Conductivity",
"Alternating Stress",
"Resistivity",
"Response Function",
"Experimental Data",
"Constant Damping Coefficient",
"Damping Factor (\u03B2)",

"Relative Permittivity",
"Electric Loss Tangent",
"Seebeck Coefficient",
"Strain-Life Parameters",

"Uniaxial Test Data",
"Uniaxial Tension Test Data",
"Uniaxial Compression Test Data",
"Biaxial Test Data",
"Shear Test Data",
"Simple Shear Test Data",
"Volumetric Test Data",
"Shock EOS",
"Polynomial EOS",

"Bulk Modulus",
"Shear Modulus",
"Plakin Special Hardening Law",
"Molecular Weight",
]

linearElasticModels = [
"Elasticity",
"Viscoelastic",
]

magneticModels = [
"Relative Permeability",
"B-H Curve",
"Demagnetization B-H Curve",
"Coercive Force & Residual Induction",
]

gasketModels = [
"Gasket Model",
]

failureModels = [
"Plastic Strain Failure",
"Principal Stress Failure",
"Principal Strain Failure",
"Stochastic Failure",
"Crack Softening Failure",
"Johnson Cook Failure",
"Grady Spall Failure",
"RHT Concrete Failure",
"Tensile Pressure Failure",
]

porosityModels = [
"Crushable Foam",
"Compaction EOS",
"P-alpha EOS",
]

plasticityModels = [
"Isotropic Hardening",
"Kinematic Hardening",
"Johnson Cook Strength",
"Cowper Symonds Strength",
"Steinberg Guinan Strength",
"Zerilli Armstrong Strength",
"Anand Viscoplasticity",
]

brittleGranularModels = [
"Drucker-Prager Strength",
"Johnson-Holmquist Strength",
"RHT Concrete Strength",
"MO Granular",
]

hyperelasticityModels = [
"Neo-Hookean",
"Mooney-Rivlin",
"Polynomial",
"Yeoh",
"Ogden",
"Gent",
"Blatz-Ko",
"Arruda-Boyce",
]

creepModels = [
"Strain Hardening",
"Time Hardening",
"Generalized Exponential",
"Generalized Graham",
"Generalized Blackburn",
"Modified Time Hardening",
"Modified Strain Hardening",
"Generalized Garofalo",
"Exponential Form",
"Norton",
"Combined Time Hardening",
"Rational Polynomial",
"Generalized Time Hardening",
]

def getMaterialBehavior(propName):
    if propName in linearElasticModels:
        return "Linear Elasticity"
    if propName in gasketModels:
        return "Gasket"
    elif propName in porosityModels:
        return "Porosity"
    elif propName in plasticityModels:
        return "Plasticity"
    elif propName in brittleGranularModels:
        return "Brittle-Granular"
    elif propName in hyperelasticityModels:
        return "Hyperelasticity"
    elif propName in creepModels:
        return "Creep"
    else:
        return None

def exportMaterialMetadata(dataContainer, path):
    f = codecs.open(path, encoding='utf-8', mode='w')
    materials = EngData.GetMaterials(EngDataContainer=dataContainer)
    for mtl in materials:
        mtlName = Project.GetEntityProperty(Entity=mtl, Name="DisplayName")
        f.write('MATERIAL_NAME: ' + mtlName + '\n')
        for prop in mtl.MaterialPropertyColl:
            propName = Project.GetEntityProperty(Entity=prop, Name='DisplayText')
            mtlBehavior = getMaterialBehavior(propName)
            # Write only interesting properties and behaviors
            if mtlBehavior is not None:
                f.write('MATERIAL_PROPERTY: ' + propName + '\n')
                f.write('MATERIAL_BEHAVIOR: ' + mtlBehavior + '\n')
    f.close()

def populateMaterialData(dataContainer, report):
    materials = EngData.GetMaterials(EngDataContainer=dataContainer)
    report.openSectionNode('Materials')
    for i in materials:
        mtlName = Project.GetEntityProperty(Entity=i, Name="DisplayName")
        report.openTableNode(mtlName + ' Properties')
        report.addColumnNode('Name')
        report.addColumnNode('Value')
        mtlProps = Project.GetEntityProperties(Entity=i)["MaterialPropertyColl"]
        for j in mtlProps:
            propDt = Project.GetEntityProperty(Entity=j, Name='DisplayText')
            mtlBehavior = getMaterialBehavior(propDt)
            if mtlBehavior is not None:
                report.addRowNode('MATERIAL_PROPERTY', propDt)
                report.addRowNode('MATERIAL_BEHAVIOR', mtlBehavior)
            propName = Project.GetEntityProperty(Entity=j, Name="TypeName")
            propDatas = Project.GetEntityProperties(Entity=j)["PropertyDataColl"]
            for k in propDatas:
                propDataName = Project.GetEntityProperty(Entity=k, Name="TypeName")
                propVars = Project.GetEntityProperties(Entity=k)["VariableColl"]
                propVals = EngData.GetValues(thisObject=k)
                nVals = len(propVals)
                nVars = 0
                for l in propVars:
                    nVars += 1
                if (nVars == nVals):
                    ii = -1
                    for l in propVars:
                        ii += 1
                        varName = Project.GetEntityProperty(Entity=l, Name="TypeName")
                        if abs(propVals[ii]) < 1e-30: continue   # 2^-100 is tiny in WB2/ED
                        report.addRowNode(varName, repr(propVals[ii]))
                else:
                        report.addRowNode(propDataName, 'Tabular Data')
        report.closeTableNode()
    report.closeSectionNode()

def exportMaterialDetails(dataContainer, path):
    report = EkmXmlReport()
    populateMaterialData(dataContainer, report)    
    f = codecs.open(path, encoding='utf-8', mode='w')
    f.write(report.getContents())
    f.close()

def exportMaterialData(outputFolder):
    for system in Project.GetAllSystems():
        if isMechanicalSystem(system):
            cont = system.GetContainer(ComponentName="Engineering Data")
            usdn = Project.GetEntityProperties(system).get('UniqueSystemDirectoryName', '')
            mtlXmlPath = os.path.join(outputFolder, usdn+'-materials.xml')
            exportMaterialDetails(cont.Name, mtlXmlPath)
            mtlMetadataPath = os.path.join(outputFolder, usdn+'-materials.txt')
            exportMaterialMetadata(cont.Name, mtlMetadataPath)
