from functools import partial

import jax.numpy as jnp
import numpy as np
import pytest
from huggingface_hub import hf_hub_download
from numpyro.distributions import constraints

from scarlet2 import init
from scarlet2.frame import Frame
from scarlet2.module import Parameter, Parameters, relative_step
from scarlet2.observation import Observation
from scarlet2.psf import ArrayPSF
from scarlet2.scene import Scene
from scarlet2.source import PointSource, Source


@pytest.fixture()
def data_file():
    """Download and load a realistic test file. This is the same data used in the
    quickstart notebook. The data will be manipulated to create invalid inputs for
    the `bad_obs` fixture."""
    filename = hf_hub_download(
        repo_id="astro-data-lab/scarlet-test-data", filename="hsc_cosmos_35.npz", repo_type="dataset"
    )
    return jnp.load(filename)


@pytest.fixture()
def good_obs(data_file):
    """Create an observation that should pass all validation checks."""
    data = jnp.asarray(data_file["images"])
    channels = [str(f) for f in data_file["filters"]]
    weights = jnp.asarray(1 / data_file["variance"])
    psf = jnp.asarray(data_file["psfs"])

    return Observation(
        data=data,
        weights=weights,
        channels=channels,
        psf=ArrayPSF(psf),
    )


@pytest.fixture()
def bad_obs(data_file):
    """Create an observation that should fail multiple validation checks."""

    data = np.asarray(data_file["images"])
    channels = [str(f) for f in data_file["filters"]]
    weights = np.asarray(1 / data_file["variance"])
    psf = np.asarray(data_file["psfs"])

    weights = weights[:-1]  # Remove the last weight to create a mismatch in dimensions
    weights[0][0] = np.inf  # Set one weight to infinity
    weights[1][0] = -1.0  # Set one weight to a negative value
    psf = psf[:-1]  # Remove the last PSF to create a mismatch in dimensions
    psf = psf[0] + 0.001

    return Observation(
        data=data,
        weights=weights,
        channels=channels,
        psf=ArrayPSF(psf),
    )


@pytest.fixture()
def good_source(good_obs, data_file):
    """Assemble a source from the good observation and the data file."""
    model_frame = Frame.from_observations(good_obs)
    with Scene(model_frame) as _:
        centers = jnp.array([(src["y"], src["x"]) for src in data_file["catalog"]])  # Note: y/x convention!
        spectrum, morph = init.from_gaussian_moments(good_obs, centers[0], min_corr=0.99)
        return Source(centers[0], spectrum, morph)


@pytest.fixture()
def bad_source(good_obs, data_file):
    """Assemble a source from the bad observation and the data file."""
    model_frame = Frame.from_observations(good_obs)
    with Scene(model_frame) as _:
        centers = jnp.array([(src["y"], src["x"]) for src in data_file["catalog"]])  # Note: y/x convention!
        spectrum, morph = init.from_gaussian_moments(good_obs, centers[0], min_corr=0.99)
        return Source(centers[0], spectrum * -1, morph)


@pytest.fixture()
def scene(good_obs, data_file):
    """Assemble a scene from the good observation and the data file."""
    model_frame = Frame.from_observations(good_obs)
    centers = jnp.array([(src["y"], src["x"]) for src in data_file["catalog"]])  # Note: y/x convention!

    with Scene(model_frame) as scene:
        for i, center in enumerate(centers):
            if i == 0:  # we know source 0 is a star
                spectrum = init.pixel_spectrum(good_obs, center, correct_psf=True)
                PointSource(center, spectrum)
            else:
                try:
                    spectrum, morph = init.from_gaussian_moments(good_obs, center, min_corr=0.99)
                except ValueError:
                    spectrum = init.pixel_spectrum(good_obs, center)
                    morph = init.compact_morphology()

                Source(center, spectrum, morph)

    # Create parameters for the scene
    spec_step = partial(relative_step, factor=0.05)
    morph_step = partial(relative_step, factor=1e-3)
    with Parameters(scene):
        for i in range(len(scene.sources)):
            Parameter(
                scene.sources[i].spectrum,
                name=f"spectrum:{i}",
                constraint=constraints.positive,
                stepsize=spec_step,
            )
    # context can be reopened
    with scene.parameters:
        for i in range(len(scene.sources)):
            if i == 0:
                Parameter(scene.sources[i].center, name=f"center:{i}", stepsize=0.1)
            else:
                Parameter(
                    scene.sources[i].morphology,
                    name=f"morph:{i}",
                    constraint=constraints.unit_interval,
                    stepsize=morph_step,
                )
    return scene
