import re
import streamlit as st
import torch
from transformers import PreTrainedTokenizerFast, BartForConditionalGeneration

# Dictionary for SU(3)/SU(2) latex representations
rep_tex_dict = {
    "SU3": {"-3": r"\bar{\textbf{3}}", "3": r"\textbf{3}"},
    "SU2": {"-2": r"\textbf{2}", "2": r"\textbf{2}", "-3": r"\textbf{3}", "3": r"\textbf{3}"},
}

def fieldobj_to_tex(obj, lor_index, pos):
    su3 = None
    su2 = None
    u1  = None
    hel = None
    sp  = None
    
    obj_mod = obj.copy()
    for tok in obj:
        if "SU3" in tok:
            su3 = tok.split("=")[-1]
            obj_mod.remove(tok)
        if "SU2" in tok:
            su2 = tok.split("=")[-1]
            obj_mod.remove(tok)
        if "U1" in tok:
            u1 = tok.split("=")[-1]
            obj_mod.remove(tok)
        if "HELICITY" in tok:
            hel = tok.split("=")[-1]
            if hel == "1":
                hel = "+1"
        if "SPIN" in tok:
            sp = tok.split("=")[-1]
    assert sp is not None

    outtex = ""
    if sp == "0": 
        outtex += r"\phi"
    if sp == "1": 
        outtex += "A" + pos + lor_index
    if sp == "1/2":
        outtex += r"\psi"

    outtex += r"_{("
    # SU(3)
    if su3 is not None:
        outtex += rep_tex_dict["SU3"].get(su3, r"\textbf{1}") + " ,"
    else:
        outtex += r"\textbf{1},"
    # SU(2)
    if su2 is not None:
        outtex += rep_tex_dict["SU2"].get(su2, r"\textbf{1}") + " ,"
    else:
        outtex += r"\textbf{1},"
    # U(1)
    if u1 is not None:
        outtex += u1 + " ,"
    else:
        outtex += r"\textbf{0},"
    # Helicity
    if hel is not None:
        outtex += "h:" + hel + " ,"
    # Finish out subscript
    if outtex[-1] == ",":
        outtex = outtex[:-1] + ")}"
    return outtex

def derobj_to_tex(obj, lor_index, pos):
    if pos == "^":
        outtex = f"D^{{{lor_index}}}_{{("
    elif pos == "_":
        outtex = f"D_{{{lor_index}}}^{{("
    else:
        raise ValueError("pos must be ^ or _")

    if "SU3" not in obj and "SU2" not in obj and "U1" not in obj:
        # Just partial derivative
        if pos == "^":
            return f"\\partial^{lor_index}"
        else:
            return f"\\partial_{lor_index}"

    if "SU3" in obj:
        outtex += "SU3,"
    if "SU2" in obj:
        outtex += "SU2,"
    if "U1" in obj:
        outtex += "U1,"
    if outtex[-1] == ",":
        outtex = outtex[:-1] + ")}"
    return outtex

def gamobj_to_tex(obj, lor_index, pos):
    return r"\sigma" + pos + lor_index

def obj_to_tex(obj, lor_index="\mu", pos="^"):
    # Convert tuple/strings to a list of tokens
    if isinstance(obj, tuple):
        obj = list(obj)
    if isinstance(obj, str):
        obj = [i for i in obj.split(" ") if i != ""]

    # Basic tokens
    if obj[0] == "+":
        return r"\quad\quad+"
    if obj[0] == "-":
        return r"\quad\quad-"
    if obj[0] == "i":
        return "i"

    # Field
    if obj[0] == "FIELD":
        return fieldobj_to_tex(obj, lor_index, pos)
    # Derivative
    if obj[0] == "DERIVATIVE":
        return derobj_to_tex(obj, lor_index, pos)
    # Sigma (gamma matrices)
    if obj[0] == "SIGMA":
        return gamobj_to_tex(obj, lor_index, pos)

    # Combined COMMUTATOR + DERIVATIVE tokens
    if obj[0] == "COMMUTATOR_ADERIVATIVE":
        new_obj = obj[:]
        new_obj[0] = "DERIVATIVE"
        return "[ " + derobj_to_tex(new_obj, lor_index, pos)
    if obj[0] == "COMMUTATOR_BDERIVATIVE":
        new_obj = obj[:]
        new_obj[0] = "DERIVATIVE"
        return ", " + derobj_to_tex(new_obj, lor_index, pos) + " ]"

    # Single COMMUTATOR tokens
    if obj[0] == "COMMUTATOR_A":
        return "[ " + derobj_to_tex(obj, lor_index, pos)
    if obj[0] == "COMMUTATOR_B":
        return ", " + derobj_to_tex(obj, lor_index, pos) + " ]"

    # Fallback for unrecognized tokens if you like:
    # return f"\\text{{Unhandled}}({obj})"
    return ""

def split_with_delimiter_preserved(string, delimiters, ignore_dots=False):
    """
    Splits a string using the given delimiters, 
    while preserving them as separate tokens.
    """
    if "." in string and not ignore_dots:
        raise ValueError("Unexpected ending to the generated Lagrangian")
    pattern = '(' + '|'.join(map(re.escape, delimiters)) + ')'
    parts = re.split(pattern, string)
    # Turn a lonely "+ " into " + "
    parts = [" + " if p == "+ " else p for p in parts]
    # Remove empty entries
    parts = [p for p in parts if p != ""]
    return parts

def clean_split(inlist, delimiters):
    """
    Merges an immediate delimiter with its next token 
    so that "FIELD " + "SPIN" -> "FIELD SPIN".
    """
    i = 0
    merged_list = []
    while i < len(inlist):
        if inlist[i] in delimiters:
            if i < len(inlist) - 1:
                merged_list.append(inlist[i] + inlist[i+1])
                i += 1
            else:
                merged_list.append(inlist[i])
        else:
            merged_list.append(inlist[i])
        i += 1
    return merged_list

def get_obj_dict(inlist):
    outdict = {}
    for iitem in inlist:
        idict = {"ID": None, "LATEX": None}
        # Find any ID=... string
        item_parts = iitem.split()
        the_ids = [x for x in item_parts if x.startswith("ID")]
        if the_ids:
            idict["ID"] = the_ids[0]
        # Always compute LATEX from obj_to_tex
        idict["LATEX"] = obj_to_tex(iitem, "\\mu", "^")
        outdict[iitem] = idict
    return outdict

def get_con_dict(inlist):
    """
    For a list of 'contractions' tokens, produce 
    a dictionary of which IDs are to be contracted 
    under LORENTZ, SU2, or SU3.
    """
    outdict = {}
    for iitem in inlist:
        tokens = iitem.split()
        tokens = [t for t in tokens if t != ""]
        sym = [t for t in tokens if ("SU" in t or "LORENTZ" in t)]
        assert len(sym) == 1, "More than one symmetry in contraction"
        ids = [t for t in tokens if ("SU" not in t and "LZ" not in t)]
        if sym[0] not in outdict:
            outdict[sym[0]] = [ids]
        else:
            outdict[sym[0]].append(ids)
    return outdict

def term_to_tex(term, verbose=True):
    """
    Converts one Lagrangian term into its LaTeX representation.
    """
    # Clean up certain strings
    term = term.replace(".", "").replace(" = ", "=").replace(" =- ", "=-")
    term = term.replace(" / ", "/")
    term = term.replace("COMMUTATOR_A DERIVATIVE", "COMMUTATOR_ADERIVATIVE")
    term = term.replace("COMMUTATOR_B DERIVATIVE", "COMMUTATOR_BDERIVATIVE")

    # Split into sub-tokens
    term = split_with_delimiter_preserved(
        term,
        [" FIELD ", " DERIVATIVE ", " SIGMA ", " COMMUTATOR_ADERIVATIVE ", " COMMUTATOR_BDERIVATIVE ", " CONTRACTIONS "]
    )
    term = clean_split(
        term,
        [" FIELD ", " DERIVATIVE ", " SIGMA ", " COMMUTATOR_ADERIVATIVE ", " COMMUTATOR_BDERIVATIVE ", " CONTRACTIONS "]
    )

    if verbose:
        print(term)

    # If it's just +, -, or i, return that token
    if term in [[" + "], [" - "], [" i "]]:
        return term[0]

    # Build dictionary for objects that aren't in "CONTRACTIONS"
    objdict = get_obj_dict([t for t in term if " CONTRACTIONS " not in t])
    if verbose:
        for k, v in objdict.items():
            print(k, "\t\t", v)
    
    # Contractions
    contractions = [t for t in term if " CONTRACTIONS " in t]
    if len(contractions) > 1:
        raise ValueError("More than one contraction in term")

    if len(contractions) == 1 and contractions != [" CONTRACTIONS "]:
        # e.g. "LORENTZ ID5 ID2", etc.
        c_str = contractions[0]
        c_str = split_with_delimiter_preserved(c_str, [" LORENTZ ", " SU2 ", " SU3 "])
        c_str = clean_split(c_str, [" LORENTZ ", " SU2 ", " SU3 "])
        c_str = [i for i in c_str if i != " CONTRACTIONS"]
        condict = get_con_dict(c_str)
        if verbose:
            print(condict)
        
        # LORENTZ contraction handling
        if "LORENTZ" in condict:
            firstlz = True
            cma = True
            for con in condict["LORENTZ"]:
                for kobj, iobj in objdict.items():
                    if iobj["ID"] is None:
                        continue
                    if iobj["ID"] in con:
                        if cma:
                            lsymb = r"\mu"
                        else:
                            lsymb = r"\nu"
                        if firstlz:
                            iobj["LATEX"] = obj_to_tex(kobj, lsymb, "^")
                            firstlz = False
                        else:
                            iobj["LATEX"] = obj_to_tex(kobj, lsymb, "_")
                            cma = False
                            firstlz = True

    # Join the final LaTeX strings
    outstr = " ".join([objdict[t]["LATEX"] for t in term if " CONTRACTIONS " not in t])
    return outstr

def str_tex(instr, num=0):
    """
    Convert list of terms into complete LaTeX lines for the Lagrangian.
    """
    if num != 0:
        instr = instr[:num]

    inlist = [term.replace(".", "") for term in instr]
    outstr = ""
    coup = 0
    mass = 0
    outstr = r"\begin{aligned}"
    for i, iterm in enumerate(inlist):
        if i == 0:
            outstr += r" \mathcal{L}= \quad \\ & "
        else:
            # Identify coupling or mass terms by counting spin-0 fields
            nqf = iterm.count("FIELD SPIN = 0")
            nD = iterm.count(" DERIVATIVE ")
            if nqf != 0 and nqf != 2 and nD == 0:
                coup += 1
                outstr += rf" \lambda_{{{coup}}} \,"
            if nqf == 2 and nD == 0:
                mass += 1
                outstr += rf" m^2_{{{mass}}} \,"
            outstr += term_to_tex(iterm, False) + r" \quad "
            if i % 4 == 0:
                outstr += r" \\ \\ & "
    return outstr

def master_str_tex(iinstr):
    """
    Master function that splits the incoming string, 
    tries to render the full Lagrangian, 
    and catches errors if the model text is truncated.
    """
    instr = split_with_delimiter_preserved(iinstr, [" + ", "+ ", " - "])
    try:
        outstr = str_tex(instr)
    except Exception as e:
        # If an error occurs, try ignoring the last token
        outstr = str_tex(instr, -1)
        outstr += "  \\cdots"
        print(e)
    outstr += r"\end{aligned}"
    return outstr

# ---------------------------------------------------------------------------------
# Model loading
device = 'cpu'
model_name = "JoseEliel/BART-Lagrangian"

@st.cache_resource
def load_model():
    model = BartForConditionalGeneration.from_pretrained(model_name).to(device)
    return model

@st.cache_resource
def load_tokenizer():
    return PreTrainedTokenizerFast.from_pretrained(model_name)

model = load_model()
hf_tokenizer = load_tokenizer()

# ---------------------------------------------------------------------------------
# Text processing wrappers
def process_input(input_text):
    # Sort fields so generation is consistent
    input_text = input_text.replace("[SOS]", "").replace("[EOS]", "").replace("FIELD", "SPLITFIELD")
    fields = input_text.split("SPLIT")[1:]
    fields = [x.strip().split(" ") for x in fields]
    fields = sorted(fields)
    fields = "[SOS] " + " ".join([" ".join(x) for x in fields]) + " [EOS]"
    return fields

def process_output(output_text):
    # Remove special tokens from model output
    return output_text.replace("[SOS]", "").replace("[EOS]", "").replace(".", "")

def reformat_expression(s):
    # e.g. turn SU2= -1 into SU2=-1, remove spaces
    return re.sub(r"(SU[23]|U1|SPIN|HEL)\s+([+-]?\s*\d+)", 
                  lambda m: f"{m.group(1)} = {m.group(2).replace(' ', '')}", 
                  s)

def generate_lagrangian(input_text):
    """
    Calls the model to produce a Lagrangian for the user-given fields.
    """
    input_text = process_input(input_text)
    inputs = hf_tokenizer([input_text], return_tensors='pt').to(device)
    with st.spinner("Generating Lagrangian..."):
        lagrangian_ids = model.generate(inputs['input_ids'], max_length=2048)
        lagrangian = hf_tokenizer.decode(lagrangian_ids[0].tolist(), skip_special_tokens=False)
        lagrangian = process_output(lagrangian)
    return lagrangian

def generate_field(sp, su2, su3, u1):
    """
    Builds a single field string with the chosen spin and gauge charges.
    """
    components = [f"FIELD SPIN={sp}"]
    # For spin = 1/2, optionally add helicity
    if sp == "1/2":
        components = [f"FIELD SPIN={sp} HEL=1/2"]

    if su2 != "$1$":
        components.append(f"SU2={su2}")
    if su3 == "$\\bar{{3}}$":
        components.append("SU3=-3")
    elif su3 != "$1$":
        components.append(f"SU3={su3.replace('$','')}")
    if u1 != "0":
        components.append(f"U1={u1}")
    return " ".join(components).replace("$", "")

# ---------------------------------------------------------------------------------
# Streamlit GUI
def main():
    st.title("$\\mathscr{L}$agrangian Generator")
    st.markdown(" ### For a set of chosen fields, this model generates the corresponding Lagrangian which encodes all interactions and dynamics of the fields.")
    
    st.markdown(" #### This is a demo of our [BART](https://arxiv.org/abs/1910.13461)-based model with ca 360M parameters")
    st.markdown(" #### Details about the model, training, and evaluation can be found in our [paper](https://arxiv.org/abs/2501.09729).")

    st.markdown(" ##### :violet[Due to computational resources, we limit the number of fields to 3. Some features in the LaTeX rendering (such as daggers) are not supported in the current version and helicity is always 1/2 (to be updated).]")
    st.markdown(" ##### Choose up to three different fields:")

    su2_options = ["$1$", "$2$", "$3$"]
    su3_options = ["$1$", "$3$", "$\\bar{3}$"]
    u1_options = ["-1", "-2/3", "-1/2", "-1/3", "0", "1/3", "1/2", "2/3", "1"]
    spin_options = ["0", "1/2"]

    if "count" not in st.session_state:
        st.session_state.count = 0
    if "field_strings" not in st.session_state:
        st.session_state.field_strings = []

    with st.form("field_selection"):
        spin_selection = st.radio("Select spin value:", spin_options)
        su2_selection = st.radio("Select SU(2) value:", su2_options)
        su3_selection = st.radio("Select SU(3) value:", su3_options)
        u1_selection  = st.radio("Select U(1) value:", u1_options)
        submitted = st.form_submit_button("Add field")
        if submitted:
            if st.session_state.count < 3:
                fs = generate_field(spin_selection, su2_selection, su3_selection, u1_selection)
                st.session_state.field_strings.append(fs)
                st.session_state.count += 1
            else:
                st.write("Maximum of 3 fields for this demo.")

    clear_fields = st.button("Clear fields")
    if clear_fields:
        st.session_state.field_strings = []
        st.session_state.count = 0

    # Display current fields
    st.write("Input Fields:")
    for i, fs in enumerate(st.session_state.field_strings, 1):
        texfield = obj_to_tex(fs)
        st.latex(r"\text{Field " + str(i) + ":} \quad " + texfield)

    # Generate Lagrangian button
    if st.button("Generate Lagrangian"):
        input_fields = " ".join(st.session_state.field_strings)
        if input_fields.strip() == "":
            st.write("Please add at least one field before generating the Lagrangian.")
            return
        
        input_fields = input_fields.replace("=", " ")
        input_fields = "[SOS] " + input_fields + " [EOS]"
        generated_lagrangian = generate_lagrangian(input_fields)
        generated_lagrangian = reformat_expression(generated_lagrangian)
        print(generated_lagrangian)
        
        # Attempt to render as LaTeX
        latex_output = master_str_tex(generated_lagrangian[1:])
        st.latex(latex_output)

    st.markdown("### Contact")
    st.markdown("For questions/suggestions, email us: [Eliel](mailto:eliel.camargo-molina@physics.uu.se) or [Yong Sheng](mailto:yongsheng.koay@physics.uu.se).")

if __name__ == "__main__":
    main()