import os
import astrophot as ap
import numpy as np
from utils import make_basic_sersic
import pytest

# torch.autograd.set_detect_anomaly(True)
######################################################################
# Model Objects
######################################################################


def test_model_sampling_modes():

    target = make_basic_sersic(90, 100)
    model = ap.Model(
        name="test_sersic",
        model_type="sersic galaxy model",
        center=[40, 41.9],
        PA=60 * np.pi / 180,
        q=0.8,
        n=0.5,
        Re=20,
        Ie=1,
        target=target,
    )

    # With subpixel integration
    model.integrate_mode = "bright"
    auto = ap.backend.to_numpy(model().data)
    model.sampling_mode = "midpoint"
    midpoint = ap.backend.to_numpy(model().data)
    midpoint_bright = midpoint.copy()
    model.sampling_mode = "simpsons"
    simpsons = ap.backend.to_numpy(model().data)
    model.sampling_mode = "upsample:5"
    upsample5 = ap.backend.to_numpy(model().data)
    model.sampling_mode = "quad:5"
    quad5 = ap.backend.to_numpy(model().data)
    assert np.allclose(midpoint, auto, rtol=1e-2), "Midpoint sampling should match auto sampling"
    assert np.allclose(midpoint, simpsons, rtol=1e-2), "Simpsons sampling should match midpoint"
    assert np.allclose(midpoint, quad5, rtol=1e-2), "Quad5 sampling should match midpoint sampling"
    assert np.allclose(
        midpoint, upsample5, rtol=1e-2
    ), "Upsample5 sampling should match midpoint sampling"
    assert np.allclose(simpsons, quad5, rtol=1e-6), "Quad5 sampling should match Simpsons sampling"

    # Without subpixel integration
    model.integrate_mode = "none"
    auto = ap.backend.to_numpy(model().data)
    model.sampling_mode = "midpoint"
    midpoint = ap.backend.to_numpy(model().data)
    model.sampling_mode = "simpsons"
    simpsons = ap.backend.to_numpy(model().data)
    model.sampling_mode = "upsample:5"
    upsample5 = ap.backend.to_numpy(model().data)
    model.sampling_mode = "quad:5"
    quad5 = ap.backend.to_numpy(model().data)
    assert np.allclose(
        midpoint, midpoint_bright, rtol=1e-2
    ), "no integrate sampling should match bright sampling"
    assert np.allclose(midpoint, auto, rtol=1e-2), "Midpoint sampling should match auto sampling"
    assert np.allclose(midpoint, simpsons, rtol=1e-2), "Simpsons sampling should match midpoint"
    assert np.allclose(midpoint, quad5, rtol=1e-2), "Quad5 sampling should match midpoint sampling"
    assert np.allclose(
        midpoint, upsample5, rtol=1e-2
    ), "Upsample5 sampling should match midpoint sampling"
    assert np.allclose(simpsons, quad5, rtol=1e-6), "Quad5 sampling should match Simpsons sampling"

    # curvature based subpixel integration
    model.integrate_mode = "curvature"
    auto = ap.backend.to_numpy(model().data)
    model.sampling_mode = "midpoint"
    midpoint = ap.backend.to_numpy(model().data)
    model.sampling_mode = "simpsons"
    simpsons = ap.backend.to_numpy(model().data)
    model.sampling_mode = "upsample:5"
    upsample5 = ap.backend.to_numpy(model().data)
    model.sampling_mode = "quad:5"
    quad5 = ap.backend.to_numpy(model().data)
    assert np.allclose(
        midpoint, midpoint_bright, rtol=1e-2
    ), "curvature integrate sampling should match bright sampling"
    assert np.allclose(midpoint, auto, rtol=1e-2), "Midpoint sampling should match auto sampling"
    assert np.allclose(midpoint, simpsons, rtol=1e-2), "Simpsons sampling should match midpoint"
    assert np.allclose(midpoint, quad5, rtol=1e-2), "Quad5 sampling should match midpoint sampling"
    assert np.allclose(
        midpoint, upsample5, rtol=1e-2
    ), "Upsample5 sampling should match midpoint sampling"
    assert np.allclose(simpsons, quad5, rtol=1e-6), "Quad5 sampling should match Simpsons sampling"

    with pytest.raises(ap.errors.SpecificationConflict):
        model.integrate_mode = "should raise"
    model.integrate_mode = "none"
    with pytest.raises(ap.errors.SpecificationConflict):
        model.sampling_mode = "should raise"
    model.sampling_mode = "midpoint"
    model.integrate_mode = "none"

    # test PSF modes
    model.psf = np.array([[0.05, 0.1, 0.05], [0.1, 0.4, 0.1], [0.05, 0.1, 0.05]])
    model.psf_convolve = True
    model()


def test_model_errors():

    # Target that is not a target image
    arr = np.zeros((10, 15))
    target = ap.image.Image(data=arr, pixelscale=1.0, zeropoint=1.0)

    with pytest.raises(ap.errors.InvalidTarget):
        ap.Model(
            name="test_model",
            model_type="sersic galaxy model",
            target=target,
        )

    # model that doesn't exist
    target = make_basic_sersic()
    with pytest.raises(ap.errors.UnrecognizedModel):
        ap.Model(
            name="test_model",
            model_type="sersic gaaxy model",
            target=target,
        )


@pytest.mark.parametrize(
    "model_type", ap.models.ComponentModel.List_Models(usable=True, types=True)
)
def test_all_model_sample(model_type):

    if ap.backend.backend == "jax" and np.random.randint(0, 3) > 0:
        pytest.skip("JAX is very slow, randomly reducing the number of tests")
    if model_type == "isothermal sech2 edgeon model" and ap.backend.backend == "jax":
        pytest.skip("JAX doesnt have bessel function k1 yet")

    if (
        model_type in ["ferrer warp galaxy model", "king warp galaxy model"]
        and ap.backend.backend == "jax"
    ):
        pytest.skip("JAX version doesnt support these models yet, difficulty with gradients")

    if any(t in model_type for t in ["warp", "fourier"]):
        pytest.skip("Warp and Fourier models are complex and slow to fit, skipping for now")
    if model_type.startswith("truncated"):
        pytest.skip("Testing truncated models is redundant")

    target = make_basic_sersic()
    target.zeropoint = 22.5
    MODEL = ap.Model(
        name="test_model",
        model_type=model_type,
        target=target,
        integrate_mode=(
            "none" if ap.backend.backend == "jax" else "bright"
        ),  # JAX JIT is reallly slow for any integration
    )
    MODEL.initialize()
    MODEL.to()
    for P in MODEL.dynamic_params:
        assert (
            P.value is not None
        ), f"Model type {model_type} parameter {P.name} should not be None after initialization"
    img = MODEL()
    assert ap.backend.all(
        ap.backend.isfinite(img.data)
    ), "Model should evaluate a real number for the full image"

    res = ap.fit.LM(MODEL, max_iter=5, verbose=1).fit()

    print(MODEL)  # test printing

    # sky has little freedom to fit, some more complex models need extra
    # attention to get a good fit so here we just check that they can improve
    if (
        "sky" in model_type
        or "king" in model_type
        or "spline" in model_type
        or model_type
        in [
            "exponential warp galaxy model",
            "ferrer warp galaxy model",
            "ferrer ray galaxy model",
            "isothermal sech2 edgeon model",
        ]
    ):
        assert res.loss_history[0] > res.loss_history[-1], (
            f"Model {model_type} should fit to the target image, but did not. "
            f"Initial loss: {res.loss_history[0]}, Final loss: {res.loss_history[-1]}"
        )
    else:  # Most models should get significantly better after just a few iterations
        assert res.loss_history[0] > (1.5 * res.loss_history[-1]), (
            f"Model {model_type} should fit to the target image, but did not. "
            f"Initial loss: {res.loss_history[0]}, Final loss: {res.loss_history[-1]}"
        )

    F = MODEL.total_flux()
    assert ap.backend.isfinite(F), "Model total flux should be finite after fitting"
    assert F > 0, "Model total flux should be positive after fitting"
    U = MODEL.total_flux_uncertainty()
    assert ap.backend.isfinite(U), "Model total flux uncertainty should be finite after fitting"
    assert U >= 0, "Model total flux uncertainty should be non-negative after fitting"
    M = MODEL.total_magnitude()
    assert ap.backend.isfinite(M), "Model total magnitude should be finite after fitting"
    U_M = MODEL.total_magnitude_uncertainty()
    assert ap.backend.isfinite(
        U_M
    ), "Model total magnitude uncertainty should be finite after fitting"
    assert U_M >= 0, "Model total magnitude uncertainty should be non-negative after fitting"

    allnames = set()
    for name in MODEL.build_params_array_names():
        assert name not in allnames, f"Duplicate parameter name found: {name}"
        allnames.add(name)


def test_sersic_save_load():

    target = make_basic_sersic()
    model = ap.Model(
        name="test_sersic",
        model_type="sersic galaxy model",
        center=[20, 20],
        PA=60 * np.pi / 180,
        q=0.5,
        n=2,
        Re=5,
        Ie=1,
        target=target,
    )

    model.initialize()
    model.save_state("test_AstroPhot_sersic.hdf5", appendable=True)
    model.center = [30, 30]
    model.PA = 30 * np.pi / 180
    model.q = 0.8
    model.n = 3
    model.Re = 10
    model.Ie = 2
    target.crtan = [1.0, 2.0]
    model.append_state("test_AstroPhot_sersic.hdf5")
    model.load_state("test_AstroPhot_sersic.hdf5", index=0)

    assert model.center.value[0].item() == 20, "Model center should be loaded correctly"
    assert model.center.value[1].item() == 20, "Model center should be loaded correctly"
    assert model.PA.value.item() == 60 * np.pi / 180, "Model PA should be loaded correctly"
    assert model.q.value.item() == 0.5, "Model q should be loaded correctly"
    assert model.n.value.item() == 2, "Model n should be loaded correctly"
    assert model.Re.value.item() == 5, "Model Re should be loaded correctly"
    assert model.Ie.value.item() == 1, "Model Ie should be loaded correctly"
    assert model.target.crtan.value[0] == 0.0, "Model target crtan should be loaded correctly"
    assert model.target.crtan.value[1] == 0.0, "Model target crtan should be loaded correctly"

    os.remove("test_AstroPhot_sersic.hdf5")


def test_batch_model(sersic):
    M = ap.Model(model_type="batch model", model=sersic)
    M.initialize()
    # Now any parameters may be given an extra dimension for the batch
    sersic.center = [[10, 10], [20, 70], [50, 50], [70, 30], [80, 80]]
    sersic.q = [0.7, 0.6, 0.5, 0.4, 0.3]
    sersic.PA = np.array([30, 60, 90, 120, 150]) * np.pi / 180

    img = M()
    assert isinstance(img, ap.image.ModelImage), "Output of batch model should be a ModelImage"
    assert img.data.shape == sersic.target.data.shape, "batched model should produce regular output"

    M.set_values(ap.backend.ones_like(M.get_values()), attribute="uncertainty")

    assert M.total_flux() > 0, "Total flux of batch model should be positive"
    assert (
        M.total_flux_uncertainty() > 0
    ), "Total flux uncertainty of batch model should be positive"


def test_batch_scene_model(sersic):
    target_batch = ap.TargetImageBatch(
        images=[
            ap.TargetImage(
                data=np.zeros((52, 50)),
                CD=ap.utils.initialize.R(np.pi / 3),
                crtan=(10, 0),
                name="image1",
            ),
            ap.TargetImage(
                data=np.zeros((52, 50)), CD=ap.utils.initialize.R(np.pi / 6), name="image2"
            ),
            ap.TargetImage(
                data=np.zeros((52, 50)),
                CD=ap.utils.initialize.R(np.pi / 16),
                crtan=(-15, 0),
                name="image3",
            ),
        ]
    )
    model = ap.Model(model_type="batch scene model", model=sersic, target=target_batch)
    model.initialize()
    sersic.Ie = [10, 20, 30]
    img = model()
    assert isinstance(
        img, ap.image.ModelImageBatch
    ), "Output of batch scene model should be a ModelImageBatch"
    assert img.data.shape[0] == 3, "Batch scene model should produce output for each target image"
    jac = model.jacobian()
    assert isinstance(
        jac, ap.image.JacobianImageBatch
    ), "Jacobian of batch scene model should be a JacobianImageBatch"
    assert (
        jac.data.shape[0] == 3
    ), "Batch scene model jacobian should produce output for each target image"
