<reponame>Radiation-Transport/mctools
"""
14/05/2021
This script reads a single MCNP output file and produces an Excel
file with information on lost particles.

This version aims at improved execution speed and usability (only one file required)

TODO: At the moment the script can only populate the list of lost particle tracks and only with the most basic info.
      Information on filler universe should be read and post-processed.
"""
from typing import List
import re
import pandas as pd


class LostParticle(object):
    """
    Each instance of this class contains all the relevant information of a lost particle.
    """
    def __init__(self):
        self.position = []
        self.direction = []
        self.position_in_log = []
        self.direction_in_log = []
        # TODO: understand which position is in global position
        #       and which in the filler position properly
        self.cell = None
        self.cell_alternative = None
        # TODO: understand what is the difference between the two cell
        #       definitions properly
        self.surface = None


def mcnp_output_to_lost_particles(mcnp_output):
    """
    This function parses the MCNP output file to extract the information of the lost particles.
    Right now it just obtains the information present in the lost particle logs at the end of the output file.
    TODO: Upgrade this function so it extracts information about the universe where the particles were lost by parsing
          the file and finding the U=?? for every lost particle cell ID.
    Returns
    -------
    List[LostParticle]
    """
    _lost_particles = []

    # Regex patterns
    pattern_float = re.compile(r'-?\d\.\d+E[+-]\d+')
    pattern_beginning_lp_log = re.compile(r'1\s+event log for lost particle')
    pattern_surface = re.compile(r'currently being tracked has reached surface')
    pattern_cell = re.compile(r'appears to be no cell on the')
    # pattern_cell_alternative = re.compile(r'side of the surface from cell')
    pattern_position = re.compile(r'x,y,z coordinates')
    pattern_direction = re.compile(r'u,v,w direction')

    # Parse the file
    with open(mcnp_output, 'r') as infile:
        line = infile.readline()
        while line != '':

            # Only interested in lost particle logs
            if pattern_beginning_lp_log.search(line):

                lost_particle = LostParticle()
                line = infile.readline()  # Step into the log

                # Parse the complete log before continuing
                while not pattern_beginning_lp_log.match(line) and line != '':

                    # Position
                    if pattern_position.search(line):
                        position_str = pattern_float.findall(line)
                        lost_particle.position = [float(x) for x in position_str]

                    # Direction
                    elif pattern_direction.search(line):
                        direction_str = pattern_float.findall(line)
                        lost_particle.direction = [float(x) for x in direction_str]

                    # Surface ID
                    elif pattern_surface.search(line):
                        surface_str = re.search(r'\d+', line).group()
                        lost_particle.surface = int(surface_str)

                    # Cell ID
                    elif pattern_cell.search(line):
                        cell_str = re.search(r'\d+', line).group()
                        lost_particle.cell = int(cell_str)
                    line = infile.readline()

                # At the end of the log we append the lost particle to the list
                _lost_particles.append(lost_particle)

            # The line is not a part of a lost particle log
            else:
                line = infile.readline()

    return _lost_particles


def generate_excel_from_lost_particles(lost_particles):
    """"
    Generates an Excel file with the location of all the particles and a summary of the most problematic cells,
    surfaces and combination cell+surface.
    TODO: Add distinction between global and local position, cell and alternative_cell etc.

    Parameters
    ----------
    lost_particles: List[LostParticle]

    """

    # Generate a DataFrame with the relevant information from the lost particles
    cells = []
    surfaces = []
    positions = []
    for lp in lost_particles:
        cells.append(lp.cell)
        surfaces.append(lp.surface)
        positions.append(lp.position)
    df_all_info = pd.DataFrame({'cellID': cells, 'surfaceID': surfaces, 'position': positions})

    # Post-process the DataFrame
    df_cell_surface_count = df_all_info.groupby(['cellID', 'surfaceID']).count()
    df_cell_surface_count = df_cell_surface_count.reset_index().set_index('position').sort_index(ascending=False)
    df_cell_surface_count.index.names = ['count']

    df_cell_count = df_all_info.groupby(['cellID']).count()
    df_cell_count = df_cell_count.reset_index().set_index('position').sort_index(ascending=False)
    df_cell_count = df_cell_count.drop('surfaceID', axis=1)
    df_cell_count.index.names = ['count']

    df_surface_count = df_all_info.groupby(['surfaceID']).count()
    df_surface_count = df_surface_count.reset_index().set_index('position').sort_index(ascending=False)
    df_surface_count = df_surface_count.drop('cellID', axis=1)
    df_surface_count.index.names = ['count']

    df_main = df_all_info.set_index(['cellID', 'surfaceID']).sort_index()

    # Write the dataframes to an Excel file
    writer = pd.ExcelWriter('LPdebug.xlsx', engine='xlsxwriter')
    workbook = writer.book
    worksheet = workbook.add_worksheet('LP Debugging')
    writer.sheets['LP Debugging'] = worksheet
    df_main.to_excel(writer, sheet_name='LP Debugging', startrow=0, startcol=0)
    worksheet.write('G1', 'Top 20 cell and surface combinations')
    worksheet.write('K1', 'Top 20 cells')
    worksheet.write('N1', 'Top 20 surfaces')
    worksheet.set_column('C:C', 25)
    df_cell_surface_count = df_cell_surface_count.head(20)
    df_cell_count = df_cell_count.head(20)
    df_surface_count = df_surface_count.head(20)
    df_cell_surface_count.to_excel(writer, sheet_name='LP Debugging', startrow=1, startcol=6)
    df_cell_count.to_excel(writer, sheet_name='LP Debugging', startrow=1, startcol=10)
    df_surface_count.to_excel(writer, sheet_name='LP Debugging', startrow=1, startcol=13)

    writer.save()
    return


def lpdebug(output_list, input_file):
    """
    Receives as argument the MCNP output file and produces the Excel file.

    Parameters
    ----------
    mcnp_output : str
    The filename or path to the MCNP output file
    """
    # Input file is not currently used for anything
    mcnp_output = output_list[0]  # Only works for a single output file

    # Create and populate a list of LostParticle instances
    print('Reading the output file...')
    lost_particles = mcnp_output_to_lost_particles(mcnp_output)

    # Post-process the lost particles to generate an Excel with the results
    print('Processing the data...')
    generate_excel_from_lost_particles(lost_particles)

    return


if __name__ == '__main__':
    "Test the script with the example files"
    lpdebug("E-Lite_IVVS_v1.0_VOID.o")
