#!/usr/bin/env python3
# -*- coding:utf-8 -*- import numpy as np import re
import argparse
import collections
import re
from natf.utils import is_blank_line, mcnp_style_str_append

comment_pattern = re.compile("^C ", re.IGNORECASE)
cont_pattern = re.compile("^      ")
comp_pattern = re.compile("^C *Component", re.IGNORECASE)
cell_range_pattern = re.compile("^C *Cell Range", re.IGNORECASE)
group_name_pattern = re.compile("^C *Group", re.IGNORECASE)
from_to_pattern = re.compile("^C.*From.*to", re.IGNORECASE)
new_comp_name_pattern = re.compile("^newcomp", re.IGNORECASE)
new_group_name_pattern = re.compile("^newgroup", re.IGNORECASE)


def is_cell_title(line):
    """Check whether this line is the first line of a cell card"""
    line_ele = line.split()
    if len(line_ele) == 0: # blank line
        return False
    # not surf title
    try:
        int(line_ele[1])
    except:
        return False
    # continue line
    if str(line_ele[0]).isdigit() and (not re.match(cont_pattern, line)):
        return True
    return False

def get_cell_cid_mid_den(line):
    """Get the cell id, mid and density information"""
    line_ele = line.split()
    cid = int(line_ele[0])
    mid = int(line_ele[1])
    den = None
    if mid > 0:
        if '(' in line_ele[2]:
            tokens = line_ele[2].split('(')
            den = float(tokens[0])
        else:
            den = float(line_ele[2])
    elif mid == 0:
        pass 
    else:
        raise ValueError(f"Wrong cell title line: {line}")
    return cid, mid, den

def cell_title_change_mat(cell_title, new_mid, new_den=-1.0):
    """Change the material in cell title"""
    line_ele = cell_title.split()
    cid = int(line_ele[0])
    mid = int(line_ele[1])
    rest = []
    if mid > 0:
        if '(' in line_ele[2]:
            tokens = line_ele[2].split('(')
            rest.append('('+tokens[1])
            rest.extend(line_ele[3:])
        else:
            rest.extend(line_ele[3:])
    elif mid == 0:
        rest.extend(line_ele[2:])
    else:
        raise ValueError(f"Wrong cell title line: {line}")
    if new_mid == 0 and new_den is None:
        return f"{cid} 0 {' '.join(rest)}"
    else:
        return f"{cid} {new_mid} {new_den} {' '.join(rest)}"

def is_comment(line):
    """Check whether this line is a comment line"""
    if re.match(comment_pattern, line):
        return True
    else:
        return False

def has_comp_name(line):
    """Check whether this line contains component name"""
    if not is_comment(line):
        return False    
    if re.match(comp_pattern, line):
        return True
    else:
        return False

def get_comp_name(line, new_comp_count=0):
    """Get the component name"""
    line_ele = line.split()
    if line_ele[-1].lower() == "component:":
        return f"newcomp{new_comp_count+1}"
    else:
        return line_ele[-1]
    
def has_cell_range(line):
    """Check whether this line contains cell range info"""
    if not is_comment(line):
        return False
    if re.match(cell_range_pattern, line):
        return True
    else:
        return False
    
def get_cell_range(line):
    """Get the cell range"""
    line_ele = line.split()
    cids = [range(int(line_ele[-3]), int(line_ele[-1]) +1)]
    return cids

def has_group_name(line):
    """Check whether this line cntains group name"""
    if re.match(group_name_pattern, line):
        return True
    else:
        return False
    
def get_group_name(line, new_group_count=0):
    line_ele = line.split()
    if line_ele[-1].lower() == "group:":
        return f"newcomp{new_group_count+1}"
    else:
        return line_ele[-1]

def has_from_to(line):
    if re.match(from_to_pattern, line):
        return True
    else:
        return False
    
def get_nonvoid_cells(inp="input"):
    """
    Read the MCNP input file generated by cosVMPT and output a list of all
    non-void cells"
    """
    nonvoid_cells = []
    with open(inp, 'r') as fin:
        cell_card_end = False
        while not cell_card_end:
            line = fin.readline()
            if is_blank_line(line): # end of cell card
                cell_card_end = True
                break
            if is_cell_title(line):
                line_ele = line.split()
                mat_id = int(line_ele[1])
                if mat_id > 0:
                    nonvoid_cells.append(int(line_ele[0]))
    return nonvoid_cells

def mcnp_tally_style(cids, sds=None, output="tally_card.txt"):
    """Convert the cell number list to a mcnp style tally description."""
    tally_card = f'C tallied cell numbers: {len(cids)}'
    tally_card = f'{tally_card}\nFC4 tally card generated via natf:mcnp_tally_style'
    tally_card = f'{tally_card}\nF4:n  '
    for i, cid in enumerate(cids):
        tally_card = mcnp_style_str_append(tally_card, str(cid))
    if sds is not None:
        tally_card = f"{tally_card}\nSD4   "
        for i, sd in enumerate(sds):
            tally_card = mcnp_style_str_append(tally_card, str(sd))
    with open(output, 'w') as fo:
        fo.write(tally_card+'\n')

def cell_vol_err_reader(inp="input"):
    """Read the cell-vol-err info file"""
    cids, vols, errs = [], [], []
    # Opening file and read data
    fin = open(inp, 'r')
    for line in fin:
        if is_blank_line(line):
            break
        line_ele = line.split()
        cids.append(int(line_ele[0]))
        vols.append(float(line_ele[1]))
        errs.append(float(line_ele[2]))
    fin.close()
    return cids, vols, errs
 

def cell_vol_to_tally(inp="input", output="tally_card.txt"):
    """
    Write the cell, vol, err info to tally card.
    """
    cell_vol_to_tally_help = ('This script read a cell-vol-err info file and\n'
                              'return a tally style string.\n')

    parser = argparse.ArgumentParser()
    parser.add_argument("-i", "--input", required=False, help="cell-vol-err in file path")
    parser.add_argument("-o", "--output", required=False, help="save the tally_card to output file")
    args = vars(parser.parse_args())

    input_file = "input"
    if args['input'] is not None:
        input_file = args['input']
    cids, vols, errs = cell_vol_err_reader(input_file)
    # save data into a tally style card
    output = "tally_card.txt"
    if args['output'] is not None:
        output = args['output']
    mcnp_tally_style(cids, sds=vols, output=output)


def nonvoid_cells_to_tally():
    nonvoid_cells_help = ('This script read a mcnp input file and return a tally style\n'
                'string of all non-void cells.\n')
    parser = argparse.ArgumentParser()
    parser.add_argument("-i", "--input", required=False, help="mcnp input file path")
    parser.add_argument("-o", "--output", required=False, help="save the string to output file")
    args = vars(parser.parse_args())

    input_file = "input"
    if args['input'] is not None:
        input_file = args['input']
    nonvoid_cells = get_nonvoid_cells(input_file)

    output = "output.txt"
    if args['output'] is not None:
        output = args['output']
    mcnp_tally_style(nonvoid_cells, output=output)

def get_part_cell_list(inp):
    """Read the components and groups input file generated by cosVMPT."""
    parts = collections.OrderedDict()
    #comps = collections.OrderedDict()  # {'name':[{group1:cids}, {group2:cids}, ...]}
    #group = collections.OrderedDict() # {'name':[cids]}
    current_comp, current_group, current_cell = None, None, None
    new_comp_count, new_group_count = 0, 0

    # read the file
    with open(inp, 'r') as fin:
        cell_card_end = False
        while not cell_card_end:
            line = fin.readline()
            if is_blank_line(line): # end of cell card
                cell_card_end = True
                break
            if has_comp_name(line):
                current_comp = get_comp_name(line, new_comp_count=new_comp_count)
                if re.match(new_comp_name_pattern, current_comp):
                    new_comp_count += 1
                #comps[current_comp] = []
                parts[current_comp] = []
                continue
            if has_group_name(line):
                current_group = get_group_name(line, new_group_count)
                if re.match(new_group_name_pattern, current_group):
                    new_group_count += 1
                #group[current_group] = []
                current_part = f"{current_comp}-{current_group}"
                parts[current_part] = []
                #comps[current_comp][current_group] = []
                continue
            if is_cell_title(line):
                line_ele = line.split()
                mat_id = int(line_ele[1])
                if mat_id > 0: # nonvoid cell
                #    comps[current_comp].append(int(line_ele[0]))
                    parts[current_comp].append(int(line_ele[0]))
                    parts[current_part].append(int(line_ele[0]))
    return parts

def format_cell_list(name, cids):
    """Format the cell list to part_cell_list style"""
    cnt = name
    for i, cid in enumerate(cids):
        cnt = mcnp_style_str_append(cnt, str(cid))
    return cnt

def save_part_cell_list(parts, output):
    """Write the components and groups cells in format of part_cell_list"""
    cnt = f"C Be careful about the content, do not enter any characters except ASCII."
    cnt = f"{cnt}\nC Lines start with `C ` is comment"
    cnt = f"{cnt}\nC Syntax of defing parts: part_name cell_list"
    cnt = f"{cnt}\nC Refer to 'NATF manual` for more details about defining parts."
    cnt = f"{cnt}\nC Generated by natf:part_cell_list"
    for key, value in parts.items():
        if len(value) > 0: # void, do not print
            cnt = f"{cnt}\n{format_cell_list(key, value)}"
    with open(output, 'w') as fo:
        fo.write(cnt)
 
def part_cell_list(inp="input", output="part_cell_list.txt"):
    """
    Read the components and groups input file generated by cosVMPT, and
    write a part_cell_list file for NATF.
    """
    list_comp_group_help = ('This script read a mcnp input file and return a\n'
                            'part_cell_list file.\n')
    parser = argparse.ArgumentParser()
    parser.add_argument("-i", "--input", required=False, help="mcnp input file path, default: input")
    parser.add_argument("-o", "--output", required=False, help="name of the part_cell_list file, default: part_cell_list.txt")
    args = vars(parser.parse_args())

    # read the file
    if args['input'] is not None:
        inp = args['input']
    parts = get_part_cell_list(inp)

    # save the content
    if args['output'] is not None:
        output = args['output']
    save_part_cell_list(parts, output=output)


