#  Copyright (c) 2024 zfit
import numpy as np
import pytest

import zfit
import zfit.z.numpy as znp
from zfit import supports


def createpdf1(obs, params):
    params = params.copy()
    yieldp = params.pop('yield')
    class SomeGradsPDF(zfit.pdf.BasePDF):

        def __init__(self, obs, a, b, c, d, e, f, name="SomeGradsPDF", **kwargs):
            params = {'a': a, 'b': b, 'c': c, 'd': d, 'e':e, 'f':f}
            super().__init__(obs=obs, params=params, name=name, autograd_params=['a', 'b', 'yield', 'e'], **kwargs)

        @supports(norm=False)
        def _pdf(self, x, norm, *, params=None):
            a = params['a']
            b = params['b']
            c = params['c']
            d = params['d']
            e = params['e']
            f = params['f']
            x = x[0]
            return (znp.exp(-a * x) + znp.exp(-b * x) + c * x + d * x + e * x + f * x) ** 2
    return SomeGradsPDF(obs=obs, **params, extended=yieldp)

def createpdfbinned1(obs, params):
    obs = obs.with_binning(43)
    pdf = createpdf1(obs, params)
    return pdf

def createsum1(obs, params):
    pdf1 = createpdf1(obs, params)
    pdf2 = zfit.pdf.Gauss(obs=obs, mu=params['a'], sigma=params['b'], extended=params['yield'])
    return zfit.pdf.SumPDF([pdf1, pdf2], fracs=params['c'], extended=params['yield'])

def createbinnedsum1(obs, params):
    obs = obs.with_binning(12)
    pdf = createsum1(obs, params)
    return pdf

def createbinnedsum2(obs, params):
    obs = obs.with_binning(12)
    pdf1 = createpdf1(obs, params)
    pdf2 = zfit.pdf.Gauss(obs=obs, mu=params['a'], sigma=params['b'], extended=params['yield'])
    return zfit.pdf.BinnedSumPDF([pdf1, pdf2], fracs=params['c'], extended=params['yield'])

def createnll1(obs,params):
    pdf = createpdf1(obs,params)
    data = pdf.sample(1000)
    nll = zfit.loss.ExtendedUnbinnedNLL(model=pdf, data=data)
    return nll

def createbinnednll1(obs,params):
    obs = obs.with_binning(12)
    pdf = createpdf1(obs,params)
    data = pdf.sample(1000)
    nll = zfit.loss.ExtendedBinnedNLL(model=pdf, data=data)
    return nll


@pytest.mark.parametrize('paramdepfactory', [
    createpdf1,
    createpdfbinned1,
    createsum1,
    createbinnedsum1,
    createbinnedsum2,
    createnll1,
    createbinnednll1

])
def test_nograd_params(paramdepfactory):


    parama = zfit.Parameter("param_a", 1.)
    paramb = zfit.Parameter("param_b", 2., floating=False)
    yieldp = zfit.Parameter("yieldp", 3.)
    parameindep = zfit.Parameter("param_eindep", 6., floating=True)
    pgrad = [parama, paramb,
             yieldp, parameindep]
    parame = zfit.ComposedParameter("param_e", lambda params: params[0] + params[1], params=[pgrad[0], pgrad[-1]])
    pgrad.append(parame)
    paramc = zfit.Parameter("param_c", 0.4, floating=True)
    paramd = zfit.Parameter("param_d", 5., floating=False)
    paramfindep = zfit.Parameter("param_findep", 3., floating=True)
    pnograd = [paramc, paramd,paramfindep]
    paramf = zfit.ComposedParameter("param_f", lambda params: params[0] + params[1],
                                       params=[pnograd[0], pnograd[-1]])
    pnograd.append(paramf)
    obs = zfit.Space('obs1', -1, 5)
    pdf = paramdepfactory(obs, {'a': parama, 'b': paramb, 'c': paramc, 'd': paramd, 'e': parame, 'f': paramf, 'yield': yieldp})
    # yieldp = pdf.get_yield()


    pgrad = set(pgrad)
    pnograd = set(pnograd)
    params = pdf.get_params(is_yield=None, autograd=None, floating=None, extract_independent=None)
    compset = set()  # add a frac composed parameter if autogenerated, supports gradients
    for p in params:
        if p not in pgrad and p not in pnograd:
            if p.name == 'Composed_autoparam_0':  # add the fraction parameter that is autocombosed
                compset = {p}
                pgrad |= compset
    assert pgrad == set(pdf.get_params(autograd=True, floating=None, extract_independent=None))
    assert pnograd == set(pdf.get_params(autograd=False, floating=None, extract_independent=None))
    assert pgrad | pnograd == params
    assert pgrad & pnograd == set()
    assert {yieldp} == pdf.get_params(is_yield=True)
    assert {yieldp} == pdf.get_params(is_yield=True, autograd=True)
    assert set() == pdf.get_params(is_yield=True, autograd=False)
    assert {parama, paramb, yieldp, parameindep, parame} | compset == set(pdf.get_params(is_yield=None, autograd=True, floating=None, extract_independent=None))
    assert {parama, paramb, yieldp, parameindep} == set(pdf.get_params(is_yield=None, autograd=True, floating=None, extract_independent=True))
    assert {parama, yieldp, parameindep} == set(pdf.get_params(is_yield=None, autograd=True, floating=True, extract_independent=True))
    assert {parama, paramb, yieldp, parame} | compset == set(pdf.get_params(is_yield=None, autograd=True, floating=None, extract_independent=False))

    assert {paramc, paramd, paramfindep} == set(pdf.get_params(is_yield=None, autograd=False, floating=None, extract_independent=True))
    assert {paramc, paramd, paramf} == set(pdf.get_params(is_yield=None, autograd=False, floating=None, extract_independent=False))
    assert {paramc, paramd, paramfindep, paramf} == set(pdf.get_params(is_yield=None, autograd=False, floating=None, extract_independent=None))
    assert set() == set(pdf.get_params(is_yield=True, autograd=False, floating=None, extract_independent=None))
    assert {paramc, paramfindep} == set(pdf.get_params(is_yield=None, autograd=False, floating=True, extract_independent=True))
    assert {paramd} == set(pdf.get_params(is_yield=None, autograd=False, floating=False, extract_independent=True))

@pytest.mark.parametrize('lossfactory', [createbinnednll1, createnll1])
def test_autograd_nogradparams_raises(lossfactory):
    parama = zfit.Parameter("param_a", 1.)
    paramb = zfit.Parameter("param_b", 2., floating=False)
    yieldp = zfit.Parameter("yieldp", 1000.)
    parameindep = zfit.Parameter("param_eindep", 6., floating=True)
    pgrad = [parama, paramb,
             yieldp, parameindep]
    parame = zfit.ComposedParameter("param_e", lambda params: params[0] * params[1], params=[pgrad[0], parameindep])
    paramc = zfit.Parameter("param_c", 0.4, floating=True)
    paramd = zfit.Parameter("param_d", 5., floating=False)
    paramfindep = zfit.Parameter("param_findep", 3., floating=True)
    pnograd = [paramc, paramd,paramfindep]
    paramf = zfit.ComposedParameter("param_f", lambda params: params[0] + params[1],
                                       params=[pnograd[0], paramfindep])
    obs = zfit.Space('obs1', -1, 5)
    loss = lossfactory(obs, {'a': parama, 'b': paramb, 'c': paramc, 'd': paramd, 'e': parame, 'f': paramf, 'yield': yieldp})
    value = loss.value()
    assert not np.isnan(value), "Test is not valid if value is nan"
    with pytest.raises(zfit.exception.AutogradNotSupported):
        loss.gradient()
    with pytest.raises(zfit.exception.AutogradNotSupported):
        loss.gradient(params=pnograd)
    with pytest.raises(zfit.exception.AutogradNotSupported):
        loss.gradient(params=pnograd[0])
    _ = loss.gradient(params=pnograd, numgrad=True)
    _ = loss.gradient(numgrad=True)
    gradient1 = loss.gradient(params=pgrad, numgrad=True)
    gradient2 = loss.gradient(params=pgrad, numgrad=False)
    assert np.allclose(gradient1, gradient2, rtol=1e-5, atol=1e-5)

    hessian1 = loss.hessian(params=pgrad, numgrad=True)
    hessian2 = loss.hessian(params=pgrad, numgrad=False)
    np.testing.assert_allclose(hessian1, hessian2, atol=1e-4, rtol=1e10)

    with pytest.raises(zfit.exception.AutogradNotSupported):
        _ = loss.hessian(params=pnograd, numgrad=False)
    hessian3 = loss.hessian(params=pnograd, numgrad=True)
    assert not np.any(np.isnan(hessian3)), f"NaN in hessian: {hessian3}"
    with pytest.raises(zfit.exception.AutogradNotSupported):
        _ = loss.hessian(params=loss.get_params(), numgrad=False)
