#=============================================================================
#
#  Copyright (c) Kitware, Inc.
#  All rights reserved.
#  See LICENSE.txt for details.
#
#  This software is distributed WITHOUT ANY WARRANTY; without even
#  the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
#  PURPOSE.  See the above copyright notice for more information.
#
#=============================================================================

""" convert_to_pyarc.py:

Helper classes for exporting a RGG model as PyARC geometry.

"""
import sys

import os
import smtk
import smtk.common
import smtk.io
import smtk.model
import rggsession
import uuid

# TODO: Use the new string format to dump out info
# '-' is a reserved keyword in pyarc. Do not use it.

def writeFile(filename, contents, mode="wt"):
    # wt = "write text"
    with open(filename, mode) as fout:
        fout.write(contents)


def prefixTabs(tabNum):
    return " " * 4 * tabNum


def rggOrderToNeamsOrder(ring, index):
    if (ring == 0):
        return 1, 1
    """Convert rgg ring index to neams ring index"""
    # AKA from 11 o'clock, counter clockwise to 3 o'clock clockwise
    if (index != 0):
        index = 6 * ring - index
    nIndex = (index + 2 * ring + 1) % (6 * ring)
    if (nIndex == 0):
        nIndex = 6 * ring
    nRing = ring + 1
    return nRing, nIndex

def surfacesToString(tabNum, surfaces):
    result = ""
    result += prefixTabs(tabNum) + "surfaces{\n"
    # Hexagon
    hexagons = sorted(list(surfaces["hexagon"]))
    for hexagon in hexagons:
        result += prefixTabs(tabNum + 1) + "hexagon ( %s ) {\n" % (hexagon[0])
        result += prefixTabs(tabNum + 2) + "orientation = %s\n" % (hexagon[1])
        result += prefixTabs(tabNum + 2) + "normal      = %s\n" % (hexagon[2])
        result += prefixTabs(tabNum + 2) + "pitch       = %f\n" % (hexagon[3])
        result += prefixTabs(tabNum + 1) + "}\n"
    # Cylinder
    #
    # TODO: these cylinders should not represent pin geometry, but rather
    #       cylindrical boundaries in the reactor. They are not used to define
    #       the reactor geometry; they are used in conjunction with planes to
    #       define a 2-dimensional lattice that is in turn used to compute
    #       neutron cross-sections.
    #
    #       From "ARC Integration into the NEAMS Workbench" (ANL/NE-17/31
    #       9/30/2017) pp. 28:
    #       How to get an equivalent RZ geometry? calculate the number of
    #       assemblies represented by each mcc3_id. Calculate the area for each
    #       region: the area of an hexagon is equal to
    #       $A=\frac{\sqrt{3}}{2}p^{2}$ with $p$ the pitch of the assembly. Find
    #       the "equivalent radius" by solving for r in $A=\pi r^{2}$.
    cylinders = sorted(list(surfaces["cylinder"]))
    for cylinder in cylinders:
        result += prefixTabs(tabNum + 1) + \
            "cylinder ( %s ) {\n" % (cylinder[0])
        result += prefixTabs(tabNum + 2) + "axis   = %s\n" % (cylinder[1])
        result += prefixTabs(tabNum + 2) + "radius = %f\n" % (cylinder[2])
        result += prefixTabs(tabNum + 1) + "}\n"
    # Plane
    planes = sorted(list(surfaces["plane"]))
    # PyARC expects the minimum plane to be at z=0, so we offset all z planes
    # to achieve this effect.
    minPlane = min([plane[1] for plane in planes])
    for plane in planes:
        result += prefixTabs(tabNum + 1) + "plane ( %s ) {\n" % (plane[0])
        result += prefixTabs(tabNum + 2) + "z = %s\n" % (plane[1] - minPlane)
        result += prefixTabs(tabNum + 1) + "}\n"
    result += prefixTabs(tabNum) + "}\n"
    return result


def materialsToString(tabNum, materialDescriptions):
    result = ""
    if materialDescriptions:
        result += prefixTabs(tabNum + 1) + 'materials{\n'
        for materialDescription in materialDescriptions:

            # ...we tokenize the description by line...
            materialDescription = materialDescription.replace("\\n","\n")
            tokenized = materialDescription.split('\n')

            # ... and do some string chicanery to get our indents right.
            for token in tokenized:
                result += prefixTabs(tabNum + 1) + token + '\n'
                if token[-1] == '{':
                    tabNum = tabNum + 1
                elif token[-1] == '}':
                    tabNum = tabNum - 1
        result += prefixTabs(tabNum + 1) + '}\n'
    return result;
class Pin(object):
    """
    A wrapper class around rgg pin for export
    """
    def __init__(self, rggPin, rggAssy, model):
        self.rggPin = rggPin
        self.rggAssy = rggAssy
        self.model = model

    def exportPinRegion(self, tabNum, indexOfAssySeg, surface):
        """
        Export a rgg pin int NEAMS son format. The assumption is that the pin
        has not frustum involved
        """
        result = ""
        materialNames = self.model.stringProperty('materials')
        label = self.rggPin.label()
        aLabel = self.rggAssy.label()
        layerMs = self.rggPin.layerMaterials()
        # Pin in NEAMS only has one radius
        radius = self.rggPin.pieces()[0].baseRadius
        result += prefixTabs(tabNum) + "pin_region ( {}_{}_{} ) {{\n".format(
            aLabel, label, indexOfAssySeg)
        for i in range(len(layerMs)):
            # sub_pin names must have their index separated by an "_"
            subPinName = label + '_' + str(i)
            materialName = materialNames[layerMs[i].subMaterialIndex]
            result += prefixTabs(
                tabNum + 1) + "sub_pin_region ( %s ) {\n" % (subPinName)
            result += prefixTabs(
                tabNum + 2) + "material   = %s\n" % (materialName)
            # Do not include inner_surf for the central region
            if (i != 0):
                currentInnerR = layerMs[i-1].normRadius * radius
                innerName = "cylinder" + str(currentInnerR)
                innerName = innerName.replace(".", "_")
                result += prefixTabs(
                    tabNum + 2) + "inner_surf = %s" % (innerName) + "\n"
                surface["cylinder"].add((innerName, "z", currentInnerR))
            currentOuterR = layerMs[i].normRadius * radius
            outerName = "cylinder" + str(currentOuterR)
            outerName = outerName.replace(".", "_")
            result += prefixTabs(
                tabNum + 2) + "outer_surf = %s" % (outerName) + "\n"
            surface["cylinder"].add((outerName, "z", currentOuterR))
            result += prefixTabs(tabNum + 1) + "}\n"
        result += prefixTabs(tabNum) + "}\n"
        return result

class Duct(object):

    def __init__(self, rggDuct, rggAssy, model):
        self.rggDuct = rggDuct
        self.rggAssy = rggAssy
        self.model = model
        self.auxGeo = 1
        self.zValues = 1
        self.zIntervals = 1
        #self.auxGeo = auxGeo
        #self.zValues = self.auxGeo.floatProperty('z values')
        #self.zIntervals = [self.zValues[0]] + \
        #    self.zValues[1:-1:2] + [self.zValues[-1]]

    def nParts(self):
        return len(self.rggDuct.segments())

    def exportRadicalRegions(self, tabNum, indexOfSubDuct, surfaces):
        """Based on the index of sub duct, generate radical_region(s) string
        and return center material for lattice string generation purpose"""
        result = ""
        materialNames = self.model.stringProperty("materials")
        coreG = smtk.model.Group(self.model.resource().findEntitiesByProperty("rggType",
        smtk.session.rgg.Core.typeDescription)[0])
        rggCore = smtk.session.rgg.Core(coreG)
        thickness = rggCore.ductThickness()[0]
        # Current segment layers: a vector of tuples
        currentSegLs = self.rggDuct.segments()[indexOfSubDuct].layers
        for i in range(len(currentSegLs)):
            # In neams, materials are exported from the outer region to inner region
            index = len(currentSegLs) - i - 1
            # do not include the innermost duct
            if (index == 0):
                continue
            # PyARC requires each radial_region to have a unique name within an
            # assembly
            result += prefixTabs(tabNum) + "radial_region ( " +\
                self.rggDuct.name() + str(
                indexOfSubDuct) + "_" + str(i) + " ) {\n"
            result += prefixTabs(tabNum + 1) + "material   = " +\
                materialNames[currentSegLs[index][0]] + "\n"
            # Do not include inner_surf for the central region
            if (index != 0):
                innerName = "hexagon" + str(currentSegLs[index-1][1]*thickness)
                innerName = innerName.replace(".", "_")
                result += prefixTabs(
                    tabNum + 1) + "inner_surf = " + innerName + "\n"
                surfaces["hexagon"].add(
                    (innerName, "y", "z", currentSegLs[index-1][1]*thickness))
            outerName = "hexagon" + str(currentSegLs[index][1]*thickness)
            outerName = outerName.replace(".", "_")
            result += prefixTabs(
                tabNum + 1) + "outer_surf = " + outerName + "\n"
            surfaces["hexagon"].add((outerName, "y", "z",
                                currentSegLs[index][1]*thickness))
            result += prefixTabs(tabNum) + "}\n"
        centerRegionMaterial = materialNames[currentSegLs[0][0]]
        return result, centerRegionMaterial

    def _radicalRegion(self, tabNum, materials, thicknessesN, indexOfSubDoct, surfaces):
        """Helper function for exportRadicalRegions"""
        materialNames = self.auxGeo.owningModel().stringProperty('materials')
        ductThickness = self.auxGeo.owningModel().floatProperty(
            'duct thickness')[0]
        thicknesses = [v * ductThickness for v in thicknessesN[::2]]
        result = ""
        for i in range(len(materials)):
            # In neams, materials are exported from the outter region to inner region
            # PyARC requires each radial_region to have a unique name within an
            # assembly
            index = len(materials) - i - 1
            result += prefixTabs(tabNum) + "radial_region ( " +\
                self.auxGeo.name() + str(
                indexOfSubDoct) + "_" + str(i) + " ) {\n"
            result += prefixTabs(tabNum + 1) + "material   = " +\
                materialNames[materials[index]] + "\n"
            # Do not include inner_surf for the central region
            if (index != 0):
                innerName = "hexagon" + str(thicknesses[index - 1])
                innerName = innerName.replace(".", "_")
                result += prefixTabs(
                    tabNum + 1) + "inner_surf = " + innerName + "\n"
                surfaces["hexagon"].add(
                    (innerName, "y", "z", thicknesses[index - 1]))
            outerName = "hexagon" + str(thicknesses[index])
            outerName = outerName.replace(".", "_")
            result += prefixTabs(
                tabNum + 1) + "outer_surf = " + outerName + "\n"
            surfaces["hexagon"].add((outerName, "y", "z", thicknesses[index]))
            result += prefixTabs(tabNum) + "}\n"
        centerRegionMaterial = materialNames[materials[0]]
        return result, centerRegionMaterial

class Assembly(object):

    def __init__(self, assy, model):
        self.rggAssy = assy
        self.model = model
        self.resource = model.resource()
        # get all pins and ducts
        self.rggDuct = smtk.session.rgg.Duct(smtk.model.AuxiliaryGeometry(
            self.resource, assy.associatedDuct()))
        self.rggPins = [smtk.session.rgg.Pin(smtk.model.AuxiliaryGeometry(
            self.resource, pinId)) for pinId in self.rggAssy.layout()]

        self.duct = Duct(self.rggDuct, self.rggAssy, self.model)
        self.pins = [Pin(rggPin, self.rggAssy, self.model) for rggPin in self.rggPins]
        # Make sure that the output order is persistant
        self.pins.sort(key = lambda x: x.rggPin.label())

        self.pitch = self.rggAssy.pitch()[0]
        self.num_ring = self.rggAssy.latticeSize()[0]

    def exportAssembly(self, tabNum, surfaces):
        """
        :param tabNum: number of tabs
        :return: the string of the assembly in son format
        """
        # QUESTION: Convert tab to space?
        # Calculate the segmentation of the assembly by z0 and z1
        assemblySegs = self.calculateAssemblySegs()
        label = self.rggAssy.label()

        result = prefixTabs(tabNum) + "assembly_hex ( " + label + " ) {\n"
        tabNum += 1
        for i in range(len(assemblySegs)):
            # QUESTION: In file it's called sub_assembly but in the documentation
            #  it's called sub_assembly_hex
            z0z1 = assemblySegs[i]
            # TODO: a meaningful name should be assigned to a subAssembly so that
            # it can be referred properly in calculation
            subAssyName = label + "_" + "_".join(map(str, z0z1))
            subAssyName = subAssyName.replace(".", "_")  # Avoid . in the name
            result += prefixTabs(tabNum) + "sub_assembly ( " + subAssyName\
                + " ) {\n"
            result += self._generateSurfaceString(tabNum + 1, z0z1, surfaces)
            _gSDSring, centerRegionMaterial =\
                self._generateSubDuctString(tabNum + 1, i, surfaces)
            result += _gSDSring
            result += self._generateLatticeString(tabNum + 1, z0z1,
                                                  i,
                                                  centerRegionMaterial,
                                                  surfaces)
            result += prefixTabs(tabNum) + "}\n"
        result += prefixTabs(tabNum - 1) + "}\n"
        return result

    def calculateAssemblySegs(self):
        """Calculate z values for segments in the duct.
        Return a list of pairs(z0,z1) in assending order"""
        ductZPairs = []
        for seg in self.rggDuct.segments():
            ductZPairs.append((seg.baseZ, seg.baseZ + seg.height))
        return sorted(ductZPairs)

    def _generateSurfaceString(self, tabNum, z0z1, surfaces):
        """
        # Helper function for exportAssembly to generate lower/upper_axial_surf
        :param tabNum: the number of tabs
        :param z0z1: current lower and upper bound along z axis
        :param surfaces: ad new objects to surfaces set if needed
        :return: result string
        """
        result = ""
        z0Name, z1Name = "z" + str(z0z1[0]), "z" + str(z0z1[1])
        z0Name, z1Name = z0Name.replace(".", "_"), z1Name.replace(".", "_")
        result += prefixTabs(tabNum) + "lower_axial_surf=" + z0Name + "\n"
        result += prefixTabs(tabNum) + "upper_axial_surf=" + z1Name + "\n"
        surfaces["plane"].add((z0Name, z0z1[0]))
        surfaces["plane"].add((z1Name, z0z1[1]))
        return result

    def _generateSubDuctString(self, tabNum, indexOfSubDuct, surfaces):
        stringResult, centerRegionMaterial =\
            self.duct.exportRadicalRegions(
                tabNum, indexOfSubDuct, surfaces)
        return stringResult, centerRegionMaterial

    def _generateLatticeString(self, tabNum, z0z1, indexofAssySeg,
                               cRMaterial, surfaces):
        result = prefixTabs(tabNum) + "assembly_hexlattice {\n"
        result += prefixTabs(tabNum + 1) + "pitch" + \
            " " * 9 + "= " + str(self.pitch) + "\n"
        result += prefixTabs(tabNum + 1) + "num_ring" + " " * 6 + "= "\
            + str(self.num_ring) + "\n"
        result += prefixTabs(tabNum + 1) + "outer" + " " * 9 + "= " + cRMaterial\
            + "\n"
        if (len(self.pins) > 0):
            result += self._generatePinLayoutString(
                tabNum + 1, z0z1, indexofAssySeg, surfaces)
        result += prefixTabs(tabNum) + "}\n"
        return result

    def _generatePinLayoutString(self, tabNum, z0z1, indexOfAssySeg, surfaces):
        result = ""
        # TODO: Use the fill list to simplify the logic
        # QUESTION: The schema planner in SMTK only supports top view(2D). If we
        # allow accumulating pins along z axis, then what should happen?

        # Replace center pin is not allowed in son format. You need to specify
        # it in the fill list. We will calculate during walk throught assy layout
        pinRegions, replaceString, centerPinLabel = "", "", "null"
        # Calculate pin regions(AKA radial_region)
        for i in range(len(self.pins)):
            pinRegions += self.pins[i].exportPinRegion(tabNum,
                                                     indexOfAssySeg,
                                                     surfaces)

        aLabel = self.rggAssy.label()
        replaceStringsL = []
        for uuid, schema in self.rggAssy.layout().items():
            rggPin = smtk.session.rgg.Pin(smtk.model.AuxiliaryGeometry(
            self.resource, uuid))
            label = rggPin.label()
            for item in schema:
                ring, index = rggOrderToNeamsOrder(item[0], item[1])
                uniquePinName = "{}_{}_{}".format(aLabel, label,
                                                    indexOfAssySeg)
                if (ring == 1 and index == 1):
                    centerPinLabel = uniquePinName
                else:
                    replaceStringsL.append((ring, index, prefixTabs(tabNum) +
                                            "replace{ ring=%d index=%d name=%s }\n"
                                            % (ring, index, uniquePinName)))
        replaceStringsL.sort(key=lambda x: (x[0], x[1]))
        for rS in replaceStringsL:
            replaceString += rS[2]

        # TODO: understand how null value is handled in son format
        result += prefixTabs(tabNum) + "fill" + " " * 10 + "= [ " +\
            "{} ".format(centerPinLabel) * \
            (self.num_ring) + "]\n"
        result += replaceString
        result += pinRegions
        return result

    def getAllSubAssyNames(self):
        # Duplicated logic in exportAssembly
        names = []
        assemblySegs = self.calculateAssemblySegs()
        label = self.rggAssy.label()

        for z0z1 in assemblySegs:
            # QUESTION: In file it's called sub_assembly but in the documentation
            #  it's called sub_assembly_hex
            # TODO: a meaningful name should be assigned to a subAssembly so that
            # it can be referred properly in calculation
            subAssyName = label + "_" + "_".join(map(str, z0z1))
            subAssyName = subAssyName.replace(".", "_")  # Avoid . in the name
            names.append(subAssyName)
        return names


class Core(object):
    """
    A wrapper class to export a rgg core into son format
    """
    def __init__(self, group):
        # core as smtk.model.group shape
        self.coreG = group
        self.resource = group.resource()
        self.model = group.owningModel()
        # core as smtk.session.rgg.Core
        self.rggCore = smtk.session.rgg.Core(self.coreG)

        assyGs = self.resource.findEntitiesByProperty("rggType", smtk.session.rgg.
        Assembly.typeDescription)
        self.rggAssys = [smtk.session.rgg.Assembly(smtk.model.Group(assyG))
         for assyG in assyGs]
        self.Assys = [Assembly(rggAssy, self.model) for rggAssy in self.rggAssys]
        # Make sure that the output order is persistant
        self.Assys.sort(key = lambda x: x.rggAssy.label())
        self.num_ring = self.rggCore.latticeSize()[0]

    def exportCore(self, tabNum, surfaces):
        result = prefixTabs(tabNum) + "regions_reactor{\n"
        result += self._calculateAxialSurfAndBoundaryCon(tabNum + 1, surfaces)
        result += self._generateCoreHexlatticeString(tabNum + 1, surfaces)
        result += self._generateRelatedAssysString(tabNum + 1, surfaces)
        result += prefixTabs(tabNum) + "}\n"
        return result

    def getAllSubAssyNames(self):
        """
        Get the names of all sub assemblies in the core.
        Ex. Let's say they are two assemblies PC and FC who are segmented as
        [0, 30, 50] and [0, 20, 50]. Then there are 4 subAssemblies as PC_0_30,
        PC_30_50, FC_0_20 and FC_20_50.
        """
        names = []
        for assy in self.Assys:
            names.extend(assy.getAllSubAssyNames())
        return names

    def _calculateAxialSurfAndBoundaryCon(self, tabNum, surfaces):
        """Helper function for exportCore to calculate axial surface and define\
        # boundary conditions. Which is also known as the z range of the core.
        """
        result = ""
        zMin = self.rggCore.zOrigin()
        zMax = zMin + self.rggCore.height()
        zMinName, zMaxName = "z" + str(zMin), "z" + str(zMax)
        zMinName, zMaxName = zMinName.replace(
            ".", "_"), zMaxName.replace(".", "_")
        surfaces["plane"].add((zMinName, zMin))
        surfaces["plane"].add((zMaxName, zMax))
        result += prefixTabs(tabNum) + "lower_axial_surf = %s\n" % (zMinName)
        result += prefixTabs(tabNum) + \
            "lower_boundary_condition = extrapolated\n\n"
        result += prefixTabs(tabNum) + "upper_axial_surf = %s\n" % (zMaxName)
        result += prefixTabs(tabNum) + \
            "upper_boundary_condition = extrapolated\n"
        return result

    def _generateCoreHexlatticeString(self, tabNum, surfaces):
        ductThickness = self.rggCore.ductThickness()[0]
        surfName = "hexagon" + str(ductThickness)
        surfName = surfName.replace(".", "_")
        surfaces["hexagon"].add((surfName, "y", "z", ductThickness))
        result = prefixTabs(tabNum) + "core_hexlattice {\n"
        # In neams workflow, assembly_surf refers to the largest duct in the core
        # In rgg workflow, duct's size is fixed per core
        result += prefixTabs(tabNum + 1) + "assembly_surf = %s\n" % (surfName)
        result += prefixTabs(tabNum + 1) + \
            "num_ring      = %d\n" % (self.num_ring)
        # Replace center assembly is not allowed in son format. You need to specify
        # it in the fill list
        replaceString, centerAssyLabel = "", "null"
        replaceStringsL = []
        for uuid, schema in self.rggCore.layout().items():
            rggAssy = smtk.session.rgg.Assembly(smtk.model.Group(
            self.resource, uuid))
            label = rggAssy.label()
            for item in schema:
                ring, index = rggOrderToNeamsOrder(item[0], item[1])
                if (ring == 1 and index == 1):
                    centerAssyLabel = label
                else:
                    replaceStringsL.append((ring, index, prefixTabs(tabNum+2) +
                                            "replace{ ring=%d index=%d name=%s }\n"
                                            % (ring, index, label)))
        replaceStringsL.sort(key=lambda x: (x[0], x[1]))
        for rS in replaceStringsL:
            replaceString += rS[2]
        result += prefixTabs(tabNum+1) + "fill" + " " * 10 + "= [ " +\
            "{} ".format(centerAssyLabel) * \
            (self.num_ring) + "]\n"
        result += replaceString
        result += prefixTabs(tabNum) + "}\n"
        return result

    def _generateRelatedAssysString(self, tabNum, surfaces):
        result = ""
        for assembly in self.Assys:
            result += assembly.exportAssembly(tabNum, surfaces)
        return result
