##############################
# Module to write MCNP input #
##############################

import logging
from pathlib import Path

import FreeCAD

from ..code_version import *
from ..utils.functions import SurfacesDict
from .functions import change_surf_sign, open_mc_surface, write_openmc_region

logger = logging.getLogger("general_logger")


class OpenmcInput:
    def __init__(self, Meta, Surfaces, options, tolerances, numeric_format):

        self.Cells = Meta
        self.options = options
        self.tolerances = tolerances
        self.numeric_format = numeric_format

        self.get_surface_table()
        self.simplify_planes(Surfaces)

        self.Surfaces = self.sorted_surfaces(Surfaces)
        self.Materials = set()

    def write_xml(self, filename):
        logger.info(f"write OpenMC xml file {filename}")
        Path(filename).parent.mkdir(parents=True, exist_ok=True)
        with open(file=filename, mode="w", encoding="utf-8") as self.inpfile:
            self.write_xml_header()
            self.inpfile.write("<geometry>\n")
            self.write_xml_cell_block()
            self.inpfile.write(" \n")
            self.write_xml_surface_block()
            self.inpfile.write("</geometry>\n")
        return

    def write_xml_header(self):
        Header = "<?xml version='1.0' encoding='UTF-8'?>\n"
        self.inpfile.write(Header)
        return

    def write_xml_cell_block(self):
        for _, cell in enumerate(self.Cells):
            if cell.MatInfo == "Graveyard":
                continue
            self.write_xml_cells(cell)
        return

    def write_xml_surface_block(self):
        for surf in self.Surfaces[:-1]:
            self.write_xml_surfaces(surf)
        self.write_xml_surfaces(self.Surfaces[-1], boundary=True)

    def write_xml_cells(self, cell):
        """Write the cell in xml OpenMC format"""
        index = cell.label
        cellName = ". ".join(cell.Comments.splitlines())
        if cell.__id__ is None:
            return

        if cell.Material == 0:
            matName = "void"
        else:
            matName = f"{cell.Material}"

        OMCcell = '  <cell id="{}" material="{}" name="{}" region="{}" universe="1"/>\n'.format(
            index,
            matName,
            cellName,
            write_openmc_region(cell.Definition, self.options, "XML"),
        )
        self.inpfile.write(OMCcell)
        return

    def write_xml_surfaces(self, surface, boundary=False):
        """Write the surfaces in xml OpenMC format"""

        surfType, coeffs = open_mc_surface(surface.Type, surface.Surf, self.tolerances, self.numeric_format)

        if not boundary:
            OMCsurf = '  <surface id="{}" type="{}" coeffs="{}" />\n'.format(surface.Index, surfType, coeffs)
        else:
            OMCsurf = '  <surface id="{}" type="{}" coeffs="{}" boundary="vacuum" />\n'.format(surface.Index, surfType, coeffs)

        self.inpfile.write(OMCsurf)
        return

    def write_py(self, filename):
        logger.info(f"write OpenMC python script {filename}")

        # get all the materials present in the model
        for cell in self.Cells:
            if cell.Material != 0:
                self.Materials.add(cell.Material)

        Path(filename).parent.mkdir(parents=True, exist_ok=True)
        with open(file=filename, mode="w", encoding="utf-8") as self.inpfile:
            self.write_py_header()

            if len(self.Materials) > 0:
                self.inpfile.write("# Materials setup\n")
                self.write_py_materials()
            self.inpfile.write("\n")

            self.inpfile.write("# Surface setup\n")
            self.write_py_surface_block()
            self.inpfile.write("\n")

            self.inpfile.write("# Cell definition \n")
            self.write_py_cell_block()
        return

    def write_py_header(self):

        Header = """\
# openMC geometry script generated by GEOUNED
import openmc

###############################################################################
# Define problem geometry
###############################################################################

"""
        self.inpfile.write(Header)
        return

    def write_py_materials(self):
        matList = tuple(sorted(self.Materials))
        strMat = []
        for m in matList:
            material = f"M{m} = openmc.Material(name='M{m}')\n"
            self.inpfile.write(material)
            strMat.append(f"M{m}")

        collect = f"materials = openmc.Materials([{', '.join(strMat)}])\n"
        self.inpfile.write(collect)
        self.inpfile.write("materials.export_to_xml()\n")

    def write_py_surface_block(self):

        for surf in self.Surfaces[:-1]:
            self.write_py_surfaces(surf)
        self.write_py_surfaces(self.Surfaces[-1], boundary=True)

    def write_py_surfaces(self, surface, boundary=False):
        """Write the surfaces in python OpenMC format"""

        surfType, coeffs = open_mc_surface(
            surface.Type,
            surface.Surf,
            self.tolerances,
            self.numeric_format,
            out_xml=False,
            quadricForm=self.options.quadricPY,
        )

        if not boundary:
            OMCsurf = f"S{surface.Index} = openmc.{surfType}({coeffs})\n"
        else:
            OMCsurf = 'S{} = openmc.{}({}, boundary_type="vacuum")\n'.format(surface.Index, surfType, coeffs)

        self.inpfile.write(OMCsurf)
        return

    def write_py_cell_block(self):

        cellNames = []
        for i, cell in enumerate(self.Cells):
            if cell.MatInfo == "Graveyard":
                continue
            self.write_py_cells(cell)
            if cell.__id__ is None:
                continue
            cellNames.append(f"C{cell.label}")

        geometry = "\ngeometry = openmc.Geometry([{}])\ngeometry.export_to_xml()\n".format(", ".join(cellNames))

        self.inpfile.write(geometry)
        return

    def write_py_cells(self, cell):
        """Write the cell in python OpenMC format"""
        index = cell.label
        cellName = ". ".join(cell.Comments.splitlines())
        if cell.__id__ is None:
            return

        if cell.Material == 0:
            OMCcell = 'C{} = openmc.Cell(name="{}", region={})\n'.format(
                index,
                cellName,
                write_openmc_region(cell.Definition, self.options, "PY"),
            )
        else:
            matName = f"M{cell.Material}"
            OMCcell = 'C{} = openmc.Cell(name="{}", fill={}, region={})\n'.format(
                index,
                cellName,
                matName,
                write_openmc_region(cell.Definition, self.options, "PY"),
            )
        self.inpfile.write(OMCcell)
        return

    def get_surface_table(self):
        self.surfaceTable = {}
        self.__solidCells__ = 0
        self.__cells__ = 0
        self.__materials__ = set()

        for i, CellObj in enumerate(self.Cells):
            if CellObj.__id__ is None:
                continue
            self.__cells__ += 1
            if CellObj.Material != 0:
                self.__materials__.add(CellObj.Material)

            surf = CellObj.Definition.get_surfaces_numbers()
            if not CellObj.Void:
                self.__solidCells__ += 1
            for index in surf:
                if index in self.surfaceTable.keys():
                    self.surfaceTable[index].add(i)
                else:
                    self.surfaceTable[index] = {i}
        return

    def simplify_planes(self, Surfaces):

        for p in Surfaces["PX"]:
            if p.Surf.Axis[0] < 0:
                p.Surf.Axis = FreeCAD.Vector(1, 0, 0)
                self.change_surf_sign(p)

        for p in Surfaces["PY"]:
            if p.Surf.Axis[1] < 0:
                p.Surf.Axis = FreeCAD.Vector(0, 1, 0)
                self.change_surf_sign(p)

        for p in Surfaces["PZ"]:
            if p.Surf.Axis[2] < 0:
                p.Surf.Axis = FreeCAD.Vector(0, 0, 1)
                self.change_surf_sign(p)
        return

    def sorted_surfaces(self, Surfaces):
        temp = SurfacesDict(Surfaces)
        surfList = []
        for ind in range(Surfaces.IndexOffset, Surfaces.surfaceNumber + Surfaces.IndexOffset):
            s = temp.get_surface(ind + 1)
            if s is not None:
                surfList.append(s)
                temp.del_surface(ind + 1)
        return surfList

    def change_surf_sign(self, p):

        if p.Index not in self.surfaceTable.keys():
            logger.info(f"{p.Type} Surface {p.Index} not used in cell definition {p.Surf.Axis} {p.Surf.Position}")
            return

        for ic in self.surfaceTable[p.Index]:
            surf = self.Cells[ic].Definition.get_surfaces_numbers()
            for s in surf:
                if s == p.Index:
                    change_surf_sign(s, self.Cells[ic].Definition)
