import json
import os
import re
import copy
import pandas as pd

## 半衰期单位转换字典
HL_UNITS = {"fs": 1e-15, "ps": 1e-12, "ns": 1e-9, "us": 1e-6, "ms": 1e-3, "s": 1, "m": 60, "h": 3600, "d": 86400, "y": 31557600, "ky": 31557600e3, "My": 31557600e6, "Gy": 31557600e9}

def nuclidesFilterZNA(nuclides_data, Z_min=None, Z_max=None, Z_oe_idx=0, N_min=None, N_max=None, N_oe_idx=0, A_min=None, A_max=None, A_oe_idx=0):
    filtered = {}
    for name, data in nuclides_data.items():
        fh = True
        if not (Z_min is None or data['z'] >= Z_min):
            fh = False
        if not (Z_max is None or data['z'] <= Z_max):
            fh = False
        if not (N_min is None or data['n'] >= N_min):
            fh = False
        if not (N_max is None or data['n'] <= N_max):
            fh = False
        if not (A_min is None or data['a'] >= A_min):
            fh = False
        if not (A_max is None or data['a'] <= A_max):
            fh = False
        
        if Z_oe_idx == 1:
            if data['z'] % 2 == 0:
                fh = False
        elif Z_oe_idx == 2:
            if data['z'] % 2 == 1:
                fh = False

        if N_oe_idx == 1:
            if data['n'] % 2 == 0:
                fh = False
        elif N_oe_idx == 2:
            if data['n'] % 2 == 1:
                fh = False

        if A_oe_idx == 1:
            if data['a'] % 2 == 0:
                fh = False
        elif A_oe_idx == 2:
            if data['a'] % 2 == 1:
                fh = False
        
        if fh:
            filtered[name] = data

    return filtered
        
def nuclidesFilterHalflife(nuclides_data, hl_min_sec=0, hl_max_sec=None):
    filtered = {}
    if hl_min_sec == None:
        if hl_max_sec == None:
            for name, data in nuclides_data.items():
                if "levels" in data:
                    for level in data["levels"]:
                        if "halflife" in level:
                            if level["halflife"]["value"] == "STABLE":
                                filtered[name] = data
                                break
    else:
        if hl_max_sec == None:
            for name, data in nuclides_data.items():
                if "levels" in data:
                    for level in data["levels"]:
                        if "halflife" in level:
                            if level["halflife"]["value"] == "STABLE":
                                filtered[name] = data
                                break
                            else:
                                if not level["halflife"]["unit"] in HL_UNITS:
                                    break
                                hl_sec = level["halflife"]["value"] * HL_UNITS[level["halflife"]["unit"]]
                                if hl_sec > hl_min_sec:
                                    filtered[name] = data
                                    break
        else:
            for name, data in nuclides_data.items():
                if "levels" in data:
                    for level in data["levels"]:
                        if "halflife" in level:
                            if level["halflife"]["value"] == "STABLE":
                                pass
                            else:
                                if not level["halflife"]["unit"] in HL_UNITS:
                                    break
                                hl_sec = level["halflife"]["value"] * HL_UNITS[level["halflife"]["unit"]]
                                if hl_sec > hl_min_sec and hl_sec < hl_max_sec:
                                    filtered[name] = data
                                    break

    return filtered

def nuclidesFilterDecayModes(nuclides_data, dm_enable_idx, decay_modes):
    filtered = {}
    nuclides_decayModes = {}
    for name, data in nuclides_data.items():
        decayModes = []
        if "levels" in data:
            for level in data["levels"]:
                if "decayModes" in level:
                    for decayData in level["decayModes"]["observed"]:
                        if not decayData["mode"] in decayModes:
                            decayModes.append(decayData["mode"])
        nuclides_decayModes[name] = copy.deepcopy(decayModes)

    decay_modes_series = pd.Series(decay_modes)

    ## nndc导出的数据里有两处写作"β⁻"了
    replace_dict = {"β⁻": "B-"}

    for name, decayModes in nuclides_decayModes.items():
        nuclide_decayModes_series = pd.Series(decayModes)
        nuclide_decayModes_series = nuclide_decayModes_series.replace(replace_dict)

        if dm_enable_idx == 1:
            if decay_modes_series.isin(nuclide_decayModes_series).all():
                filtered[name] = nuclides_data[name]
        elif dm_enable_idx == 2:
            if decay_modes_series.isin(nuclide_decayModes_series).any():
                filtered[name] = nuclides_data[name]

    return filtered

def nuclidesSearchingNom(nuclides_data, nuclide_nom):
    nuclide = None
    element = None
    nuclide_A = None
    if not re.fullmatch(rf"([A-Za-z]+)([-_|]*)([0-9]+)", nuclide_nom) == None:
        match00 = re.fullmatch(rf"([A-Za-z]+)([-_|]*)([0-9]+)", nuclide_nom)
        element = match00.group(1)
        nuclide_A = match00.group(3)
    elif not re.fullmatch(rf"([0-9]+)([-_|]*)([A-Za-z]+)", nuclide_nom) == None:
        match00 = re.fullmatch(rf"([0-9]+)([-_|]*)([A-Za-z]+)", nuclide_nom)
        element = match00.group(3)
        nuclide_A = match00.group(1)

    if not (element == None or nuclide_A == None):
        if not element in ("n", "N"):
            if len(element) == 1:
                element = element.upper()
            else:
                element = element[:1].upper() + element[1:].lower()
        while nuclide_A.startswith("0"):
            nuclide_A = nuclide_A[1:]
        nom = nuclide_A + element
        if nom in nuclides_data:
            nuclide = nuclides_data[nom]
    
    return nuclide

def nuclidesSearchingZN(nuclides_data, Z_in=0, N_in=0):
    nuclide = None
    for nom, data in nuclides_data.items():
        if data["z"] == Z_in and data["n"] == N_in:
            nuclide = data
            break

    return nuclide

def nuclidesSearchingZA(nuclides_data, Z_in=0, A_in=0):
    nuclide = None
    for nom, data in nuclides_data.items():
        if data["z"] == Z_in and data["a"] == A_in:
            nuclide = data
            break

    return nuclide

def nuclidesSearchingNA(nuclides_data, N_in=0, A_in=0):
    nuclide = None
    for nom, data in nuclides_data.items():
        if data["n"] == N_in and data["a"] == A_in:
            nuclide = data
            break

    return nuclide

def nuclideData_dict2dataframe(nuclideData_dict):
    nom = nuclideData_dict["name"]
    z = nuclideData_dict["z"]
    n = nuclideData_dict["n"]
    a = nuclideData_dict["a"]
    data = []
    for level in nuclideData_dict["levels"]:
        if level["halflife"]["value"] == "STABLE":
            data.append((nom, z, n, a, level["energy"]["value"], level["energy"]["unit"], level["spinParity"], level["massExcess"]["value"], level["massExcess"]["unit"], level["massExcess"]["uncertainty"], None, None, None, None, None))
        else:
            hl = level["halflife"]["value"]
            hl_unit = level["halflife"]["unit"]
            hl_unc = level["halflife"]["uncertainty"]["value"]
            for decayMode in level["decayModes"]["observed"]:
                data.append((nom, z, n, a, level["energy"]["value"], level["energy"]["unit"], level["spinParity"], level["massExcess"]["value"], level["massExcess"]["unit"], level["massExcess"]["uncertainty"], hl, hl_unit, hl_unc, decayMode["mode"], decayMode["value"]))

    nuclideDataframe = pd.DataFrame(data, columns=["Nuclide", "Z", "N", "A", "E(level)", "E(level) unit", "Spin Parity", "Mass Excess", "Mass Excess unit", "Mass Excess uncertainty", "Halflife", "Halflife unit", "Halflife uncertainty", "Decay Mode", "Branch Ratio"])

    return nuclideDataframe

def nuclideData_dict2dataframeCompact(nuclideData_dict):
    data = []
    for level in nuclideData_dict["levels"]:
        if level["halflife"]["value"] == "STABLE":
            data.append((str(level["energy"]["value"])+" "+level["energy"]["unit"], level["spinParity"], level["massExcess"]["formats"]["NDS"]+" "+level["massExcess"]["unit"], "STABLE", None))
        else:
            decayModes = []
            for decayMode in level["decayModes"]["observed"]:
                decayModes.append((decayMode["mode"], decayMode["value"]))
            decayModeDF = pd.DataFrame(decayModes, columns=["Decay Mode", "Branch Ratio"])
            data.append((str(level["energy"]["value"])+" "+level["energy"]["unit"], level["spinParity"], level["massExcess"]["formats"]["NDS"]+" "+level["massExcess"]["unit"], level["halflife"]["formats"]["NDS"]+" "+level["halflife"]["unit"], decayModeDF))

    nuclideDataframe = pd.DataFrame(data, columns=["E(level)", "Spin Parity", "Mass Excess", "Halflife", "Decay Modes"])

    return nuclideDataframe

###
#nuclidesFilter(1)