'''
Created on 12-Jul-2021

@author: rishi_priyansh

# pip install radioactivedecay
# https://radioactivedecay.github.io/overview.html#introduction

'''
import random
import inspect
import traceback

import networkx as nx
from collections import deque
import numpy as np
from radioactivedecay import Nuclide

from xv.util import UtilClass

class Radioactivity_func_copied(UtilClass):
    
    
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        pass
    
    @classmethod
    def _parse_nuclide_label(cls, nuclide: str) -> str:
        """
        Format a nuclide string to mass number, meta-stable state character in
        superscript, then element symbol. Output is used on node labels in decay
        chain plots.
    
        Parameters
        ----------
        nuclide : str
            Nuclide string in element-mass format.
    
        Returns
        -------
        str
            Parsed string for node label in ^{mass}element format.
    
        """
    
        if nuclide == "SF":
            return "various"
    
        nuclide_conversion = {
            "0": "\N{SUPERSCRIPT ZERO}",
            "1": "\N{SUPERSCRIPT ONE}",
            "2": "\N{SUPERSCRIPT TWO}",
            "3": "\N{SUPERSCRIPT THREE}",
            "4": "\N{SUPERSCRIPT FOUR}",
            "5": "\N{SUPERSCRIPT FIVE}",
            "6": "\N{SUPERSCRIPT SIX}",
            "7": "\N{SUPERSCRIPT SEVEN}",
            "8": "\N{SUPERSCRIPT EIGHT}",
            "9": "\N{SUPERSCRIPT NINE}",
            "m": "\N{MODIFIER LETTER SMALL M}",
            "n": "\N{SUPERSCRIPT LATIN SMALL LETTER N}",
            "p": "\N{MODIFIER LETTER SMALL P}",
            "q": "\N{LATIN SMALL LETTER Q}",  # Unicode has no superscript q
            "r": "\N{MODIFIER LETTER SMALL R}",
            "x": "\N{MODIFIER LETTER SMALL X}",
        }
    
        element, isotope = nuclide.split("-")
        return "".join(map(lambda char: nuclide_conversion[char], list(isotope))) + element
    
    @classmethod
    def _parse_decay_mode_label(cls, mode: str) -> str:
        """
        Format a decay mode string for edge label on decay chain plot.
    
        Parameters
        ----------
        mode : str
            Decay mode string.
    
        Returns
        -------
        str
            Formatted decay mode string for use in an edge label.
    
        """
    
        mode_conversion = {
            "α": "\N{GREEK SMALL LETTER ALPHA}",
            "β": "\N{GREEK SMALL LETTER BETA}",
            "ε": "\N{GREEK SMALL LETTER EPSILON}",
            "+": "\N{SUPERSCRIPT PLUS SIGN}",
            "-": "\N{SUPERSCRIPT MINUS}",
            "12C": "\N{SUPERSCRIPT ONE}\N{SUPERSCRIPT TWO}C",
            "14C": "\N{SUPERSCRIPT ONE}\N{SUPERSCRIPT FOUR}C",
            "20O": "\N{SUPERSCRIPT TWO}\N{SUPERSCRIPT ZERO}O",
            "23F": "\N{SUPERSCRIPT TWO}\N{SUPERSCRIPT THREE}F",
            "22Ne": "\N{SUPERSCRIPT TWO}\N{SUPERSCRIPT TWO}Ne",
            "24Ne": "\N{SUPERSCRIPT TWO}\N{SUPERSCRIPT FOUR}Ne",
            "25Ne": "\N{SUPERSCRIPT TWO}\N{SUPERSCRIPT FIVE}Ne",
            "26Ne": "\N{SUPERSCRIPT TWO}\N{SUPERSCRIPT SIX}Ne",
            "28Mg": "\N{SUPERSCRIPT TWO}\N{SUPERSCRIPT EIGHT}Mg",
            "29Mg": "\N{SUPERSCRIPT TWO}\N{SUPERSCRIPT NINE}Mg",
            "30Mg": "\N{SUPERSCRIPT THREE}\N{SUPERSCRIPT ZERO}Mg",
            "32Si": "\N{SUPERSCRIPT THREE}\N{SUPERSCRIPT TWO}Si",
            "34Si": "\N{SUPERSCRIPT THREE}\N{SUPERSCRIPT FOUR}Si",
        }
    
        for unformatted, formatted in mode_conversion.items():
            mode = mode.replace(unformatted, formatted)
        return mode
    
    @classmethod
    def _build_decay_digraph(cls, 
        parent: Nuclide,
        digraph: nx.classes.digraph.DiGraph,
    ) -> nx.classes.digraph.DiGraph:
        """
        Build a networkx DiGraph for the decay chain of this nuclide.
    
        Parameters
        ----------
        parent : Nuclide
            Nuclide instance of the parent nuclide of the decay chain.
        digraph : networkx.classes.digraph.DiGraph
            DiGraph for the decay chain.
    
        Returns
        -------
        digraph : networkx.classes.digraph.DiGraph
            DiGraph of the decay chain.
        max_generation : int
            Number of generations of progeny in the decay chain.
        max_xpos : int
            Maximum number of progeny within any one generation of the decay chain.
    
        """
        
        generation_max_xpos = {0: 0}
    
        dequeue = deque([parent.nuclide])
        generations = deque([0])
        xpositions = deque([0])
        node_label = (
            f"{cls._parse_nuclide_label(parent.nuclide)}\n{parent.half_life('readable')}"
        )
        digraph.add_node(parent.nuclide, generation=0, xpos=0, label=node_label)
        seen = {parent.nuclide}
    
        while len(dequeue) > 0:
            parent_name = dequeue.popleft()
            generation = generations.popleft() + 1
            xpos = xpositions.popleft()
            if generation not in generation_max_xpos:
                generation_max_xpos[generation] = -1
            parent = Nuclide(parent_name, parent.decay_data)
    
            progeny = parent.progeny()
            branching_fractions = parent.branching_fractions()
            decay_modes = parent.decay_modes()
    
            if xpos < generation_max_xpos[generation] + 1:
                xpos = generation_max_xpos[generation] + 1
            xcounter = 0
            for idx, prog in enumerate(progeny):
                if prog not in seen:
                    node_label = cls._parse_nuclide_label(prog)
                    if prog in parent.decay_data.nuclide_dict:
                        node_label += f"\n{parent.decay_data.half_life(prog, 'readable')}"
                        if np.isfinite(parent.decay_data.half_life(prog)):
                            dequeue.append(prog)
                            generations.append(generation)
                            xpositions.append(xpos + xcounter)
                    if prog == "SF":
                        prog = parent.nuclide + "_SF"
    
                    digraph.add_node(
                        prog,
                        generation=generation,
                        xpos=xpos + xcounter,
                        label=node_label,
                    )
                    seen.add(prog)
    
                    if xpos + xcounter > generation_max_xpos[generation]:
                        generation_max_xpos[generation] = xpos + xcounter
                    xcounter += 1
    
                edge_label = (
                    cls._parse_decay_mode_label(decay_modes[idx])
                    + "\n"
                    + str(branching_fractions[idx])
                )
                digraph.add_edge(parent.nuclide, prog, label=edge_label)
    
        for node in digraph:
            digraph.nodes[node]["pos"] = (
                digraph.nodes[node]["xpos"],
                digraph.nodes[node]["generation"] * -1,
            )
    
        return digraph, max(generation_max_xpos), max(generation_max_xpos.values())
    pass
