
# This file was generated by decaytreeedit-backend
# Copyright © 2025 Kai Habermann

# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

# The Author(s) shall be credited in any publication or presentation that uses this software, including but not limited to conference presentations, journal articles, and theses. The Authors shall be credited by reference to the decayangle and decayamplitude repositories, as well as the citation of the original paper: https://doi.org/10.1103/PhysRevD.111.056015

# THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 

from decayamplitude.resonance import Resonance, ResonanceDict
from decayamplitude.rotation import QN
from decayamplitude.chain import MultiChain
from decayamplitude.combiner import ChainCombiner
from decayamplitude.backend import numpy as np
from decayangle.decay_topology import Topology, Node
from decayamplitude.particle import Particle, DecaySetup

from decayangle.config import config as decayangle_config
from decayamplitude.kinematics_helpers import mass_from_node
decayangle_config.backend = "jax"
decayangle_config.sorting = "value"  # this makes sure, that the order of entries in a node is handeled internaly and we do not have to worry about it

from jax import jit, grad

def constant_lineshape(L, S, *args):
    #This is a dummy lineshape function, it does nothing, but is needed for the code to work
    # It is suited to model a mother resonance for example
    return 1.0

def BW_lineshape(mass):
    # usually the function is called for a fixed set of points in the phasespace. 
    # Thus we can give the mass values as fixed
    # This allows the compiler to optimize the function better
    # decayamplitude will provide l and s values per default for the resonances, but not the mass values
    def BW(l,s, m0, gamma):
        return 1 / (mass**2 - m0**2 + 1j * mass * gamma) 
    return BW

        
final_state_qn = {
    1: Particle(quantum_numbers=QN(2, -1), name="J/psi(1S)"), # name for orientation: J/psi(1S)
    2: Particle(quantum_numbers=QN(0, -1), name="pi+"), # name for orientation: pi+
    3: Particle(quantum_numbers=QN(0, -1), name="pi-"), # name for orientation: pi-
    4: Particle(quantum_numbers=QN(0, -1), name="pi+"), # name for orientation: pi+
    5: Particle(quantum_numbers=QN(0, -1), name="pi-"), # name for orientation: pi-
}

setup = DecaySetup(final_state_particles=final_state_qn)

topologies = setup.topologies  # This ensures correct symmetrization

def amplitude(momenta):
    # momenta has the form of 
    # {
    #  1: np.array(shape=(...,4)), with indices 0,1,2,3 for p_x,p_y,p_z,E
    #  2: np.array(shape=(...,4)), with indices 0,1,2,3 for p_x,p_y,p_z,E
    #  ...
    #}

    resonances = ResonanceDict({
        0:    [
            Resonance(Node(0), quantum_numbers=QN(0, -1), lineshape=constant_lineshape, argnames=[], preserve_partity=False, name='B0') 
            # Mother particle, set the QN to what they need to be for your decay. Usually a weak initial decay is assumed, so parity is not preserved. If this is not the case set it to True here or in the frontend.
                    # Replace the lineshape with a function of your liking, the call signatrue has to be f(L, S, *args), where the args will be exposed under the names you give in argnames
                    # The mass points at which to evaluate the lineshape are not given at runtime. Look at the bw_lineshpe for details
            ],
        
        (2, 3, 4, 5): [
        Resonance(Node((2, 3, 4, 5)), quantum_numbers=QN(6, -1), lineshape=BW_lineshape(mass_from_node(Node((2, 3, 4, 5)), momenta)), argnames=['rho_316900_mass', 'rho_316900_width'], preserve_partity=True, name='rho_3(1690)0')
        ],

        (4, 5): [
        Resonance(Node((4, 5)), quantum_numbers=QN(2, -1), lineshape=BW_lineshape(mass_from_node(Node((4, 5)), momenta)), argnames=['rho7700_mass', 'rho7700_width'], preserve_partity=True, name='rho(770)0')
        ],

        (2, 3): [
        Resonance(Node((2, 3)), quantum_numbers=QN(2, -1), lineshape=BW_lineshape(mass_from_node(Node((2, 3)), momenta)), argnames=['rho7700_mass', 'rho7700_width'], preserve_partity=True, name='rho(770)0')
        ],

        (1, 2, 3): [
        Resonance(Node((1, 2, 3)), quantum_numbers=QN(2, -1), lineshape=BW_lineshape(mass_from_node(Node((1, 2, 3)), momenta)), argnames=['psi2S_mass', 'psi2S_width'], preserve_partity=True, name='psi(2S)')
        ]
    })

    topologies = setup.filled_topologies(
        resonances=resonances
    )
    print(f"TOPOLOGIES: {topologies}")
    chains = [
        MultiChain(
            topology=topology,
            resonances = resonances,
            momenta = momenta,
            final_state_qn = final_state_qn
        ) for topology in topologies
    ]
    full = ChainCombiner(chains)

    # The unpolarized amplitude is the simplest one, and the default case in LHCb
    unpolarized, argnames = full.unpolarized_amplitude(
        full.generate_couplings() # This is a helper function to generate the couplings for the hadronic system, if you want to restrict them, you will have to do it manually.
                                    # Alternatively you can also restrict the couplings in the fitter later.      
        )
    # argnames are the names of the arguments of the function, which are the masses and widths of the resonances and the couplings
    # The order of argnames is the order of the arguments in the function
    # The function can be called with positional arguments, or with keyword arguments
    # so unpolarized(*[1, 2, 3, 4, 5, 6]) is the same as unpolarized(mass_resonance_1=1, width_resonance_1=2, mass_resonance_2=3, width_resonance_2=4, mass_resonance_3=5, width_resonance_3=6), omitting the couplings
    # print(argnames)
    for name in argnames:
        print(name)

    # an issue with jax, where the internal caching structure needs to be prewarmed, so that in the compilation step the correct types are inferred
    print(unpolarized(*([1.0] * len(argnames))))

    # we can now jit the function, to make it faster after the compile
    unpolarized = jit(unpolarized) 
    print(unpolarized(*([1.0] * len(argnames))))

    # for the gradient calculation we need to define a log likelihood function or something, that produces a single value
    def LL(*args):
        return np.sum(
            np.log(unpolarized(*args))
                )
    # we can calc the gradient of the log likelihood function
    unpolarized_grad = jit(grad(LL, argnums=[i for i in range(len(argnames))]))

    # and a test call (may take quite some time)
    # print(unpolarized_grad(*([1.0] * len(argnames))))

    # Further calls will be much faster, since we dont need to compile again
    # print(unpolarized_grad(*([1.0] * len(argnames))))
    # print(unpolarized_grad(*([2.0] * len(argnames))))


    # Other options for amplitudes, one might be interested in
    # polarized, lambdas ,polarized_argnames = full.polarized_amplitude(full.generate_couplings())
    # print(lambdas)
    # print(polarized(*lambda_values,*([1] * len(polarized_argnames))) )

    # matrix_function, matrix_argnames = full.matrix_function(full.generate_couplings())
    # print(matrix_argnames)
    # print(matrix_function(0, *([1] * len(argnames))) )

    return full

if __name__ == "__main__":
    mother_mass = 1
    daughter_masses = [0.1] * len(final_state_qn)
    from decayangle.kinematics import from_mass_and_momentum
    import random 

    random_angles = [
        (random.uniform(0, 2 * np.pi) , random.uniform(0, np.pi)) 
         for _ in range(len(final_state_qn))
         ]

    three_momenta = [
        np.array([0.1,0,0]) * np.array([np.cos(theta) * np.sin(phi), np.sin(theta) * np.sin(phi), np.cos(phi)])
        for theta, phi in random_angles
    ]

    momenta = {
        i: from_mass_and_momentum(m, p)
            for i, (m, p) in enumerate(zip(daughter_masses, three_momenta), start=1)
    }
    momenta = topologies[0].to_rest_frame(momenta)
    full = amplitude(momenta)